From 9e5e0e23b2ec5032670b04f80220d1f51f8f19be Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sat, 24 Feb 2024 15:26:27 -0800
Subject: [PATCH] perf: `mkSplitterProof`

---
 src/Lean/Meta/Match/MatchEqs.lean    | 31 ++++++++++++++++++++++++++++
 tests/lean/run/bv_math_lit_perf.lean |  8 +------
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/src/Lean/Meta/Match/MatchEqs.lean b/src/Lean/Meta/Match/MatchEqs.lean
index 4839c1a8e3..2caecddd3b 100644
--- a/src/Lean/Meta/Match/MatchEqs.lean
+++ b/src/Lean/Meta/Match/MatchEqs.lean
@@ -15,6 +15,35 @@ import Lean.Meta.Tactic.Contradiction
 
 namespace Lean.Meta
 
+/--
+A custom, approximated, and quick `contradiction` tactic.
+It only finds contradictions of the form `(h₁ : p)` and `(h₂ : ¬ p)` where
+`p`s are structurally equal. The procedure is not quadratic like `contradiction`.
+
+We use it to improve the performance of `proveSubgoalLoop` at `mkSplitterProof`.
+We will eventually have to write more efficient proof automation for this module.
+The new proof automation should exploit the structure of the generated goals and avoid general purpose tactics
+such as `contradiction`.
+-/
+private def _root_.Lean.MVarId.contradictionQuick (mvarId : MVarId) : MetaM Bool := do
+  mvarId.withContext do
+    let mut posMap : HashMap Expr FVarId := {}
+    let mut negMap : HashMap Expr FVarId := {}
+    for localDecl in (← getLCtx) do
+      unless localDecl.isImplementationDetail do
+        if let some p ← matchNot? localDecl.type then
+          if let some pFVarId := posMap.find? p then
+            mvarId.assign (← mkAbsurd (← mvarId.getType) (mkFVar pFVarId) localDecl.toExpr)
+            return true
+          negMap := negMap.insert p localDecl.fvarId
+        if (← isProp localDecl.type) then
+          if let some nFVarId := negMap.find? localDecl.type then
+            mvarId.assign (← mkAbsurd (← mvarId.getType) localDecl.toExpr (mkFVar nFVarId))
+            return true
+          posMap := posMap.insert localDecl.type localDecl.fvarId
+      pure ()
+    return false
+
 /--
   Helper method for `proveCondEqThm`. Given a goal of the form `C.rec ... xMajor = rhs`,
   apply `cases xMajor`. -/
@@ -567,6 +596,8 @@ where
 
   proveSubgoalLoop (mvarId : MVarId) : MetaM Unit := do
     trace[Meta.Match.matchEqs] "proveSubgoalLoop\n{mvarId}"
+    if (← mvarId.contradictionQuick) then
+      return ()
     match (← injectionAny mvarId) with
     | InjectionAnyResult.solved => return ()
     | InjectionAnyResult.failed =>
diff --git a/tests/lean/run/bv_math_lit_perf.lean b/tests/lean/run/bv_math_lit_perf.lean
index 5bfe08d4f7..8a43e7c284 100644
--- a/tests/lean/run/bv_math_lit_perf.lean
+++ b/tests/lean/run/bv_math_lit_perf.lean
@@ -16,13 +16,7 @@ def f (x : BitVec 32) : Nat :=
   | 920#32 => 12
   | _      => 1000
 
--- TODO
-theorem ex1 : f 500#32 = x := by
-  set_option trace.Meta.Tactic.simp true in
-  simp [f]
-  sorry
-
-set_option maxHeartbeats 500
+set_option maxHeartbeats 2800
 example : f 500#32 = x := by
   simp [f]
   sorry