feat: top level variable back reference infra (kcl-lang#1284)

Signed-off-by: peefy <xpf6677@163.com>
Signed-off-by: JeevaRamanathan <jeevaramanathan.m@infosys.com>

kclvm/compiler/src/codegen/llvm/context.rs

@@ -85,6 +85,7 @@ pub struct LLVMCodeGenContext<'ctx> {
     pub program: &'ctx ast::Program,
     pub functions: RefCell<Vec<Rc<FunctionValue<'ctx>>>>,
     pub imported: RefCell<HashSet<String>>,
+    pub setter_keys: RefCell<HashSet<String>>,
     pub schema_stack: RefCell<Vec<value::SchemaType>>,
     pub lambda_stack: RefCell<Vec<usize>>,
     pub schema_expr_stack: RefCell<Vec<()>>,
@@ -365,6 +366,20 @@ impl<'ctx> BuilderMethods for LLVMCodeGenContext<'ctx> {
             self.module.add_function(name, fn_ty, None)
         }
     }
+
+    /// Add a setter function named `name`.
+    fn add_setter_function(&self, name: &str) -> Self::Function {
+        let fn_ty = self.setter_func_type();
+        if self.no_link {
+            let pkgpath = self.current_pkgpath();
+            let msg = format!("pkgpath {} is not found", pkgpath);
+            let modules = self.modules.borrow_mut();
+            let module = modules.get(&pkgpath).expect(&msg).borrow_mut();
+            module.inner.add_function(name, fn_ty, None)
+        } else {
+            self.module.add_function(name, fn_ty, None)
+        }
+    }
 }
 
 /* Value methods */
@@ -604,6 +619,16 @@ impl<'ctx> ValueMethods for LLVMCodeGenContext<'ctx> {
             .get_first_param()
             .expect(kcl_error::CONTEXT_VAR_NOT_FOUND_MSG)
     }
+    /// Get the global evaluation scope pointer.
+    fn current_scope_ptr(&self) -> Self::Value {
+        self.builder
+            .get_insert_block()
+            .unwrap()
+            .get_parent()
+            .unwrap()
+            .get_nth_param(1)
+            .expect(kcl_error::CONTEXT_VAR_NOT_FOUND_MSG)
+    }
 }
 
 impl<'ctx> ValueCalculationMethods for LLVMCodeGenContext<'ctx> {
@@ -1255,6 +1280,7 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
             functions: RefCell::new(vec![]),
             imported: RefCell::new(HashSet::new()),
             local_vars: RefCell::new(HashSet::new()),
+            setter_keys: RefCell::new(HashSet::new()),
             schema_stack: RefCell::new(vec![]),
             // 1 denotes the top global main function lambda and 0 denotes the builtin scope.
             // Any user-defined lambda scope greater than 1.
@@ -1714,7 +1740,12 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
     }
 
     /// Append a variable or update the existed variable
-    pub fn add_or_update_global_variable(&self, name: &str, value: BasicValueEnum<'ctx>) {
+    pub fn add_or_update_global_variable(
+        &self,
+        name: &str,
+        value: BasicValueEnum<'ctx>,
+        save_scope: bool,
+    ) {
         // Find argument name in the scope
         let current_pkgpath = self.current_pkgpath();
         let mut pkg_scopes = self.pkg_scopes.borrow_mut();

kclvm/compiler/src/codegen/llvm/module.rs

@@ -1,11 +1,17 @@
 // Copyright The KCL Authors. All rights reserved.
 
+use indexmap::IndexMap;
+use inkwell::values::FunctionValue;
+use inkwell::AddressSpace;
 use kclvm_ast::ast;
 use kclvm_ast::walker::TypedResultWalker;
+use kclvm_runtime::ApiFunc;
+use kclvm_sema::pkgpath_without_prefix;
 
-use super::context::LLVMCodeGenContext;
-use crate::codegen::error as kcl_error;
-use crate::codegen::traits::ValueMethods;
+use super::context::{BacktrackMeta, LLVMCodeGenContext};
+use crate::codegen::traits::{BuilderMethods, ProgramCodeGen, ValueMethods};
+use crate::codegen::{error as kcl_error, ENTRY_NAME};
+use crate::value;
 use std::str;
 
 impl<'ctx> LLVMCodeGenContext<'ctx> {
@@ -31,6 +37,8 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
                 _ => {}
             };
         }
+        // Pre define global variables with setter functions.
+        // self.predefine_global_setters(module);
     }
 
     pub fn predefine_global_types(&self, name: &str) {
@@ -49,13 +57,81 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
         self.emit_global_vars(&module.body);
     }
 
