Fix if codegen bugs

compiler/erg_common/opcode.rs

@@ -149,6 +149,10 @@ impl CommonOpcode {
     pub const fn take_arg(&self) -> bool {
         90 <= (*self as u8) && (*self as u8) < 220
     }
+
+    pub fn is_jump_op(op: u8) -> bool {
+        [93, 110, 111, 112, 113, 114, 115, 140, 143, 175, 176].contains(&op)
+    }
 }
 
 impl_u8_enum! {CompareOp;

compiler/erg_compiler/codegen.rs

@@ -303,6 +303,9 @@ impl PyCodeGenerator {
         } else {
             jump_to
         };
+        if !CommonOpcode::is_jump_op(*self.cur_block_codeobj().code.get(idx - 1).unwrap()) {
+            self.crash(&format!("calc_edit_jump: not jump op: {idx} {jump_to}"));
+        }
         self.edit_code(idx, arg);
     }
 
@@ -1509,6 +1512,7 @@ impl PyCodeGenerator {
         let cond = args.remove(0);
         self.emit_expr(cond);
         let idx_pop_jump_if_false = self.lasti();
+        // Opcode310::POP_JUMP_IF_FALSE == Opcode311::POP_JUMP_FORWARD_IF_FALSE
         self.write_instr(Opcode310::POP_JUMP_IF_FALSE);
         // cannot detect where to jump to at this moment, so put as 0
         self.write_arg(0);
@@ -1523,10 +1527,15 @@ impl PyCodeGenerator {
             }
         }
         if args.get(0).is_some() {
+            let idx_jump_forward = self.lasti();
             self.write_instr(JUMP_FORWARD); // jump to end
             self.write_arg(0);
             // else block
-            let idx_else_begin = self.lasti();
+            let idx_else_begin = if self.py_version.minor >= Some(11) {
+                self.lasti() - idx_pop_jump_if_false - 2
+            } else {
+                self.lasti()
+            };
             self.calc_edit_jump(idx_pop_jump_if_false + 1, idx_else_begin);
             match args.remove(0) {
                 Expr::Lambda(lambda) => {
@@ -1537,7 +1546,6 @@ impl PyCodeGenerator {
                     self.emit_expr(other);
                 }
             }
-            let idx_jump_forward = idx_else_begin - 2;
             let idx_end = self.lasti();
             self.calc_edit_jump(idx_jump_forward + 1, idx_end - idx_jump_forward - 2);
             // FIXME: this is a hack to make sure the stack is balanced
@@ -1545,6 +1553,8 @@ impl PyCodeGenerator {
                 self.stack_dec();
             }
         } else {
+            self.write_instr(JUMP_FORWARD);
+            self.write_arg(1);
             // no else block
             let idx_end = if self.py_version.minor >= Some(11) {
                 self.lasti() - idx_pop_jump_if_false - 1

compiler/erg_compiler/context/compare.rs

@@ -452,7 +452,7 @@ impl Context {
             (Type, Poly { name, params }) | (Poly { name, params }, Type)
                 if &name[..] == "Tuple" =>
             {
-                let tps = Vec::try_from(params[0].clone()).unwrap();
+                if let Ok(tps) = Vec::try_from(params[0].clone()) {
                     for tp in tps {
                         let Ok(t) = self.convert_tp_into_ty(tp) else {
                             return false;
@@ -461,6 +461,7 @@ impl Context {
                             return false;
                         }
                     }
+                }
                 false
             }
             (Type, Poly { name, params }) | (Poly { name, params }, Type)

compiler/erg_compiler/context/eval.rs

@@ -446,33 +446,20 @@ impl Context {
             self.py_mod_cache.clone(),
             self.clone(),
         );
-        let return_t = lambda_ctx.eval_const_block(&lambda.body)?;
-        // FIXME: lambda: i: Int -> Int
-        // => sig_t: (i: Type) -> Type
-        // => as_type: (i: Int) -> Int
+        let return_t = v_enum(set! {lambda_ctx.eval_const_block(&lambda.body)?});
         let sig_t = subr_t(
             SubrKind::from(lambda.op.kind),
             non_default_params.clone(),
-            var_params.clone(),
+            var_params,
             default_params.clone(),
-            v_enum(set![return_t.clone()]),
+            return_t,
         );
         let sig_t = self.generalize_t(sig_t);
-        let as_type = subr_t(
-            SubrKind::from(lambda.op.kind),
-            non_default_params,
-            var_params,
-            default_params,
-            // TODO: unwrap
-            return_t.as_type().unwrap().into_typ(),
-        );
-        let as_type = self.generalize_t(as_type);
         let subr = ConstSubr::User(UserConstSubr::new(
             Str::ever("<lambda>"),
             lambda.sig.params.clone(),
             lambda.body.clone(),
             sig_t,
-            Some(as_type),
         ));
         Ok(ValueObj::Subr(subr))
     }

compiler/erg_compiler/ty/mod.rs

@@ -16,6 +16,8 @@ use std::path::PathBuf;
 
 use constructors::dict_t;
 use erg_common::dict::Dict;
+#[allow(unused_imports)]
+use erg_common::log;
 use erg_common::set::Set;
 use erg_common::traits::LimitedDisplay;
 use erg_common::vis::Field;
@@ -24,7 +26,7 @@ use erg_common::{enum_unwrap, fmt_option, fmt_set_split_with, set, Str};
 use erg_parser::ast::{Block, Params};
 use erg_parser::token::TokenKind;
 
-use self::constructors::{int_interval, mono};
+use self::constructors::{int_interval, mono, subr_t};
 use self::free::{
     fresh_varname, CanbeFree, Constraint, Free, FreeKind, FreeTyVar, HasLevel, Level, GENERIC_LEVEL,
 };
@@ -141,23 +143,15 @@ pub struct UserConstSubr {
     params: Params,
     block: Block,
     sig_t: Type,
-    as_type: Option<Type>,
 }
 
 impl UserConstSubr {
-    pub const fn new(
-        name: Str,
-        params: Params,
-        block: Block,
-        sig_t: Type,
-        as_type: Option<Type>,
-    ) -> Self {
+    pub const fn new(name: Str, params: Params, block: Block, sig_t: Type) -> Self {
         Self {
             name,
             params,
             block,
             sig_t,
-            as_type,
         }
     }
 }
@@ -265,10 +259,30 @@ impl ConstSubr {
         }
     }
 
-    pub fn as_type(&self) -> Option<&Type> {
+    /// ConstSubr{sig_t: Int -> {Int}, ..}.as_type() == Int -> Int
+    pub fn as_type(&self) -> Option<Type> {
         match self {
-            ConstSubr::User(user) => user.as_type.as_ref(),
-            ConstSubr::Builtin(builtin) => builtin.as_type.as_ref(),
+            ConstSubr::User(user) => {
+                let Type::Subr(subr) = &user.sig_t else { return None };
+                if let Type::Refinement(refine) = subr.return_t.as_ref() {
+                    if refine.preds.len() == 1 {
+                        let pred = refine.preds.iter().next().unwrap().clone();
+                        if let Predicate::Equal { rhs, .. } = pred {
+                            let return_t = Type::try_from(rhs).ok()?;
+                            let var_params = subr.var_params.as_ref().map(|t| t.as_ref());
+                            return Some(subr_t(
+                                subr.kind,
+                                subr.non_default_params.clone(),
+                                var_params.cloned(),
+                                subr.default_params.clone(),
+                                return_t,
+                            ));
+                        }
+                    }
+                }
+                None
+            }
+            ConstSubr::Builtin(builtin) => builtin.as_type.clone(),
         }
     }
 }

compiler/erg_compiler/ty/typaram.rs

@@ -496,6 +496,21 @@ impl<'a> TryFrom<&'a TyParam> for &'a Type {
     }
 }
 
+impl TryFrom<TyParam> for Type {
+    type Error = ();
+    fn try_from(tp: TyParam) -> Result<Type, ()> {
+        match tp {
+            TyParam::FreeVar(fv) if fv.is_linked() => {
+                Type::try_from(fv.forced_as_ref().linked().unwrap().clone()).map_err(|_| ())
+            }
+            TyParam::Type(t) => Ok(*t),
+            TyParam::Value(v) => Type::try_from(v),
+            // TODO: Array, Dict, Set
+            _ => Err(()),
+        }
+    }
+}
+
 impl HasLevel for TyParam {
     fn level(&self) -> Option<Level> {
         match self {

compiler/erg_compiler/ty/value.rs

@@ -1139,7 +1139,7 @@ impl ValueObj {
                 }
                 Some(TypeObj::Builtin(Type::Record(attr_ts)))
             }
-            Self::Subr(subr) => Some(TypeObj::Builtin(subr.as_type().unwrap().clone())),
+            Self::Subr(subr) => subr.as_type().map(TypeObj::Builtin),
             Self::Array(_) | Self::Tuple(_) | Self::Dict(_) => todo!(),
             _other => None,
         }

compiler/erg_parser/parse.rs

@@ -706,11 +706,6 @@ impl Parser {
         debug_call_info!(self);
         match self.peek() {
             Some(t) if t.is(Symbol) => {
-                if &t.inspect()[..] == "do" || &t.inspect()[..] == "do!" {
-                    let lambda = self.try_reduce_do_block().map_err(|_| self.stack_dec())?;
-                    self.level -= 1;
-                    return Ok(PosOrKwArg::Pos(PosArg::new(Expr::Lambda(lambda))));
-                }
                 if self.nth_is(1, Walrus) {
                     let acc = self.try_reduce_acc_lhs().map_err(|_| self.stack_dec())?;
                     debug_power_assert!(self.cur_is(Walrus));
@@ -1410,6 +1405,11 @@ impl Parser {
     fn try_reduce_bin_lhs(&mut self, in_type_args: bool, in_brace: bool) -> ParseResult<Expr> {
         debug_call_info!(self);
         match self.peek() {
+            Some(t) if &t.inspect()[..] == "do" || &t.inspect()[..] == "do!" => {
+                let lambda = self.try_reduce_do_block().map_err(|_| self.stack_dec())?;
+                self.level -= 1;
+                Ok(Expr::Lambda(lambda))
+            }
             Some(t) if t.category_is(TC::Literal) => {
                 // TODO: 10.times ...などメソッド呼び出しもある
                 let lit = self.try_reduce_lit().map_err(|_| self.stack_dec())?;

tests/should_ok/if.er

@@ -0,0 +1,16 @@
+a = if True:
+    do 1
+    do 2
+b = if False:
+    do 1
+    do 2
+assert a == 1
+assert b == 2
+
+c = if True:
+    do 1
+d = if False:
+    do 1
+
+assert c == 1
+assert d == None

tests/test.rs

@@ -54,6 +54,11 @@ fn exec_helloworld() -> Result<(), ()> {
     }
 }
 
+#[test]
+fn exec_if() -> Result<(), ()> {
+    expect_success("tests/should_ok/if.er")
+}
+
 #[test]
 fn exec_impl() -> Result<(), ()> {
     expect_success("examples/impl.er")
@@ -89,6 +94,11 @@ fn exec_raw_ident() -> Result<(), ()> {
     expect_success("examples/raw_ident.er")
 }
 
+#[test]
+fn exec_rec() -> Result<(), ()> {
+    expect_success("tests/should_ok/rec.er")
+}
+
 #[test]
 fn exec_record() -> Result<(), ()> {
     expect_success("examples/record.er")
@@ -148,12 +158,6 @@ fn exec_pyimport() -> Result<(), ()> {
     }
 }
 
-#[test]
-fn exec_rec() -> Result<(), ()> {
-    // this script is valid but the current code generating process has a bug.
-    expect_end_with("tests/should_err/rec.er", 1)
-}
-
 #[test]
 fn exec_set() -> Result<(), ()> {
     expect_failure("examples/set.er", 1)