fix(ssa refactor): fix array element propagation through constant fol…

…ding and DIE (#1674)

crates/nargo_cli/tests/test_data_ssa_refactor/simple_add_and_ret_arr/Nargo.toml

@@ -0,0 +1,5 @@
+[package]
+authors = [""]
+compiler_version = "0.6.0"
+
+[dependencies]

crates/nargo_cli/tests/test_data_ssa_refactor/simple_add_and_ret_arr/Prover.toml

@@ -0,0 +1 @@
+x = 1

crates/nargo_cli/tests/test_data_ssa_refactor/simple_add_and_ret_arr/src/main.nr

@@ -0,0 +1,8 @@
+// A simple program to test that SSA array values elements 
+// aren't disconnected from their instruction results, and
+// that dead instruction elemination looks inside of arrays
+// when deciding whether of not an instruction should be
+// retained.
+fn main(x: Field) -> pub [Field; 1] {
+    [x + 1]
+}

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

@@ -1,11 +1,11 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::HashSet;
 
 use iter_extended::vecmap;
 
 use crate::ssa_refactor::{
     ir::{
         basic_block::BasicBlockId, dfg::InsertInstructionResult, function::Function,
-        instruction::InstructionId, value::ValueId,
+        instruction::InstructionId,
     },
     ssa_gen::Ssa,
 };
@@ -44,7 +44,6 @@ fn constant_fold(function: &mut Function) {
 #[derive(Default)]
 struct Context {
     /// Maps pre-folded ValueIds to the new ValueIds obtained by re-inserting the instruction.
-    values: HashMap<ValueId, ValueId>,
     visited_blocks: HashSet<BasicBlockId>,
     block_queue: Vec<BasicBlockId>,
 }
@@ -56,57 +55,31 @@ 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(function.dfg.resolve(value)));
-
-        function.dfg.set_block_terminator(block, terminator);
         self.block_queue.extend(function.dfg[block].successors());
     }
 
-    fn get_value(&self, value: ValueId) -> ValueId {
-        self.values.get(&value).copied().unwrap_or(value)
-    }
-
     fn push_instruction(
         &mut self,
         function: &mut Function,
         block: BasicBlockId,
         id: InstructionId,
     ) {
-        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 instruction = function.dfg[id].clone();
+        let old_results = function.dfg.instruction_results(id).to_vec();
 
         let ctrl_typevars = instruction
             .requires_ctrl_typevars()
-            .then(|| vecmap(&results, |result| function.dfg.type_of_value(*result)));
+            .then(|| vecmap(&old_results, |result| function.dfg.type_of_value(*result)));
 
         let new_results =
-            function.dfg.insert_instruction_and_results(instruction, block, ctrl_typevars);
-
-        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
-    /// ValueId (from the source_function) and its new ValueId in the destination function.
-    fn insert_new_instruction_results(
-        values: &mut HashMap<ValueId, ValueId>,
-        old_results: &[ValueId],
-        new_results: InsertInstructionResult,
-    ) {
+            match function.dfg.insert_instruction_and_results(instruction, block, ctrl_typevars) {
+                InsertInstructionResult::SimplifiedTo(new_result) => vec![new_result],
+                InsertInstructionResult::Results(new_results) => new_results.to_vec(),
+                InsertInstructionResult::InstructionRemoved => vec![],
+            };
         assert_eq!(old_results.len(), new_results.len());
-
-        match new_results {
-            InsertInstructionResult::SimplifiedTo(new_result) => {
-                values.insert(old_results[0], new_result);
-            }
-            InsertInstructionResult::Results(new_results) => {
-                for (old_result, new_result) in old_results.iter().zip(new_results) {
-                    values.insert(*old_result, *new_result);
-                }
-            }
-            InsertInstructionResult::InstructionRemoved => (),
+        for (old_result, new_result) in old_results.iter().zip(new_results) {
+            function.dfg.set_value_from_id(*old_result, new_result);
         }
     }
 }
@@ -119,6 +92,7 @@ mod test {
             instruction::{BinaryOp, TerminatorInstruction},
             map::Id,
             types::Type,
+            value::Value,
         },
         ssa_builder::FunctionBuilder,
     };
