Apply latest fixes

Author: asterite

noir/noir-repo/compiler/noirc_evaluator/src/ssa/ir/dom.rs

@@ -119,6 +119,29 @@ impl DominatorTree {
         }
     }
 
+    /// Walk up the dominator tree until we find a block for which `f` returns `Some` value.
+    /// Otherwise return `None` when we reach the top.
+    ///
+    /// Similar to `Iterator::filter_map` but only returns the first hit.
+    pub(crate) fn find_map_dominator<T>(
+        &self,
+        mut block_id: BasicBlockId,
+        f: impl Fn(BasicBlockId) -> Option<T>,
+    ) -> Option<T> {
+        if !self.is_reachable(block_id) {
+            return None;
+        }
+        loop {
+            if let Some(value) = f(block_id) {
+                return Some(value);
+            }
+            block_id = match self.immediate_dominator(block_id) {
+                Some(immediate_dominator) => immediate_dominator,
+                None => return None,
+            }
+        }
+    }
+
     /// Allocate and compute a dominator tree from a pre-computed control flow graph and
     /// post-order counterpart.
     pub(crate) fn with_cfg_and_post_order(cfg: &ControlFlowGraph, post_order: &PostOrder) -> Self {

noir/noir-repo/compiler/noirc_evaluator/src/ssa/ir/instruction.rs

@@ -319,6 +319,44 @@ impl Instruction {
         matches!(self.result_type(), InstructionResultType::Unknown)
     }
 
+    /// Indicates if the instruction has a side effect, ie. it can fail, or it interacts with memory.
+    ///
+    /// This is similar to `can_be_deduplicated`, but it doesn't depend on whether the caller takes
+    /// constraints into account, because it might not use it to isolate the side effects across branches.
+    pub(crate) fn has_side_effects(&self, dfg: &DataFlowGraph) -> bool {
+        use Instruction::*;
+
+        match self {
+            // These either have side-effects or interact with memory
+            EnableSideEffectsIf { .. }
+            | Allocate
+            | Load { .. }
+            | Store { .. }
+            | IncrementRc { .. }
+            | DecrementRc { .. } => true,
+
+            Call { func, .. } => match dfg[*func] {
+                Value::Intrinsic(intrinsic) => intrinsic.has_side_effects(),
+                _ => true, // Be conservative and assume other functions can have side effects.
+            },
+
+            // These can fail.
+            Constrain(..) | RangeCheck { .. } => true,
+
+            // This should never be side-effectful
+            MakeArray { .. } => false,
+
+            // These can have different behavior depending on the EnableSideEffectsIf context.
+            Binary(_)
+            | Cast(_, _)
+            | Not(_)
+            | Truncate { .. }
+            | IfElse { .. }
+            | ArrayGet { .. }
+            | ArraySet { .. } => self.requires_acir_gen_predicate(dfg),
+        }
+    }
+
     /// Indicates if the instruction can be safely replaced with the results of another instruction with the same inputs.
     /// If `deduplicate_with_predicate` is set, we assume we're deduplicating with the instruction
     /// and its predicate, rather than just the instruction. Setting this means instructions that

noir/noir-repo/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs

@@ -68,16 +68,16 @@ impl Function {
     /// The structure of this pass is simple:
     /// Go through each block and re-insert all instructions.
     pub(crate) fn constant_fold(&mut self, use_constraint_info: bool) {
-        let mut context = Context::new(self, use_constraint_info);
-        context.block_queue.push_back(self.entry_block());
+        let mut context = Context::new(use_constraint_info);
+        let mut dom = DominatorTree::with_function(self);
 
         while let Some(block) = context.block_queue.pop_front() {
             if context.visited_blocks.contains(&block) {
                 continue;
             }
 
             context.visited_blocks.insert(block);
-            context.fold_constants_in_block(self, block);
+            context.fold_constants_in_block(&mut self.dfg, &mut dom, block);
         }
     }
 }
@@ -99,15 +99,56 @@ struct Context {
     /// We also keep track of how a value was simplified to other values per block. That is,
     /// a same ValueId could have been simplified to one value in one block and to another value
     /// in another block.
-    constraint_simplification_mappings:
-        HashMap<ValueId, HashMap<ValueId, Vec<(BasicBlockId, ValueId)>>>,
+    constraint_simplification_mappings: ConstraintSimplificationCache,
 
     // Cache of instructions without any side-effects along with their outputs.
     cached_instruction_results: InstructionResultCache,
+}
 
-    dom: DominatorTree,
+/// Records a simplified equivalents of an [`Instruction`] in the blocks
+/// where the constraint that advised the simplification has been encountered.
+///
+/// For more information see [`ConstraintSimplificationCache`].
+#[derive(Default)]
+struct SimplificationCache {
+    /// Simplified expressions where we found them.
+    ///
+    /// It will always have at least one value because `add` is called
+    /// after the default is constructed.
+    simplifications: HashMap<BasicBlockId, ValueId>,
 }
 
+impl SimplificationCache {
+    /// Called with a newly encountered simplification.
+    fn add(&mut self, dfg: &DataFlowGraph, simple: ValueId, block: BasicBlockId) {
+        self.simplifications
+            .entry(block)
+            .and_modify(|existing| {
+                // `SimplificationCache` may already hold a simplification in this block
+                // so we check whether `simple` is a better simplification than the current one.
+                if let Some((_, simpler)) = simplify(dfg, *existing, simple) {
+                    *existing = simpler;
+                };
+            })
+            .or_insert(simple);
+    }
+
+    /// Try to find a simplification in a visible block.
+    fn get(&self, block: BasicBlockId, dom: &DominatorTree) -> Option<ValueId> {
+        // Deterministically walk up the dominator chain until we encounter a block that contains a simplification.
+        dom.find_map_dominator(block, |b| self.simplifications.get(&b).cloned())
+    }
+}
+
+/// HashMap from `(side_effects_enabled_var, Instruction)` to a simplified expression that it can
+/// be replaced with based on constraints that testify to their equivalence, stored together
+/// with the set of blocks at which this constraint has been observed.
+///
+/// Only blocks dominated by one in the cache should have access to this information, otherwise
+/// we create a sort of time paradox where we replace an instruction with a constant we believe
+/// it _should_ equal to, without ever actually producing and asserting the value.
+type ConstraintSimplificationCache = HashMap<ValueId, HashMap<ValueId, SimplificationCache>>;
+
 /// HashMap from (Instruction, side_effects_enabled_var) to the results of the instruction.
 /// Stored as a two-level map to avoid cloning Instructions during the `.get` call.
 ///
@@ -124,55 +165,55 @@ struct ResultCache {
 }
 
 impl Context {
-    fn new(function: &Function, use_constraint_info: bool) -> Self {
+    fn new(use_constraint_info: bool) -> Self {
         Self {
             use_constraint_info,
             visited_blocks: Default::default(),
             block_queue: Default::default(),
             constraint_simplification_mappings: Default::default(),
             cached_instruction_results: Default::default(),
-            dom: DominatorTree::with_function(function),
         }
     }
 
-    fn fold_constants_in_block(&mut self, function: &mut Function, block: BasicBlockId) {
-        let instructions = function.dfg[block].take_instructions();
+    fn fold_constants_in_block(
+        &mut self,
+        dfg: &mut DataFlowGraph,
+        dom: &mut DominatorTree,
+        block: BasicBlockId,
+    ) {
+        let instructions = dfg[block].take_instructions();
 
-        let mut side_effects_enabled_var =
-            function.dfg.make_constant(FieldElement::one(), Type::bool());
+        // Default side effect condition variable with an enabled state.
+        let mut side_effects_enabled_var = dfg.make_constant(FieldElement::one(), Type::bool());
 
         for instruction_id in instructions {
             self.fold_constants_into_instruction(
-                &mut function.dfg,
+                dfg,
+                dom,
                 block,
                 instruction_id,
                 &mut side_effects_enabled_var,
             );
         }
-        self.block_queue.extend(function.dfg[block].successors());
+        self.block_queue.extend(dfg[block].successors());
     }
 
     fn fold_constants_into_instruction(
         &mut self,
         dfg: &mut DataFlowGraph,
+        dom: &mut DominatorTree,
         block: BasicBlockId,
         id: InstructionId,
         side_effects_enabled_var: &mut ValueId,
     ) {
-        let constraint_simplification_mapping =
-            self.constraint_simplification_mappings.get(side_effects_enabled_var);
-        let instruction = Self::resolve_instruction(
-            id,
-            block,
-            dfg,
-            &mut self.dom,
-            constraint_simplification_mapping,
-        );
+        let constraint_simplification_mapping = self.get_constraint_map(*side_effects_enabled_var);
+        let instruction =
+            Self::resolve_instruction(id, block, dfg, dom, constraint_simplification_mapping);
         let old_results = dfg.instruction_results(id).to_vec();
 
         // If a copy of this instruction exists earlier in the block, then reuse the previous results.
         if let Some(cached_results) =
-            self.get_cached(dfg, &instruction, *side_effects_enabled_var, block)
+            self.get_cached(dfg, dom, &instruction, *side_effects_enabled_var, block)
         {
             Self::replace_result_ids(dfg, &old_results, cached_results);
             return;
@@ -204,7 +245,7 @@ impl Context {
         block: BasicBlockId,
         dfg: &DataFlowGraph,
         dom: &mut DominatorTree,
-        constraint_simplification_mapping: Option<&HashMap<ValueId, Vec<(BasicBlockId, ValueId)>>>,
+        constraint_simplification_mapping: &HashMap<ValueId, SimplificationCache>,
     ) -> Instruction {
         let instruction = dfg[instruction_id].clone();
 
@@ -214,30 +255,28 @@ impl Context {
         // This allows us to reach a stable final `ValueId` for each instruction input as we add more
         // constraints to the cache.
         fn resolve_cache(
+            block: BasicBlockId,
             dfg: &DataFlowGraph,
             dom: &mut DominatorTree,
-            cache: Option<&HashMap<ValueId, Vec<(BasicBlockId, ValueId)>>>,
+            cache: &HashMap<ValueId, SimplificationCache>,
             value_id: ValueId,
-            block: BasicBlockId,
         ) -> ValueId {
             let resolved_id = dfg.resolve(value_id);
-            let Some(cached_values) = cache.and_then(|cache| cache.get(&resolved_id)) else {
-                return resolved_id;
-            };
-
-            for (cached_block, cached_value) in cached_values {
-                // We can only use the simplified value if it was simplified in a block that dominates the current one
-                if dom.dominates(*cached_block, block) {
-                    return resolve_cache(dfg, dom, cache, *cached_value, block);
+            match cache.get(&resolved_id) {
+                Some(simplification_cache) => {
+                    if let Some(simplified) = simplification_cache.get(block, dom) {
+                        resolve_cache(block, dfg, dom, cache, simplified)
+                    } else {
+                        resolved_id
+                    }
                 }
+                None => resolved_id,
             }
-
-            resolved_id
         }
 
         // Resolve any inputs to ensure that we're comparing like-for-like instructions.
         instruction.map_values(|value_id| {
-            resolve_cache(dfg, dom, constraint_simplification_mapping, value_id, block)
+            resolve_cache(block, dfg, dom, constraint_simplification_mapping, value_id)
         })
     }
 
@@ -288,7 +327,7 @@ impl Context {
                     self.get_constraint_map(side_effects_enabled_var)
                         .entry(complex)
                         .or_default()
-                        .push((block, simple));
+                        .add(dfg, simple, block);
                 }
             }
         }
@@ -312,7 +351,7 @@ impl Context {
     fn get_constraint_map(
         &mut self,
         side_effects_enabled_var: ValueId,
-    ) -> &mut HashMap<ValueId, Vec<(BasicBlockId, ValueId)>> {
+    ) -> &mut HashMap<ValueId, SimplificationCache> {
         self.constraint_simplification_mappings.entry(side_effects_enabled_var).or_default()
     }
 
@@ -327,19 +366,21 @@ impl Context {
         }
     }
 
+    /// Get a cached result if it can be used in this context.
     fn get_cached<'a>(
-        &'a mut self,
+        &self,
         dfg: &DataFlowGraph,
+        dom: &mut DominatorTree,
         instruction: &Instruction,
         side_effects_enabled_var: ValueId,
         block: BasicBlockId,
-    ) -> Option<&'a [ValueId]> {
+    ) -> Option<&[ValueId]> {
         let results_for_instruction = self.cached_instruction_results.get(instruction)?;
 
         let predicate = self.use_constraint_info && instruction.requires_acir_gen_predicate(dfg);
         let predicate = predicate.then_some(side_effects_enabled_var);
 
-        results_for_instruction.get(&predicate)?.get(block, &mut self.dom)
+        results_for_instruction.get(&predicate)?.get(block, dom)
     }
 }