Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: capture outside variables in lambda. #548

Merged
merged 1 commit into from
May 15, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 30 additions & 25 deletions kclvm/compiler/src/codegen/llvm/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -1735,7 +1735,6 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
&[closure_map, string_ptr_value],
)
}
// Nest comp for
None => self.builder.build_load(*var, name),
}
} else {
Expand Down Expand Up @@ -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
Expand Down
9 changes: 5 additions & 4 deletions kclvm/compiler/src/codegen/llvm/node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions kclvm/compiler/src/codegen/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
6 changes: 6 additions & 0 deletions test/grammar/lambda/closure_0/main.k
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
x = lambda {
a = 1
lambda {
a + 1
}()
}()
1 change: 1 addition & 0 deletions test/grammar/lambda/closure_0/stdout.golden
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
x: 2
7 changes: 7 additions & 0 deletions test/grammar/lambda/closure_1/main.k
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
x = lambda {
a = 1
b = 2
lambda x {
a + b + x
}(3)
}()
1 change: 1 addition & 0 deletions test/grammar/lambda/closure_1/stdout.golden
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
x: 6
10 changes: 10 additions & 0 deletions test/grammar/lambda/closure_2/main.k
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
x = lambda {
a = 1
lambda {
b = 2
c = 3
lambda a {
a + b + c
}(a)
}()
}()
1 change: 1 addition & 0 deletions test/grammar/lambda/closure_2/stdout.golden
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
x: 6
9 changes: 9 additions & 0 deletions test/grammar/lambda/closure_3/main.k
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
x = lambda {
a = 1
lambda {
b = 2
lambda a {
b + a
}(a)
}()
}()
1 change: 1 addition & 0 deletions test/grammar/lambda/closure_3/stdout.golden
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
x: 3
9 changes: 9 additions & 0 deletions test/grammar/lambda/closure_4/main.k
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
x = lambda {
a = 1
lambda {
b = a + 1
lambda {
c = b + 1
}()
}()
}()
1 change: 1 addition & 0 deletions test/grammar/lambda/closure_4/stdout.golden
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
x: 3
12 changes: 12 additions & 0 deletions test/grammar/lambda/closure_5/main.k
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
x = lambda {
a = 1
lambda {
b = a + 1
lambda {
c = b + 1
lambda {
d = c + 1
}()
}()
}()
}()
1 change: 1 addition & 0 deletions test/grammar/lambda/closure_5/stdout.golden
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
x: 4