if else type analysis

README.md

@@ -5,12 +5,29 @@ A toy language I'm making to learn LLVM and compilers.
 Syntax is subject to change any time right now. It has a rusty style for now.
 
 ```
-fn add(a: i64, b: i64) -> i64 {
-    return a + b;
+struct Hello {
+    x: i32,
+    y: i32,
 }
 
-fn main() {
-    let x = 2 + 3;
-    let y = add(x, 4);
+fn test(x: Hello) {
+    return;
 }
+
+fn works(x: i64) -> i64 {
+    let z = 0;
+    if 2 == x {
+        z = x * 2;
+    } else {
+        z = x * 3;
+    }
+    return z;
+}
+
+fn main() -> i64 {
+    let y: i64 = 2;
+    let z = y;
+    return works(z);
+}
+
 ```

simple.ed

@@ -8,8 +8,8 @@ fn test(x: Hello) {
 }
 
 fn works(x: i64) -> i64 {
-    let z: i64 = 0;
-    if x == 2 {
+    let z = 0;
+    if 2 == x {
         z = x * 2;
     } else {
         z = x * 3;

src/codegen.rs

@@ -489,10 +489,11 @@ impl<'ctx> CodeGen<'ctx> {
         let (lhs, lhs_type) = self
             .compile_expression(lhs, variables, types)?
             .expect("should have result");
-        let (rhs, _rhs_type) = self
+        let (rhs, rhs_type) = self
             .compile_expression(rhs, variables, types)?
             .expect("should have result");
 
+        assert_eq!(lhs_type, rhs_type);
         let lhs = lhs.into_int_value();
         let rhs = rhs.into_int_value();
 

src/type_analysis.rs

@@ -10,6 +10,10 @@ pub fn type_inference(ast: &mut ast::Program) {
             let ret_type = function.return_type.clone();
             let mut var_cache: HashMap<String, TypeExp> = HashMap::new();
 
+            for arg in &function.params {
+                var_cache.insert(arg.ident.clone(), arg.type_exp.clone());
+            }
+
             if let Some(ret_type) = &ret_type {
                 let ret_type_exp = fn_return_type(function);
 
@@ -175,6 +179,7 @@ fn set_exp_types_from_cache(
             if let Some(value_type) = value_type {
                 // todo: check types matches?
                 var_cache.insert(name, value_type.clone());
+                *env = Some(value_type.clone());
             } else if var_cache.contains_key(&name) {
                 *value_type = var_cache.get(&name).cloned();
                 if env.is_none() {
@@ -183,10 +188,16 @@ fn set_exp_types_from_cache(
             }
         }
         Expression::BinaryOp(lhs, op, rhs) => match op {
-            ast::OpCode::Eq | ast::OpCode::Ne => {}
+            ast::OpCode::Eq | ast::OpCode::Ne => {
+                set_exp_types_from_cache(lhs, var_cache, env);
+                set_exp_types_from_cache(rhs, var_cache, env);
+                set_exp_types_from_cache(lhs, var_cache, env);
+                *env = Some(TypeExp::Boolean);
+            }
             _ => {
                 set_exp_types_from_cache(lhs, var_cache, env);
                 set_exp_types_from_cache(rhs, var_cache, env);
+                set_exp_types_from_cache(lhs, var_cache, env); // needed in case 2 == x
             }
         },
         Expression::Literal(lit) => match lit {
@@ -241,7 +252,7 @@ fn set_expression_type(
             }
         }
         Expression::BinaryOp(lhs, op, rhs) => match op {
-            ast::OpCode::Eq | ast::OpCode::Ne => {}
+            // ast::OpCode::Eq | ast::OpCode::Ne => {}
             _ => {
                 set_expression_type(lhs, expected_type, var_cache);
                 set_expression_type(rhs, expected_type, var_cache);