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>>, composer: TopLevelComposer, top_level: Option>, - to_be_registered: Vec, primitive_ids: PrimitivePythonId, global_value_ids: Arc>>, 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(()) }