feat(ssa): Mark transitively dead parameters during DIE

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

@@ -801,6 +801,26 @@ impl TerminatorInstruction {
         }
     }
 
+    /// Apply a function to each value along with its index
+    pub(crate) fn for_eachi_value<T>(&self, mut f: impl FnMut(usize, ValueId) -> T) {
+        use TerminatorInstruction::*;
+        match self {
+            JmpIf { condition, .. } => {
+                f(0, *condition);
+            }
+            Jmp { arguments, .. } => {
+                for (index, argument) in arguments.iter().enumerate() {
+                    f(index, *argument);
+                }
+            }
+            Return { return_values, .. } => {
+                for (index, return_value) in return_values.iter().enumerate() {
+                    f(index, *return_value);
+                }
+            }
+        }
+    }
+
     /// Mutate each BlockId to a new BlockId specified by the given mapping function.
     pub(crate) fn mutate_blocks(&mut self, mut f: impl FnMut(BasicBlockId) -> BasicBlockId) {
         use TerminatorInstruction::*;

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

@@ -12,7 +12,7 @@ use crate::ssa::{
         basic_block::{BasicBlock, BasicBlockId},
         dfg::DataFlowGraph,
         function::{Function, FunctionId},
-        instruction::{BinaryOp, Instruction, InstructionId, Intrinsic},
+        instruction::{BinaryOp, Instruction, InstructionId, Intrinsic, TerminatorInstruction},
         post_order::PostOrder,
         types::Type,
         value::{Value, ValueId},
@@ -41,9 +41,45 @@ impl Ssa {
         self.dead_instruction_elimination_with_pruning(true, true)
     }
 
-    fn dead_instruction_elimination_with_pruning(self, flattened: bool, skip_brillig: bool) -> Ssa {
-        let (ssa, result) = self.dead_instruction_elimination_inner(flattened, skip_brillig);
-        ssa.prune_dead_parameters(result.unused_parameters)
+    /// The elimination of certain unused instructions assumes that the DIE pass runs after
+    /// the flattening of the CFG, but if that's not the case then we should not eliminate
+    /// them just yet.
+    #[tracing::instrument(level = "trace", skip(self))]
+    pub(crate) fn dead_instruction_elimination_pre_flattening(self) -> Ssa {
+        self.dead_instruction_elimination_with_pruning(false, false)
+    }
+
+    fn dead_instruction_elimination_with_pruning(
+        mut self,
+        flattened: bool,
+        skip_brillig: bool,
+    ) -> Ssa {
+        let mut previous_unused_params = None;
+        loop {
+            let (new_ssa, result) =
+                self.dead_instruction_elimination_inner(flattened, skip_brillig);
+
+            // Determine whether we have any unused variables
+            let has_unused = result
+                .unused_parameters
+                .values()
+                .any(|block_map| block_map.values().any(|params| !params.is_empty()));
+            // If there are no unused parameters, return early
+            if !has_unused {
+                return new_ssa;
+            }
+
+            if let Some(previous) = &previous_unused_params {
+                // If no changes to dead parameters occurred, return early
+                if previous == &result.unused_parameters {
+                    return new_ssa;
+                }
+            }
+
+            // Prune unused parameters and repeat
+            self = new_ssa.prune_dead_parameters(&result.unused_parameters);
+            previous_unused_params = Some(result.unused_parameters);
+        }
     }
 
     fn dead_instruction_elimination_inner(
@@ -97,6 +133,12 @@ impl Ssa {
 impl Function {
     /// Removes any unused instructions in the reachable blocks of the given function.
     ///
+    /// This method is designed to be run within the context of the full SSA, not in isolation.
+    /// Running DIE on a single function may cause inconsistencies, such as leaving dangling unused parameters.
+    /// The pruning of block parameters depends on the full SSA context.
+    /// Therefore, this method must remain private, and DIE should run over the entire SSA,
+    /// ensuring proper tracking of unused parameters across all blocks.
+    ///
     /// The blocks of the function are iterated in post order, such that any blocks containing
     /// instructions that reference results from an instruction in another block are evaluated first.
     /// If we did not iterate blocks in this order we could not safely say whether or not the results
@@ -107,7 +149,7 @@ impl Function {
     ///   After processing all functions, the union of these sets enables determining the unused globals.
     /// - A mapping of (block id -> unused parameters) for the given function.
     ///   This can be used by follow-up passes to prune unused parameters from blocks.
-    pub(crate) fn dead_instruction_elimination(
+    fn dead_instruction_elimination(
         &mut self,
         flattened: bool,
         skip_brillig: bool,
@@ -129,13 +171,20 @@ impl Function {
         for block in blocks.as_slice() {
             context.remove_unused_instructions_in_block(self, *block);
 
-            let unused_params = self.dfg[*block]
-                .parameters()
+            let parameters = self.dfg[*block].parameters();
+            let mut keep_list = Vec::with_capacity(parameters.len());
+            let unused_params = parameters
                 .iter()
-                .filter(|value| !context.used_values.contains(value))
+                .filter(|value| {
+                    let keep = context.used_values.contains(value);
+                    keep_list.push(keep);
+                    !keep
+                })
                 .copied()
                 .collect::<Vec<_>>();
+
             unused_params_per_block.insert(*block, unused_params);
+            context.parameter_keep_list.insert(*block, keep_list);
         }
 
         context.remove_rc_instructions(&mut self.dfg);
@@ -170,6 +219,17 @@ struct Context {
 
     /// Track IncrementRc instructions per block to determine whether they are useless.
     rc_tracker: RcTracker,
+
+    /// A per-block list indicating which block parameters are still considered alive.
+    ///
+    /// Each entry maps a [BasicBlockId] to a `Vec<bool>`, where the `i`th boolean corresponds to
+    /// the `i`th parameter of that block. A value of `true` means the parameter is used and should
+    /// be preserved. A value of `false` means it is unused and can be pruned.
+    ///
+    /// This keep list is used during terminator analysis to avoid incorrectly marking values as used
+    /// simply because they appear as terminator arguments. Only parameters marked as live here
+    /// should result in values being marked as used in terminator arguments.
+    parameter_keep_list: HashMap<BasicBlockId, Vec<bool>>,
 }
 
 impl Context {
@@ -246,8 +306,19 @@ impl Context {
 
     /// Adds values referenced by the terminator to the set of used values.
     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);
+        let terminator = block.unwrap_terminator();
+        let jmp_destination = if let TerminatorInstruction::Jmp { destination, .. } = terminator {
+            Some(*destination)
+        } else {
+            None
+        };
+
+        block.unwrap_terminator().for_eachi_value(|index, value| {
+            let keep_list = jmp_destination.and_then(|dest| self.parameter_keep_list.get(&dest));
+            let should_keep = keep_list.is_none_or(|list| list[index]);
+            if should_keep {
+                self.mark_used_instruction_results(&function.dfg, value);
+            }
         });
     }
 
@@ -575,7 +646,7 @@ mod test {
                 store Field 1 at v8
                 v9 = load v8 -> Field
                 v10 = add v9, Field 1
-                v11 = add v9, Field 2
+                v11 = add v1, Field 2
                 v13 = add v9, Field 3
                 v14 = add v13, v13
                 call assert_constant(v10)
@@ -595,7 +666,7 @@ mod test {
             store Field 1 at v4
             v6 = load v4 -> Field
             v7 = add v6, Field 1
-            v8 = add v6, Field 2
+            v8 = add v1, Field 2
             call assert_constant(v7)
             return v8
         }

compiler/noirc_evaluator/src/ssa/opt/die/prune_dead_parameters.rs

@@ -58,10 +58,10 @@ impl Ssa {
     #[tracing::instrument(level = "trace", skip(self))]
     pub(crate) fn prune_dead_parameters(
         mut self,
-        unused_parameters: HashMap<FunctionId, HashMap<BasicBlockId, Vec<ValueId>>>,
+        unused_parameters: &HashMap<FunctionId, HashMap<BasicBlockId, Vec<ValueId>>>,
     ) -> Self {
         for (func_id, unused_parameters) in unused_parameters {
-            let function = self.functions.get_mut(&func_id).expect("ICE: Function should exist");
+            let function = self.functions.get_mut(func_id).expect("ICE: Function should exist");
             function.prune_dead_parameters(unused_parameters);
         }
         self
@@ -70,7 +70,7 @@ impl Ssa {
 
 impl Function {
     /// See [`prune_dead_parameters`][self] module for more information
-    fn prune_dead_parameters(&mut self, unused_params: HashMap<BasicBlockId, Vec<ValueId>>) {
+    fn prune_dead_parameters(&mut self, unused_params: &HashMap<BasicBlockId, Vec<ValueId>>) {
         let cfg = ControlFlowGraph::with_function(self);
         let post_order = PostOrder::with_cfg(&cfg);
 
@@ -169,7 +169,7 @@ mod tests {
         assert_eq!(b1_unused[0].to_u32(), 0);
         assert_eq!(b1_unused[1].to_u32(), 2);
 
-        let ssa = ssa.prune_dead_parameters(die_result.unused_parameters);
+        let ssa = ssa.prune_dead_parameters(&die_result.unused_parameters);
 
         assert_ssa_snapshot!(ssa, @r#"
         brillig(inline) fn test f0 {
@@ -226,7 +226,8 @@ mod tests {
         assert_eq!(b3_unused.len(), 1);
         assert_eq!(b3_unused[0].to_u32(), 2);
 
-        let ssa = ssa.prune_dead_parameters(die_result.unused_parameters);
+        let ssa = ssa.prune_dead_parameters(&die_result.unused_parameters);
+        let (ssa, _) = ssa.dead_instruction_elimination_inner(false, false);
 
         // We expect b3 to have no parameters anymore and both predecessors (b1 and b2)
         // should no longer pass any arguments to their terminator (which jumps to b3).
@@ -240,7 +241,6 @@ mod tests {
             jmpif v5 then: b1, else: b2
           b1():
             v7 = mul u32 601072115, u32 2825334515
-            v8 = cast v7 as u64
             jmp b3()
           b2():
             jmp b3()
@@ -275,7 +275,7 @@ mod tests {
         let b1_unused = function.get(&Id::test_new(1)).expect("Should have unused parameters");
         assert!(b1_unused.is_empty());
 
-        let ssa = ssa.prune_dead_parameters(die_result.unused_parameters);
+        let ssa = ssa.prune_dead_parameters(&die_result.unused_parameters);
 
         // b0 still has both parameters even though v0 is unused
         // as b0 is the entry block which would also change the function signature.
@@ -338,6 +338,7 @@ mod tests {
         }
         "#;
         let ssa = Ssa::from_str(src).unwrap();
+
         // DIE is necessary to fetch the block parameters liveness information
         let (ssa, die_result) = ssa.dead_instruction_elimination_inner(false, false);
 
@@ -350,39 +351,83 @@ mod tests {
             if block_id.to_u32() == 9 {
                 assert!(unused_params.len() == 1);
                 assert_eq!(unused_params[0].to_u32(), 3);
+            } else if block_id.to_u32() == 5 {
+                assert!(unused_params.len() == 1);
+                assert_eq!(unused_params[0].to_u32(), 1);
+            } else if block_id.to_u32() == 6 {
+                assert!(unused_params.len() == 1);
+                assert_eq!(unused_params[0].to_u32(), 2);
             } else {
                 assert!(unused_params.is_empty());
             }
         }
 
-        let ssa = ssa.prune_dead_parameters(die_result.unused_parameters);
+        let ssa = ssa.prune_dead_parameters(&die_result.unused_parameters);
+
+        let (ssa, die_result) = ssa.dead_instruction_elimination_inner(false, false);
+
+        assert!(die_result.unused_parameters.len() == 1);
+        let function = die_result
+            .unused_parameters
+            .get(&Id::test_new(0))
+            .expect("Should have unused parameters");
+        for unused_params in function.values() {
+            assert!(unused_params.is_empty());
+        }
 
-        // We only expect b9 to have v3 pruned.
-        // Only v1 in b5 is marked as used with a single DIE pass.
-        // Another DIE pass on the resulting SSA below would be necessary to now see v1 in b5 as unused.
         assert_ssa_snapshot!(ssa, @r#"
         brillig(inline) predicate_pure fn main f0 {
           b0(v0: i16):
-            v4 = lt i16 3, v0
-            jmpif v4 then: b1, else: b2
+            v2 = lt i16 3, v0
+            jmpif v2 then: b1, else: b2
           b1():
-            v7 = lt i16 4, v0
-            jmpif v7 then: b3, else: b4
+            v4 = lt i16 4, v0
+            jmpif v4 then: b3, else: b4
           b2():
-            jmp b5(Field 3)
+            jmp b5()
           b3():
-            jmp b6(Field 1)
+            jmp b6()
           b4():
-            jmp b6(Field 2)
-          b5(v1: Field):
-            v11 = lt i16 5, v0
-            jmpif v11 then: b7, else: b8
-          b6(v2: Field):
-            jmp b5(v2)
+            jmp b6()
+          b5():
+            v6 = lt i16 5, v0
+            jmpif v6 then: b7, else: b8
+          b6():
+            jmp b5()
+          b7():
+            jmp b9()
+          b8():
+            jmp b9()
+          b9():
+            return
+        }
+        "#);
+
+        // Now check that calling the DIE -> parameter pruning feedback loop produces the same result
+        let ssa = Ssa::from_str(src).unwrap();
+        let ssa = ssa.dead_instruction_elimination_with_pruning(false, false);
+        assert_ssa_snapshot!(ssa, @r#"
+        brillig(inline) predicate_pure fn main f0 {
+          b0(v0: i16):
+            v2 = lt i16 3, v0
+            jmpif v2 then: b1, else: b2
+          b1():
+            v4 = lt i16 4, v0
+            jmpif v4 then: b3, else: b4
+          b2():
+            jmp b5()
+          b3():
+            jmp b6()
+          b4():
+            jmp b6()
+          b5():
+            v6 = lt i16 5, v0
+            jmpif v6 then: b7, else: b8
+          b6():
+            jmp b5()
           b7():
             jmp b9()
           b8():
-            v12 = add v1, Field 1
             jmp b9()
           b9():
             return