@@ -178,4 +152,46 @@ mod test {
             _ => unreachable!("b0 should have a return terminator"),
         }
     }
+
+    #[test]
+    fn arrays_elements_are_updated() {
+        // fn main f0 {
+        //   b0(v0: Field):
+        //     v1 = add v0, Field 1
+        //     return [v1]
+        // }
+        //
+        // After constructing this IR, we run constant folding with no expected benefit, but to
+        // ensure that all new values ids are correctly propagated.
+        let main_id = Id::test_new(0);
+
+        // Compiling main
+        let mut builder = FunctionBuilder::new("main".into(), main_id, RuntimeType::Acir);
+        let v0 = builder.add_parameter(Type::field());
+        let one = builder.field_constant(1u128);
+        let v1 = builder.insert_binary(v0, BinaryOp::Add, one);
+        let arr =
+            builder.current_function.dfg.make_array(vec![v1].into(), vec![Type::field()].into());
+        builder.terminate_with_return(vec![arr]);
+
+        let ssa = builder.finish().fold_constants();
+        let main = ssa.main();
+        let entry_block_id = main.entry_block();
+        let entry_block = &main.dfg[entry_block_id];
+        assert_eq!(entry_block.instructions().len(), 1);
+        let new_add_instr = entry_block.instructions().first().unwrap();
+        let new_add_instr_result = main.dfg.instruction_results(*new_add_instr)[0];
+        assert_ne!(new_add_instr_result, v1);
+
+        let return_value_id = match entry_block.unwrap_terminator() {
+            TerminatorInstruction::Return { return_values } => return_values[0],
+            _ => unreachable!(),
+        };
+        let return_element = match &main.dfg[return_value_id] {
+            Value::Array { array, .. } => array[0],
+            _ => unreachable!(),
+        };
+        // The return element is expected to refer to the new add instruction result.
+        assert_eq!(main.dfg.resolve(new_add_instr_result), main.dfg.resolve(return_element));
+    }
 }

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

@@ -5,10 +5,11 @@ use std::collections::HashSet;
 use crate::ssa_refactor::{
     ir::{
         basic_block::{BasicBlock, BasicBlockId},
+        dfg::DataFlowGraph,
         function::Function,
         instruction::{Instruction, InstructionId},
         post_order::PostOrder,
-        value::ValueId,
+        value::{Value, ValueId},
     },
     ssa_gen::Ssa,
 };
@@ -64,14 +65,16 @@ impl Context {
         block_id: BasicBlockId,
     ) {
         let block = &function.dfg[block_id];
-        self.mark_terminator_values_as_used(block);
+        self.mark_terminator_values_as_used(function, block);
 
         for instruction in block.instructions().iter().rev() {
             if self.is_unused(*instruction, function) {
                 self.instructions_to_remove.insert(*instruction);
             } else {
                 let instruction = &function.dfg[*instruction];
-                instruction.for_each_value(|value| self.used_values.insert(value));
+                instruction.for_each_value(|value| {
+                    self.mark_used_instruction_results(&function.dfg, value);
+                });
             }
         }
 
@@ -99,8 +102,29 @@ impl Context {
     }
 
     /// Adds values referenced by the terminator to the set of used values.
-    fn mark_terminator_values_as_used(&mut self, block: &BasicBlock) {
-        block.unwrap_terminator().for_each_value(|value| self.used_values.insert(value));
+    fn mark_terminator_values_as_used(&mut self, function: &Function, block: &BasicBlock) {
+        block.unwrap_terminator().for_each_value(|value| {
+            self.mark_used_instruction_results(&function.dfg, value);
+        });
+    }
+
+    /// Inspects a value recursively (as it could be an array) and marks all comprised instruction
+    /// results as used.
+    fn mark_used_instruction_results(&mut self, dfg: &DataFlowGraph, value_id: ValueId) {
+        let value_id = dfg.resolve(value_id);
+        match &dfg[value_id] {
+            Value::Instruction { .. } => {
+                self.used_values.insert(value_id);
+            }
+            Value::Array { array, .. } => {
+                for elem in array {
+                    self.mark_used_instruction_results(dfg, *elem);
+                }
+            }
+            _ => {
+                // Does not comprise of any instruction results
+            }
+        }
     }
 }