From 107303aa98f5e1c8278c633739af519edc2ba65d Mon Sep 17 00:00:00 2001
From: sofia <sofia@teascade.net>
Date: Mon, 25 Aug 2025 21:12:34 +0300
Subject: [PATCH] Update generic types for let-statements

---
 examples/generics.reid   |  5 +++--
 reid/src/mir/generics.rs | 35 +++++++++++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/examples/generics.reid b/examples/generics.reid
index b7b2a78..3ddce32 100644
--- a/examples/generics.reid
+++ b/examples/generics.reid
@@ -1,9 +1,10 @@
 // Arithmetic, function calls and imports!
 
 fn test<T>(value: T) -> T {
-  return value;
+  let b: T = value;
+  return b;
 }
 
 fn main() -> u32 {
-  return test<u32>(15);
+  return test<u64>(15) as u32 + test<u32>(5);
 }
\ No newline at end of file
diff --git a/reid/src/mir/generics.rs b/reid/src/mir/generics.rs
index 606351b..1ce412c 100644
--- a/reid/src/mir/generics.rs
+++ b/reid/src/mir/generics.rs
@@ -29,10 +29,15 @@ pub struct Functions {
     assoc_calls: HashMap<(TypeKind, String), Calls>,
 }
 
-type GenericsPassState<'map, 'st, 'sc> = PassState<'st, 'sc, (), ErrorKind>;
+#[derive(Default, Clone)]
+pub struct GenericsPassData {
+    generic_types: HashMap<String, TypeKind>,
+}
+
+type GenericsPassState<'map, 'st, 'sc> = PassState<'st, 'sc, GenericsPassData, ErrorKind>;
 
 impl Pass for GenericsPass {
-    type Data = ();
+    type Data = GenericsPassData;
     type TError = ErrorKind;
 
     fn context(&mut self, context: &mut mir::Context, mut _state: PassState<Self::Data, Self::TError>) -> PassResult {
@@ -142,6 +147,32 @@ impl Pass for GenericsPass {
         }
         Ok(())
     }
+
+    fn function(
+        &mut self,
+        func: &mut FunctionDefinition,
+        mut state: PassState<Self::Data, Self::TError>,
+    ) -> PassResult {
+        for (name, ty) in &func.generics {
+            state.scope.data.generic_types.insert(name.clone(), ty.clone());
+        }
+        Ok(())
+    }
+
+    fn stmt(&mut self, stmt: &mut mir::Statement, mut state: PassState<Self::Data, Self::TError>) -> PassResult {
+        match &mut stmt.0 {
+            mir::StmtKind::Let(var_ref, _, _) => match var_ref.0.clone() {
+                TypeKind::CustomType(custom_type_key) => {
+                    if let Some(ty) = state.scope.data.generic_types.get(&custom_type_key.0) {
+                        var_ref.0 = ty.clone();
+                    }
+                }
+                _ => {}
+            },
+            _ => {}
+        }
+        Ok(())
+    }
 }
 
 impl mir::Block {