+    /// Predefine all global variables.
+    pub fn predefine_global_setters(&self, module: &'ctx ast::Module) {
+        // New a function block to the global setter constrcution process.
+        let global_setter_block = self.append_block("");
+        self.br(global_setter_block);
+        self.builder.position_at_end(global_setter_block);
+        let mut place_holder_map: IndexMap<String, Vec<FunctionValue<'ctx>>> = IndexMap::new();
+        let mut body_map: IndexMap<String, Vec<&ast::Node<ast::Stmt>>> = IndexMap::new();
+        let pkgpath = &self.current_pkgpath();
+        // Setter function name format: "$set.<pkg_path>.$<var_name>"
+        self.emit_global_setters(
+            &module.body,
+            &pkgpath,
+            false,
+            &mut place_holder_map,
+            &mut body_map,
+            &mut vec![],
+        );
+        // Build global attribute backtrack functions.
+        {
+            for (k, functions) in &place_holder_map {
+                if k == kclvm_runtime::CAL_MAP_INDEX_SIGNATURE {
+                    continue;
+                }
+                let stmt_list = body_map.get(k).expect(kcl_error::INTERNAL_ERROR_MSG);
+                let mut if_level = 0;
+                for (attr_func, stmt) in functions.iter().zip(stmt_list) {
+                    let function = *attr_func;
+                    let name = function
+                        .get_name()
+                        .to_str()
+                        .expect(kcl_error::INTERNAL_ERROR_MSG);
+                    // Get attribute function from the module.
+                    let function = self.lookup_function(name);
+                    self.push_function(function);
+                    let attr_block = self.append_block(ENTRY_NAME);
+                    self.builder.position_at_end(attr_block);
+                    // Backtrack meta begin
+                    if matches!(&stmt.node, ast::Stmt::If(..)) {
+                        if_level += 1;
+                        *self.backtrack_meta.borrow_mut() = Some(BacktrackMeta {
+                            target: k.clone(),
+                            level: if_level,
+                            count: 0,
+                            stop: false,
+                        });
+                    } else {
+                        if_level = 0;
+                    }
+                    let result = self.walk_stmt(stmt).expect(kcl_error::COMPILE_ERROR_MSG);
+                    // Backtrack meta end
+                    if matches!(&stmt.node, ast::Stmt::If(..)) {
+                        *self.backtrack_meta.borrow_mut() = None
+                    }
+                    // Build return
+                    self.builder.build_return(Some(&result));
+                    // Position at global main function block
+                    self.builder.position_at_end(global_setter_block);
+                    self.pop_function();
+                }
+            }
+        }
+    }
+
     fn emit_global_vars(&self, body: &'ctx [Box<ast::Node<ast::Stmt>>]) {
         for stmt in body {
             match &stmt.node {
                 ast::Stmt::Unification(unification_stmt) => {
                     let names = &unification_stmt.target.node.names;
                     if names.len() == 1 {
-                        self.add_or_update_global_variable(&names[0].node, self.undefined_value());
+                        self.add_or_update_global_variable(
+                            &names[0].node,
+                            self.undefined_value(),
+                            false,
+                        );
                     }
                 }
                 ast::Stmt::Assign(assign_stmt) => {
@@ -65,6 +141,7 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
                             self.add_or_update_global_variable(
                                 &names[0].node,
                                 self.undefined_value(),
+                                false,
                             );
                         }
                     }
@@ -122,4 +199,134 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
             }
         }
     }
