diff --git a/nac3artiq/demo/min_artiq.py b/nac3artiq/demo/min_artiq.py index 62d32cc..10d7b88 100644 --- a/nac3artiq/demo/min_artiq.py +++ b/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 from embedding_map import EmbeddingMap @@ -17,7 +18,8 @@ __all__ = [ "rpc", "ms", "us", "ns", "print_int32", "print_int64", "Core", "TTLOut", - "parallel", "sequential" + "parallel", "sequential", + "StringWrapper" ] @@ -91,6 +93,7 @@ artiq_builtins = { "virtual": virtual, "_ConstGenericMarker": _ConstGenericMarker, "Option": Option, + "str": builtins.str, } compiler = nac3artiq.NAC3(core_arguments["target"], artiq_builtins) allow_registration = True @@ -204,6 +207,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, set()) @@ -291,6 +296,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""" diff --git a/nac3artiq/src/lib.rs b/nac3artiq/src/lib.rs index 6ab9f45..01d40d7 100644 --- a/nac3artiq/src/lib.rs +++ b/nac3artiq/src/lib.rs @@ -379,16 +379,16 @@ impl Nac3 { link_fn: &dyn Fn(&Module) -> PyResult, ) -> PyResult { Python::with_gil(|_py| { - let mut string_store = self.string_store.write(); - for arg in &args { - if let Ok(s) = arg.extract::() { - if !string_store.contains_key(&s) { - let next_id = i32::try_from(string_store.len()).unwrap(); - string_store.insert(s.clone(), next_id); - } - } + 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(); let (mut composer, mut builtins_def, mut builtins_ty) = TopLevelComposer::new( @@ -565,10 +565,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}"); - module.add(&name, arg)?; - name_to_pyid.insert(name.clone().into(), id_fun.call1((arg,))?.extract()?); - arg_names.push(name); + if let Ok(st) = arg.extract::() { + let literal = format!("{st:?}"); + arg_names.push(literal); + } 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(", ")) @@ -834,20 +840,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::>(); - 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) } @@ -1107,7 +1099,6 @@ impl Nac3 { fs::write(working_directory.path().join("kernel.ld"), include_bytes!("kernel.ld")).unwrap(); let mut string_store: HashMap = HashMap::default(); - string_store.insert(String::new(), 0); // Keep this list of exceptions in sync with `EXCEPTION_ID_LOOKUP` in `artiq::firmware::ksupport::eh_artiq` // The exceptions declared here must be defined in `artiq.coredevice.exceptions`