From 916b2a653e094b427287e7b24646297fbaac10b0 Mon Sep 17 00:00:00 2001
From: mars <mars@tebibyte.media>
Date: Sat, 26 Mar 2022 16:45:22 -0600
Subject: [PATCH] Mutable locals

---
 src/jit/engine.rs    | 28 ++++++++++++++++
 src/jit/translate.rs | 78 ++++++++++++++++++++++++++++++++++----------
 2 files changed, 88 insertions(+), 18 deletions(-)

diff --git a/src/jit/engine.rs b/src/jit/engine.rs
index 0c9c130..7740a9e 100644
--- a/src/jit/engine.rs
+++ b/src/jit/engine.rs
@@ -101,4 +101,32 @@ mod tests {
         let code_fn = jit_fn::<(i64, i64), i64>(source);
         assert_eq!(code_fn((2, 3)), 15);
     }
+
+    #[test]
+    fn let_shadowing() {
+        let source = r#"
+            let_shadowing(i64 a, i64 b) i64 {
+                let c = a + b;
+                let c = (c * b);
+                let c = (c + a);
+                let c = (c * a);
+                c
+            }"#;
+        let code_fn = jit_fn::<(i64, i64), i64>(source);
+        assert_eq!(code_fn((2, 3)), 34);
+    }
+
+    #[test]
+    fn let_mutable() {
+        let source = r#"
+            let_mutable(i64 a, i64 b) i64 {
+                let mut c = a + b;
+                c = (c * b);
+                c = (c + a);
+                c = (c * a);
+                c
+            }"#;
+        let code_fn = jit_fn::<(i64, i64), i64>(source);
+        assert_eq!(code_fn((2, 3)), 34);
+    }
 }
diff --git a/src/jit/translate.rs b/src/jit/translate.rs
index 7e42128..0458543 100644
--- a/src/jit/translate.rs
+++ b/src/jit/translate.rs
@@ -7,7 +7,14 @@ use cranelift::prelude::*;
 pub struct FunctionTranslator<'a> {
     pub int: types::Type,
     pub builder: FunctionBuilder<'a>,
-    pub locals: Vec<(String, Value)>,
+    pub locals: Vec<Local>,
+    pub variable_index: usize,
+}
+
+pub struct Local {
+    pub name: String,
+    pub var: Variable,
+    pub mutable: bool,
 }
 
 impl<'a> FunctionTranslator<'a> {
@@ -16,6 +23,7 @@ impl<'a> FunctionTranslator<'a> {
             int,
             builder,
             locals: Vec::new(),
+            variable_index: 0,
         }
     }
 
@@ -62,7 +70,8 @@ impl<'a> FunctionTranslator<'a> {
         for (i, arg) in signature.args.iter().enumerate() {
             if let Some(name) = arg.name {
                 let val = self.builder.block_params(entry_block)[i];
-                self.locals.push((name.to_string(), val));
+                let var = self.add_local(self.int, name, false);
+                self.builder.def_var(var, val);
             }
         }
 
@@ -86,12 +95,9 @@ impl<'a> FunctionTranslator<'a> {
                 self.translate_expr(expr);
             }
             Let { var, mutable, expr } => {
-                if *mutable {
-                    unimplemented!("Mutable variables")
-                }
-
                 let val = self.translate_expr(expr);
-                self.locals.push((var.to_string(), val));
+                let var = self.add_local(self.int, var, *mutable);
+                self.builder.def_var(var, val);
             }
             _ => unimplemented!(),
         }
@@ -100,26 +106,53 @@ impl<'a> FunctionTranslator<'a> {
     pub fn translate_expr(&mut self, expr: &ast::Expr) -> Value {
         use ast::Expr::*;
         match expr {
-            Local(name) => self
-                .get_local(name)
-                .expect(&format!("Unrecognized local {}", name)),
+            Local(name) => {
+                let var = self
+                    .get_local(name)
+                    .expect(&format!("Unrecognized local {}", name))
+                    .var;
+                self.builder.use_var(var)
+            }
             Literal(ast::Literal::DecimalInteger(literal)) => {
                 // TODO parse integers while building ast so that codegen doesn't have to
                 let val: i64 = literal.parse().unwrap();
                 self.builder.ins().iconst(self.int, val)
             }
-            BinaryOp(op, terms) => {
-                let lhs = self.translate_expr(&terms.0);
-                let rhs = self.translate_expr(&terms.1);
-                self.translate_binary_op(op, lhs, rhs)
-            }
+            BinaryOp(op, terms) => self.translate_binary_op(op, &terms.0, &terms.1),
             // TODO the AST doesn't need this either
             Group(expr) => self.translate_expr(expr),
             _ => unimplemented!("Expression: {:#?}", expr),
         }
     }
 
-    pub fn translate_binary_op(&mut self, op: &ast::BinaryOp, lhs: Value, rhs: Value) -> Value {
+    pub fn translate_binary_op(
+        &mut self,
+        op: &ast::BinaryOp,
+        lhs: &ast::Expr,
+        rhs: &ast::Expr,
+    ) -> Value {
+        let rhs = self.translate_expr(rhs);
+
+        if *op == ast::BinaryOp::Assign {
+            if let ast::Expr::Local(var) = lhs {
+                let local = self
+                    .get_local(var)
+                    .expect(&format!("Unrecognized local {}", var));
+
+                if !local.mutable {
+                    panic!("Attempted to assign to immutable variable {}", var);
+                }
+
+                let var = local.var;
+                self.builder.def_var(var, rhs);
+                return rhs;
+            } else {
+                unimplemented!("Assign to non-local lhs");
+            }
+        }
+
+        let lhs = self.translate_expr(lhs);
+
         use ast::BinaryOp::*;
         let ins = self.builder.ins();
         match op {
@@ -131,7 +164,16 @@ impl<'a> FunctionTranslator<'a> {
         }
     }
 
-    pub fn get_local(&self, name: &str) -> Option<Value> {
-        self.locals.iter().rev().find(|l| &l.0 == name).map(|l| l.1)
+    pub fn get_local(&self, name: &str) -> Option<&Local> {
+        self.locals.iter().rev().find(|l| &l.name == name)
+    }
+
+    pub fn add_local(&mut self, var_type: types::Type, name: &str, mutable: bool) -> Variable {
+        let name = name.to_string();
+        let var = Variable::new(self.variable_index);
+        self.builder.declare_var(var, var_type);
+        self.locals.push(Local { name, var, mutable });
+        self.variable_index += 1;
+        var
     }
 }