diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs index 9b713eee06e..ab2018b1df8 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs +++ b/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) { + 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>, diff --git a/crates/noirc_evaluator/src/ssa_refactor/ssa_builder/function_builder.rs b/crates/noirc_evaluator/src/ssa_refactor/ssa_builder/function_builder.rs index c0a94be6f80..d11e9a763cd 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ssa_builder/function_builder.rs +++ b/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) { 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); + } } diff --git a/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/context.rs b/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/context.rs index 48175ebb52b..10206e28c2d 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/context.rs +++ b/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, + 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 diff --git a/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/mod.rs b/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/mod.rs index f8faf8eeeb4..a7880032d42 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/mod.rs +++ b/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 { diff --git a/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/value.rs b/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/value.rs index 31a93374940..52ff52d75f2 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/value.rs +++ b/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 { Branch(Vec>), 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 Tree { Tree::Leaf(value) => f(value), } } + + pub(super) fn map_mut(&mut self, mut f: impl FnMut(&T) -> Tree) { + self.map_mut_helper(&mut f); + } + + fn map_mut_helper(&mut self, f: &mut impl FnMut(&T) -> Tree) { + 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(self, mut f: impl FnMut(T) -> Tree) -> Tree { + self.map_helper(&mut f) + } + + fn map_helper(self, f: &mut impl FnMut(T) -> Tree) -> Tree { + 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 for Values { @@ -82,7 +122,7 @@ impl Tree { impl Tree { /// Flattens and evaluates this Tree into a list of ir values /// for return statements, branching instructions, or function parameters. - pub(super) fn into_value_list(self) -> Vec { - vecmap(self.flatten(), Value::eval) + pub(super) fn into_value_list(self, ctx: &mut FunctionContext) -> Vec { + vecmap(self.flatten(), |value| value.eval(ctx)) } }