From ae09a0d444026af2a7ff4db3a8cfdc34d9960165 Mon Sep 17 00:00:00 2001
From: occheung <dc@m-labs.hk>
Date: Tue, 3 Dec 2024 16:45:05 +0800
Subject: [PATCH] exceptions: preallocate in NAC3 instead

---
 nac3artiq/demo/embedding_map.py | 27 ---------------------
 nac3artiq/src/lib.rs            | 43 ++++++++++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 28 deletions(-)

diff --git a/nac3artiq/demo/embedding_map.py b/nac3artiq/demo/embedding_map.py
index 2116c69e..a43af698 100644
--- a/nac3artiq/demo/embedding_map.py
+++ b/nac3artiq/demo/embedding_map.py
@@ -7,33 +7,6 @@ class EmbeddingMap:
         self.function_map = {}
         self.attributes_writeback = []
 
-        # preallocate exception names
-        self.preallocate_runtime_exception_names(["RuntimeError",
-                                          "RTIOUnderflow",
-                                          "RTIOOverflow",
-                                          "RTIODestinationUnreachable",
-                                          "DMAError",
-                                          "I2CError",
-                                          "CacheError",
-                                          "SPIError",
-                                          "0:ZeroDivisionError",
-                                          "0:IndexError",
-                                          "0:ValueError",
-                                          "0:RuntimeError",
-                                          "0:AssertionError",
-                                          "0:KeyError",
-                                          "0:NotImplementedError",
-                                          "0:OverflowError",
-                                          "0:IOError",
-                                          "0:UnwrapNoneError"])
-
-    def preallocate_runtime_exception_names(self, names):
-        for i, name in enumerate(names):
-            if ":" not in name:
-                name = "0:artiq.coredevice.exceptions." + name
-            exn_id = self.store_str(name)
-            assert exn_id == i
-
     def store_function(self, key, fun):
         self.function_map[key] = fun
         return key
diff --git a/nac3artiq/src/lib.rs b/nac3artiq/src/lib.rs
index b86c2534..f9cf7aa9 100644
--- a/nac3artiq/src/lib.rs
+++ b/nac3artiq/src/lib.rs
@@ -1071,6 +1071,47 @@ impl Nac3 {
         let working_directory = tempfile::Builder::new().prefix("nac3-").tempdir().unwrap();
         fs::write(working_directory.path().join("kernel.ld"), include_bytes!("kernel.ld")).unwrap();
 
+        let mut string_store: HashMap<String, i32> = Default::default();
+
+        // 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`
+        // Verify synchronization by running the test cases in `artiq.test.coredevice.test_exceptions`
+        let runtime_exception_names = [
+            "RTIOUnderflow",
+            "RTIOOverflow",
+            "RTIODestinationUnreachable",
+            "DMAError",
+            "I2CError",
+            "CacheError",
+            "SPIError",
+            "SubkernelError",
+            "0:AssertionError",
+            "0:AttributeError",
+            "0:IndexError",
+            "0:IOError",
+            "0:KeyError",
+            "0:NotImplementedError",
+            "0:OverflowError",
+            "0:RuntimeError",
+            "0:TimeoutError",
+            "0:TypeError",
+            "0:ValueError",
+            "0:ZeroDivisionError",
+            "0:LinAlgError",
+            "UnwrapNoneError",
+        ];
+
+        // Preallocate runtime exception names
+        for (i, name) in runtime_exception_names.iter().enumerate() {
+            let exn_name = if name.find(':').is_none() {
+                format!("0:artiq.coredevice.exceptions.{name}")
+            } else {
+                name.to_string()
+            };
+
+            string_store.insert(exn_name, i as i32);
+        }
+
         Ok(Nac3 {
             isa,
             time_fns,
@@ -1080,7 +1121,7 @@ impl Nac3 {
             top_levels: Vec::default(),
             pyid_to_def: Arc::default(),
             working_directory,
-            string_store: Arc::default(),
+            string_store: Arc::new(string_store.into()),
             exception_ids: Arc::default(),
             deferred_eval_store: DeferredEvaluationStore::new(),
             llvm_options: CodeGenLLVMOptions {