chore: don't use set_value_from_id in remove_truncate_after_range_checks

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

@@ -383,6 +383,50 @@ impl DataFlowGraph {
         self.instructions[id] = instruction;
     }
 
+    /// Replaces all values in the given blocks with the values in the given map.
+    ///
+    /// This method should be preferred over `set_value_from_id` which might eventually be removed.
+    pub(crate) fn replace_values_in_blocks(
+        &mut self,
+        blocks: impl Iterator<Item = BasicBlockId>,
+        values_to_replace: &HashMap<ValueId, ValueId>,
+    ) {
+        if values_to_replace.is_empty() {
+            return;
+        }
+
+        let replacement_fn = |value_id| {
+            if let Some(replacement_id) = values_to_replace.get(&value_id) {
+                *replacement_id
+            } else {
+                value_id
+            }
+        };
+
+        for block in blocks {
+            // Replace in all the block's instructions
+            for instruction_id in self.blocks[block].instructions() {
+                let instruction = &mut self.instructions[*instruction_id];
+                instruction.map_values_mut(replacement_fn);
+
+                // Make sure we also replace the instruction results
+                let results = self.results.get_mut(instruction_id);
+                if let Some(results) = results {
+                    for result in results {
+                        if let Some(replacement_id) = values_to_replace.get(result) {
+                            *result = *replacement_id;
+                        }
+                    }
+                }
+            }
+
+            // Finally, the value might show up in a terminator
+            if self[block].terminator().is_some() {
+                self[block].unwrap_terminator_mut().map_values_mut(replacement_fn);
+            }
+        }
+    }
+
     /// 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

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

@@ -1,4 +1,5 @@
-use im::{HashMap, HashSet};
+use fxhash::FxHashMap as HashMap;
+use fxhash::FxHashSet as HashSet;
 
 use crate::ssa::{
     ir::{function::Function, instruction::Instruction, value::ValueId},
@@ -20,11 +21,14 @@ impl Ssa {
 
 impl Function {
     pub(crate) fn remove_truncate_after_range_check(&mut self) {
-        for block in self.reachable_blocks() {
-            // Keeps the minimum bit size a value was range-checked against
-            let mut range_checks: HashMap<ValueId, u32> = HashMap::new();
-            let mut instructions_to_remove = HashSet::new();
-            let mut values_to_replace = HashMap::<ValueId, ValueId>::new();
+        let mut values_to_replace = HashMap::<ValueId, ValueId>::default();
+        // Keeps the minimum bit size a value was range-checked against
+        let mut range_checks: HashMap<ValueId, u32> = HashMap::default();
+
+        let blocks = self.reachable_blocks();
+        for block in &blocks {
+            let block = *block;
+            let mut instructions_to_remove = HashSet::default();
 
             for instruction_id in self.dfg[block].instructions() {
                 let instruction = &self.dfg[*instruction_id];
@@ -69,11 +73,9 @@ impl Function {
             self.dfg[block]
                 .instructions_mut()
                 .retain(|instruction| !instructions_to_remove.contains(instruction));
-
-            for (old_value, new_value) in values_to_replace {
-                self.dfg.set_value_from_id(old_value, new_value);
-            }
         }
+
+        self.dfg.replace_values_in_blocks(blocks.into_iter(), &values_to_replace);
     }
 }
 
@@ -91,6 +93,8 @@ mod tests {
           b0(v0: Field):
             range_check v0 to 64 bits // This is to make sure we keep the smallest one
             range_check v0 to 32 bits
+            jmp b1() // Make sure the optimization is applied across blocks
+          b1():
             v1 = truncate v0 to 32 bits, max_bit_size: 254
             return v1
         }
@@ -103,6 +107,8 @@ mod tests {
           b0(v0: Field):
             range_check v0 to 64 bits
             range_check v0 to 32 bits
+            jmp b1()
+          b1():
             return v0
         }
         ");