Refactor Toplevel composer

flake.nix

@@ -180,7 +180,9 @@
           clippy
           pre-commit
           rustfmt
+          rust-analyzer
         ];
+        RUST_SRC_PATH = "${pkgs.rust.packages.stable.rustPlatform.rustLibSrc}";
         shellHook =
           ''
           export DEMO_LINALG_STUB=${packages.x86_64-linux.demo-linalg-stub}/lib/liblinalg.a

nac3artiq/demo/demo.py

@@ -1,26 +1,87 @@
 from min_artiq import *
+from numpy import int32
 
 
+# @nac3
+# class A:
+#     a: int32
+#     core: KernelInvariant[Core]
+
+#     def __init__(self, a: int32):
+#         self.core = Core()
+#         self.a = a
+    
+#     @kernel
+#     def output_all_fields(self):
+#         #print(self.a)
+#         pass
+    
+#     @kernel
+#     def set_a(self, a: int32):
+#         self.a = a
+
+# @nac3
+# class B(A):
+#     b: int32
+
+#     def __init__(self, b: int32):
+#         # A.__init__(self, b + 1)
+#         self.core = Core()
+#         self.a = b
+#         self.b = b
+#         self.set_b(b)
+    
+#     @kernel
+#     def output_parent_fields(self):
+#         # A.output_all_fields(self)
+#         pass
+
+#     @kernel
+#     def output_all_fields(self):
+#         # A.output_all_fields(self)
+#         pass
+#         #print(self.b)
+    
+#     @kernel
+#     def set_b(self, b: int32):
+#         self.b = b
+
 @nac3
-class Demo:
+class C:
+    c: Kernel[int32]
+    a: Kernel[int32]
+    b: Kernel[int32]
     core: KernelInvariant[Core]
-    led0: KernelInvariant[TTLOut]
-    led1: KernelInvariant[TTLOut]
 
-    def __init__(self):
+    def __init__(self, c: int32):
+        # B.__init__(self, c + 1)
         self.core = Core()
-        self.led0 = TTLOut(self.core, 18)
+        self.a = c
-        self.led1 = TTLOut(self.core, 19)
+        self.b = c
+        self.c = c
+    
+    @kernel
+    def output_parent_fields(self):
+        # B.output_all_fields(self)
+        pass
 
     @kernel
-    def run(self):
+    def output_all_fields(self):
-        self.core.reset()
+        # B.output_all_fields(self)
-        while True:
+        #print(self.c)
-            with parallel:
+        pass
-                self.led0.pulse(100.*ms)
-                self.led1.pulse(100.*ms)
-            self.core.delay(100.*ms)
 
+    @kernel
+    def set_c(self, c: int32):
+        self.c = c
+    
+    @kernel
+    def run(self):
+        self.output_all_fields()
+        # self.set_a(1)
+        # self.set_b(2)
+        self.set_c(3)
+        self.output_all_fields()
 
 if __name__ == "__main__":
-    Demo().run()
+    C(10).run()

nac3core/src/toplevel/composer.rs

@@ -5,7 +5,7 @@ use crate::{
     codegen::{expr::get_subst_key, stmt::exn_constructor},
     symbol_resolver::SymbolValue,
     typecheck::{
-        type_inferencer::{FunctionData, Inferencer},
+        type_inferencer::{report_error, FunctionData, Inferencer},
         typedef::{TypeVar, VarMap},
     },
 };
