diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs index 690bb121634..0c7427dc228 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs +++ b/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, + values_to_replace: &HashMap, + ) { + 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 diff --git a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs index 8b8d6a05dc8..9273453d6cb 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs +++ b/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 = HashMap::new(); - let mut instructions_to_remove = HashSet::new(); - let mut values_to_replace = HashMap::::new(); + let mut values_to_replace = HashMap::::default(); + // Keeps the minimum bit size a value was range-checked against + let mut range_checks: HashMap = 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 } ");