fix: closure bug

GreasySlug · Nov 18, 2023 · 962e010 · 962e010
1 parent 439ba1d
commit 962e010
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 0 deletions.
diff --git a/crates/erg_compiler/context/mod.rs b/crates/erg_compiler/context/mod.rs
@@ -1382,6 +1382,16 @@ impl Context {
         None
     }
 
+    pub fn current_function_ctx(&self) -> Option<&Context> {
+        if self.kind.is_subr() {
+            Some(self)
+        } else if let Some(outer) = self.get_outer() {
+            outer.current_function_ctx()
+        } else {
+            None
+        }
+    }
+
     pub(crate) fn check_types(&self) {
         if DEBUG_MODE {
             for (_, ctx) in self.poly_types.iter() {

diff --git a/crates/erg_compiler/lower.rs b/crates/erg_compiler/lower.rs
@@ -932,6 +932,11 @@ impl<A: ASTBuildable> GenericASTLowerer<A> {
         if !ident.vi.is_toplevel()
             && ident.vi.def_namespace() != &self.module.context.name
             && ident.vi.kind.can_capture()
+            && self
+                .module
+                .context
+                .current_function_ctx()
+                .is_some_and(|ctx| ctx.control_kind().is_none())
         {
             self.module.context.captured_names.push(ident.clone());
         }

diff --git a/tests/should_ok/closure.er b/tests/should_ok/closure.er
@@ -12,3 +12,11 @@ func2! version: Int =
 
 arr = func2!(1)
 assert arr[0] == 1
+
+for! [1], _ =>
+    result = !""
+    push! left, right =
+        result.push! "| \{left} | \{right} |\n"
+
+    push! "a", "b"
+    assert result == "| a | b |\n"