From 863e9c073ba864ee2463ff80e201f2ce6e7ef4bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20B=C3=B6ving?= <hargonix@gmail.com>
Date: Tue, 1 Oct 2024 17:28:39 +0200
Subject: [PATCH] feat: generalize the `bv_normalize` pipeline to support more
 general preprocessing passes (#5568)

Beyond what's in the title this also fixes: #5543
---
 .../Tactic/BVDecide/Frontend/BVCheck.lean     |  6 +-
 .../Tactic/BVDecide/Frontend/BVDecide.lean    | 10 +--
 .../Tactic/BVDecide/Frontend/Normalize.lean   | 89 +++++++++++++------
 tests/lean/run/bv_decide_rewriter.lean        | 10 +++
 4 files changed, 78 insertions(+), 37 deletions(-)

diff --git a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVCheck.lean b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVCheck.lean
index 7decc61f33c6..0a32d89b0ab2 100644
--- a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVCheck.lean
+++ b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVCheck.lean
@@ -55,9 +55,9 @@ def evalBvCheck : Tactic := fun
   | `(tactic| bv_check%$tk $path:str) => do
     let cfg ← BVDecide.Frontend.BVCheck.mkContext path.getString
     liftMetaFinishingTactic fun g => do
-      let res ← Normalize.bvNormalize g
-      match res.goal with
-      | some g => bvCheck g cfg
+      let g'? ← Normalize.bvNormalize g
+      match g'? with
+      | some g' => bvCheck g' cfg
       | none =>
         let bvNormalizeStx ← `(tactic| bv_normalize)
         TryThis.addSuggestion tk bvNormalizeStx (origSpan? := ← getRef)
diff --git a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide.lean b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide.lean
index 6d2dae182676..1dac032f678d 100644
--- a/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide.lean
+++ b/src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide.lean
@@ -265,10 +265,6 @@ def bvUnsat (g : MVarId) (cfg : TacticContext) : MetaM (Except CounterExample Lr
 The result of calling `bv_decide`.
 -/
 structure Result where
-  /--
-  Trace of the `simp` used in `bv_decide`'s normalization procedure.
-  -/
-  simpTrace : Simp.Stats
   /--
   If the normalization step was not enough to solve the goal this contains the LRAT proof
   certificate.
@@ -280,10 +276,10 @@ Try to close `g` using a bitblaster. Return either a `CounterExample` if one is
 if `g` is proven.
 -/
 def bvDecide' (g : MVarId) (cfg : TacticContext) : MetaM (Except CounterExample Result) := do
-  let ⟨g?, simpTrace⟩ ← Normalize.bvNormalize g
-  let some g := g? | return .ok ⟨simpTrace, none⟩
+  let g? ← Normalize.bvNormalize g
+  let some g := g? | return .ok ⟨none⟩
   match ← bvUnsat g cfg with
-  | .ok lratCert => return .ok ⟨simpTrace, some lratCert⟩
+  | .ok lratCert => return .ok ⟨some lratCert⟩
   | .error counterExample => return .error counterExample
 
 /--
diff --git a/src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize.lean b/src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize.lean
index 261a6f1c9da7..11ed833d6fe1 100644
--- a/src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize.lean
+++ b/src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize.lean
@@ -37,40 +37,75 @@ builtin_simproc [bv_normalize] eqToBEq (((_ : Bool) = (_ : Bool))) := fun e => d
     let proof := mkApp2 (mkConst ``Bool.eq_to_beq) lhs rhs
     return .done { expr := new, proof? := some proof }
 
-structure Result where
-  goal : Option MVarId := none
-  stats : Simp.Stats := {}
+/--
+A pass in the normalization pipeline. Takes the current goal and produces a refined one or closes
+the goal fully, indicated by returning `none`.
+-/
+abbrev Pass := MVarId → MetaM (Option MVarId)
+
+namespace Pass
+
+/--
+Repeatedly run a list of `Pass` until they either close the goal or an iteration doesn't change
+the goal anymore.
+-/
+partial def fixpointPipeline (passes : List Pass) (goal : MVarId) : MetaM (Option MVarId) := do
+  let runPass (goal? : Option MVarId) (pass : Pass) : MetaM (Option MVarId) := do
+    let some goal := goal? | return none
+    pass goal
+
+  let some newGoal := ← passes.foldlM (init := some goal) runPass | return none
+  if goal != newGoal then
+    trace[Meta.Tactic.bv] m!"Rerunning pipeline on:\n{newGoal}"
+    fixpointPipeline passes newGoal
+  else
+    trace[Meta.Tactic.bv] "Pipeline reached a fixpoint"
+    return newGoal
+
+/--
+Responsible for applying the Bitwuzla style rewrite rules.
+-/
+def rewriteRulesPass : Pass := fun goal => do
+  let bvThms ← bvNormalizeExt.getTheorems
+  let bvSimprocs ← bvNormalizeSimprocExt.getSimprocs
+  let sevalThms ← getSEvalTheorems
+  let sevalSimprocs ← Simp.getSEvalSimprocs
+
+  let simpCtx : Simp.Context := {
+    config := { failIfUnchanged := false, zetaDelta := true }
+    simpTheorems := #[bvThms, sevalThms]
+    congrTheorems := (← getSimpCongrTheorems)
+  }
+
+  let hyps ← goal.getNondepPropHyps
+  let ⟨result?, _⟩ ← simpGoal goal
+    (ctx := simpCtx)
+    (simprocs := #[bvSimprocs, sevalSimprocs])
+    (fvarIdsToSimp := hyps)
+  let some (_, newGoal) := result? | return none
+  return newGoal
+
+/--
+The normalization passes used by `bv_normalize` and thus `bv_decide`.
+-/
+def defaultPipeline : List Pass := [rewriteRulesPass]
+
+end Pass
 
-def bvNormalize (g : MVarId) : MetaM Result := do
+def bvNormalize (g : MVarId) : MetaM (Option MVarId) := do
   withTraceNode `bv (fun _ => return "Normalizing goal") do
     -- Contradiction proof
-    let some g ← g.falseOrByContra | return {}
-
-    -- Normalization by simp
-    let bvThms ← bvNormalizeExt.getTheorems
-    let bvSimprocs ← bvNormalizeSimprocExt.getSimprocs
-    let sevalThms ← getSEvalTheorems
-    let sevalSimprocs ← Simp.getSEvalSimprocs
-
-    let simpCtx : Simp.Context := {
-      config := { failIfUnchanged := false, zetaDelta := true }
-      simpTheorems := #[bvThms, sevalThms]
-      congrTheorems := (← getSimpCongrTheorems)
-    }
-
-    let hyps ← g.getNondepPropHyps
-    let ⟨result?, stats⟩ ← simpGoal g
-      (ctx := simpCtx)
-      (simprocs := #[bvSimprocs, sevalSimprocs])
-      (fvarIdsToSimp := hyps)
-    let some (_, g) := result? | return ⟨none, stats⟩
-    return ⟨some g, stats⟩
+    let some g ← g.falseOrByContra | return none
+    trace[Meta.Tactic.bv] m!"Running preprocessing pipeline on:\n{g}"
+    Pass.fixpointPipeline Pass.defaultPipeline g
 
 @[builtin_tactic Lean.Parser.Tactic.bvNormalize]
 def evalBVNormalize : Tactic := fun
   | `(tactic| bv_normalize) => do
-    liftMetaFinishingTactic fun g => do
-      discard <| bvNormalize g
+    let g ← getMainGoal
+    match ← bvNormalize g with
+    | some newGoal => setGoals [newGoal]
+    | none => setGoals []
   | _ => throwUnsupportedSyntax
 
 end Frontend.Normalize
diff --git a/tests/lean/run/bv_decide_rewriter.lean b/tests/lean/run/bv_decide_rewriter.lean
index 3d696fc60588..cd3ed7920b0d 100644
--- a/tests/lean/run/bv_decide_rewriter.lean
+++ b/tests/lean/run/bv_decide_rewriter.lean
@@ -16,3 +16,13 @@ example :
 
 example (x y z : BitVec 8) (h1 : x = z → False) (h2 : x = y) (h3 : y = z) : False := by
   bv_decide
+
+def mem_subset (a1 a2 b1 b2 : BitVec 64) : Bool :=
+  (b2 - b1 = BitVec.ofNat 64 (2^64 - 1)) ||
+  ((a2 - b1 <= b2 - b1 && a1 - b1 <= a2 - b1))
+
+-- Show that bv_normalize yields the preprocessed goal
+theorem mem_subset_refl : mem_subset a1 a2 a1 a2 := by
+  unfold mem_subset
+  bv_normalize
+  sorry