From c7cca11ad140b8ff013cf22beaaab58ebc0abedb Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Thu, 11 Nov 2021 16:32:37 +0800
Subject: [PATCH] update NAC3

---
 artiq/coredevice/core.py |  6 +++---
 artiq/language/core.py   | 27 ++++++++++++++++-----------
 flake.lock               |  8 ++++----
 3 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/artiq/coredevice/core.py b/artiq/coredevice/core.py
index 766ff0cef..fa332858e 100644
--- a/artiq/coredevice/core.py
+++ b/artiq/coredevice/core.py
@@ -60,9 +60,9 @@ class Core:
         self.comm.close()
 
     def compile(self, method, args, kwargs, file_output=None):
-        if core_language._allow_module_registration:
-            self.compiler.analyze_modules(core_language._registered_modules)
-            core_language._allow_module_registration = False
+        if core_language._allow_registration:
+            self.compiler.analyze(core_language._registered_functions, core_language._registered_classes)
+            core_language._allow_registration = False
 
         if hasattr(method, "__self__"):
             obj = method.__self__
diff --git a/artiq/language/core.py b/artiq/language/core.py
index 36d8b459e..c6a96ecbe 100644
--- a/artiq/language/core.py
+++ b/artiq/language/core.py
@@ -4,7 +4,7 @@ Core ARTIQ extensions to the Python language.
 
 from typing import Generic, TypeVar
 from functools import wraps
-from inspect import getfullargspec, getmodule
+from inspect import getfullargspec
 from types import SimpleNamespace
 
 
@@ -25,24 +25,29 @@ def round64(x):
     return round(x)
 
 
-_allow_module_registration = True
-_registered_modules = set()
+_allow_registration = True
+# Delay NAC3 analysis until all referenced variables are supposed to exist on the CPython side.
+_registered_functions = set()
+_registered_classes = set()
 
-def _register_module_of(obj):
-    assert _allow_module_registration
-    # Delay NAC3 analysis until all referenced variables are supposed to exist on the CPython side.
-    _registered_modules.add(getmodule(obj))
+def _register_function(fun):
+    assert _allow_registration
+    _registered_functions.add(fun)
+
+def _register_class(cls):
+    assert _allow_registration
+    _registered_classes.add(cls)
 
 
 def extern(function):
     """Decorates a function declaration defined by the core device runtime."""
-    _register_module_of(function)
+    _register_function(function)
     return function
 
 
 def kernel(function_or_method):
     """Decorates a function or method to be executed on the core device."""
-    _register_module_of(function_or_method)
+    _register_function(function_or_method)
     argspec = getfullargspec(function_or_method)
     if argspec.args and argspec.args[0] == "self":
         @wraps(function_or_method)
@@ -59,7 +64,7 @@ def kernel(function_or_method):
 
 def portable(function):
     """Decorates a function or method to be executed on the same device (host/core device) as the caller."""
-    _register_module_of(function)
+    _register_function(function)
     return function
 
 
@@ -68,7 +73,7 @@ def nac3(cls):
     Decorates a class to be analyzed by NAC3.
     All classes containing kernels or portable methods must use this decorator.
     """
-    _register_module_of(cls)
+    _register_class(cls)
     return cls
 
 
diff --git a/flake.lock b/flake.lock
index 36927591b..5619ab884 100644
--- a/flake.lock
+++ b/flake.lock
@@ -82,11 +82,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1636252154,
-        "narHash": "sha256-YPRjk9M9GdUiAZircxUl5BDU8rN5t0B0ZkIau+dMqmg=",
+        "lastModified": 1636618109,
+        "narHash": "sha256-dJu9Tw+8mVL3ZMIlcvM9g6WT3hYrKDaVceZJV+BxjTA=",
         "ref": "master",
-        "rev": "50f1aca1aa7ecb4412aa4b8a0daa1cc644d8ca89",
-        "revCount": 412,
+        "rev": "c004da85f705fef8952885c1493407610a05b61d",
+        "revCount": 415,
         "type": "git",
         "url": "https://git.m-labs.hk/M-Labs/nac3.git"
       },