Walk up the dominator tree to find the first simplification

Author: aakoshh

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

@@ -119,6 +119,26 @@ impl DominatorTree {
         }
     }
 
+    /// Walk up the dominator tree until we find one that to 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> {
+        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 {

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

@@ -218,19 +218,8 @@ impl SimplificationCache {
 
     /// Try to find a simplification in a visible block.
     fn get(&self, block: BasicBlockId, dom: &mut DominatorTree) -> Option<ValueId> {
-        // See if we have a direct simplification in this block.
-        if let Some(value) = self.simplifications.get(&block) {
-            return Some(*value);
-        }
         // Check if there is a dominating block we can take a simplification from.
-        // Going backwards so that we find a constraint closest to what we have already processed
-        // (assuming block IDs of blocks further down in the SSA are larger).
-        for (constraining_block, value) in self.simplifications.iter().rev() {
-            if dom.dominates(*constraining_block, block) {
-                return Some(*value);
-            }
-        }
-        None
+        dom.find_map_dominator(block, |b| self.simplifications.get(&b).cloned())
     }
 }