From 50f1aca1aa7ecb4412aa4b8a0daa1cc644d8ca89 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Sun, 7 Nov 2021 10:29:14 +0800
Subject: [PATCH] nac3artiq: move module registration list to CPython side

In ARTIQ, we cannot create a global NAC3 object because we do not
know the ISA in advance.
---
 nac3artiq/demo/min_artiq.py | 12 ++++--------
 nac3artiq/src/lib.rs        | 15 ++++-----------
 2 files changed, 8 insertions(+), 19 deletions(-)

diff --git a/nac3artiq/demo/min_artiq.py b/nac3artiq/demo/min_artiq.py
index 800dc65d..5057e0df 100644
--- a/nac3artiq/demo/min_artiq.py
+++ b/nac3artiq/demo/min_artiq.py
@@ -17,7 +17,7 @@ core_arguments = device_db.device_db["core"]["arguments"]
 
 compiler = nac3artiq.NAC3(core_arguments["target"])
 allow_module_registration = True
-registered_ids = set()
+registered_modules = set()
 
 
 T = TypeVar('T')
@@ -26,13 +26,9 @@ class KernelInvariant(Generic[T]):
 
 
 def register_module_of(obj):
-    global registered_ids
     assert allow_module_registration
-    module = getmodule(obj)
-    module_id = id(module)
-    if module_id not in registered_ids:
-        compiler.register_module(module)
-        registered_ids.add(module_id)
+    # Delay NAC3 analysis until all referenced variables are supposed to exist on the CPython side.
+    registered_modules.add(getmodule(obj))
 
 
 def extern(function):
@@ -111,7 +107,7 @@ class Core:
     def run(self, method, *args, **kwargs):
         global allow_module_registration
         if allow_module_registration:
-            compiler.analyze()
+            compiler.analyze_modules(registered_modules)
             allow_module_registration = False
 
         if hasattr(method, "__self__"):
diff --git a/nac3artiq/src/lib.rs b/nac3artiq/src/lib.rs
index 7ac86c10..bfb319a8 100644
--- a/nac3artiq/src/lib.rs
+++ b/nac3artiq/src/lib.rs
@@ -9,7 +9,7 @@ use inkwell::{
     OptimizationLevel,
 };
 use pyo3::prelude::*;
-use pyo3::{exceptions, types::PyList, types::PyBytes};
+use pyo3::{exceptions, types::PyList, types::PySet, types::PyBytes};
 use nac3parser::{
     ast::{self, StrRef},
     parser::{self, parse_program},
@@ -66,7 +66,6 @@ struct Nac3 {
     pyid_to_type: Arc<RwLock<HashMap<u64, Type>>>,
     composer: TopLevelComposer,
     top_level: Option<Arc<TopLevelContext>>,
-    to_be_registered: Vec<PyObject>,
     primitive_ids: PrimitivePythonId,
     global_value_ids: Arc<Mutex<HashSet<u64>>>,
     working_directory: TempDir,
@@ -299,20 +298,14 @@ impl Nac3 {
             top_level: None,
             pyid_to_def: Default::default(),
             pyid_to_type: Default::default(),
-            to_be_registered: Default::default(),
             global_value_ids: Default::default(),
             working_directory
         })
     }
 
-    fn register_module(&mut self, obj: PyObject) {
-        // Delay registration until all referenced variables are supposed to exist on the CPython side
-        self.to_be_registered.push(obj);
-    }
-
-    fn analyze(&mut self) -> PyResult<()> {
-        for obj in std::mem::take(&mut self.to_be_registered).into_iter() {
-            self.register_module_impl(obj)?;
+    fn analyze_modules(&mut self, modules: &PySet) -> PyResult<()> {
+        for obj in modules.iter() {
+            self.register_module_impl(obj.into())?;
         }
         Ok(())
     }