From f8ec9c19d6de1bf8cf7855574c3717c843427fd7 Mon Sep 17 00:00:00 2001
From: mars <mars@tebibyte.media>
Date: Sat, 26 Mar 2022 23:27:02 -0600
Subject: [PATCH] Translate if statements

---
 src/jit/engine.rs    | 19 +++++++++++++++++++
 src/jit/translate.rs | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/src/jit/engine.rs b/src/jit/engine.rs
index 0f29d61..7305fdf 100644
--- a/src/jit/engine.rs
+++ b/src/jit/engine.rs
@@ -171,4 +171,23 @@ mod tests {
         assert_eq!(code_fn((2, 4)), 22); // Test 1
         assert_eq!(code_fn((7, 3)), 65); // Test 2
     }
+
+    #[test]
+    fn if_statement() {
+        let source = r#"
+            if_statement(i64 a, i64 b) i64 {
+                let mut c = a + b;
+
+                if a < b {
+                    c = (c * a);
+                } else {
+                    c = (c * b);
+                }
+
+                c
+            }"#;
+        let code_fn = jit_fn::<(i64, i64), i64>(source);
+        assert_eq!(code_fn((2, 4)), 12);
+        assert_eq!(code_fn((7, 3)), 30);
+    }
 }
diff --git a/src/jit/translate.rs b/src/jit/translate.rs
index fc747c3..9f35b5e 100644
--- a/src/jit/translate.rs
+++ b/src/jit/translate.rs
@@ -99,6 +99,7 @@ impl<'a> FunctionTranslator<'a> {
                 let var = self.add_local(self.int, var, *mutable);
                 self.builder.def_var(var, val);
             }
+            If(if_stmt) => self.translate_if_stmt(if_stmt),
             _ => unimplemented!(),
         }
     }
@@ -160,6 +161,41 @@ impl<'a> FunctionTranslator<'a> {
         self.builder.block_params(merge_block)[0]
     }
 
+    pub fn translate_if_stmt(&mut self, if_stmt: &ast::IfStmt) {
+        let ast::IfStmt {
+            test_expr,
+            then_body,
+            else_body,
+        } = if_stmt;
+
+        let test_val = self.translate_expr(test_expr);
+
+        let then_block = self.builder.create_block();
+        let else_info = else_body.as_ref().map(|b| (b, self.builder.create_block()));
+        let merge_block = self.builder.create_block();
+
+        let else_jump = else_info.map(|b| b.1).unwrap_or(merge_block);
+        self.builder.ins().brz(test_val, else_jump, &[]);
+        self.builder.ins().jump(then_block, &[]);
+
+        let mut translate_branch = |block, body| {
+            self.builder.switch_to_block(block);
+            self.builder.seal_block(block);
+            let val = self.translate_branch_body(body);
+            assert_eq!(val, None, "If statement branch has tail expression");
+            self.builder.ins().jump(merge_block, &[]);
+        };
+
+        translate_branch(then_block, then_body);
+
+        if let Some((else_body, else_block)) = else_info {
+            translate_branch(else_block, else_body);
+        }
+
+        self.builder.switch_to_block(merge_block);
+        self.builder.seal_block(merge_block);
+    }
+
     pub fn translate_branch_body(&mut self, branch_body: &ast::BranchBody) -> Option<Value> {
         let scope_size = self.locals.len();