[WIP] Merge Globals into Module

flake.nix

@@ -180,9 +180,12 @@
           clippy
           pre-commit
           rustfmt
+          rust-analyzer
         ];
+        RUST_SRC_PATH = "${pkgs.rust.packages.stable.rustPlatform.rustLibSrc}";
         shellHook =
           ''
+          export PYTHONPATH=/home/abdul/nac3/pyo3:$PYTHONPATH
           export DEMO_LINALG_STUB=${packages.x86_64-linux.demo-linalg-stub}/lib/liblinalg.a
           export DEMO_LINALG_STUB32=${packages.x86_64-linux.demo-linalg-stub32}/lib/liblinalg.a
           '';

nac3artiq/demo/demo.py

@@ -1,5 +1,6 @@
 from min_artiq import *
 
+attr1: Kernel[str] = "ss"
 
 @nac3
 class Demo:
@@ -14,6 +15,8 @@ class Demo:
 
     @kernel
     def run(self):
+        global attr1
+        # attr1 = "2"
         self.core.reset()
         while True:
             with parallel:

nac3artiq/demo/nac3artiq.so

@@ -1 +0,0 @@
-../../target/release/libnac3artiq.so

nac3artiq/src/codegen.rs

@@ -8,8 +8,7 @@ use std::{
 
 use itertools::Itertools;
 use pyo3::{
-    types::{PyDict, PyList},
-    PyObject, PyResult, Python,
+    types::{PyDict, PyList}, PyObject, PyResult, Python
 };
 
 use super::{symbol_resolver::InnerResolver, timeline::TimeFns};
@@ -38,7 +37,7 @@ use nac3core::{
     toplevel::{
         helper::{extract_ndims, PrimDef},
         numpy::unpack_ndarray_var_tys,
-        DefinitionId, GenCall,
+        DefinitionId, GenCall, TopLevelDef,
     },
     typecheck::typedef::{iter_type_vars, FunSignature, FuncArg, Type, TypeEnum, VarMap},
 };
@@ -982,7 +981,8 @@ pub fn attributes_writeback<'ctx>(
             values.push((ty, obj.to_basic_value_enum(ctx, generator, ty).unwrap()));
         }
 
-        for val in (*globals).values() {
+        // For now the global variables are just values so completely useless (14.0 stored as a float object in globals)
+        for (gloabl_id, val) in &*globals {
             let val = val.as_ref(py);
             let ty = inner_resolver.get_obj_type(
                 py,
@@ -995,6 +995,33 @@ pub fn attributes_writeback<'ctx>(
                 return Ok(Err(ty));
             }
             let ty = ty.unwrap();
+
+            if let Some(def_id) = inner_resolver.pyid_to_def.read().get(gloabl_id) {
+                if let TopLevelDef::Variable { name, simple_name, .. } = &*top_levels[def_id.0].read() {
+                    // println!("[+] Varaible with type: {:?}\n{:?}\n", ctx.unifier.stringify(*ty), ctx.unifier.get_ty(*ty));
+                    println!("Sending Value of {:?}", val.to_string());
+                    if gen_rpc_tag(ctx, ty, &mut scratch_buffer).is_ok() {
+                        // let Some(val) = ctx.module.get_global(simple_name.to_string().as_str()) else {continue;};
+                        // let val = val.as_pointer_value();
+                        let pydict = PyDict::new(py);
+                        pydict.set_item("global", val)?;
+                        pydict.set_item("name", name)?;
+                        host_attributes.append(pydict)?;
+
+                        values.push((
+                            ty,
+                            // ctx.build_gep_and_load(
+                            //     val,
+                            //     &[zero, int32.const_int(0, false)],
+                            //     None,
+                            // ),
+                            inner_resolver.get_obj_value(py, val, ctx, generator, ty)?.unwrap(),
+                        ));
+                    }
+                    continue;
+                }
+            }
+            
             match &*ctx.unifier.get_ty(ty) {
                 TypeEnum::TObj { fields, obj_id, .. }
                     if *obj_id != ctx.primitives.option.obj_id(&ctx.unifier).unwrap() =>
@@ -1003,6 +1030,15 @@ pub fn attributes_writeback<'ctx>(
                     // for non-primitive attributes, they should be in another global
                     let mut attributes = Vec::new();
                     let obj = inner_resolver.get_obj_value(py, val, ctx, generator, ty)?.unwrap();
+                    
+                    if !obj.is_pointer_value() && gen_rpc_tag(ctx, ty, &mut scratch_buffer).is_ok() {
+                        println!("[-] Other function skipped");
+                    //     values.push((ty, obj));
+                    //     let pydict = PyDict::new(py);
+                    //     pydict.set_item("global", val)?;
+                    //     host_attributes.append(pydict)?;
+                    //     continue;
+                    }
                     for (name, (field_ty, is_mutable)) in fields {
                         if !is_mutable {
                             continue;

nac3artiq/src/lib.rs

@@ -228,6 +228,17 @@ impl Nac3 {
                         }
                     })
                 }
+                // Allow global declaration with Kernel[X] to pass across
+                // These are kept in sync between kernel and host
+                // KernelInvariants are constants, and hence not included here
+                StmtKind::AnnAssign { ref annotation, .. } => {
+                    match &annotation.node {
+                        ExprKind::Subscript { value, .. } 
+                        if matches!(&value.node, ExprKind::Name { id, .. } 
+                        if id.to_string().as_str() == "Kernel") => true,
+                        _ => false
+                    }
+                },
                 _ => false,
             };
 
@@ -544,6 +555,15 @@ impl Nac3 {
                     pyid_to_ty.insert(id, ty);
                 }
             }
