refactor: capture outside variables in lambda. (#548)

kcl-lang · May 15, 2023 · b536376 · b536376
1 parent 74b1794
commit b536376
Show file tree

Hide file tree

Showing 15 changed files with 96 additions and 31 deletions.
diff --git a/kclvm/compiler/src/codegen/llvm/context.rs b/kclvm/compiler/src/codegen/llvm/context.rs
@@ -29,7 +29,7 @@ use kclvm_sema::builtin;
 use kclvm_sema::plugin;
 
 use crate::codegen::abi::Align;
-use crate::codegen::{error as kcl_error, EmitOptions};
+use crate::codegen::{error as kcl_error, EmitOptions, INNER_LEVEL};
 use crate::codegen::{
     traits::*, ENTRY_NAME, GLOBAL_VAL_ALIGNMENT, KCL_CONTEXT_VAR_NAME, MODULE_NAME,
     PKG_INIT_FUNCTION_SUFFIX,
@@ -1735,7 +1735,6 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
                                     &[closure_map, string_ptr_value],
                                 )
                             }
-                            // Nest comp for
                             None => self.builder.build_load(*var, name),
                         }
                     } else {
@@ -1815,48 +1814,54 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
         Ok(value)
     }
 
-    /// Get closure dict in the current scope.
-    pub(crate) fn get_closure_dict_in_current_scope(&self) -> BasicValueEnum<'ctx> {
-        let is_in_schema = self.schema_stack.borrow().len() > 0;
-        // Get closures in the current scope
-        let dict_value = self.dict_value();
+    /// Get closure map in the current scope.
+    pub(crate) fn get_closure_map(&self) -> BasicValueEnum<'ctx> {
+        // Get closures in the current scope.
+        let closure_map = self.dict_value();
         {
             let pkgpath = self.current_pkgpath();
             let pkgpath = if !pkgpath.starts_with(PKG_PATH_PREFIX) && pkgpath != MAIN_PKG_PATH {
                 format!("{}{}", PKG_PATH_PREFIX, pkgpath)
             } else {
                 pkgpath
             };
-            let pkg_scopes = self.pkg_scopes.borrow_mut();
+            let pkg_scopes = self.pkg_scopes.borrow();
             let scopes = pkg_scopes
                 .get(&pkgpath)
                 .unwrap_or_else(|| panic!("package {} is not found", pkgpath));
-            let closures_mut = scopes
-                .last()
-                .expect(kcl_error::INTERNAL_ERROR_MSG)
-                .closures
-                .borrow_mut();
-            let variables_mut = scopes
-                .last()
-                .expect(kcl_error::INTERNAL_ERROR_MSG)
-                .variables
-                .borrow_mut();
-            for (key, ptr) in &*closures_mut {
-                if variables_mut.contains_key(key) {
-                    let value = self.builder.build_load(*ptr, "");
-                    self.dict_insert_override_item(dict_value, key.as_str(), value);
-                }
+            // Clouure variable must be inner of the global scope.
+            if scopes.len() > INNER_LEVEL {
+                let closures = scopes
+                    .last()
+                    .expect(kcl_error::INTERNAL_ERROR_MSG)
+                    .closures
+                    .borrow();
+                // Curret scope vaiable.
+                let variables = scopes
+                    .get(scopes.len() - INNER_LEVEL)
+                    .expect(kcl_error::INTERNAL_ERROR_MSG)
+                    .variables
+                    .borrow();
+                // Transverse all scope and capture closures except the builtin amd global scope.
+                for (key, ptr) in &*closures {
+                    if variables.contains_key(key) {
+                        let value = self.builder.build_load(*ptr, "");
+                        self.dict_insert_override_item(closure_map, key.as_str(), value);
+                    }
+                } // Curret scope vaiable.
             }
         }
+        // Capture schema `self` closure.
+        let is_in_schema = self.schema_stack.borrow().len() > 0;
         if is_in_schema {
             for shcmea_closure_name in value::SCHEMA_VARIABLE_LIST {
                 let value = self
                     .get_variable(shcmea_closure_name)
                     .expect(kcl_error::INTERNAL_ERROR_MSG);
-                self.dict_insert_override_item(dict_value, shcmea_closure_name, value);
+                self.dict_insert_override_item(closure_map, shcmea_closure_name, value);
             }
         }
-        dict_value
+        closure_map
     }
 
     /// Push a function call frame into the function stack

diff --git a/kclvm/compiler/src/codegen/llvm/node.rs b/kclvm/compiler/src/codegen/llvm/node.rs
@@ -16,7 +16,7 @@ use crate::codegen::error as kcl_error;
 use crate::codegen::llvm::context::BacktrackMeta;
 use crate::codegen::llvm::utils;
 use crate::codegen::traits::*;