@@ -389,7 +389,26 @@ impl TopLevelComposer {
     }
 
     pub fn start_analysis(&mut self, inference: bool) -> Result<(), HashSet<String>> {
-        self.analyze_top_level_class_type_var()?;
+        let unifier = self.unifier.borrow_mut();
+        let primitives_store = &self.primitives_ty;
+        let temp_def_list = self.extract_def_list();
+
+        // Separate class definitions
+        let def_list = &self.definition_ast_list;
+        let class_def_list = def_list
+            .iter()
+            .skip(self.builtin_num)
+            .filter(|def| def.1.is_some() && matches!(&*def.0.read(), TopLevelDef::Class { .. }))
+            .collect::<Vec<_>>();
+
+        // Step 1. Analyze type variables within class definitions
+        Self::analyze_top_level_class_type_var(
+            class_def_list,
+            temp_def_list.clone(),
+            unifier,
+            primitives_store,
+        )?;
+
         self.analyze_top_level_class_bases()?;
         self.analyze_top_level_class_fields_methods()?;
         self.analyze_top_level_function()?;
@@ -399,116 +418,117 @@ impl TopLevelComposer {
         Ok(())
     }
 
-    /// step 1, analyze the type vars associated with top level class
+    fn analyze_bases(
-    fn analyze_top_level_class_type_var(&mut self) -> Result<(), HashSet<String>> {
+        class_def: &Arc<RwLock<TopLevelDef>>,
-        let def_list = &self.definition_ast_list;
+        class_ast: &Option<Stmt>,
-        let temp_def_list = self.extract_def_list();
+        temp_def_list: &Vec<Arc<RwLock<TopLevelDef>>>,
-        let unifier = self.unifier.borrow_mut();
+        unifier: &mut Unifier,
-        let primitives_store = &self.primitives_ty;
+        primitives_store: &PrimitiveStore,
-
+    ) -> Result<(), HashSet<String>> {
-        let mut analyze = |class_def: &Arc<RwLock<TopLevelDef>>, class_ast: &Option<Stmt>| {
+        let mut class_def = class_def.write();
-            // only deal with class def here
+        let (class_def_id, class_ancestors, class_bases_ast, class_type_vars, class_resolver) = {
-            let mut class_def = class_def.write();
+            let TopLevelDef::Class { object_id, ancestors, type_vars, resolver, .. } =
-            let (class_bases_ast, class_def_type_vars, class_resolver) = {
+                &mut *class_def
-                if let TopLevelDef::Class { type_vars, resolver, .. } = &mut *class_def {
+            else {
-                    let Some(ast::Located { node: ast::StmtKind::ClassDef { bases, .. }, .. }) =
+                unreachable!()
-                        class_ast
-                    else {
-                        unreachable!()
-                    };
-
-                    (bases, type_vars, resolver)
-                } else {
-                    return Ok(());
-                }
             };
-            let class_resolver = class_resolver.as_ref().unwrap();
+            let Some(ast::Located { node: ast::StmtKind::ClassDef { bases, .. }, .. }) = class_ast
-            let class_resolver = &**class_resolver;
+            else {
+                unreachable!()
+            };
+            (object_id, ancestors, bases, type_vars, resolver.as_ref().unwrap().as_ref())
+        };
 
-            let mut is_generic = false;
+        let mut is_generic = false;
-            for b in class_bases_ast {
+        let mut has_base = false;
-                match &b.node {
+        // Check class bases for typevars
-                    // analyze typevars bounded to the class,
+        for b in class_bases_ast {
-                    // only support things like `class A(Generic[T, V])`,
+            match &b.node {
-                    // things like `class A(Generic[T, V, ImportedModule.T])` is not supported
+                // analyze typevars bounded to the class,
-                    // i.e. only simple names are allowed in the subscript
+                // only support things like `class A(Generic[T, V])`,
-                    // should update the TopLevelDef::Class.typevars and the TypeEnum::TObj.params
+                // things like `class A(Generic[T, V, ImportedModule.T])` is not supported
-                    ast::ExprKind::Subscript { value, slice, .. }
+                // i.e. only simple names are allowed in the subscript
-                        if {
+                // should update the TopLevelDef::Class.typevars and the TypeEnum::TObj.params
-                            matches!(
+                ast::ExprKind::Subscript { value, slice, .. } if matches!(&value.node, ast::ExprKind::Name { id, .. } if id == &"Generic".into()) =>
-                                &value.node,
+                {
-                                ast::ExprKind::Name { id, .. } if id == &"Generic".into()
+                    if is_generic {
-                            )
+                        return report_error("only single Generic[...] is allowed", b.location);
-                        } =>
+                    }
-                    {
+                    is_generic = true;
-                        if is_generic {
-                            return Err(HashSet::from([format!(
-                                "only single Generic[...] is allowed (at {})",
-                                b.location
-                            )]));
-                        }
-                        is_generic = true;
 
-                        let type_var_list: Vec<&ast::Expr<()>>;
+                    let type_var_list: Vec<&ast::Expr<()>>;
-                        // if `class A(Generic[T, V, G])`
+                    // if `class A(Generic[T, V, G])`
-                        if let ast::ExprKind::Tuple { elts, .. } = &slice.node {
+                    if let ast::ExprKind::Tuple { elts, .. } = &slice.node {
-                            type_var_list = elts.iter().collect_vec();
+                        type_var_list = elts.iter().collect_vec();
-                        // `class A(Generic[T])`
+                    // `class A(Generic[T])`
-                        } else {
+                    } else {
-                            type_var_list = vec![&**slice];
+                        type_var_list = vec![&**slice];
-                        }
-
-                        // parse the type vars
-                        let type_vars = type_var_list
-                            .into_iter()
-                            .map(|e| {
-                                class_resolver.parse_type_annotation(
-                                    &temp_def_list,
-                                    unifier,
-                                    primitives_store,
-                                    e,
-                                )
-                            })
-                            .collect::<Result<Vec<_>, _>>()?;
-
-                        // check if all are unique type vars
-                        let all_unique_type_var = {
-                            let mut occurred_type_var_id: HashSet<TypeVarId> = HashSet::new();
-                            type_vars.iter().all(|x| {
-                                let ty = unifier.get_ty(*x);
-                                if let TypeEnum::TVar { id, .. } = ty.as_ref() {
-                                    occurred_type_var_id.insert(*id)
-                                } else {
-                                    false
-                                }
-                            })
-                        };
-                        if !all_unique_type_var {
-                            return Err(HashSet::from([format!(
-                                "duplicate type variable occurs (at {})",
-                                slice.location
-                            )]));
-                        }
-
-                        // add to TopLevelDef
-                        class_def_type_vars.extend(type_vars);
                     }
 
-                    // if others, do nothing in this function
+                    let type_vars = type_var_list
-                    _ => continue,
+                        .into_iter()
+                        .map(|e| {
+                            class_resolver.parse_type_annotation(
+                                &temp_def_list,
+                                unifier,
+                                primitives_store,
+                                e,
+                            )
+                        })
+                        .collect::<Result<Vec<_>, _>>()?;
+
+                    class_type_vars.extend(type_vars);
                 }
+                ast::ExprKind::Subscript { .. } => {
+                    if has_base {
+                        return report_error("a class definition can only have at most one base class declaration and one generic declaration", b.location);
+                    }
+                    has_base = true;
+
+                    // the function parse_ast_to make sure that no type var occurred in
+                    // bast_ty if it is a CustomClassKind
+                    let base_ty = parse_ast_to_type_annotation_kinds(
+                        class_resolver,
+                        &temp_def_list,
+                        unifier,
+                        &primitives_store,
+                        b,
+                        vec![(*class_def_id, class_type_vars.clone())]
+                            .into_iter()
+                            .collect::<HashMap<_, _>>(),
+                    )?;
+
+                    if let TypeAnnotation::CustomClass { .. } = &base_ty {
+                        class_ancestors.push(base_ty);
+                    } else {
+                        return report_error(
+                            "class base declaration can only be custom class",
+                            b.location,
+                        );
+                    }
+                }
+                // TODO: Report Error here
+                _ => unreachable!(),
             }
-            Ok(())
+        }
-        };
+
+        Ok(())
+    }
+
+    /// step 1, analyze the type vars associated with top level class
+    fn analyze_top_level_class_type_var(
+        def_list: Vec<&DefAst>,
+        temp_def_list: Vec<Arc<RwLock<TopLevelDef>>>,
+        unifier: &mut Unifier,
+        primitives_store: &PrimitiveStore,
+    ) -> Result<(), HashSet<String>> {
         let mut errors = HashSet::new();
-        for (class_def, class_ast) in def_list.iter().skip(self.builtin_num) {
+        for (class_def, class_ast) in def_list.iter() {
-            if class_ast.is_none() {
+            if let Err(e) = Self::analyze_bases(class_def, class_ast, &temp_def_list, unifier, primitives_store) {
-                continue;
-            }
-            if let Err(e) = analyze(class_def, class_ast) {
                 errors.extend(e);
             }
         }
+
         if !errors.is_empty() {
             return Err(errors);
         }

nac3core/src/typecheck/type_inferencer/mod.rs

@@ -112,7 +112,7 @@ impl Fold<()> for NaiveFolder {
     }
 }
 
-fn report_error<T>(msg: &str, location: Location) -> Result<T, InferenceError> {
+pub fn report_error<T>(msg: &str, location: Location) -> Result<T, InferenceError> {
     Err(HashSet::from([format!("{msg} at {location}")]))
 }