Merge upstream into feature/invk

nac3artiq/demo/min_artiq.py

@@ -4,6 +4,7 @@ from types import SimpleNamespace
 from numpy import int32, int64
 from typing import Generic, TypeVar
 from math import floor, ceil
+import builtins
 
 import nac3artiq
 
@@ -16,7 +17,8 @@ __all__ = [
     "rpc", "ms", "us", "ns",
     "print_int32", "print_int64",
     "Core", "TTLOut",
-    "parallel", "legacy_parallel", "sequential"
+    "parallel", "legacy_parallel", "sequential",
+    "StringWrapper"
 ]
 
 
@@ -90,6 +92,7 @@ artiq_builtins = {
     "virtual": virtual,
     "_ConstGenericMarker": _ConstGenericMarker,
     "Option": Option,
+    "str": builtins.str, 
 }
 compiler = nac3artiq.NAC3(core_arguments["target"], artiq_builtins)
 allow_registration = True
@@ -243,6 +246,8 @@ class Core:
         global allow_registration
 
         embedding = EmbeddingMap()
+        for value, str_id in sorted(string_store.items(), key=lambda x: x[1]):
+            embedding.string_map[value] = str_id
 
         if allow_registration:
             compiler.analyze(registered_functions, registered_classes, special_ids, set())
@@ -330,6 +335,36 @@ class KernelContextManager:
     def __exit__(self):
         pass
 
+@nac3
+class StringWrapper:
+    """Wrapper for Python strings in NAC3"""
+    artiq_builtin = True
+    _value: str
+    _id: int 
+
+    def __init__(self, value: str):
+        global next_string_id
+        self._value = value
+        if value not in string_store:
+            string_store[value] = next_string_id
+            next_string_id += 1
+        self._id = string_store[value]
+
+    def __str__(self):
+        return self._value
+
+    def get_identifier(self) -> int:
+        return self._id
+    
+string_store = {}
+NAC3_INTERNAL_STRINGS = {
+    "0:artiq.coredevice.exceptions.RTIOUnderflow": 0,
+    "": 1,
+}
+
+for s, id in NAC3_INTERNAL_STRINGS.items():
+    string_store[s] = id
+next_string_id = max(NAC3_INTERNAL_STRINGS.values()) + 1
 @nac3
 class UnwrapNoneError(Exception):
     """raised when unwrapping a none value"""

nac3artiq/demo/nac3artiq.so

@@ -1 +0,0 @@
-../../target/release/libnac3artiq.so

nac3artiq/demo/test_cache.py

@@ -0,0 +1,20 @@
+from min_artiq import *
+from numpy import int32
+
+@nac3
+class StringListTest:
+    core: KernelInvariant[Core]
+
+    def __init__(self):
+        self.core = Core()
+
+    @kernel 
+    def test_string_and_list(self, key: str, value: list[int32]):
+        print_int32(int32(42))
+
+def test_params():
+    exp = StringListTest()
+    exp.test_string_and_list("x4", [])
+
+if __name__ == "__main__":
+    test_params()

nac3artiq/src/lib.rs

