fix(ssa refactor): resolve replaced value ids for printing (#1535)

* fix(ssa refactor): resolve replaced value ids for printing

* fix(ssa refactor): Expand PR #1535 to resolve ValueIds in all SSA passes (#1642)

* Expand PR

* chore(ssa refactor): more value id resolving

* chore(ssa refactor): another value id resolve

---------

Co-authored-by: Joss <joss@aztecprotocol.com>

---------

Co-authored-by: jfecher <jake@aztecprotocol.com>

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

@@ -325,6 +325,7 @@ impl Context {
     /// involving such values are evaluated via a separate path and stored in
     /// `ssa_value_to_array_address` instead.
     fn convert_value(&mut self, value_id: ValueId, dfg: &DataFlowGraph) -> AcirValue {
+        let value_id = dfg.resolve(value_id);
         let value = &dfg[value_id];
         if let Some(acir_value) = self.ssa_values.get(&value_id) {
             return acir_value.clone();

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

@@ -59,6 +59,11 @@ pub(crate) struct DataFlowGraph {
 
     /// All blocks in a function
     blocks: DenseMap<BasicBlock>,
+
+    /// Debugging information about which `ValueId`s have had their underlying `Value` substituted
+    /// for that of another. This information is purely used for printing the SSA, and has no
+    /// material effect on the SSA itself.
+    replaced_value_ids: HashMap<ValueId, ValueId>,
 }
 
 impl DataFlowGraph {
@@ -146,19 +151,28 @@ impl DataFlowGraph {
         self.values.insert(value)
     }
 
-    /// Replaces the value specified by the given ValueId with a new Value.
+    /// Set the value of value_to_replace to refer to the value referred to by new_value.
     ///
     /// This is the preferred method to call for optimizations simplifying
     /// values since other instructions referring to the same ValueId need
     /// not be modified to refer to a new ValueId.
-    pub(crate) fn set_value(&mut self, value_id: ValueId, new_value: Value) {
-        self.values[value_id] = new_value;
+    pub(crate) fn set_value_from_id(&mut self, value_to_replace: ValueId, new_value: ValueId) {
+        if value_to_replace != new_value {
+            self.replaced_value_ids.insert(value_to_replace, self.resolve(new_value));
+            let new_value = self.values[new_value].clone();
+            self.values[value_to_replace] = new_value;
+        }
     }
 
-    /// Set the value of value_to_replace to refer to the value referred to by new_value.
-    pub(crate) fn set_value_from_id(&mut self, value_to_replace: ValueId, new_value: ValueId) {
-        let new_value = self.values[new_value].clone();
-        self.values[value_to_replace] = new_value;
+    /// If `original_value_id`'s underlying `Value` has been substituted for that of another
+    /// `ValueId`, this function will return the `ValueId` from which the substitution was taken.
+    /// If `original_value_id`'s underlying `Value` has not been substituted, the same `ValueId`
+    /// is returned.
+    pub(crate) fn resolve(&self, original_value_id: ValueId) -> ValueId {
+        match self.replaced_value_ids.get(&original_value_id) {
+            Some(id) => self.resolve(*id),
+            None => original_value_id,
+        }
     }
 
     /// Creates a new constant value, or returns the Id to an existing one if

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

@@ -61,6 +61,7 @@ pub(crate) fn display_block(
 /// constant or a function we print those directly.
 fn value(function: &Function, id: ValueId) -> String {
     use super::value::Value;
+    let id = function.dfg.resolve(id);
     match &function.dfg[id] {
         Value::NumericConstant { constant, typ } => {
             format!("{typ} {constant}")

crates/noirc_evaluator/src/ssa_refactor/opt/constant_folding.rs

@@ -56,9 +56,9 @@ impl Context {
         for instruction in instructions {
             self.push_instruction(function, block, instruction);
         }
-
-        let terminator =
-            function.dfg[block].unwrap_terminator().map_values(|value| self.get_value(value));
+        let terminator = function.dfg[block]
+            .unwrap_terminator()
+            .map_values(|value| self.get_value(function.dfg.resolve(value)));
 
         function.dfg.set_block_terminator(block, terminator);
         self.block_queue.extend(function.dfg[block].successors());
@@ -74,8 +74,9 @@ impl Context {
         block: BasicBlockId,
         id: InstructionId,
     ) {
-        let instruction = function.dfg[id].map_values(|id| self.get_value(id));
-        let results = function.dfg.instruction_results(id).to_vec();
+        let instruction =
+            function.dfg[id].map_values(|id| self.get_value(function.dfg.resolve(id)));
+        let results = vecmap(function.dfg.instruction_results(id), |id| function.dfg.resolve(*id));
 
         let ctrl_typevars = instruction
             .requires_ctrl_typevars()

crates/noirc_evaluator/src/ssa_refactor/opt/flatten_cfg.rs

@@ -534,7 +534,8 @@ impl<'f> Context<'f> {
     fn push_instruction(&mut self, id: InstructionId) {
         let instruction = self.function.dfg[id].map_values(|id| self.translate_value(id));
         let instruction = self.handle_instruction_side_effects(instruction);
-        let results = self.function.dfg.instruction_results(id).to_vec();
+        let results = self.function.dfg.instruction_results(id);
+        let results = vecmap(results, |id| self.function.dfg.resolve(*id));
 
         let ctrl_typevars = instruction
             .requires_ctrl_typevars()

crates/noirc_evaluator/src/ssa_refactor/opt/inlining.rs

@@ -362,13 +362,14 @@ impl<'function> PerFunctionContext<'function> {
     fn push_instruction(&mut self, id: InstructionId) {
         let instruction = self.source_function.dfg[id].map_values(|id| self.translate_value(id));
         let results = self.source_function.dfg.instruction_results(id);
+        let results = vecmap(results, |id| self.source_function.dfg.resolve(*id));
 
         let ctrl_typevars = instruction
             .requires_ctrl_typevars()
-            .then(|| vecmap(results, |result| self.source_function.dfg.type_of_value(*result)));
+            .then(|| vecmap(&results, |result| self.source_function.dfg.type_of_value(*result)));
 
         let new_results = self.context.builder.insert_instruction(instruction, ctrl_typevars);
-        Self::insert_new_instruction_results(&mut self.values, results, new_results);
+        Self::insert_new_instruction_results(&mut self.values, &results, new_results);
     }
 
     /// Modify the values HashMap to remember the mapping between an instruction result's previous

crates/noirc_evaluator/src/ssa_refactor/opt/mem2reg.rs

@@ -103,8 +103,7 @@ impl PerBlockContext {
                 Instruction::Load { address } => {
                     if let Some(address) = self.try_const_address(*address, dfg) {
                         if let Some(last_value) = self.last_stores.get(&address) {
-                            let last_value = dfg[*last_value].clone();
-                            loads_to_substitute.push((*instruction_id, last_value));
+                            loads_to_substitute.push((*instruction_id, *last_value));
                         } else {
                             self.failed_substitutes.insert(address);
                         }
@@ -138,7 +137,7 @@ impl PerBlockContext {
                 .instruction_results(*instruction_id)
                 .first()
                 .expect("ICE: Load instructions should have single result");
-            dfg.set_value(result_value, new_value.clone());
+            dfg.set_value_from_id(result_value, *new_value);
         }
 
         // Delete load instructions

crates/noirc_evaluator/src/ssa_refactor/opt/unrolling.rs

@@ -441,7 +441,8 @@ impl<'f> LoopIteration<'f> {
 
     fn push_instruction(&mut self, id: InstructionId) {
         let instruction = self.function.dfg[id].map_values(|id| self.get_value(id));
-        let results = self.function.dfg.instruction_results(id).to_vec();
+        let results = self.function.dfg.instruction_results(id);
+        let results = vecmap(results, |id| self.function.dfg.resolve(*id));
 
         let ctrl_typevars = instruction
             .requires_ctrl_typevars()