fix: closure codegen bug

crates/erg_compiler/codegen.rs

@@ -320,6 +320,10 @@ impl PyCodeGenerator {
         self.units.last().unwrap()
     }
 
+    fn is_toplevel(&self) -> bool {
+        self.cur_block() == self.toplevel_block()
+    }
+
     fn captured_vars(&self) -> Vec<&str> {
         let mut caps = vec![];
         for unit in self.units.iter() {
@@ -634,7 +638,7 @@ impl PyCodeGenerator {
             })
     }
 
-    fn local_search(&self, name: &str, _acc_kind: AccessKind) -> Option<Name> {
+    fn local_search(&self, name: &str, acc_kind: AccessKind) -> Option<Name> {
         if self.py_version.minor < Some(11) {
             if let Some(idx) = self
                 .cur_block_codeobj()
@@ -645,34 +649,47 @@ impl PyCodeGenerator {
                 return Some(Name::deref(idx));
             }
         }
-        if let Some(idx) = self
+        match acc_kind {
-            .cur_block_codeobj()
+            AccessKind::Name => {
-            .varnames
+                if let Some(idx) = self
-            .iter()
+                    .cur_block_codeobj()
-            .position(|v| &**v == name)
+                    .freevars
-        {
+                    .iter()
-            if self.captured_vars().contains(&name) {
+                    .position(|f| &**f == name)
-                Some(Name::deref(idx))
+                {
-            } else {
+                    Some(Name::deref(idx))
-                Some(Name::fast(idx))
+                } else if let Some(idx) = self
+                    .cur_block_codeobj()
+                    .varnames
+                    .iter()
+                    .position(|v| &**v == name)
+                {
+                    if self.captured_vars().contains(&name) {
+                        None
+                    } else {
+                        Some(Name::fast(idx))
+                    }
+                } else if let Some(idx) = self
+                    .cur_block_codeobj()
+                    .names
+                    .iter()
+                    .position(|n| &**n == name)
+                {
+                    if !self.is_toplevel() {
+                        None
+                    } else {
+                        Some(Name::local(idx))
+                    }
+                } else {
+                    None
+                }
             }
-        } else if let Some(idx) = self
+            _ => self
-            .cur_block_codeobj()
+                .cur_block_codeobj()
-            .names
+                .names
-            .iter()
-            .position(|n| &**n == name)
-        {
-            if self.py_version.minor >= Some(11) && self.captured_vars().contains(&name) {
-                None
-            } else {
-                Some(Name::local(idx))
-            }
-        } else {
-            self.cur_block_codeobj()
-                .freevars
                 .iter()
-                .position(|f| &**f == name)
+                .position(|n| &**n == name)
-                .map(Name::deref)
+                .map(Name::local),
         }
     }
 
@@ -702,7 +719,7 @@ impl PyCodeGenerator {
     }
 
     fn register_name(&mut self, name: Str, kind: RegisterNameKind) -> Name {
-        let current_is_toplevel = self.cur_block() == self.toplevel_block();
+        let current_is_toplevel = self.is_toplevel();
         match self.rec_search(&name) {
             Some(st @ (StoreLoadKind::Local | StoreLoadKind::Global)) => {
                 if kind.is_fast() {
@@ -843,6 +860,7 @@ impl PyCodeGenerator {
         self.write_instr(instr);
         self.write_arg(name.idx);
         self.stack_inc();
+        self.mut_cur_block_codeobj().stacksize += 2;
         if instr == LOAD_GLOBAL as u8 && self.py_version.minor >= Some(11) {
             self.write_bytes(&[0; 2]);
             self.write_bytes(&[0; 8]);
@@ -3719,11 +3737,11 @@ impl PyCodeGenerator {
             0
         };
         let mut cells = vec![];
-        if self.py_version.minor >= Some(11) {
+        for captured in captured_names {
-            for captured in captured_names {
+            self.mut_cur_block()
-                self.mut_cur_block()
+                .captured_vars
-                    .captured_vars
+                .push(captured.inspect().clone());
-                    .push(captured.inspect().clone());
+            if self.py_version.minor >= Some(11) {
                 self.write_instr(Opcode311::MAKE_CELL);
                 cells.push((captured, self.lasti()));
                 self.write_arg(0);

crates/erg_compiler/ty/codeobj.rs

@@ -629,10 +629,17 @@ impl CodeObj {
         }
         match op308 {
             Opcode308::STORE_DEREF | Opcode308::LOAD_DEREF => {
-                write!(instrs, "{arg} ({})", self.freevars.get(arg).unwrap()).unwrap();
+                write!(
+                    instrs,
+                    "{arg} ({})",
+                    self.freevars
+                        .get(arg)
+                        .unwrap_or_else(|| &self.cellvars[arg])
+                )
+                .unwrap();
             }
             Opcode308::LOAD_CLOSURE => {
-                write!(instrs, "{arg} ({})", self.cellvars.get(arg).unwrap()).unwrap();
+                write!(instrs, "{arg} ({})", self.cellvars[arg]).unwrap();
             }
             Opcode308::JUMP_ABSOLUTE => {
                 write!(instrs, "{arg} (to {})", arg).unwrap();
@@ -670,10 +677,17 @@ impl CodeObj {
         }
         match op309 {
             Opcode309::STORE_DEREF | Opcode309::LOAD_DEREF => {
-                write!(instrs, "{arg} ({})", self.freevars.get(arg).unwrap()).unwrap();
+                write!(
+                    instrs,
+                    "{arg} ({})",
+                    self.freevars
+                        .get(arg)
+                        .unwrap_or_else(|| &self.cellvars[arg])
+                )
+                .unwrap();
             }
             Opcode309::LOAD_CLOSURE => {
-                write!(instrs, "{arg} ({})", self.cellvars.get(arg).unwrap()).unwrap();
+                write!(instrs, "{arg} ({})", self.cellvars[arg]).unwrap();
             }
             Opcode309::JUMP_ABSOLUTE => {
                 write!(instrs, "{arg} (to {})", arg).unwrap();
@@ -716,13 +730,12 @@ impl CodeObj {
                     "{arg} ({})",
                     self.freevars
                         .get(arg)
-                        .or_else(|| self.varnames.get(arg))
+                        .unwrap_or_else(|| &self.cellvars[arg])
-                        .unwrap()
                 )
                 .unwrap();
             }
             Opcode310::LOAD_CLOSURE => {
-                write!(instrs, "{arg} ({})", self.cellvars.get(arg).unwrap()).unwrap();
+                write!(instrs, "{arg} ({})", self.cellvars[arg]).unwrap();
             }
             Opcode310::JUMP_ABSOLUTE => {
                 write!(instrs, "{arg} (to {})", arg * 2).unwrap();
@@ -762,10 +775,10 @@ impl CodeObj {
         }
         match op311 {
             Opcode311::STORE_DEREF | Opcode311::LOAD_DEREF => {
-                write!(instrs, "{arg} ({})", self.varnames.get(arg).unwrap()).unwrap();
+                write!(instrs, "{arg} ({})", self.varnames[arg]).unwrap();
             }
             Opcode311::MAKE_CELL | Opcode311::LOAD_CLOSURE => {
-                write!(instrs, "{arg} ({})", self.varnames.get(arg).unwrap()).unwrap();
+                write!(instrs, "{arg} ({})", self.varnames[arg]).unwrap();
             }
             Opcode311::POP_JUMP_FORWARD_IF_FALSE | Opcode311::POP_JUMP_FORWARD_IF_TRUE => {
                 write!(instrs, "{arg} (to {})", idx + arg * 2 + 2).unwrap();
@@ -787,7 +800,7 @@ impl CodeObj {
                 write!(instrs, "{arg}").unwrap();
             }
             Opcode311::KW_NAMES => {
-                write!(instrs, "{arg} ({})", self.consts.get(arg).unwrap()).unwrap();
+                write!(instrs, "{arg} ({})", self.consts[arg]).unwrap();
             }
             Opcode311::BINARY_OP => {
                 write!(