chore: Add Instruction::map_values_mut (#6756)

Co-authored-by: Tom French <[email protected]>
Co-authored-by: Ary Borenszweig <[email protected]>

compiler/noirc_evaluator/src/ssa/function_builder/data_bus.rs

@@ -94,6 +94,20 @@ impl DataBus {
         DataBus { call_data, return_data: self.return_data.map(&mut f) }
     }
 
+    /// Updates the databus values in place with the provided function
+    pub(crate) fn map_values_mut(&mut self, mut f: impl FnMut(ValueId) -> ValueId) {
+        for cd in self.call_data.iter_mut() {
+            cd.array_id = f(cd.array_id);
+
+            // Can't mutate a hashmap's keys so we need to collect into a new one.
+            cd.index_map = cd.index_map.iter().map(|(k, v)| (f(*k), *v)).collect();
+        }
+
+        if let Some(data) = self.return_data.as_mut() {
+            *data = f(*data);
+        }
+    }
+
     pub(crate) fn call_data_array(&self) -> Vec<(u32, ValueId)> {
         self.call_data.iter().map(|cd| (cd.call_data_id, cd.array_id)).collect()
     }

compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs

@@ -73,25 +73,24 @@ impl<'f> FunctionInserter<'f> {
 
     /// Get an instruction and make sure all the values in it are freshly resolved.
     pub(crate) fn map_instruction(&mut self, id: InstructionId) -> (Instruction, CallStack) {
-        (
-            self.function.dfg[id].clone().map_values(|id| self.resolve(id)),
-            self.function.dfg.get_call_stack(id),
-        )
+        let mut instruction = self.function.dfg[id].clone();
+        instruction.map_values_mut(|id| self.resolve(id));
+        (instruction, self.function.dfg.get_call_stack(id))
     }
 
     /// Maps a terminator in place, replacing any ValueId in the terminator with the
     /// resolved version of that value id from this FunctionInserter's internal value mapping.
     pub(crate) fn map_terminator_in_place(&mut self, block: BasicBlockId) {
         let mut terminator = self.function.dfg[block].take_terminator();
-        terminator.mutate_values(|value| self.resolve(value));
+        terminator.map_values_mut(|value| self.resolve(value));
         self.function.dfg[block].set_terminator(terminator);
     }
 
     /// Maps the data bus in place, replacing any ValueId in the data bus with the
     /// resolved version of that value id from this FunctionInserter's internal value mapping.
     pub(crate) fn map_data_bus_in_place(&mut self) {
-        let data_bus = self.function.dfg.data_bus.clone();
-        let data_bus = data_bus.map_values(|value| self.resolve(value));
+        let mut data_bus = self.function.dfg.data_bus.clone();
+        data_bus.map_values_mut(|value| self.resolve(value));
         self.function.dfg.data_bus = data_bus;
     }
 

compiler/noirc_evaluator/src/ssa/ir/instruction.rs

@@ -669,6 +669,70 @@ impl Instruction {
         }
     }
 
+    /// Maps each ValueId inside this instruction to a new ValueId in place.
+    pub(crate) fn map_values_mut(&mut self, mut f: impl FnMut(ValueId) -> ValueId) {
+        match self {
+            Instruction::Binary(binary) => {
+                binary.lhs = f(binary.lhs);
+                binary.rhs = f(binary.rhs);
+            }
+            Instruction::Cast(value, _) => *value = f(*value),
+            Instruction::Not(value) => *value = f(*value),
+            Instruction::Truncate { value, bit_size: _, max_bit_size: _ } => {
+                *value = f(*value);
+            }
+            Instruction::Constrain(lhs, rhs, assert_message) => {
+                *lhs = f(*lhs);
+                *rhs = f(*rhs);
+                if let Some(ConstrainError::Dynamic(_, _, payload_values)) = assert_message {
+                    for value in payload_values {
+                        *value = f(*value);
+                    }
+                }
+            }
+            Instruction::Call { func, arguments } => {
+                *func = f(*func);
+                for argument in arguments {
+                    *argument = f(*argument);
+                }
+            }
+            Instruction::Allocate => (),
+            Instruction::Load { address } => *address = f(*address),
+            Instruction::Store { address, value } => {
+                *address = f(*address);
+                *value = f(*value);
+            }
+            Instruction::EnableSideEffectsIf { condition } => {
+                *condition = f(*condition);
+            }
+            Instruction::ArrayGet { array, index } => {
+                *array = f(*array);
+                *index = f(*index);
+            }
+            Instruction::ArraySet { array, index, value, mutable: _ } => {
+                *array = f(*array);
+                *index = f(*index);
+                *value = f(*value);
+            }
+            Instruction::IncrementRc { value } => *value = f(*value),
+            Instruction::DecrementRc { value } => *value = f(*value),
+            Instruction::RangeCheck { value, max_bit_size: _, assert_message: _ } => {
+                *value = f(*value);
+            }
+            Instruction::IfElse { then_condition, then_value, else_condition, else_value } => {
+                *then_condition = f(*then_condition);
+                *then_value = f(*then_value);
+                *else_condition = f(*else_condition);
+                *else_value = f(*else_value);
+            }
+            Instruction::MakeArray { elements, typ: _ } => {
+                for element in elements.iter_mut() {
+                    *element = f(*element);
+                }
+            }
+        }
+    }
+
     /// Applies a function to each input value this instruction holds.
     pub(crate) fn for_each_value<T>(&self, mut f: impl FnMut(ValueId) -> T) {
         match self {
@@ -1193,7 +1257,7 @@ impl TerminatorInstruction {
     }
 
     /// Mutate each ValueId to a new ValueId using the given mapping function
-    pub(crate) fn mutate_values(&mut self, mut f: impl FnMut(ValueId) -> ValueId) {
+    pub(crate) fn map_values_mut(&mut self, mut f: impl FnMut(ValueId) -> ValueId) {
         use TerminatorInstruction::*;
         match self {
             JmpIf { condition, .. } => {

compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs

@@ -372,7 +372,7 @@ impl<'brillig> Context<'brillig> {
         dom: &mut DominatorTree,
         constraint_simplification_mapping: &HashMap<ValueId, SimplificationCache>,
     ) -> Instruction {
-        let instruction = dfg[instruction_id].clone();
+        let mut instruction = dfg[instruction_id].clone();
 
         // Alternate between resolving `value_id` in the `dfg` and checking to see if the resolved value
         // has been constrained to be equal to some simpler value in the current block.
@@ -400,9 +400,10 @@ impl<'brillig> Context<'brillig> {
         }
 
         // Resolve any inputs to ensure that we're comparing like-for-like instructions.
-        instruction.map_values(|value_id| {
+        instruction.map_values_mut(|value_id| {
             resolve_cache(block, dfg, dom, constraint_simplification_mapping, value_id)
-        })
+        });
+        instruction
     }
 
     /// Pushes a new [`Instruction`] into the [`DataFlowGraph`] which applies any optimizations

compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs

@@ -1207,20 +1207,15 @@ mod test {
     ) -> Vec<u128> {
         match dfg[value] {
             Value::Instruction { instruction, .. } => {
-                let mut values = vec![];
-                dfg[instruction].map_values(|value| {
-                    values.push(value);
-                    value
-                });
+                let mut constants = vec![];
 
-                let mut values: Vec<_> = values
-                    .into_iter()
-                    .flat_map(|value| get_all_constants_reachable_from_instruction(dfg, value))
-                    .collect();
+                dfg[instruction].for_each_value(|value| {
+                    constants.extend(get_all_constants_reachable_from_instruction(dfg, value));
+                });
 
-                values.sort();
-                values.dedup();
-                values
+                constants.sort();
+                constants.dedup();
+                constants
             }
             Value::NumericConstant { constant, .. } => vec![constant.to_u128()],
             _ => Vec::new(),

compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs

@@ -599,8 +599,9 @@ impl<'f> PerFunctionContext<'f> {
     }
 
     fn update_data_bus(&mut self) {
-        let databus = self.inserter.function.dfg.data_bus.clone();
-        self.inserter.function.dfg.data_bus = databus.map_values(|t| self.inserter.resolve(t));
+        let mut databus = self.inserter.function.dfg.data_bus.clone();
+        databus.map_values_mut(|t| self.inserter.resolve(t));
+        self.inserter.function.dfg.data_bus = databus;
     }
 
     fn handle_terminator(&mut self, block: BasicBlockId, references: &mut Block) {

compiler/noirc_evaluator/src/ssa/opt/normalize_value_ids.rs

@@ -109,9 +109,10 @@ impl Context {
             }
 
             let old_block = &mut old_function.dfg[old_block_id];
-            let mut terminator = old_block
-                .take_terminator()
-                .map_values(|value| self.new_ids.map_value(new_function, old_function, value));
+            let mut terminator = old_block.take_terminator();
+            terminator
+                .map_values_mut(|value| self.new_ids.map_value(new_function, old_function, value));
+
             terminator.mutate_blocks(|old_block| self.new_ids.blocks[&old_block]);
             new_function.dfg.set_block_terminator(new_block_id, terminator);
         }

compiler/noirc_evaluator/src/ssa/opt/unrolling.rs

@@ -938,10 +938,9 @@ impl<'f> LoopIteration<'f> {
             }
             self.inserter.push_instruction(instruction, self.insert_block);
         }
-        let mut terminator = self.dfg()[self.source_block]
-            .unwrap_terminator()
-            .clone()
-            .map_values(|value| self.inserter.resolve(value));
+        let mut terminator = self.dfg()[self.source_block].unwrap_terminator().clone();
+
+        terminator.map_values_mut(|value| self.inserter.resolve(value));
 
         // Replace the blocks in the terminator with fresh one with the same parameters,
         // while remembering which were the original block IDs.