+
+    /// Emit setter functions for global variables.
+    pub(crate) fn emit_global_setters(
+        &self,
+        body: &'ctx [Box<ast::Node<ast::Stmt>>],
+        pkgpath: &str,
+        is_in_if: bool,
+        place_holder_map: &mut IndexMap<String, Vec<FunctionValue<'ctx>>>,
+        body_map: &mut IndexMap<String, Vec<&'ctx ast::Node<ast::Stmt>>>,
+        in_if_names: &mut Vec<String>,
+    ) {
+        let add_stmt =
+            |name: &str,
+             stmt: &'ctx ast::Node<ast::Stmt>,
+             place_holder_map: &mut IndexMap<String, Vec<FunctionValue<'ctx>>>,
+             body_map: &mut IndexMap<String, Vec<&'ctx ast::Node<ast::Stmt>>>| {
+                // The function form e.g., $set.__main__.a(&Context, &LazyScope, &ValueRef, &ValueRef)
+                let var_key = format!("{}.{name}", pkgpath_without_prefix!(pkgpath));
+                let function =
+                    self.add_setter_function(&format!("{}.{}", value::GLOBAL_SETTER, var_key));
+                let lambda_fn_ptr = self.builder.build_bitcast(
+                    function.as_global_value().as_pointer_value(),
+                    self.context.i64_type().ptr_type(AddressSpace::default()),
+                    "",
+                );
+                if !place_holder_map.contains_key(name) {
+                    place_holder_map.insert(name.to_string(), vec![]);
+                }
+                let name_vec = place_holder_map
+                    .get_mut(name)
+                    .expect(kcl_error::INTERNAL_ERROR_MSG);
+                name_vec.push(function);
+                self.build_void_call(
+                    &ApiFunc::kclvm_scope_add_setter.name(),
+                    &[
+                        self.current_runtime_ctx_ptr(),
+                        self.current_scope_ptr(),
+                        self.native_global_string(pkgpath, "").into(),
+                        self.native_global_string(name, "").into(),
+                        lambda_fn_ptr,
+                    ],
+                );
+                let key = format!("{}.{name}", pkgpath_without_prefix!(pkgpath));
+                self.setter_keys.borrow_mut().insert(key);
+                if !body_map.contains_key(name) {
+                    body_map.insert(name.to_string(), vec![]);
+                }
+                let body_vec = body_map.get_mut(name).expect(kcl_error::INTERNAL_ERROR_MSG);
+                body_vec.push(stmt);
+            };
+        for stmt in body {
+            match &stmt.node {
+                ast::Stmt::Unification(unification_stmt) => {
+                    let name = &unification_stmt.target.node.names[0].node;
+                    if is_in_if {
+                        in_if_names.push(name.to_string());
+                    } else {
+                        add_stmt(name, stmt, place_holder_map, body_map);
+                    }
+                }
+                ast::Stmt::Assign(assign_stmt) => {
+                    for target in &assign_stmt.targets {
+                        let name = &target.node.names[0].node;
+                        if is_in_if {
+                            in_if_names.push(name.to_string());
+                        } else {
+                            add_stmt(name, stmt, place_holder_map, body_map);
+                        }
+                    }
+                }
+                ast::Stmt::AugAssign(aug_assign_stmt) => {
+                    let target = &aug_assign_stmt.target;
+                    let name = &target.node.names[0].node;
+                    if is_in_if {
+                        in_if_names.push(name.to_string());
+                    } else {
+                        add_stmt(name, stmt, place_holder_map, body_map);
+                    }
+                }
+                ast::Stmt::If(if_stmt) => {
+                    let mut names: Vec<String> = vec![];
+                    self.emit_global_setters(
+                        &if_stmt.body,
+                        pkgpath,
+                        true,
+                        place_holder_map,
+                        body_map,
+                        &mut names,
+                    );
+                    if is_in_if {
+                        for name in &names {
+                            in_if_names.push(name.to_string());
+                        }
+                    } else {
+                        for name in &names {
+                            add_stmt(name, stmt, place_holder_map, body_map);
+                        }
+                        names.clear();
+                    }
+                    self.emit_global_setters(
+                        &if_stmt.orelse,
+                        pkgpath,
+                        true,
+                        place_holder_map,
+                        body_map,
+                        &mut names,
+                    );
+                    if is_in_if {
+                        for name in &names {
+                            in_if_names.push(name.to_string());
+                        }
+                    } else {
+                        for name in &names {
+                            add_stmt(name, stmt, place_holder_map, body_map);
+                        }
+                        names.clear();
+                    }
+                }
+                ast::Stmt::SchemaAttr(schema_attr) => {
+                    let name = schema_attr.name.node.as_str();
+                    if is_in_if {
+                        in_if_names.push(name.to_string());
+                    } else {
+                        add_stmt(name, stmt, place_holder_map, body_map);
+                    }
+                }
+                _ => {}
+            }
+        }
+    }
 }

