Fix codegen bug

compiler/erg_common/python_util.rs

@@ -6,7 +6,7 @@ use std::process::Command;
 use crate::serialize::get_magic_num_from_bytes;
 
 #[cfg(unix)]
-pub const BUILTIN_PYTHON_MODS: [&str; 17] = [
+pub const BUILTIN_PYTHON_MODS: [&str; 18] = [
     "glob",
     "http",
     "importlib",
@@ -17,6 +17,7 @@ pub const BUILTIN_PYTHON_MODS: [&str; 17] = [
     "posix",
     "random",
     "re",
+    "shutil",
     "socket",
     "subprocess",
     "sys",
@@ -26,7 +27,7 @@ pub const BUILTIN_PYTHON_MODS: [&str; 17] = [
     "zipfile",
 ];
 #[cfg(not(unix))]
-pub const BUILTIN_PYTHON_MODS: [&str; 16] = [
+pub const BUILTIN_PYTHON_MODS: [&str; 17] = [
     "glob",
     "http",
     "importlib",
@@ -36,6 +37,7 @@ pub const BUILTIN_PYTHON_MODS: [&str; 16] = [
     "os",
     "random",
     "re",
+    "shutil",
     "socket",
     "subprocess",
     "sys",

compiler/erg_compiler/codegen.rs

@@ -947,9 +947,17 @@ impl CodeGenerator {
         match &bin.op.kind {
             // l..<r == range(l, r)
             TokenKind::RightOpen => {
-                self.emit_load_name_instr(Identifier::public("range"));
+                self.emit_load_name_instr(Identifier::public("RightOpenRange"));
+            }
+            TokenKind::LeftOpen => {
+                self.emit_load_name_instr(Identifier::public("LeftOpenRange"));
+            }
+            TokenKind::Closed => {
+                self.emit_load_name_instr(Identifier::public("ClosedRange"));
+            }
+            TokenKind::Open => {
+                self.emit_load_name_instr(Identifier::public("OpenRange"));
             }
-            TokenKind::LeftOpen | TokenKind::Closed | TokenKind::Open => todo!(),
             TokenKind::InOp => {
                 // if no-std, always `x in y == True`
                 if self.cfg.no_std {
@@ -1056,6 +1064,7 @@ impl CodeGenerator {
 
     fn emit_if_instr(&mut self, mut args: Args) {
         log!(info "entered {}", fn_name!());
+        let init_stack_len = self.cur_block().stack_len;
         let cond = args.remove(0);
         self.emit_expr(cond);
         let idx_pop_jump_if_false = self.cur_block().lasti;
@@ -1065,8 +1074,8 @@ impl CodeGenerator {
         match args.remove(0) {
             // then block
             Expr::Lambda(lambda) => {
-                let params = self.gen_param_names(&lambda.params);
+                // let params = self.gen_param_names(&lambda.params);
-                self.emit_frameless_block(lambda.body, params);
+                self.emit_frameless_block(lambda.body, vec![]);
             }
             other => {
                 self.emit_expr(other);
@@ -1080,8 +1089,8 @@ impl CodeGenerator {
             self.edit_code(idx_pop_jump_if_false + 1, idx_else_begin / 2);
             match args.remove(0) {
                 Expr::Lambda(lambda) => {
-                    let params = self.gen_param_names(&lambda.params);
+                    // let params = self.gen_param_names(&lambda.params);
-                    self.emit_frameless_block(lambda.body, params);
+                    self.emit_frameless_block(lambda.body, vec![]);
                 }
                 other => {
                     self.emit_expr(other);
@@ -1090,15 +1099,18 @@ impl CodeGenerator {
             let idx_jump_forward = idx_else_begin - 2;
             let idx_end = self.cur_block().lasti;
             self.edit_code(idx_jump_forward + 1, (idx_end - idx_jump_forward - 2) / 2);
+            // FIXME: this is a hack to make sure the stack is balanced
+            while self.cur_block().stack_len != init_stack_len + 1 {
                 self.stack_dec();
-            // self.stack_dec();
+            }
         } else {
             // no else block
             let idx_end = self.cur_block().lasti;
             self.edit_code(idx_pop_jump_if_false + 1, idx_end / 2);
             self.emit_load_const(ValueObj::None);
+            while self.cur_block().stack_len != init_stack_len + 1 {
                 self.stack_dec();
-            self.stack_dec();
+            }
         }
     }
 
@@ -1110,18 +1122,24 @@ impl CodeGenerator {
         self.write_arg(0);
         let idx_for_iter = self.cur_block().lasti;
         self.write_instr(FOR_ITER);
+        self.stack_inc();
         // FOR_ITER pushes a value onto the stack, but we can't know how many
         // but after executing this instruction, stack_len should be 1
         // cannot detect where to jump to at this moment, so put as 0
         self.write_arg(0);
         let lambda = enum_unwrap!(args.remove(0), Expr::Lambda);
+        let init_stack_len = self.cur_block().stack_len;
         let params = self.gen_param_names(&lambda.params);
-        self.emit_frameless_block(lambda.body, params); // ここでPOPされる
+        self.emit_frameless_block(lambda.body, params);
+        if self.cur_block().stack_len >= init_stack_len {
+            self.emit_pop_top();
+        }
         self.write_instr(JUMP_ABSOLUTE);
         self.write_arg((idx_for_iter / 2) as u8);
         let idx_end = self.cur_block().lasti;
         self.edit_code(idx_for_iter + 1, (idx_end - idx_for_iter - 2) / 2);
-        self.emit_pop_top();
+        // self.emit_pop_top();
+        self.stack_dec();
         self.emit_load_const(ValueObj::None);
     }
 
@@ -1661,11 +1679,10 @@ impl CodeGenerator {
                 Name,
             );
         }
+        let init_stack_len = self.cur_block().stack_len;
         for expr in block.into_iter() {
             self.emit_expr(expr);
-            // TODO: discard
+            if self.cur_block().stack_len > init_stack_len {
-            // 最終的に帳尻を合わせる(コード生成の順番的にスタックの整合性が一時的に崩れる場合がある)
-            if self.cur_block().stack_len == 1 {
                 self.emit_pop_top();
             }
         }
@@ -1681,10 +1698,11 @@ impl CodeGenerator {
                 Name,
             );
         }
+        let init_stack_len = self.cur_block().stack_len;
         for expr in block.into_iter() {
             self.emit_expr(expr);
             // __exit__, __enter__() are on the stack
-            if self.cur_block().stack_len != 2 {
+            if self.cur_block().stack_len > init_stack_len {
                 self.emit_pop_top();
             }
         }
@@ -1852,19 +1870,20 @@ impl CodeGenerator {
             &name,
             firstlineno,
         ));
+        let init_stack_len = self.cur_block().stack_len;
         for expr in block.into_iter() {
             self.emit_expr(expr);
             // NOTE: 各行のトップレベルでは0個または1個のオブジェクトが残っている
             // Pythonの場合使わなかったオブジェクトはそのまま捨てられるが、Ergではdiscardを使う必要がある
             // TODO: discard
-            if self.cur_block().stack_len == 1 {
+            if self.cur_block().stack_len > init_stack_len {
                 self.emit_pop_top();
             }
         }
         self.cancel_pop_top(); // 最後の値は戻り値として取っておく
-        if self.cur_block().stack_len == 0 {
+        if self.cur_block().stack_len == init_stack_len {
             self.emit_load_const(ValueObj::None);
-        } else if self.cur_block().stack_len > 1 {
+        } else if self.cur_block().stack_len > init_stack_len + 1 {
             let block_id = self.cur_block().id;
             let stack_len = self.cur_block().stack_len;
             CompileError::stack_bug(

compiler/erg_compiler/context/initialize/mod.rs

@@ -1283,6 +1283,8 @@ impl Context {
         str_iterator.register_superclass(Obj, &obj);
         let mut array_iterator = Self::builtin_poly_class("ArrayIterator", vec![PS::t_nd("T")], 1);
         array_iterator.register_superclass(Obj, &obj);
+        let mut range_iterator = Self::builtin_poly_class("RangeIterator", vec![PS::t_nd("T")], 1);
+        range_iterator.register_superclass(Obj, &obj);
         /* Float_mut */
         let mut float_mut = Self::builtin_mono_class("Float!", 2);
         float_mut.register_superclass(Float, &float);
@@ -1523,6 +1525,15 @@ impl Context {
             Public,
         );
         range.register_trait(range_t.clone(), range_eq);
+        let mut range_iterable =
+            Self::builtin_methods(Some(poly("Iterable", vec![ty_tp(mono_q("T"))])), 2);
+        range_iterable.register_builtin_impl(
+            "iter",
+            fn0_met(Str, mono("RangeIterator")),
+            Immutable,
+            Public,
+        );
+        range.register_trait(range_t.clone(), range_iterable);
         /* Proc */
         let mut proc = Self::builtin_mono_class("Proc", 2);
         proc.register_superclass(Obj, &obj);
@@ -1589,6 +1600,13 @@ impl Context {
             Const,
             Some("array_iterator"),
         );
+        self.register_builtin_type(
+            poly("RangeIterator", vec![ty_tp(mono_q("T"))]),
+            range_iterator,
+            Private,
+            Const,
+            Some("RangeIterator"),
+        );
         self.register_builtin_type(mono("Int!"), int_mut, Private, Const, Some("int"));
         self.register_builtin_type(mono("Nat!"), nat_mut, Private, Const, Some("Nat"));
         self.register_builtin_type(mono("Float!"), float_mut, Private, Const, Some("float"));

compiler/erg_compiler/lib/pystd/os.d/path.d.er

@@ -1 +1 @@
-.expanduser!: (path: PathLike,) => NoneType
+.expanduser!: (path: PathLike,) => Str

compiler/erg_compiler/lib/pystd/shutil.d.er

@@ -0,0 +1,6 @@
+.copyfile!: (src: PathLike, dst: PathLike,) => NoneType
+.copy!: (src: PathLike, dst: PathLike,) => NoneType
+.copytree!: (src: PathLike, dst: PathLike,) => NoneType
+.rmtree!: (path: PathLike,) => NoneType
+.move!: (src: PathLike, dst: PathLike) => NoneType
+.which!: (cmd: Str,) => Str or NoneType

compiler/erg_compiler/lib/pystd/zipfile.d.er

@@ -1,7 +1,8 @@
 .ZipFile!: ClassType
 .ZipFile! <: FileLike!
 
-.ZipFile!.open!: (path: PathLike or FileLike!, mode := Str) => .ZipFile!
+.open! = 'ZipFile': (path: PathLike or FileLike!, mode := Str) => .ZipFile!
+.ZipFile!.open!: (name: PathLike, mode := Str) => .ZipFile!
 .ZipFile!.add!: (self: RefMut(.ZipFile!), name: PathLike, arcname: PathLike or NoneType := NoneType, recursive := Bool) => NoneType
 .ZipFile!.close!: (self: .ZipFile!,) => NoneType
 .ZipFile!.extractall!: (self: RefMut(.ZipFile!), path := PathLike, members: [Str; _] or NoneType := NoneType, numeric_owner := Bool) => NoneType

compiler/erg_compiler/lib/std/_erg_std_prelude.py

@@ -67,7 +67,6 @@ class Bool(Nat):
 
 class Str(str):
     def __instancecheck__(cls, obj):
-        print(cls, obj)
         return obj == Str or obj == str
 
     def try_new(s: str): # -> Result[Nat]
@@ -133,10 +132,10 @@ class RangeIterator:
     def __init__(self, rng):
         self.rng = rng
         self.needle = self.rng.start
-        if type(self.rng.start) == int:
+        if issubclass(Nat, type(self.rng.start)):
             if not(self.needle in self.rng):
                 self.needle += 1
-        elif type(self.rng.start) == str:
+        elif issubclass(Str, type(self.rng.start)):
             if not(self.needle in self.rng):
                 self.needle = chr(ord(self.needle) + 1)
         else:
@@ -147,12 +146,12 @@ class RangeIterator:
         return self
 
     def __next__(self):
-        if type(self.rng.start) == int:
+        if issubclass(Nat, type(self.rng.start)):
             if self.needle in self.rng:
                 result = self.needle
                 self.needle += 1
                 return result
-        elif type(self.rng.start) == str:
+        elif issubclass(Str, type(self.rng.start)):
             if self.needle in self.rng:
                 result = self.needle
                 self.needle = chr(ord(self.needle) + 1)