From d24dfa10313b64968915745448f5ca15de2dcaaa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20B=C3=B6ving?= <hargonix@gmail.com>
Date: Sun, 23 Mar 2025 14:56:00 +0100
Subject: [PATCH] perf: add a cache to bv_decide's reflection procedure (#7644)

This PR adds a cache to the reflection procedure of bv_decide.

This was motivated by the following profile on QF_BV SMTLIB problem
`sage/app12/bench_3564.smt2`: https://share.firefox.dev/4iTG8KX. After
this change we roughly get a 10x speedup and `simp` is the bottleneck
again: https://share.firefox.dev/4iuezYT
---
 .../BVDecide/Frontend/BVDecide/Reflect.lean   | 42 +++++++++++++++++++
 .../BVDecide/Frontend/BVDecide/Reify.lean     | 23 +++++-----
 2 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reflect.lean b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reflect.lean
index ca07f84c33..4fb8683e2d 100644
--- a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reflect.lean
+++ b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reflect.lean
@@ -300,6 +300,18 @@ structure LemmaState where
   The list of top level lemmas that got created on the fly during reflection.
   -/
   lemmas : Array SatAtBVLogical := #[]
+  /--
+  Cache for reification of `BVExpr`.
+  -/
+  bvExprCache : Std.HashMap Expr (Option ReifiedBVExpr) := {}
+  /--
+  Cache for reification of `BVPred`.
+  -/
+  bvPredCache : Std.HashMap Expr (Option ReifiedBVPred) := {}
+  /--
+  Cache for reification of `BVLogicalExpr`.
+  -/
+  bvLogicalCache : Std.HashMap Expr (Option ReifiedBVLogical) := {}
 
 /--
 The lemma reflection monad. It extends the usual reflection monad `M` by adding the ability to
@@ -319,6 +331,36 @@ Add another top level lemma.
 def addLemma (lemma : SatAtBVLogical) : LemmaM Unit := do
   modify fun s => { s with lemmas := s.lemmas.push lemma }
 
+@[specialize]
+def withBVExprCache (e : Expr) (f : Expr → LemmaM (Option ReifiedBVExpr)) :
+    LemmaM (Option ReifiedBVExpr) := do
+  match (← get).bvExprCache[e]? with
+  | some hit => return hit
+  | none =>
+    let res ← f e
+    modify fun s => { s with bvExprCache := s.bvExprCache.insert e res }
+    return res
+
+@[specialize]
+def withBVPredCache (e : Expr) (f : Expr → LemmaM (Option ReifiedBVPred)) :
+    LemmaM (Option ReifiedBVPred) := do
+  match (← get).bvPredCache[e]? with
+  | some hit => return hit
+  | none =>
+    let res ← f e
+    modify fun s => { s with bvPredCache := s.bvPredCache.insert e res }
+    return res
+
+@[specialize]
+def withBVLogicalCache (e : Expr) (f : Expr → LemmaM (Option ReifiedBVLogical)) :
+    LemmaM (Option ReifiedBVLogical) := do
+  match (← get).bvLogicalCache[e]? with
+  | some hit => return hit
+  | none =>
+    let res ← f e
+    modify fun s => { s with bvLogicalCache := s.bvLogicalCache.insert e res }
+    return res
+
 end LemmaM
 
 end Frontend
diff --git a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reify.lean b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reify.lean
index 4a9a698eaa..d5132acd05 100644
--- a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reify.lean
+++ b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reify.lean
@@ -184,10 +184,11 @@ where
   to return `some`.
   -/
   goOrAtom (x : Expr) : LemmaM (Option ReifiedBVExpr) := do
-    let res ← go x
-    match res with
-    | some exp => return some exp
-    | none => ReifiedBVExpr.bitVecAtom x false
+    LemmaM.withBVExprCache x fun x => do
+      let res ← go x
+      match res with
+      | some exp => return some exp
+      | none => ReifiedBVExpr.bitVecAtom x false
 
   shiftConstLikeReflection (distance : Nat) (innerExpr : Expr) (shiftOp : Nat → BVUnOp)
       (shiftOpName : Name) (congrThm : Name) :
@@ -283,9 +284,10 @@ Reify an `Expr` that is a predicate about `BitVec`.
 Unless this function is called on something that is not a `Bool` it is always going to return `some`.
 -/
 partial def ReifiedBVPred.of (t : Expr) : LemmaM (Option ReifiedBVPred) := do
-  match ← go t with
-  | some pred => return some pred
-  | none => ReifiedBVPred.boolAtom t
+  LemmaM.withBVPredCache t fun t => do
+    match ← go t with
+    | some pred => return some pred
+    | none => ReifiedBVPred.boolAtom t
 where
   /--
   Reify `t`, returns `none` if the reification procedure failed.
@@ -344,9 +346,10 @@ where
   Unless this function is called on something that is not a `Bool` it is always going to return `some`.
   -/
   goOrAtom (t : Expr) : LemmaM (Option ReifiedBVLogical) := do
-    match ← go t with
-    | some boolExpr => return some boolExpr
-    | none => ReifiedBVLogical.boolAtom t
+    LemmaM.withBVLogicalCache t fun t => do
+      match ← go t with
+      | some boolExpr => return some boolExpr
+      | none => ReifiedBVLogical.boolAtom t
 
   gateReflection (lhsExpr rhsExpr : Expr) (gate : Gate) :
       LemmaM (Option ReifiedBVLogical) := do