-use crate::codegen::{ENTRY_NAME, GLOBAL_LEVEL, PKG_INIT_FUNCTION_SUFFIX, SCHEMA_LEVEL};
+use crate::codegen::{ENTRY_NAME, GLOBAL_LEVEL, INNER_LEVEL, PKG_INIT_FUNCTION_SUFFIX};
 use crate::{check_backtrack_stop, pkgpath_without_prefix};
 
 use super::context::{CompileResult, LLVMCodeGenContext};
@@ -2046,8 +2046,9 @@ impl<'ctx> TypedResultWalker<'ctx> for LLVMCodeGenContext<'ctx> {
         // Exist the function
         self.builder.position_at_end(func_before_block);
         let closure = self.list_value();
-        let dict_value = self.get_closure_dict_in_current_scope();
-        self.list_append(closure, dict_value);
+        // Use closure map in the laste scope to construct curret closure map.
+        // The default value of the closure map is `{}`.
+        self.list_append(closure, self.get_closure_map());
         let function = self.closure_value(function, closure);
         self.leave_scope();
         self.pop_function();
@@ -2359,7 +2360,7 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
                                 (&config_entry, then_block),
                             ]);
                             let config_value = phi.as_basic_value();
-                            if self.scope_level() >= SCHEMA_LEVEL && !is_local_var {
+                            if self.scope_level() >= INNER_LEVEL && !is_local_var {
                                 self.dict_merge(schema_value, name, value, 1, -1);
                                 self.value_union(schema_value, config_value);
                                 let cal_map = self

diff --git a/kclvm/compiler/src/codegen/mod.rs b/kclvm/compiler/src/codegen/mod.rs
@@ -19,8 +19,8 @@ pub(crate) const KCL_CONTEXT_VAR_NAME: &str = "context";
 pub(crate) const PKG_INIT_FUNCTION_SUFFIX: &str = "init";
 /// Global level
 pub(crate) const GLOBAL_LEVEL: usize = 1;
-/// Schema level
-pub(crate) const SCHEMA_LEVEL: usize = 2;
+/// Inner level
+pub(crate) const INNER_LEVEL: usize = 2;
 /// Global variable alignment
 pub(crate) const GLOBAL_VAL_ALIGNMENT: u32 = 8;
 

diff --git a/test/grammar/lambda/closure_0/main.k b/test/grammar/lambda/closure_0/main.k
@@ -0,0 +1,6 @@
+x = lambda {
+    a = 1
+    lambda {
+        a + 1
+    }()
+}()
diff --git a/test/grammar/lambda/closure_0/stdout.golden b/test/grammar/lambda/closure_0/stdout.golden
@@ -0,0 +1 @@
+x: 2
diff --git a/test/grammar/lambda/closure_1/main.k b/test/grammar/lambda/closure_1/main.k
@@ -0,0 +1,7 @@
+x = lambda {
+    a = 1
+    b = 2
+    lambda x {
+        a + b + x
+    }(3)
+}()
diff --git a/test/grammar/lambda/closure_1/stdout.golden b/test/grammar/lambda/closure_1/stdout.golden
@@ -0,0 +1 @@
+x: 6
diff --git a/test/grammar/lambda/closure_2/main.k b/test/grammar/lambda/closure_2/main.k
@@ -0,0 +1,10 @@
+x = lambda {
+    a = 1
+    lambda {
+        b = 2
+        c = 3
+        lambda a {
+            a + b + c
+        }(a)
+    }()
+}()
diff --git a/test/grammar/lambda/closure_2/stdout.golden b/test/grammar/lambda/closure_2/stdout.golden
@@ -0,0 +1 @@
+x: 6
diff --git a/test/grammar/lambda/closure_3/main.k b/test/grammar/lambda/closure_3/main.k
@@ -0,0 +1,9 @@
+x = lambda {
+    a = 1
+    lambda {
+        b = 2
+        lambda a {
+            b + a
+        }(a)
+    }()
+}()
diff --git a/test/grammar/lambda/closure_3/stdout.golden b/test/grammar/lambda/closure_3/stdout.golden
@@ -0,0 +1 @@
+x: 3
diff --git a/test/grammar/lambda/closure_4/main.k b/test/grammar/lambda/closure_4/main.k
@@ -0,0 +1,9 @@
+x = lambda {
+    a = 1
+    lambda {
+        b = a + 1
+        lambda {
+            c = b + 1
+        }()
+    }()
+}()
diff --git a/test/grammar/lambda/closure_4/stdout.golden b/test/grammar/lambda/closure_4/stdout.golden
@@ -0,0 +1 @@
+x: 3
diff --git a/test/grammar/lambda/closure_5/main.k b/test/grammar/lambda/closure_5/main.k
@@ -0,0 +1,12 @@
+x = lambda {
+    a = 1
+    lambda {
+        b = a + 1
+        lambda {
+            c = b + 1
+            lambda {
+                d = c + 1
+            }()
+        }()
+    }()
+}()
diff --git a/test/grammar/lambda/closure_5/stdout.golden b/test/grammar/lambda/closure_5/stdout.golden
@@ -0,0 +1 @@
+x: 4