From ba6f253e4f4e440436aafe4a2333ecde8010ea25 Mon Sep 17 00:00:00 2001
From: Siddharth Bhat <siddu.druid@gmail.com>
Date: Wed, 30 Oct 2024 13:50:59 -0500
Subject: [PATCH] feat: Switch bvOmegaBench to use MetaM [3/?]

---
 Tactics/BvOmegaBench.lean | 102 +++++++++++++++++++++++++++++---------
 1 file changed, 79 insertions(+), 23 deletions(-)

diff --git a/Tactics/BvOmegaBench.lean b/Tactics/BvOmegaBench.lean
index 1505335c..8a159eab 100644
--- a/Tactics/BvOmegaBench.lean
+++ b/Tactics/BvOmegaBench.lean
@@ -10,41 +10,96 @@ and allows us to send bug reports to the Lean developers.
 -/
 import Tactics.Attr
 import Lean
-open Lean Elab Meta Tactic
+import Lean.Elab.Tactic.Omega.Frontend
+import Lean.Elab.Tactic.Omega
+import Tactics.Simp
+
+open Lean Elab Meta Tactic Omega
 
 namespace BvOmegaBench
 
+
+-- Adapted mkSimpContext:
+-- from https://github.com/leanprover/lean4/blob/master/src/Lean/Elab/Tactic/Simp.lean#L287
+def bvOmegaSimpCtx : MetaM (Simp.Context ×  Array Simp.Simprocs) := do
+  let mut simprocs := #[]
+  let mut simpTheorems := #[]
+
+  let some ext ← (getSimpExtension? `bv_toNat)
+    | throwError m!"[omega] Error: unable to find `bv_toNat"
+  simpTheorems := simpTheorems.push (← ext.getTheorems)
+
+  if let some ext ← (Simp.getSimprocExtension? `bv_toNat) then
+    let s ← ext.getSimprocs
+    simprocs := simprocs.push s
+
+  let congrTheorems ← Meta.getSimpCongrTheorems
+  let config : Simp.Config := { failIfUnchanged := false }
+  let ctx : Simp.Context := { config, simpTheorems, congrTheorems }
+  return (ctx, simprocs)
+
 /--
 Run bv_omega, gather the results, and then store them at the value that is given by the option.
 By default, it's stored at `pwd`, with filename `omega-bench`. The file is appended to,
 with the goal state that is being run, and the time elapsed to solve the goal is written.
+
+Code adapted from: 
+- https://github.com/leanprover/lean4/blob/master/src/Lean/Elab/Tactic/Simp.lean#L406
+- https://github.com/leanprover/lean4/blob/master/src/Lean/Elab/Tactic/Omega/Frontend.lean#L685
 -/
-def run : TacticM Unit := do
+def run (g : MVarId) : MetaM Unit := do
   let minMs ← getBvOmegaBenchMinMs
-  let goal ← getMainGoal
-  let goalStr ← ppGoal goal
+  let goalStr ← ppGoal g
   let startTime ← IO.monoMsNow
+  let filePath ← getBvOmegaBenchFilePath
   try
-    withMainContext do
-      withoutRecover do
-        evalTactic (← `(tactic| bv_omega))
-        if !(← getBvOmegaBenchIsEnabled) then
-          return ()
-        let endTime ← IO.monoMsNow
-        let delta := endTime - startTime
-        let filePath ← getBvOmegaBenchFilePath
-        IO.FS.withFile filePath IO.FS.Mode.append fun h => do
-          if delta >= minMs then
-            h.putStrLn "\n---\n"
-            h.putStrLn s!"time"
-            h.putStrLn s!"{delta}"
-            h.putStrLn s!"endtime"
-            h.putStrLn s!"goal"
-            h.putStrLn goalStr.pretty
-            h.putStrLn s!"endgoal"
+      g.withContext do
+        let (bvToNatSimpCtx, bvToNatSimprocs) ← bvOmegaSimpCtx
+        let nonDepHyps ← g.getNondepPropHyps
+        let mut g := g
+
+        try
+          let (result?, _stats) ← 
+            simpGoal g bvToNatSimpCtx bvToNatSimprocs
+              (simplifyTarget := true) (discharge? := .none) nonDepHyps
+          let .some (_, g') := result? | return ()
+          g := g'
+        catch e => 
+          trace[simp_mem.info] "in BvOmega, ran `simp only [bv_toNat]` and got error: {indentD e.toMessageData}"
+          throw e
+
+        g.withContext do 
+          let some g ← g.falseOrByContra | return ()
+          g.withContext do
+            let hyps := (← getLocalHyps).toList
+            omega hyps g {}
+            let endTime ← IO.monoMsNow
+            let delta := endTime - startTime
+            if (← getBvOmegaBenchIsEnabled) && delta ≥ minMs then 
+              IO.FS.withFile filePath IO.FS.Mode.append fun h => do
+                  h.putStrLn "\n---\n"
+                  h.putStrLn s!"goal"
+                  h.putStrLn goalStr.pretty
+                  h.putStrLn s!"endgoal"
+                  h.putStrLn s!"time"
+                  h.putStrLn s!"{delta}"
+                  h.putStrLn s!"endtime"
   catch e =>
+    let endTime ← IO.monoMsNow
+    let delta := endTime - startTime
+    if (← getBvOmegaBenchIsEnabled) && delta ≥ minMs then 
+      IO.FS.withFile filePath IO.FS.Mode.append fun h => do
+          h.putStrLn "\n---\n"
+          h.putStrLn s!"goal"
+          h.putStrLn goalStr.pretty
+          h.putStrLn s!"endgoal"
+          h.putStrLn s!"time"
+          h.putStrLn s!"{delta}"
+          h.putStrLn s!"endtime"
+          h.putStrLn s!"error"
+          h.putStrLn s!"{← e.toMessageData.toString}"
+          h.putStrLn s!"enderror"
     throw e
-  return ()
 
 end BvOmegaBench
 
@@ -53,5 +108,6 @@ syntax (name := bvOmegaBenchTac) "bv_omega_bench" : tactic
 @[tactic bvOmegaBenchTac]
 def bvOmegaBenchImpl : Tactic
 | `(tactic| bv_omega_bench) =>
-   BvOmegaBench.run
+    liftMetaFinishingTactic fun g => do
+      BvOmegaBench.run g
 | _ => throwUnsupportedSyntax