From 5f692debd8bd77e0dd5febdbf070055466f25196 Mon Sep 17 00:00:00 2001
From: David Mak <csdavidmak@ust.hk>
Date: Wed, 13 Dec 2023 18:23:32 +0800
Subject: [PATCH] core: Add PrimitiveStore into Unifier

This will be used during unification between a const generic variable
and a `Literal`.
---
 nac3core/src/codegen/mod.rs                    |  1 +
 nac3core/src/toplevel/helper.rs                |  1 +
 nac3core/src/typecheck/type_inferencer/test.rs |  3 +++
 nac3core/src/typecheck/typedef/mod.rs          | 15 ++++++++++++++-
 4 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/nac3core/src/codegen/mod.rs b/nac3core/src/codegen/mod.rs
index 0917d8c15..51496e699 100644
--- a/nac3core/src/codegen/mod.rs
+++ b/nac3core/src/codegen/mod.rs
@@ -580,6 +580,7 @@ pub fn gen_func_impl<'ctx, G: CodeGenerator, F: FnOnce(&mut G, &mut CodeGenConte
         let (unifier, primitives) = &top_level_ctx.unifiers.read()[task.unifier_index];
         (Unifier::from_shared_unifier(unifier), *primitives)
     };
+    unifier.put_primitive_store(&primitives);
     unifier.top_level = Some(top_level_ctx.clone());
 
     let mut cache = HashMap::new();
diff --git a/nac3core/src/toplevel/helper.rs b/nac3core/src/toplevel/helper.rs
index e11de9986..0f4d37558 100644
--- a/nac3core/src/toplevel/helper.rs
+++ b/nac3core/src/toplevel/helper.rs
@@ -145,6 +145,7 @@ impl TopLevelComposer {
             exception,
             option,
         };
+        unifier.put_primitive_store(&primitives);
         crate::typecheck::magic_methods::set_primitives_magic_methods(&primitives, &mut unifier);
         (primitives, unifier)
     }
diff --git a/nac3core/src/typecheck/type_inferencer/test.rs b/nac3core/src/typecheck/type_inferencer/test.rs
index 4589c839f..f3ac2d08d 100644
--- a/nac3core/src/typecheck/type_inferencer/test.rs
+++ b/nac3core/src/typecheck/type_inferencer/test.rs
@@ -148,6 +148,7 @@ impl TestEnvironment {
             uint64,
             option,
         };
+        unifier.put_primitive_store(&primitives);
         set_primitives_magic_methods(&primitives, &mut unifier);
 
         let id_to_name = [
@@ -296,6 +297,8 @@ impl TestEnvironment {
             option,
         };
 
+        unifier.put_primitive_store(&primitives);
+
         let (v0, id) = unifier.get_dummy_var();
 
         let foo_ty = unifier.add_ty(TypeEnum::TObj {
diff --git a/nac3core/src/typecheck/typedef/mod.rs b/nac3core/src/typecheck/typedef/mod.rs
index 85dfedc26..03802e186 100644
--- a/nac3core/src/typecheck/typedef/mod.rs
+++ b/nac3core/src/typecheck/typedef/mod.rs
@@ -13,6 +13,7 @@ use super::type_error::{TypeError, TypeErrorKind};
 use super::unification_table::{UnificationKey, UnificationTable};
 use crate::symbol_resolver::SymbolValue;
 use crate::toplevel::{DefinitionId, TopLevelContext, TopLevelDef};
+use crate::typecheck::type_inferencer::PrimitiveStore;
 
 #[cfg(test)]
 mod test;
@@ -211,7 +212,8 @@ pub struct Unifier {
     pub(crate) calls: Vec<Rc<Call>>,
     var_id: u32,
     unify_cache: HashSet<(Type, Type)>,
-    snapshot: Option<(usize, u32)>
+    snapshot: Option<(usize, u32)>,
+    primitive_store: Option<PrimitiveStore>,
 }
 
 impl Default for Unifier {
@@ -231,9 +233,19 @@ impl Unifier {
             unify_cache: HashSet::new(),
             top_level: None,
             snapshot: None,
+            primitive_store: None,
         }
     }
 
+    /// Sets the [PrimitiveStore] instance within this `Unifier`.
+    ///
+    /// This function can only be invoked once. Any subsequent invocations will result in an
+    /// assertion error..
+    pub fn put_primitive_store(&mut self, primitives: &PrimitiveStore) {
+        assert!(self.primitive_store.is_none());
+        self.primitive_store.replace(primitives.clone());
+    }
+
     pub unsafe fn get_unification_table(&mut self) -> &mut UnificationTable<Rc<TypeEnum>> {
         &mut self.unification_table
     }
@@ -252,6 +264,7 @@ impl Unifier {
             top_level: None,
             unify_cache: HashSet::new(),
             snapshot: None,
+            primitive_store: None,
         }
     }