Add string registration in NAC

nac3artiq/src/lib.rs

@@ -378,6 +378,18 @@ impl Nac3 {
         py: Python,
         link_fn: &dyn Fn(&Module) -> PyResult<T>,
     ) -> PyResult<T> {
+        // Pre-register string arguments in string store
+        Python::with_gil(|py| {
+            let mut string_store = self.string_store.write();
+            for arg in &args {
+                if let Ok(s) = arg.extract::<String>() {
+                    if !string_store.contains_key(&s) {
+                        string_store.insert(s.clone(), i32::try_from(string_store.len()).unwrap());
+                    }
+                }
+            }
+        });
+
         let size_t = self.isa.get_size_type();
         let (mut composer, mut builtins_def, mut builtins_ty) = TopLevelComposer::new(
             self.builtins.clone(),
@@ -993,7 +1005,8 @@ impl Nac3 {
             Isa::CortexA9 | Isa::Host => &timeline::EXTERN_TIME_FNS,
         };
         let (primitive, _) = TopLevelComposer::make_primitives(isa.get_size_type());
-        let builtins = vec![
+        let builtins = {
+            let mut b = vec![
                 (
                     "now_mu".into(),
                     FunSignature { args: vec![], ret: primitive.int64, vars: VarMap::new() },
@@ -1042,6 +1055,14 @@ impl Nac3 {
                     }))),
                 ),
             ];
+            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();
@@ -1086,6 +1107,7 @@ impl Nac3 {
         fs::write(working_directory.path().join("kernel.ld"), include_bytes!("kernel.ld")).unwrap();
 
         let mut string_store: HashMap<String, i32> = 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`

nac3artiq/src/symbol_resolver.rs

@@ -679,6 +679,11 @@ impl InnerResolver {
             return Ok(Ok(ty));
         }
 
+        // Early check for string type
+        if let Ok(_s) = obj.extract::<String>() {
+            return Ok(Ok(primitives.str));
+        }
+
         let (extracted_ty, inst_check) = match self.get_pyty_obj_type(
             py,
             {
@@ -686,6 +691,7 @@ impl InnerResolver {
                     self.primitive_ids.typevar,
                     self.primitive_ids.generic_alias.0,
                     self.primitive_ids.generic_alias.1,
+                    self.primitive_ids.string, // Add string type to the check
                 ]
                 .contains(&self.helper.id_fn.call1(py, (ty.clone(),))?.extract::<u64>(py)?)
                 {
@@ -943,12 +949,19 @@ impl InnerResolver {
                         |_| Ok(Err(format!("{obj} is not in the range of float64"))),
                         |_| Ok(Ok(extracted_ty)),
                     )
+                } else if let Ok(s) = obj.extract::<String>() {
+                        if unifier.unioned(extracted_ty, primitives.str) {
+                            Ok(Ok(primitives.str))
+                        } else {
+                            Ok(Err(format!("expected str, got {s}")))
+                        }
                 } else {
                     Ok(Ok(extracted_ty))
                 }
                 }
             }
         }
+    }
 
     pub fn get_obj_value<'ctx>(
         &self,
@@ -988,7 +1001,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));
@@ -1480,7 +1493,6 @@ impl InnerResolver {
             Err("only primitives values, option and tuple can be default parameter value".into())
         })
     }
-}
 
 impl SymbolResolver for Resolver {
     fn get_default_param_value(&self, expr: &ast::Expr) -> Option<SymbolValue> {