artiq: register modules

nac3artiq/src/lib.rs

@@ -229,6 +229,8 @@ impl Nac3 {
                         }
                     })
                 }
+                // Required to capture alias used within module
+                StmtKind::Import { .. } => false,
                 _ => false,
             };
 
@@ -422,12 +424,14 @@ impl Nac3 {
         ];
         add_exceptions(&mut composer, &mut builtins_def, &mut builtins_ty, &exception_names);
 
+        // Stores a mapping from module id to attributes
         let mut module_to_resolver_cache: HashMap<u64, _> = HashMap::new();
 
         let mut rpc_ids = vec![];
         for (stmt, path, module) in &self.top_levels {
             let py_module: &PyAny = module.extract(py)?;
             let module_id: u64 = id_fn.call1((py_module,))?.extract()?;
+            let module_name: String = py_module.getattr("__name__")?.extract()?;
             let helper = helper.clone();
             let class_obj;
             if let StmtKind::ClassDef { name, .. } = &stmt.node {
@@ -442,9 +446,10 @@ impl Nac3 {
             } else {
                 class_obj = None;
             }
-            let (name_to_pyid, resolver) =
+            let (name_to_pyid, resolver, _) =
                 module_to_resolver_cache.get(&module_id).cloned().unwrap_or_else(|| {
                     let mut name_to_pyid: HashMap<StrRef, u64> = HashMap::new();
+                    name_to_pyid.insert(module_name.clone().into(), module_id);
                     let members: &PyDict =
                         py_module.getattr("__dict__").unwrap().downcast().unwrap();
                     for (key, val) in members {
@@ -470,12 +475,14 @@ impl Nac3 {
                         deferred_eval_store: self.deferred_eval_store.clone(),
                     })))
                         as Arc<dyn SymbolResolver + Send + Sync>;
+                    println!("Name_to_PyID for {:?}: {:?}({module_id})", py_module.to_string(), name_to_pyid);
                     let name_to_pyid = Rc::new(name_to_pyid);
+                    
                     module_to_resolver_cache
-                        .insert(module_id, (name_to_pyid.clone(), resolver.clone()));
+                        .insert(module_id, (name_to_pyid.clone(), resolver.clone(), module_name.clone()));
-                    (name_to_pyid, resolver)
+                    (name_to_pyid, resolver, module_name)
                 });
-
+            
             let (name, def_id, ty) = composer
                 .register_top_level(stmt.clone(), Some(resolver.clone()), path, false)
                 .map_err(|e| {
@@ -547,6 +554,19 @@ impl Nac3 {
             }
         }
 
+        // Adding top level module definitions
+        for (module_id, (module_name_to_pyid, module_resolver, module_name)) in module_to_resolver_cache.into_iter() {
+            let def_id= composer.register_top_level_module(
+                module_name,
+                module_name_to_pyid,
+                module_resolver
+            ).map_err(|e| {
+                CompileError::new_err(format!("compilation failed\n----------\n{e}"))
+            })?;
+
+            self.pyid_to_def.write().insert(module_id, def_id);
+        }
+
         let id_fun = PyModule::import(py, "builtins")?.getattr("id")?;
         let mut name_to_pyid: HashMap<StrRef, u64> = HashMap::new();
         let module = PyModule::new(py, "tmp")?;

nac3core/src/toplevel/composer.rs

@@ -202,6 +202,34 @@ impl TopLevelComposer {
         self.definition_ast_list.iter().map(|(def, ..)| def.clone()).collect_vec()
     }
 
+    /// register top level modules
+    pub fn register_top_level_module(
+        &mut self,
+        module_name: String,
+        name_to_pyid: Rc<HashMap<StrRef, u64>>,
+        resolver: Arc<dyn SymbolResolver + Send + Sync>,
+    ) -> Result<DefinitionId, String> {
+        // Construct methods field
+        let mut methods: HashMap<StrRef, DefinitionId> = HashMap::new();
+        for (name, _) in name_to_pyid.iter() {
+            if let Ok(def_id) = resolver.get_identifier_def(*name) {
+                methods.insert(*name, def_id);
+            };
+        }
+        let module_def = TopLevelDef::Module { 
+            name: module_name.clone().into(), 
+            alias: module_name.into(), 
+            module_id: DefinitionId(self.definition_ast_list.len()), 
+            methods, 
+            resolver: Some(resolver), 
+            loc: None
+        };
+
+        println!("{:?}", module_def.to_string(&mut self.unifier));
+        self.definition_ast_list.push((Arc::new(RwLock::new(module_def)).into(), None));
+        Ok(DefinitionId(self.definition_ast_list.len() - 1))
+    }
+
     /// register, just remember the names of top level classes/function
     /// and check duplicate class/method/function definition
     pub fn register_top_level(