From 7f82e61969718480e549f884998bd396a7fc1fc6 Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Wed, 27 Nov 2024 10:23:30 +0000
Subject: [PATCH] Walk up the dominator tree to find the first simplification

---
 compiler/noirc_evaluator/src/ssa/ir/dom.rs    | 20 +++++++++++++++++++
 .../src/ssa/opt/constant_folding.rs           | 13 +-----------
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/ir/dom.rs b/compiler/noirc_evaluator/src/ssa/ir/dom.rs
index c1a7f14e0d1..085789fb67d 100644
--- a/compiler/noirc_evaluator/src/ssa/ir/dom.rs
+++ b/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 {
diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs
index b53cbab858a..e758d039c94 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs
+++ b/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())
     }
 }