perf: improve bv_decide preprocessing based on Bitwuzla optimisations (…

…#6641)

This PR implements several optimisation tricks from Bitwuzla's
preprocessing passes into the Lean equivalent in `bv_decide`. Note that
these changes are mostly geared towards large proof states as for
example seen in SMT-Lib.

src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize.lean

@@ -40,7 +40,7 @@ def passPipeline : PreProcessM (List Pass) := do
 
 def bvNormalize (g : MVarId) (cfg : BVDecideConfig) : MetaM (Option MVarId) := do
   withTraceNode `bv (fun _ => return "Preprocessing goal") do
-    (go g).run cfg
+    (go g).run cfg g
 where
   go (g : MVarId) : PreProcessM (Option MVarId) := do
     let some g ← g.falseOrByContra | return none

src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/AndFlatten.lean

@@ -19,6 +19,11 @@ namespace Frontend.Normalize
 
 open Lean.Meta
 
+structure AndFlattenState where
+  hypsToDelete : Array FVarId := #[]
+  hypsToAdd : Array Hypothesis := #[]
+  cache : Std.HashSet Expr := {}
+
 /--
 Flatten out ands. That is look for hypotheses of the form `h : (x && y) = true` and replace them
 with `h.left : x = true` and `h.right : y = true`. This can enable more fine grained substitutions
@@ -27,59 +32,67 @@ in embedded constraint substitution.
 partial def andFlatteningPass : Pass where
   name := `andFlattening
   run' goal := do
+    let (_, { hypsToDelete, hypsToAdd, .. }) ← processGoal goal |>.run {}
+    if hypsToAdd.isEmpty then
+      return goal
+    else
+      let (_, goal) ← goal.assertHypotheses hypsToAdd
+      -- Given that we collected the hypotheses in the correct order above the invariant is given
+      let goal ← goal.tryClearMany hypsToDelete
+      return goal
+where
+  processGoal (goal : MVarId) : StateRefT AndFlattenState MetaM Unit := do
     goal.withContext do
       let hyps ← goal.getNondepPropHyps
-      let mut newHyps := #[]
-      let mut oldHyps := #[]
-      for fvar in hyps do
-        let hyp : Hypothesis := {
-          userName := (← fvar.getDecl).userName
-          type := ← fvar.getType
-          value := mkFVar fvar
-        }
-        let sizeBefore := newHyps.size
-        newHyps ← splitAnds hyp newHyps
-        if newHyps.size > sizeBefore then
-          oldHyps := oldHyps.push fvar
-      if newHyps.size == 0 then
-        return goal
-      else
-        let (_, goal) ← goal.assertHypotheses newHyps
-        -- Given that we collected the hypotheses in the correct order above the invariant is given
-        let goal ← goal.tryClearMany oldHyps
-        return goal
-where
-  splitAnds (hyp : Hypothesis) (hyps : Array Hypothesis) (first : Bool := true) :
-      MetaM (Array Hypothesis) := do
-    match ← trySplit hyp with
-    | some (left, right) =>
-      let hyps ← splitAnds left hyps false
-      splitAnds right hyps false
-    | none =>
-      if first then
-        return hyps
-      else
-        return hyps.push hyp
+      hyps.forM processFVar
+
+  processFVar (fvar : FVarId) : StateRefT AndFlattenState MetaM Unit := do
+    let type ← fvar.getType
+    if (← get).cache.contains type then
+      modify (fun s => { s with hypsToDelete := s.hypsToDelete.push fvar })
+    else
+      let hyp := {
+        userName := (← fvar.getDecl).userName
+        type := type
+        value := mkFVar fvar
+      }
+      let some (lhs, rhs) ← trySplit hyp | return ()
+      modify (fun s => { s with hypsToDelete := s.hypsToDelete.push fvar })
+      splitAnds [lhs, rhs]
+
+  splitAnds (worklist : List Hypothesis) : StateRefT AndFlattenState MetaM Unit := do
+    match worklist with
+    | [] => return ()
+    | hyp :: worklist =>
+      match ← trySplit hyp with
+      | some (left, right) => splitAnds <| left :: right :: worklist
+      | none =>
+        modify (fun s => { s with hypsToAdd := s.hypsToAdd.push hyp })
+        splitAnds worklist
 
-  trySplit (hyp : Hypothesis) : MetaM (Option (Hypothesis × Hypothesis)) := do
+  trySplit (hyp : Hypothesis) :
+      StateRefT AndFlattenState MetaM (Option (Hypothesis × Hypothesis)) := do
     let typ := hyp.type
-    let_expr Eq α eqLhs eqRhs := typ | return none
-    let_expr Bool.and lhs rhs := eqLhs | return none
-    let_expr Bool.true := eqRhs | return none
-    let_expr Bool := α | return none
-    let mkEqTrue (lhs : Expr) : Expr :=
-      mkApp3 (mkConst ``Eq [1]) (mkConst ``Bool) lhs (mkConst ``Bool.true)
-    let leftHyp : Hypothesis := {
-      userName := hyp.userName,
-      type := mkEqTrue lhs,
-      value := mkApp3 (mkConst ``Std.Tactic.BVDecide.Normalize.Bool.and_left) lhs rhs hyp.value
-    }
-    let rightHyp : Hypothesis := {
-      userName := hyp.userName,
-      type := mkEqTrue rhs,
-      value := mkApp3 (mkConst ``Std.Tactic.BVDecide.Normalize.Bool.and_right) lhs rhs hyp.value
-    }
-    return some (leftHyp, rightHyp)
+    if (← get).cache.contains typ then
+      return none
+    else
+      modify (fun s => { s with cache := s.cache.insert typ })
+      let_expr Eq _ eqLhs eqRhs := typ | return none
+      let_expr Bool.and lhs rhs := eqLhs | return none
+      let_expr Bool.true := eqRhs | return none
+      let mkEqTrue (lhs : Expr) : Expr :=
+        mkApp3 (mkConst ``Eq [1]) (mkConst ``Bool) lhs (mkConst ``Bool.true)
+      let leftHyp : Hypothesis := {
+        userName := hyp.userName,
+        type := mkEqTrue lhs,
+        value := mkApp3 (mkConst ``Std.Tactic.BVDecide.Normalize.Bool.and_left) lhs rhs hyp.value
+      }
+      let rightHyp : Hypothesis := {
+        userName := hyp.userName,
+        type := mkEqTrue rhs,
+        value := mkApp3 (mkConst ``Std.Tactic.BVDecide.Normalize.Bool.and_right) lhs rhs hyp.value
+      }
+      return some (leftHyp, rightHyp)
 
 
 end Frontend.Normalize

src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/Basic.lean

@@ -16,14 +16,33 @@ namespace Frontend.Normalize
 
 open Lean.Meta
 
-abbrev PreProcessM : Type → Type := ReaderT BVDecideConfig MetaM
+structure PreProcessState where
+  /--
+  Contains `FVarId` that we already know are in `bv_normalize` simp normal form and thus don't
+  need to be processed again when we visit the next time.
+  -/
+  rewriteCache : Std.HashSet FVarId := {}
+
+abbrev PreProcessM : Type → Type := ReaderT BVDecideConfig <| StateRefT PreProcessState MetaM
 
 namespace PreProcessM
 
 def getConfig : PreProcessM BVDecideConfig := read
 
-def run (cfg : BVDecideConfig) (x : PreProcessM α) : MetaM α :=
-  ReaderT.run x cfg
+@[inline]
+def checkRewritten (fvar : FVarId) : PreProcessM Bool := do
+  let val := (← get).rewriteCache.contains fvar
+  trace[Meta.Tactic.bv] m!"{mkFVar fvar} was already rewritten? {val}"
+  return val
+
+@[inline]
+def rewriteFinished (fvar : FVarId) : PreProcessM Unit := do
+  trace[Meta.Tactic.bv] m!"Adding {mkFVar fvar} to the rewritten set"
+  modify (fun s => { s with rewriteCache := s.rewriteCache.insert fvar })
+
+def run (cfg : BVDecideConfig) (goal : MVarId) (x : PreProcessM α) : MetaM α := do
+  let hyps ← goal.getNondepPropHyps
+  ReaderT.run x cfg |>.run' { rewriteCache := Std.HashSet.empty hyps.size }
 
 end PreProcessM
 

src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/EmbeddedConstraint.lean

@@ -33,26 +33,26 @@ def embeddedConstraintPass : Pass where
       let mut duplicates : Array FVarId := #[]
       for hyp in hyps do
         let typ ← hyp.getType
-        let_expr Eq α lhs rhs := typ | continue
+        let_expr Eq _ lhs rhs := typ | continue
         let_expr Bool.true := rhs | continue
-        let_expr Bool := α | continue
         if seen.contains lhs then
-          -- collect and later remove duplicates on the fly
           duplicates := duplicates.push hyp
         else
           seen := seen.insert lhs
           let localDecl ← hyp.getDecl
-          let proof  := localDecl.toExpr
+          let proof := localDecl.toExpr
           relevantHyps ← relevantHyps.addTheorem (.fvar hyp) proof
 
       let goal ← goal.tryClearMany duplicates
-      let cfg ← PreProcessM.getConfig
 
+      if relevantHyps.isEmpty then
+        return goal
+
+      let cfg ← PreProcessM.getConfig
       let simpCtx ← Simp.mkContext
         (config := { failIfUnchanged := false, maxSteps := cfg.maxSteps })
         (simpTheorems := relevantHyps)
         (congrTheorems := (← getSimpCongrTheorems))
-
       let ⟨result?, _⟩ ← simpGoal goal (ctx := simpCtx) (fvarIdsToSimp := ← goal.getNondepPropHyps)
       let some (_, newGoal) := result? | return none
       return newGoal

src/Lean/Elab/Tactic/BVDecide/Frontend/Normalize/Rewrite.lean

@@ -35,13 +35,27 @@ def rewriteRulesPass : Pass where
       (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
+    let hyps ← getHyps goal
+    if hyps.isEmpty then
+      return goal
+    else
+      let ⟨result?, _⟩ ← simpGoal goal
+        (ctx := simpCtx)
+        (simprocs := #[bvSimprocs, sevalSimprocs])
+        (fvarIdsToSimp := hyps)
+
+      let some (_, newGoal) := result? | return none
+      newGoal.withContext do
+        (← newGoal.getNondepPropHyps).forM PreProcessM.rewriteFinished
+      return newGoal
+where
+  getHyps (goal : MVarId) : PreProcessM (Array FVarId) := do
+    goal.withContext do
+      let mut hyps ← goal.getNondepPropHyps
+      let filter hyp := do
+        return !(← PreProcessM.checkRewritten hyp)
+      hyps.filterM filter
+
 
 end Frontend.Normalize
 end Lean.Elab.Tactic.BVDecide

tests/lean/run/bv_decide_bool.lean

@@ -1,9 +1,7 @@
 import Std.Tactic.BVDecide
 
-example (h : true = false) : False := by bv_decide
 example {x y z : Bool} (_ : (x && y) = z) (_ : x = !y) (_ : z = true) : False := by bv_decide
 example {a b c d e f : Bool} (_ : (a && b) = c) (_ : (b && c) = d) (_ : (c && d) = e) (_ : (d && e) = f)
     (_ : a = false) (_ : f = true) : False := by bv_decide
 
-example (h : true = false) : False := by bv_decide
 example (h : x = false) : false = x := by bv_decide