@@ -440,6 +440,18 @@ impl Nac3 {
         py: Python,
         link_fn: &dyn Fn(&Module) -> PyResult<T>,
     ) -> PyResult<T> {
+        Python::with_gil(|_py| {
+            let string_map_py = embedding_map.getattr("string_map")?;
+            let reverse_map_py = embedding_map.getattr("string_reverse_map")?;
+
+            let string_store = self.string_store.read();
+            for (s, key) in string_store.iter() {
+                string_map_py.set_item(s, key)?;
+                reverse_map_py.set_item(key, s)?;
+            }
+            Ok::<_, PyErr>(())
+        })?;
+
         let size_t = self.isa.get_size_type(&Context::create());
         let (mut composer, mut builtins_def, mut builtins_ty) = TopLevelComposer::new(
             self.builtins.clone(),
@@ -643,10 +655,16 @@ impl Nac3 {
         name_to_pyid.insert("base".into(), id_fun.call1((obj,))?.extract()?);
         let mut arg_names = vec![];
         for (i, arg) in args.into_iter().enumerate() {
-            let name = format!("tmp{i}");
+            if let Ok(st) = arg.extract::<String>() {
-            module.add(&name, arg)?;
+                let literal = format!("{st:?}");
-            name_to_pyid.insert(name.clone().into(), id_fun.call1((arg,))?.extract()?);
+                arg_names.push(literal);
-            arg_names.push(name);
+            } else {
+                let tmp_name = format!("tmp{i}");
+                module.add(&tmp_name, arg)?;
+                let pyid_val: u64 = id_fun.call1((arg,))?.extract()?;
+                name_to_pyid.insert(tmp_name.clone().into(), pyid_val);
+                arg_names.push(tmp_name);
+            }
         }
         let synthesized = if method_name.is_empty() {
             format!("def __modinit__():\n    base({})", arg_names.join(", "))
@@ -926,20 +944,6 @@ impl Nac3 {
             panic!("Failed to run optimization for module `main`: {}", err.to_string());
         }
 
-        Python::with_gil(|py| {
-            let string_store = self.string_store.read();
-            let mut string_store_vec = string_store.iter().collect::<Vec<_>>();
-            string_store_vec.sort_by(|(_s1, key1), (_s2, key2)| key1.cmp(key2));
-            for (s, key) in string_store_vec {
-                let embed_key: i32 = helper.store_str.call1(py, (s,)).unwrap().extract(py).unwrap();
-                assert_eq!(
-                    embed_key, *key,
-                    "string {s} is out of sync between embedding map (key={embed_key}) and \
-                    the internal string store (key={key})"
-                );
-            }
-        });
-
         link_fn(&main)
     }
 
@@ -1056,55 +1060,64 @@ impl Nac3 {
         };
         let (primitive, _) =
             TopLevelComposer::make_primitives(isa.get_size_type(&Context::create()));
-        let builtins = vec![
+        let builtins = {
-            (
+            let mut b = vec![
-                "now_mu".into(),
+                (
-                FunSignature { args: vec![], ret: primitive.int64, vars: VarMap::new() },
+                    "now_mu".into(),
-                Arc::new(GenCall::new(Box::new(move |ctx, _, _, _, _| {
+                    FunSignature { args: vec![], ret: primitive.int64, vars: VarMap::new() },
-                    Ok(Some(time_fns.emit_now_mu(ctx)))
+                    Arc::new(GenCall::new(Box::new(move |ctx, _, _, _, _| {
-                }))),
+                        Ok(Some(time_fns.emit_now_mu(ctx)))
-            ),
+                    }))),
-            (
+                ),
-                "at_mu".into(),
+                (
-                FunSignature {
+                    "at_mu".into(),
-                    args: vec![FuncArg {
+                    FunSignature {
-                        name: "t".into(),
+                        args: vec![FuncArg {
-                        ty: primitive.int64,
+                            name: "t".into(),
-                        default_value: None,
+                            ty: primitive.int64,
-                        is_vararg: false,
+                            default_value: None,
-                    }],
+                            is_vararg: false,
-                    ret: primitive.none,
+                        }],
-                    vars: VarMap::new(),
+                        ret: primitive.none,
-                },
+                        vars: VarMap::new(),
-                Arc::new(GenCall::new(Box::new(move |ctx, _, fun, args, generator| {
+                    },
-                    let arg_ty = fun.0.args[0].ty;
+                    Arc::new(GenCall::new(Box::new(move |ctx, _, fun, args, generator| {
-                    let arg =
+                        let arg_ty = fun.0.args[0].ty;
-                        args[0].1.clone().to_basic_value_enum(ctx, generator, arg_ty).unwrap();
+                        let arg =
-                    time_fns.emit_at_mu(ctx, arg);
+                            args[0].1.clone().to_basic_value_enum(ctx, generator, arg_ty).unwrap();
-                    Ok(None)
+                        time_fns.emit_at_mu(ctx, arg);
-                }))),
+                        Ok(None)
-            ),
+                    }))),
-            (
+                ),
-                "delay_mu".into(),
+                (
-                FunSignature {
+                    "delay_mu".into(),
-                    args: vec![FuncArg {
+                    FunSignature {
-                        name: "dt".into(),
+                        args: vec![FuncArg {
-                        ty: primitive.int64,
+                            name: "dt".into(),
-                        default_value: None,
+                            ty: primitive.int64,
-                        is_vararg: false,
+                            default_value: None,
-                    }],
+                            is_vararg: false,
-                    ret: primitive.none,
+                        }],
-                    vars: VarMap::new(),
+                        ret: primitive.none,
-                },
+                        vars: VarMap::new(),
-                Arc::new(GenCall::new(Box::new(move |ctx, _, fun, args, generator| {
+                    },
-                    let arg_ty = fun.0.args[0].ty;
+                    Arc::new(GenCall::new(Box::new(move |ctx, _, fun, args, generator| {
-                    let arg =
+                        let arg_ty = fun.0.args[0].ty;
-                        args[0].1.clone().to_basic_value_enum(ctx, generator, arg_ty).unwrap();
+                        let arg =
-                    time_fns.emit_delay_mu(ctx, arg);
+                            args[0].1.clone().to_basic_value_enum(ctx, generator, arg_ty).unwrap();
-                    Ok(None)
+                        time_fns.emit_delay_mu(ctx, arg);
-                }))),
+                        Ok(None)
-            ),
+                    }))),
-        ];
+                ),
+            ];
+            b.push((
+                "str".into(),
+                FunSignature { args: vec![], ret: primitive.str, vars: VarMap::new() },
+                Arc::new(GenCall::new(Box::new(|_, _, _, _, _| Ok(None)))),
+            ));
+
+            b
+        };
 
         let builtins_mod = PyModule::import(py, "builtins").unwrap();
         let id_fn = builtins_mod.getattr("id").unwrap();

nac3artiq/src/symbol_resolver.rs

@@ -720,7 +720,6 @@ impl InnerResolver {
             self.pyid_to_type.write().insert(py_obj_id, ty);
             return Ok(Ok(ty));
         }
-
         let (extracted_ty, inst_check) = match self.get_pyty_obj_type(
             py,
             {
@@ -728,6 +727,7 @@ impl InnerResolver {
                     self.primitive_ids.typevar,
                     self.primitive_ids.generic_alias.0,
                     self.primitive_ids.generic_alias.1,
+                    self.primitive_ids.string,
                 ]
                 .contains(&self.helper.id_fn.call1(py, (ty.clone(),))?.extract::<u64>(py)?)
                 {
@@ -985,6 +985,11 @@ impl InnerResolver {
                         |_| Ok(Err(format!("{obj} is not in the range of float64"))),
                         |_| Ok(Ok(extracted_ty)),
                     )
+                } else if unifier.unioned(extracted_ty, primitives.str) {
+                    obj.extract::<String>().map_or_else(
+                        |_| Ok(Err(format!("{obj} is not a valid string"))),
+                        |_| Ok(Ok(primitives.str)),
+                    )
                 } else {
                     Ok(Ok(extracted_ty))
                 }
@@ -1030,7 +1035,7 @@ impl InnerResolver {
         } else if ty_id == self.primitive_ids.string || ty_id == self.primitive_ids.np_str_ {
             let val: String = obj.extract().unwrap();
             self.id_to_primitive.write().insert(id, PrimitiveValue::Str(val.clone()));
-            Ok(Some(ctx.gen_string(generator, val).into()))
+            return Ok(Some(ctx.gen_string(generator, val).into()));
         } else if ty_id == self.primitive_ids.float || ty_id == self.primitive_ids.float64 {
             let val: f64 = obj.extract().unwrap();
             self.id_to_primitive.write().insert(id, PrimitiveValue::F64(val));