chore(ssa refactor): Implement mutable and immutable variables

crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs

@@ -126,6 +126,17 @@ impl DataFlowGraph {
         id
     }
 
+    /// Replace an instruction id with another.
+    ///
+    /// This function should generally be avoided if possible in favor of inserting new
+    /// instructions since it does not check whether the instruction results of the removed
+    /// instruction are still in use. Users of this function thus need to ensure the old
+    /// instruction's results are no longer in use or are otherwise compatible with the
+    /// new instruction's result count and types.
+    pub(crate) fn replace_instruction(&mut self, id: Id<Instruction>, instruction: Instruction) {
+        self.instructions[id] = instruction;
+    }
+
     /// Insert a value into the dfg's storage and return an id to reference it.
     /// Until the value is used in an instruction it is unreachable.
     pub(crate) fn make_value(&mut self, value: Value) -> ValueId {
@@ -141,8 +152,11 @@ impl DataFlowGraph {
 
     /// Attaches results to the instruction, clearing any previous results.
     ///
+    /// This does not normally need to be called manually as it is called within
+    /// make_instruction automatically.
+    ///
     /// Returns the results of the instruction
-    fn make_instruction_results(
+    pub(crate) fn make_instruction_results(
         &mut self,
         instruction_id: InstructionId,
         ctrl_typevars: Option<Vec<Type>>,

crates/noirc_evaluator/src/ssa_refactor/ssa_builder/function_builder.rs

@@ -5,7 +5,7 @@ use crate::ssa_refactor::ir::{
     function::{Function, FunctionId},
     instruction::{Binary, BinaryOp, Instruction, TerminatorInstruction},
     types::Type,
-    value::ValueId,
+    value::{Value, ValueId},
 };
 
 use super::SharedBuilderContext;
@@ -204,4 +204,30 @@ impl<'ssa> FunctionBuilder<'ssa> {
     pub(crate) fn terminate_with_return(&mut self, return_values: Vec<ValueId>) {
         self.terminate_block_with(TerminatorInstruction::Return { return_values });
     }
+
+    /// Mutates a load instruction into a store instruction.
+    ///
+    /// This function is used while generating ssa-form for assignments currently.
+    /// To re-use most of the expression infrastructure, the lvalue of an assignment
+    /// is compiled as an expression and to assign to it we replace the final load
+    /// (which should always be present to load a mutable value) with a store of the
+    /// assigned value.
+    pub(crate) fn mutate_load_into_store(&mut self, load_result: ValueId, value_to_store: ValueId) {
+        let (instruction, address) = match &self.current_function.dfg[load_result] {
+            Value::Instruction { instruction, .. } => {
+                match &self.current_function.dfg[*instruction] {
+                    Instruction::Load { address } => (*instruction, *address),
+                    other => {
+                        panic!("mutate_load_into_store: Expected Load instruction, found {other:?}")
+                    }
+                }
+            }
+            other => panic!("mutate_load_into_store: Expected Load instruction, found {other:?}"),
+        };
+
+        let store = Instruction::Store { address, value: value_to_store };
+        self.current_function.dfg.replace_instruction(instruction, store);
+        // Clear the results of the previous load for safety
+        self.current_function.dfg.make_instruction_results(instruction, None);
+    }
 }

crates/noirc_evaluator/src/ssa_refactor/ssa_gen/context.rs

@@ -21,6 +21,7 @@ type FunctionQueue = Vec<(ast::FuncId, IrFunctionId)>;
 
 pub(super) struct FunctionContext<'a> {
     definitions: HashMap<LocalId, Values>,
+
     pub(super) builder: FunctionBuilder<'a>,
     shared_context: &'a SharedContext,
 }
@@ -165,10 +166,55 @@ impl<'a> FunctionContext<'a> {
         address
     }
 
+    /// Define a local variable to be some Values that can later be retrieved
+    /// by calling self.lookup(id)
     pub(super) fn define(&mut self, id: LocalId, value: Values) {
         let existing = self.definitions.insert(id, value);
         assert!(existing.is_none(), "Variable {id:?} was defined twice in ssa-gen pass");
     }
+
+    /// Looks up the value of a given local variable. Expects the variable to have
+    /// been previously defined or panics otherwise.
+    pub(super) fn lookup(&self, id: LocalId) -> Values {
+        self.definitions.get(&id).expect("lookup: variable not defined").clone()
+    }
+
+    /// Extract the given field of the tuple. Panics if the given Values is not
+    /// a Tree::Branch or does not have enough fields.
+    pub(super) fn get_field(tuple: Values, field_index: usize) -> Values {
+        match tuple {
+            Tree::Branch(mut trees) => trees.remove(field_index),
+            Tree::Leaf(value) => {
+                unreachable!("Tried to extract tuple index {field_index} from non-tuple {value:?}")
+            }
+        }
+    }
+
+    /// Mutate lhs to equal rhs
+    pub(crate) fn assign(&mut self, lhs: Values, rhs: Values) {
+        match (lhs, rhs) {
+            (Tree::Branch(lhs_branches), Tree::Branch(rhs_branches)) => {
+                assert_eq!(lhs_branches.len(), rhs_branches.len());
+
+                for (lhs, rhs) in lhs_branches.into_iter().zip(rhs_branches) {
+                    self.assign(lhs, rhs);
+                }
+            }
+            (Tree::Leaf(lhs), Tree::Leaf(rhs)) => {
+                // Re-evaluating these should have no effect
+                let (lhs, rhs) = (lhs.eval(self), rhs.eval(self));
+
+                // Expect lhs to be previously evaluated. If it is a load we need to undo
+                // the load to get the address to store to.
+                self.builder.mutate_load_into_store(lhs, rhs);
+            }
+            (lhs, rhs) => {
+                unreachable!(
+                    "assign: Expected lhs and rhs values to match but found {lhs:?} and {rhs:?}"
+                )
+            }
+        }
+    }
 }
 
 /// True if the given operator cannot be encoded directly and needs

crates/noirc_evaluator/src/ssa_refactor/ssa_gen/mod.rs

@@ -63,16 +63,16 @@ impl<'a> FunctionContext<'a> {
     /// Codegen any non-tuple expression so that we can unwrap the Values
     /// tree to return a single value for use with most SSA instructions.
     fn codegen_non_tuple_expression(&mut self, expr: &Expression) -> ValueId {
-        match self.codegen_expression(expr) {
-            Tree::Branch(branches) => {
-                panic!("codegen_non_tuple_expression called on tuple {branches:?}")
-            }
-            Tree::Leaf(value) => value.eval(),
-        }
+        self.codegen_expression(expr).into_leaf().eval(self)
     }
 
-    fn codegen_ident(&mut self, _ident: &ast::Ident) -> Values {
-        todo!()
+    fn codegen_ident(&mut self, ident: &ast::Ident) -> Values {
+        match &ident.definition {
+            ast::Definition::Local(id) => self.lookup(*id).map(|value| value.eval(self).into()),
+            ast::Definition::Function(_) => todo!(),
+            ast::Definition::Builtin(_) => todo!(),
+            ast::Definition::LowLevel(_) => todo!(),
+        }
     }
 
     fn codegen_literal(&mut self, literal: &ast::Literal) -> Values {
@@ -107,9 +107,10 @@ impl<'a> FunctionContext<'a> {
         // Now we must manually store all the elements into the array
         let mut i = 0u128;
         for element in elements {
-            element.for_each(|value| {
+            element.for_each(|element| {
                 let address = self.make_offset(array, i);
-                self.builder.insert_store(address, value.eval());
+                let element = element.eval(self);
+                self.builder.insert_store(address, element);
                 i += 1;
             });
         }
@@ -145,15 +146,26 @@ impl<'a> FunctionContext<'a> {
 
     fn codegen_index(&mut self, index: &ast::Index) -> Values {
         let array = self.codegen_non_tuple_expression(&index.collection);
-        let base_offset = self.codegen_non_tuple_expression(&index.index);
+        self.codegen_array_index(array, &index.index, &index.element_type)
+    }
+
+    /// This is broken off from codegen_index so that it can also be
+    /// used to codegen a LValue::Index
+    fn codegen_array_index(
+        &mut self,
+        array: super::ir::value::ValueId,
+        index: &ast::Expression,
+        element_type: &ast::Type,
+    ) -> Values {
+        let base_offset = self.codegen_non_tuple_expression(index);
 
         // base_index = base_offset * type_size
-        let type_size = Self::convert_type(&index.element_type).size_of_type();
+        let type_size = Self::convert_type(element_type).size_of_type();
         let type_size = self.builder.field_constant(type_size as u128);
         let base_index = self.builder.insert_binary(base_offset, BinaryOp::Mul, type_size);
 
         let mut field_index = 0u128;
-        self.map_type(&index.element_type, |ctx, typ| {
+        self.map_type(element_type, |ctx, typ| {
             let offset = ctx.make_offset(base_index, field_index);
             field_index += 1;
             ctx.builder.insert_load(array, offset, typ).into()
@@ -221,11 +233,13 @@ impl<'a> FunctionContext<'a> {
                 ctx.builder.add_block_parameter(end_block, typ).into()
             });
 
-            self.builder.terminate_with_jmp(end_block, else_value.into_value_list());
+            let else_values = else_value.into_value_list(self);
+            self.builder.terminate_with_jmp(end_block, else_values);
 
             // Must also set the then block to jmp to the end now
             self.builder.switch_to_block(then_block);
-            self.builder.terminate_with_jmp(end_block, then_value.into_value_list());
+            let then_values = then_value.into_value_list(self);
+            self.builder.terminate_with_jmp(end_block, then_values);
             self.builder.switch_to_block(end_block);
         } else {
             // In the case we have no 'else', the 'else' block is actually the end block.
@@ -240,21 +254,30 @@ impl<'a> FunctionContext<'a> {
         Tree::Branch(vecmap(tuple, |expr| self.codegen_expression(expr)))
     }
 
-    fn codegen_extract_tuple_field(&mut self, tuple: &Expression, index: usize) -> Values {
-        match self.codegen_expression(tuple) {
-            Tree::Branch(mut trees) => trees.remove(index),
-            Tree::Leaf(value) => {
-                unreachable!("Tried to extract tuple index {index} from non-tuple {value:?}")
-            }
-        }
+    fn codegen_extract_tuple_field(&mut self, tuple: &Expression, field_index: usize) -> Values {
+        let tuple = self.codegen_expression(tuple);
+        Self::get_field(tuple, field_index)
     }
 
     fn codegen_call(&mut self, _call: &ast::Call) -> Values {
         todo!()
     }
 
-    fn codegen_let(&mut self, _let_expr: &ast::Let) -> Values {
-        todo!()
+    fn codegen_let(&mut self, let_expr: &ast::Let) -> Values {
+        let mut values = self.codegen_expression(&let_expr.expression);
+
+        if let_expr.mutable {
+            values.map_mut(|value| {
+                let value = value.eval(self);
+                // Size is always 1 here since we're recursively unpacking tuples
+                let alloc = self.builder.insert_allocate(1);
+                self.builder.insert_store(alloc, value);
+                alloc.into()
+            });
+        }
+
+        self.define(let_expr.id, values);
+        self.unit_value()
     }
 
     fn codegen_constrain(&mut self, expr: &Expression, _location: Location) -> Values {
@@ -263,8 +286,25 @@ impl<'a> FunctionContext<'a> {
         self.unit_value()
     }
 
-    fn codegen_assign(&mut self, _assign: &ast::Assign) -> Values {
-        todo!()
+    fn codegen_assign(&mut self, assign: &ast::Assign) -> Values {
+        let lhs = self.codegen_lvalue(&assign.lvalue);
+        let rhs = self.codegen_expression(&assign.expression);
+        self.assign(lhs, rhs);
+        self.unit_value()
+    }
+
+    fn codegen_lvalue(&mut self, lvalue: &ast::LValue) -> Values {
+        match lvalue {
+            ast::LValue::Ident(ident) => self.codegen_ident(ident),
+            ast::LValue::Index { array, index, element_type, location: _ } => {
+                let array = self.codegen_lvalue(array).into_leaf().eval(self);
+                self.codegen_array_index(array, index, element_type)
+            }
+            ast::LValue::MemberAccess { object, field_index } => {
+                let object = self.codegen_lvalue(object);
+                Self::get_field(object, *field_index)
+            }
+        }
     }
 
     fn codegen_semi(&mut self, expr: &Expression) -> Values {

crates/noirc_evaluator/src/ssa_refactor/ssa_gen/value.rs

@@ -4,25 +4,34 @@ use crate::ssa_refactor::ir::function::FunctionId as IrFunctionId;
 use crate::ssa_refactor::ir::types::Type;
 use crate::ssa_refactor::ir::value::ValueId as IrValueId;
 
-#[derive(Debug)]
+use super::context::FunctionContext;
+
+#[derive(Debug, Clone)]
 pub(super) enum Tree<T> {
     Branch(Vec<Tree<T>>),
     Leaf(T),
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Copy, Clone)]
 pub(super) enum Value {
     Normal(IrValueId),
     Function(IrFunctionId),
+
+    /// A mutable variable that must be loaded as the given type before being used
+    Mutable(IrValueId, Type),
 }
 
 impl Value {
     /// Evaluate a value, returning an IrValue from it.
-    /// This has no effect on Value::Normal, but any variables will be updated with their latest
-    /// use.
-    pub(super) fn eval(self) -> IrValueId {
+    /// This has no effect on Value::Normal, but any variables will
+    /// need to be loaded from memory
+    pub(super) fn eval(self, ctx: &mut FunctionContext) -> IrValueId {
         match self {
             Value::Normal(value) => value,
+            Value::Mutable(address, typ) => {
+                let offset = ctx.builder.field_constant(0u128);
+                ctx.builder.insert_load(address, offset, typ)
+            }
             Value::Function(_) => panic!("Tried to evaluate a function value"),
         }
     }
@@ -56,6 +65,37 @@ impl<T> Tree<T> {
             Tree::Leaf(value) => f(value),
         }
     }
+
+    pub(super) fn map_mut(&mut self, mut f: impl FnMut(&T) -> Tree<T>) {
+        self.map_mut_helper(&mut f);
+    }
+
+    fn map_mut_helper(&mut self, f: &mut impl FnMut(&T) -> Tree<T>) {
+        match self {
+            Tree::Branch(trees) => trees.iter_mut().for_each(|tree| tree.map_mut_helper(f)),
+            Tree::Leaf(value) => *self = f(value),
+        }
+    }
+
+    pub(super) fn map<U>(self, mut f: impl FnMut(T) -> Tree<U>) -> Tree<U> {
+        self.map_helper(&mut f)
+    }
+
+    fn map_helper<U>(self, f: &mut impl FnMut(T) -> Tree<U>) -> Tree<U> {
+        match self {
+            Tree::Branch(trees) => Tree::Branch(vecmap(trees, |tree| tree.map_helper(f))),
+            Tree::Leaf(value) => f(value),
+        }
+    }
+
+    /// Unwraps this Tree into the value of the leaf node. Panics if
+    /// this Tree is a Branch
+    pub(super) fn into_leaf(self) -> T {
+        match self {
+            Tree::Branch(_) => panic!("into_leaf called on a Tree::Branch"),
+            Tree::Leaf(value) => value,
+        }
+    }
 }
 
 impl From<IrValueId> for Values {
@@ -82,7 +122,7 @@ impl Tree<Type> {
 impl Tree<Value> {
     /// Flattens and evaluates this Tree<Value> into a list of ir values
     /// for return statements, branching instructions, or function parameters.
-    pub(super) fn into_value_list(self) -> Vec<IrValueId> {
-        vecmap(self.flatten(), Value::eval)
+    pub(super) fn into_value_list(self, ctx: &mut FunctionContext) -> Vec<IrValueId> {
+        vecmap(self.flatten(), |value| value.eval(ctx))
     }
 }