+            if let StmtKind::AnnAssign { target, .. } = &stmt.node {
+                let ExprKind::Name { id: name, .. } = target.node else { unreachable!() };
+                
+                global_value_ids
+                    .write()
+                    // .insert(id, py_module.getattr(name.to_string().as_str()).unwrap().into());
+                    .insert(id, module.as_ref(py).into());
+
+            }
         }
 
         let id_fun = PyModule::import(py, "builtins")?.getattr("id")?;

nac3artiq/src/symbol_resolver.rs

@@ -23,7 +23,7 @@ use nac3core::{
     inkwell::{
         module::Linkage,
         types::{BasicType, BasicTypeEnum},
-        values::BasicValueEnum,
+        values::{BasicValue, BasicValueEnum},
         AddressSpace,
     },
     nac3parser::ast::{self, StrRef},
@@ -1555,9 +1555,40 @@ impl SymbolResolver for Resolver {
     fn get_symbol_value<'ctx>(
         &self,
         id: StrRef,
-        _: &mut CodeGenContext<'ctx, '_>,
-        _: &mut dyn CodeGenerator,
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        generator: &mut dyn CodeGenerator,
     ) -> Option<ValueEnum<'ctx>> {
+        if let Some(global_value) = self
+            .0
+            .name_to_pyid
+            .get(&id)
+            .and_then(|pyid| self.0.global_value_ids.read().get(pyid).cloned())
+        {
+            let val = ctx.module.get_global(id.to_string().as_str()).unwrap_or_else(|| {
+                let v = Python::with_gil(|py| -> PyResult<SymbolValue> {
+                    Ok(self
+                        .0
+                        .get_default_param_obj_value(py, global_value.as_ref(py))
+                        .unwrap()
+                        .unwrap())
+                })
+                .unwrap();
+
+                let ty = v.get_type(&ctx.primitives, &mut ctx.unifier);
+
+                let init_val = ctx.gen_symbol_val(generator, &v, ty);
+                let llvm_ty = init_val.get_type();
+
+                println!("Adding {id}");
+                let global = ctx.module.add_global(llvm_ty, None, &id.to_string());
+                global.set_linkage(Linkage::LinkOnceAny);
+                global.set_initializer(&init_val);
+
+                global
+            });
+            return Some(val.as_basic_value_enum().into());
+        }
+
         let sym_value = {
             let id_to_val = self.0.id_to_pyval.read();
             id_to_val.get(&id).cloned()

nac3core/src/toplevel/composer.rs

@@ -434,7 +434,7 @@ impl TopLevelComposer {
         location: Location,
     ) -> Result<(StrRef, DefinitionId, Option<Type>), String> {
         if self.keyword_list.contains(&name) {
-            return Err(format!("cannot use keyword `{name}` as a class name (at {location})"));
+            return Err(format!("cannot use keyword `{name}` as a variable name (at {location})"));
         }
 
         let global_var_name =
@@ -469,10 +469,10 @@ impl TopLevelComposer {
         self.analyze_top_level_class_definition()?;
         self.analyze_top_level_class_fields_methods()?;
         self.analyze_top_level_function()?;
+        self.analyze_top_level_variables()?;
         if inference {
             self.analyze_function_instance()?;
         }
-        self.analyze_top_level_variables()?;
         Ok(())
     }
 
@@ -1410,7 +1410,7 @@ impl TopLevelComposer {
         Ok(())
     }
 
-    /// step 4, analyze and call type inferencer to fill the `instance_to_stmt` of
+    /// step 5, analyze and call type inferencer to fill the `instance_to_stmt` of
     /// [`TopLevelDef::Function`]
     fn analyze_function_instance(&mut self) -> Result<(), HashSet<String>> {
         // first get the class constructor type correct for the following type check in function body
@@ -1941,12 +1941,109 @@ impl TopLevelComposer {
         Ok(())
     }
 
-    /// Step 5. Analyze and populate the types of global variables.
+    /// Step 4. Analyze and populate the types of global variables.
     fn analyze_top_level_variables(&mut self) -> Result<(), HashSet<String>> {
         let def_list = &self.definition_ast_list;
         let temp_def_list = self.extract_def_list();
         let unifier = &mut self.unifier;
         let primitives_store = &self.primitives_ty;
+        
+
+        // let dummy_field_type = unifier.get_dummy_var().ty;
+
+        // let annotation = match value {
+        //     None => {
+        //         // handle Kernel[T], KernelInvariant[T]
+        //         let (annotation, mutable) = match &annotation.node {
+        //             ExprKind::Subscript { value, slice, .. }
+        //                 if matches!(
+        //                     &value.node,
+        //                     ast::ExprKind::Name { id, .. } if id == &core_config.kernel_invariant_ann.into()
+        //                 ) =>
+        //             {
+        //                 (slice, false)
+        //             }
+        //             ExprKind::Subscript { value, slice, .. }
+        //                 if matches!(
+        //                     &value.node,
+        //                     ast::ExprKind::Name { id, .. } if core_config.kernel_ann.map_or(false, |c| id == &c.into())
+        //                 ) =>
+        //             {
+        //                 (slice, true)
+        //             }
+        //             _ if core_config.kernel_ann.is_none() => (annotation, true),
+        //             _ => continue, // ignore fields annotated otherwise
+        //         };
+        //         class_fields_def.push((*attr, dummy_field_type, mutable));
+        //         annotation
+        //     }
+        //     // Supporting Class Attributes
+        //     Some(boxed_expr) => {
+        //         // Class attributes are set as immutable regardless
+        //         let (annotation, _) = match &annotation.node {
+        //             ExprKind::Subscript { slice, .. } => (slice, false),
+        //             _ if core_config.kernel_ann.is_none() => (annotation, false),
+        //             _ => continue,
+        //         };
+
+        //         match &**boxed_expr {
+        //             ast::Located {location: _, custom: (), node: ExprKind::Constant { value: v, kind: _ }} => {
+        //                 // Restricting the types allowed to be defined as class attributes
+        //                 match v {
+        //                     ast::Constant::Bool(_) | ast::Constant::Str(_) | ast::Constant::Int(_) | ast::Constant::Float(_) => {}
+        //                     _ => {
+        //                         return Err(HashSet::from([
+        //                         format!(
+        //                             "unsupported statement in class definition body (at {})",
+        //                             b.location
+        //                         ),
+        //                     ]))
+        //                     }
+        //                 }
+        //                 class_attributes_def.push((*attr, dummy_field_type, v.clone()));
+        //             }
+        //             _ => {
+        //                 return Err(HashSet::from([
+        //                     format!(
+        //                         "unsupported statement in class definition body (at {})",
+        //                         b.location
+        //                     ),
+        //                 ]))
+        //             }
+        //         }
+        //         annotation
+        //     }
+        // };
+        // let parsed_annotation = parse_ast_to_type_annotation_kinds(
+        //     class_resolver,
+        //     temp_def_list,
+        //     unifier,
+        //     primitives,
+        //     annotation.as_ref(),
+        //     vec![(class_id, class_type_vars_def.clone())]
+        //         .into_iter()
+        //         .collect::<HashMap<_, _>>(),
+        // )?;
+        // // find type vars within this return type annotation
+        // let type_vars_within =
+        //     get_type_var_contained_in_type_annotation(&parsed_annotation);
+        // // handle the class type var and the method type var
+        // for type_var_within in type_vars_within {
+        //     let TypeAnnotation::TypeVar(t) = type_var_within else {
+        //         unreachable!("must be type var annotation")
+        //     };
+
+        //     if !class_type_vars_def.contains(&t){
+        //         return Err(HashSet::from([
+        //             format!(
+        //                 "class fields can only use type \
+        //                     vars over which the class is generic (at {})",
+        //                 annotation.location
+        //             ),
+        //         ]))
+        //     }
+        // }
+        // type_var_to_concrete_def.insert(dummy_field_type, parsed_annotation);
 
         let mut analyze = |variable_def: &Arc<RwLock<TopLevelDef>>| -> Result<_, HashSet<String>> {
             let TopLevelDef::Variable { ty: dummy_ty, ty_decl, resolver, loc, .. } =
@@ -1955,10 +2052,23 @@ impl TopLevelComposer {
                 // not top level variable def, skip
                 return Ok(());
             };
-
+            
             let resolver = &**resolver.as_ref().unwrap();
 
             if let Some(ty_decl) = ty_decl {
+                let ty_decl = match &ty_decl.node {
+                    ExprKind::Subscript { value, slice, .. }
+                        if matches!(
+                            &value.node,
+                            ast::ExprKind::Name { id, .. } if self.core_config.kernel_ann.map_or(false, |c| id == &c.into())
+                        ) =>
+                    {
+                        slice 
+                    }
+                    _ if self.core_config.kernel_ann.is_none() => ty_decl,
+                    _ => unreachable!("Global variables should be annotated with Kernel[]") // ignore fields annotated otherwise
+                };
+
                 let ty_annotation = parse_ast_to_type_annotation_kinds(
                     resolver,
                     &temp_def_list,