From 964945d244f09a37f55d476f9675a616442a425a Mon Sep 17 00:00:00 2001
From: occheung <dc@m-labs.hk>
Date: Tue, 3 Dec 2024 16:45:46 +0800
Subject: [PATCH] string_store: update embedding map after compilation

---
 nac3artiq/src/lib.rs             | 16 +++++++++++++++-
 nac3artiq/src/symbol_resolver.rs |  5 +----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/nac3artiq/src/lib.rs b/nac3artiq/src/lib.rs
index f9cf7aa9..a9d8fa1d 100644
--- a/nac3artiq/src/lib.rs
+++ b/nac3artiq/src/lib.rs
@@ -577,7 +577,7 @@ impl Nac3 {
             field_to_val: RwLock::default(),
             name_to_pyid,
             module: module.to_object(py),
-            helper,
+            helper: helper.clone(),
             string_store: self.string_store.clone(),
             exception_ids: self.exception_ids.clone(),
             deferred_eval_store: self.deferred_eval_store.clone(),
@@ -822,6 +822,20 @@ 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)
     }
 
diff --git a/nac3artiq/src/symbol_resolver.rs b/nac3artiq/src/symbol_resolver.rs
index 9e29acf8..fa57755d 100644
--- a/nac3artiq/src/symbol_resolver.rs
+++ b/nac3artiq/src/symbol_resolver.rs
@@ -1536,10 +1536,7 @@ impl SymbolResolver for Resolver {
         if let Some(id) = string_store.get(s) {
             *id
         } else {
-            let id = Python::with_gil(|py| -> PyResult<i32> {
-                self.0.helper.store_str.call1(py, (s,))?.extract(py)
-            })
-            .unwrap();
+            let id = string_store.len() as i32;
             string_store.insert(s.into(), id);
             id
         }