Update generic types for let-statements

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);
 }

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 {