kclvm/compiler/src/codegen/llvm/node.rs

@@ -2402,6 +2402,7 @@ impl<'ctx> LLVMCodeGenContext<'ctx> {
                         self.add_or_update_global_variable(
                             name,
                             right_value.expect(kcl_error::INTERNAL_ERROR_MSG),
+                            true,
                         );
                     // Lambda local variables.
                     } else if self.is_in_lambda() {

kclvm/compiler/src/codegen/mod.rs

@@ -17,6 +17,8 @@ pub const ENTRY_NAME: &str = "entry";
 pub const VALUE_TYPE_NAME: &str = "kclvm_value_ref_t";
 /// The kclvm runtime context type name.
 pub const CONTEXT_TYPE_NAME: &str = "kclvm_context_t";
+/// The kclvm runtime evaluation type name.
+pub const SCOPE_TYPE_NAME: &str = "kclvm_eval_scope_t";
 /// Package init function name suffix
 pub const PKG_INIT_FUNCTION_SUFFIX: &str = "init";
 /// Global level

kclvm/compiler/src/codegen/traits/builder.rs

@@ -82,4 +82,6 @@ pub trait BuilderMethods: BackendTypes {
     fn lookup_function(&self, name: &str) -> Self::Function;
     /// Add a function named `name`.
     fn add_function(&self, name: &str) -> Self::Function;
+    /// Add a setter function named `name`.
+    fn add_setter_function(&self, name: &str) -> Self::Function;
 }

kclvm/compiler/src/codegen/traits/type.rs

@@ -1,7 +1,7 @@
 //! Copyright The KCL Authors. All rights reserved.
 
 use crate::codegen::abi::AddressSpace;
-use crate::codegen::{CONTEXT_TYPE_NAME, VALUE_TYPE_NAME};
+use crate::codegen::{CONTEXT_TYPE_NAME, SCOPE_TYPE_NAME, VALUE_TYPE_NAME};
 
 use super::BackendTypes;
 
@@ -53,6 +53,10 @@ pub trait DerivedTypeMethods: BaseTypeMethods {
     fn context_ptr_type(&self) -> Self::Type {
         self.ptr_type_to(self.get_intrinsic_type(CONTEXT_TYPE_NAME))
     }
+    /// Get the context pointer type.
+    fn scope_ptr_type(&self) -> Self::Type {
+        self.ptr_type_to(self.get_intrinsic_type(SCOPE_TYPE_NAME))
+    }
     /// Get the function type.
     fn function_type(&self) -> Self::FunctionLet {
         let value_ptr_type = self.value_ptr_type();
@@ -62,6 +66,13 @@ pub trait DerivedTypeMethods: BaseTypeMethods {
             value_ptr_type,
         )
     }
+    /// Get the setter function type.
+    fn setter_func_type(&self) -> Self::FunctionLet {
+        let context_ptr_type = self.context_ptr_type();
+        let scope_ptr_type = self.scope_ptr_type();
+        let value_ptr_type = self.value_ptr_type();
+        self.function_let(&[context_ptr_type, scope_ptr_type], value_ptr_type)
+    }
 }
 
 /// TypeCodeGen defines all type APIs.

kclvm/compiler/src/codegen/traits/value.rs

@@ -43,6 +43,8 @@ pub trait ValueMethods: BackendTypes {
     fn global_value_ptr(&self, name: &str) -> Self::Value;
     /// Get current runtime context pointer.
     fn current_runtime_ctx_ptr(&self) -> Self::Value;
+    /// Get the global evaluation scope pointer.
+    fn current_scope_ptr(&self) -> Self::Value;
 }
 
 /// DerivedValueCalculationMethods defines all value base calculation APIs.