From ad1b4f1b9bba39e7018c04bba7d8e477a8b6f860 Mon Sep 17 00:00:00 2001
From: abdul124 <ar@m-labs.hk>
Date: Wed, 19 Jun 2024 17:53:55 +0800
Subject: [PATCH] nac3artiq: allow class attribute access without init function

---
 nac3artiq/demo/support_class_attr_issue102.py | 40 +++++++++++++++++++
 nac3artiq/src/symbol_resolver.rs              | 37 ++++++++++++++---
 2 files changed, 71 insertions(+), 6 deletions(-)
 create mode 100644 nac3artiq/demo/support_class_attr_issue102.py

diff --git a/nac3artiq/demo/support_class_attr_issue102.py b/nac3artiq/demo/support_class_attr_issue102.py
new file mode 100644
index 0000000..1b93144
--- /dev/null
+++ b/nac3artiq/demo/support_class_attr_issue102.py
@@ -0,0 +1,40 @@
+from min_artiq import *
+from numpy import int32
+
+
+@nac3
+class Demo:
+    attr1: KernelInvariant[int32] = 2
+    attr2: int32 = 4
+    attr3: Kernel[int32]
+
+    @kernel
+    def __init__(self):
+        self.attr3 = 8
+
+
+@nac3
+class NAC3Devices:
+    core: KernelInvariant[Core]
+    attr4: KernelInvariant[int32] = 16
+    
+    def __init__(self):
+        self.core = Core()
+
+    @kernel
+    def run(self):
+        Demo.attr1  # Supported
+        # Demo.attr2  # Field not accessible on Kernel
+        # Demo.attr3  # Only attributes can be accessed in this way
+        # Demo.attr1 = 2  # Attributes are immutable
+
+        self.attr4  # Attributes can be accessed within class
+
+        obj = Demo()
+        obj.attr1  # Attributes can be accessed by class objects
+
+        NAC3Devices.attr4  # Attributes accessible for classes without __init__
+
+
+if __name__ == "__main__":
+    NAC3Devices().run()
diff --git a/nac3artiq/src/symbol_resolver.rs b/nac3artiq/src/symbol_resolver.rs
index b9fb6ba..62597b8 100644
--- a/nac3artiq/src/symbol_resolver.rs
+++ b/nac3artiq/src/symbol_resolver.rs
@@ -627,12 +627,15 @@ impl InnerResolver {
         let pyid_to_def = self.pyid_to_def.read();
         let constructor_ty = pyid_to_def.get(&py_obj_id).and_then(|def_id| {
             defs.iter().find_map(|def| {
-                if let TopLevelDef::Class { object_id, methods, constructor, .. } = &*def.read() {
-                    if object_id == def_id
-                        && constructor.is_some()
-                        && methods.iter().any(|(s, _, _)| s == &"__init__".into())
+                if let Some(rear_guard) = def.try_read() {
+                    if let TopLevelDef::Class { object_id, methods, constructor, .. } = &*rear_guard
                     {
-                        return *constructor;
+                        if object_id == def_id
+                            && constructor.is_some()
+                            && methods.iter().any(|(s, _, _)| s == &"__init__".into())
+                        {
+                            return *constructor;
+                        }
                     }
                 }
                 None
@@ -664,7 +667,29 @@ impl InnerResolver {
             primitives,
         )? {
             Ok(s) => s,
-            Err(e) => return Ok(Err(e)),
+            Err(e) => {
+                // Allow access to Class Attributes of Classes without having to initialize Objects
+                if self.pyid_to_def.read().contains_key(&py_obj_id) {
+                    if let Some(def_id) = self.pyid_to_def.read().get(&py_obj_id).copied() {
+                        let def = defs[def_id.0].read();
+                        let TopLevelDef::Class { object_id, .. } = &*def else {
+                            // only object is supported, functions are not supported
+                            unreachable!("function type is not supported, should not be queried")
+                        };
+
+                        let ty = TypeEnum::TObj {
+                            obj_id: *object_id,
+                            params: VarMap::new(),
+                            fields: HashMap::new(),
+                        };
+                        (unifier.add_ty(ty), true)
+                    } else {
+                        return Ok(Err(e));
+                    }
+                } else {
+                    return Ok(Err(e));
+                }
+            }
         };
         match (&*unifier.get_ty(extracted_ty), inst_check) {
             // do the instantiation for these four types