perf: bv_decide uses rfl in reflection if possible (#6286)

This PR ensure `bv_decide` uses definitional equality in its reflection
procedure as much as possible. Previously it would build up explicit
congruence proofs for the kernel to check. This reduces the size of
proof terms passed to kernel speeds up checking of large reflection
proofs.

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

@@ -157,9 +157,9 @@ structure ReifiedBVExpr where
   -/
   bvExpr : BVExpr width
   /--
-  A proof that `bvExpr.eval atomsAssignment = originalBVExpr`.
+  A proof that `bvExpr.eval atomsAssignment = originalBVExpr`, none if it holds by `rfl`.
   -/
-  evalsAtAtoms : M Expr
+  evalsAtAtoms : M (Option Expr)
   /--
   A cache for `toExpr bvExpr`.
   -/
@@ -174,9 +174,9 @@ structure ReifiedBVPred where
   -/
   bvPred : BVPred
   /--
-  A proof that `bvPred.eval atomsAssignment = originalBVPredExpr`.
+  A proof that `bvPred.eval atomsAssignment = originalBVPredExpr`, none if it holds by `rfl`.
   -/
-  evalsAtAtoms : M Expr
+  evalsAtAtoms : M (Option Expr)
   /--
   A cache for `toExpr bvPred`
   -/
@@ -191,9 +191,9 @@ structure ReifiedBVLogical where
   -/
   bvExpr : BVLogicalExpr
   /--
-  A proof that `bvExpr.eval atomsAssignment = originalBVLogicalExpr`.
+  A proof that `bvExpr.eval atomsAssignment = originalBVLogicalExpr`, none if it holds by `rfl`.
   -/
-  evalsAtAtoms : M Expr
+  evalsAtAtoms : M (Option Expr)
   /--
   A cache for `toExpr bvExpr`
   -/
@@ -263,6 +263,29 @@ where
     let newAtomsAssignment ← mkListLit packedType packed
     modify fun s => { s with atomsAssignmentCache := newAtomsAssignment }
 
+@[specialize]
+def simplifyBinaryProof' (mkFRefl : Expr → Expr) (fst : Expr) (fproof : Option Expr)
+    (mkSRefl : Expr → Expr) (snd : Expr) (sproof : Option Expr) : Option (Expr × Expr) := do
+  match fproof, sproof with
+  | some fproof, some sproof => some (fproof, sproof)
+  | some fproof, none => some (fproof, mkSRefl snd)
+  | none, some sproof => some (mkFRefl fst, sproof)
+  | none, none => none
+
+@[specialize]
+def simplifyBinaryProof (mkRefl : Expr → Expr) (fst : Expr) (fproof : Option Expr) (snd : Expr)
+    (sproof : Option Expr) : Option (Expr × Expr) := do
+  simplifyBinaryProof' mkRefl fst fproof mkRefl snd sproof
+
+@[specialize]
+def simplifyTernaryProof (mkRefl : Expr → Expr) (fst : Expr) (fproof : Option Expr) (snd : Expr)
+    (sproof : Option Expr) (thd : Expr) (tproof : Option Expr) : Option (Expr × Expr × Expr) := do
+  match fproof, simplifyBinaryProof mkRefl snd sproof thd tproof with
+  | some fproof, some stproof => some (fproof, stproof)
+  | some fproof, none => some (fproof, mkRefl snd, mkRefl thd)
+  | none, some stproof => some (mkRefl fst, stproof)
+  | none, none => none
+
 end M
 
 /--

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

@@ -37,9 +37,8 @@ Register `e` as an atom of `width` that might potentially be `synthetic`.
 def mkAtom (e : Expr) (width : Nat) (synthetic : Bool) : M ReifiedBVExpr := do
   let ident ← M.lookup e width synthetic
   let expr := mkApp2 (mkConst ``BVExpr.var) (toExpr width) (toExpr ident)
-  let proof := do
-    let evalExpr ← mkEvalExpr width expr
-    return mkBVRefl width evalExpr
+  -- This is safe because this proof always holds definitionally.
+  let proof := pure none
   return ⟨width, .var ident, proof, expr⟩
 
 /--
@@ -70,9 +69,8 @@ Build a reified version of the constant `val`.
 def mkBVConst (val : BitVec w) : M ReifiedBVExpr := do
   let bvExpr : BVExpr w := .const val
   let expr := mkApp2 (mkConst ``BVExpr.const) (toExpr w) (toExpr val)
-  let proof := do
-    let evalExpr ← ReifiedBVExpr.mkEvalExpr w expr
-    return ReifiedBVExpr.mkBVRefl w evalExpr
+  -- This is safe because this proof always holds definitionally.
+  let proof := pure none
   return ⟨w, bvExpr, proof, expr⟩
 
 end ReifiedBVExpr

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

@@ -49,7 +49,8 @@ Build a reified version of the constant `val`.
 def mkBoolConst (val : Bool) : M ReifiedBVLogical := do
   let boolExpr := .const val
   let expr := mkApp2 (mkConst ``BoolExpr.const) (mkConst ``BVPred) (toExpr val)
-  let proof := pure <| ReifiedBVLogical.mkRefl (toExpr val)
+  -- This is safe because this proof always holds definitionally.
+  let proof := pure none
   return ⟨boolExpr, proof, expr⟩
 
 /--
@@ -71,8 +72,13 @@ def mkGate (lhs rhs : ReifiedBVLogical) (lhsExpr rhsExpr : Expr) (gate : Gate) :
   let proof := do
     let lhsEvalExpr ← ReifiedBVLogical.mkEvalExpr lhs.expr
     let rhsEvalExpr ← ReifiedBVLogical.mkEvalExpr rhs.expr
-    let lhsProof ← lhs.evalsAtAtoms
-    let rhsProof ← rhs.evalsAtAtoms
+    let lhsProof? ← lhs.evalsAtAtoms
+    let rhsProof? ← rhs.evalsAtAtoms
+    let some (lhsProof, rhsProof) :=
+      M.simplifyBinaryProof
+        ReifiedBVLogical.mkRefl
+        lhsEvalExpr lhsProof?
+        rhsEvalExpr rhsProof? | return none
     return mkApp6
       (mkConst congrThm)
       lhsExpr rhsExpr
@@ -95,8 +101,9 @@ def mkNot (sub : ReifiedBVLogical) (subExpr : Expr) : M ReifiedBVLogical := do
   let boolExpr := .not sub.bvExpr
   let expr := mkApp2 (mkConst ``BoolExpr.not) (mkConst ``BVPred) sub.expr
   let proof := do
+    -- This is safe as `not_congr` holds definitionally if the arguments are defeq.
+    let some subProof ← sub.evalsAtAtoms | return none
     let subEvalExpr ← ReifiedBVLogical.mkEvalExpr sub.expr
-    let subProof ← sub.evalsAtAtoms
     return mkApp3 (mkConst ``Std.Tactic.BVDecide.Reflect.Bool.not_congr) subExpr subEvalExpr subProof
   return ⟨boolExpr, proof, expr⟩
 
@@ -119,9 +126,15 @@ def mkIte (discr lhs rhs : ReifiedBVLogical) (discrExpr lhsExpr rhsExpr : Expr)
     let discrEvalExpr ← ReifiedBVLogical.mkEvalExpr discr.expr
     let lhsEvalExpr ← ReifiedBVLogical.mkEvalExpr lhs.expr
     let rhsEvalExpr ← ReifiedBVLogical.mkEvalExpr rhs.expr
-    let discrProof ← discr.evalsAtAtoms
-    let lhsProof ← lhs.evalsAtAtoms
-    let rhsProof ← rhs.evalsAtAtoms
+    let discrProof? ← discr.evalsAtAtoms
+    let lhsProof? ← lhs.evalsAtAtoms
+    let rhsProof? ← rhs.evalsAtAtoms
+    let some (discrProof, lhsProof, rhsProof) :=
+      M.simplifyTernaryProof
+        ReifiedBVLogical.mkRefl
+        discrEvalExpr discrProof?
+        lhsEvalExpr lhsProof?
+        rhsEvalExpr rhsProof? | return none
     return mkApp9
       (mkConst ``Std.Tactic.BVDecide.Reflect.Bool.ite_congr)
       discrExpr lhsExpr rhsExpr

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

@@ -35,8 +35,10 @@ def boolAtom (t : Expr) : M (Option ReifiedBVPred) := do
   let bvExpr : BVPred := .getLsbD atom.bvExpr 0
   let expr := mkApp3 (mkConst ``BVPred.getLsbD) (toExpr 1) atom.expr (toExpr 0)
   let proof := do
+    -- ofBool_congr does not hold definitionally, if this ever becomes an issue we need to find
+    -- a more clever encoding for boolean atoms
     let atomEval ← ReifiedBVExpr.mkEvalExpr atom.width atom.expr
-    let atomProof ← atom.evalsAtAtoms
+    let atomProof := (← atom.evalsAtAtoms).getD (ReifiedBVExpr.mkBVRefl atom.width atomEval)
     return mkApp3
       (mkConst ``Std.Tactic.BVDecide.Reflect.BitVec.ofBool_congr)
       t
@@ -63,9 +65,14 @@ def mkBinPred (lhs rhs : ReifiedBVExpr) (lhsExpr rhsExpr : Expr) (pred : BVBinPr
         rhs.expr
     let proof := do
       let lhsEval ← ReifiedBVExpr.mkEvalExpr lhs.width lhs.expr
-      let lhsProof ← lhs.evalsAtAtoms
       let rhsEval ← ReifiedBVExpr.mkEvalExpr rhs.width rhs.expr
-      let rhsProof ← rhs.evalsAtAtoms
+      let lhsProof? ← lhs.evalsAtAtoms
+      let rhsProof? ← rhs.evalsAtAtoms
+      let some (lhsProof, rhsProof) :=
+        M.simplifyBinaryProof
+          (ReifiedBVExpr.mkBVRefl lhs.width)
+          lhsEval lhsProof?
+          rhsEval rhsProof? | return none
       return mkApp7
         (mkConst congrThm)
         (toExpr lhs.width)
@@ -90,8 +97,9 @@ def mkGetLsbD (sub : ReifiedBVExpr) (subExpr : Expr) (idx : Nat) : M ReifiedBVPr
   let idxExpr := toExpr idx
   let expr := mkApp3 (mkConst ``BVPred.getLsbD) (toExpr sub.width) sub.expr idxExpr
   let proof := do
+    -- This is safe as `getLsbD_congr` holds definitionally if the arguments are defeq.
+    let some subProof ← sub.evalsAtAtoms | return none
     let subEval ← ReifiedBVExpr.mkEvalExpr sub.width sub.expr
-    let subProof ← sub.evalsAtAtoms
     return mkApp5
       (mkConst ``Std.Tactic.BVDecide.Reflect.BitVec.getLsbD_congr)
       idxExpr

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

@@ -62,7 +62,7 @@ where
 
     let proof := do
       let evalExpr ← ReifiedBVLogical.mkEvalExpr imp.expr
-      let congrProof ← imp.evalsAtAtoms
+      let congrProof := (← imp.evalsAtAtoms).getD (ReifiedBVLogical.mkRefl evalExpr)
       let lemmaProof := mkApp4 (mkConst lemmaName) (toExpr lhs.width) discrExpr lhsExpr rhsExpr
 
       let trueExpr := mkConst ``Bool.true

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

@@ -112,7 +112,8 @@ where
           inner.expr
       let proof := do
         let innerEval ← ReifiedBVExpr.mkEvalExpr inner.width inner.expr
-        let innerProof ← inner.evalsAtAtoms
+        -- This is safe as `zeroExtend_congr` holds definitionally if the arguments are defeq.
+        let some innerProof ← inner.evalsAtAtoms | return none
         return mkApp5 (mkConst ``Std.Tactic.BVDecide.Reflect.BitVec.zeroExtend_congr)
           newWidthExpr
           (toExpr inner.width)
@@ -132,7 +133,8 @@ where
           inner.expr
       let proof := do
         let innerEval ← ReifiedBVExpr.mkEvalExpr inner.width inner.expr
-        let innerProof ← inner.evalsAtAtoms
+        -- This is safe as `zeroExtend_congr` holds definitionally if the arguments are defeq.
+        let some innerProof ← inner.evalsAtAtoms | return none
         return mkApp5 (mkConst ``Std.Tactic.BVDecide.Reflect.BitVec.signExtend_congr)
           newWidthExpr
           (toExpr inner.width)
@@ -150,9 +152,13 @@ where
         lhs.expr rhs.expr
       let proof := do
         let lhsEval ← ReifiedBVExpr.mkEvalExpr lhs.width lhs.expr
-        let lhsProof ← lhs.evalsAtAtoms
-        let rhsProof ← rhs.evalsAtAtoms
         let rhsEval ← ReifiedBVExpr.mkEvalExpr rhs.width rhs.expr
+        let lhsProof? ← lhs.evalsAtAtoms
+        let rhsProof? ← rhs.evalsAtAtoms
+        let some (lhsProof, rhsProof) :=
+          M.simplifyBinaryProof'
+            (ReifiedBVExpr.mkBVRefl lhs.width) lhsEval lhsProof?
+            (ReifiedBVExpr.mkBVRefl rhs.width) rhsEval rhsProof? | return none
         return mkApp8 (mkConst ``Std.Tactic.BVDecide.Reflect.BitVec.append_congr)
           (toExpr lhs.width) (toExpr rhs.width)
           lhsExpr lhsEval
@@ -169,7 +175,8 @@ where
         inner.expr
       let proof := do
         let innerEval ← ReifiedBVExpr.mkEvalExpr inner.width inner.expr
-        let innerProof ← inner.evalsAtAtoms
+        -- This is safe as `zeroExtend_congr` holds definitionally if the arguments are defeq.
+        let some innerProof ← inner.evalsAtAtoms | return none
         return mkApp5 (mkConst ``Std.Tactic.BVDecide.Reflect.BitVec.replicate_congr)
           (toExpr n)
           (toExpr inner.width)
@@ -189,7 +196,8 @@ where
         inner.expr
       let proof := do
         let innerEval ← ReifiedBVExpr.mkEvalExpr inner.width inner.expr
-        let innerProof ← inner.evalsAtAtoms
+        -- This is safe as `zeroExtend_congr` holds definitionally if the arguments are defeq.
+        let some innerProof ← inner.evalsAtAtoms | return none
         return mkApp6 (mkConst ``Std.Tactic.BVDecide.Reflect.BitVec.extract_congr)
           startExpr
           lenExpr
@@ -301,11 +309,16 @@ where
       return none
 
   binaryCongrProof (lhs rhs : ReifiedBVExpr) (lhsExpr rhsExpr : Expr) (congrThm : Expr) :
-      M Expr := do
+      M (Option Expr) := do
     let lhsEval ← ReifiedBVExpr.mkEvalExpr lhs.width lhs.expr
-    let lhsProof ← lhs.evalsAtAtoms
-    let rhsProof ← rhs.evalsAtAtoms
     let rhsEval ← ReifiedBVExpr.mkEvalExpr rhs.width rhs.expr
+    let lhsProof? ← lhs.evalsAtAtoms
+    let rhsProof? ← rhs.evalsAtAtoms
+    let some (lhsProof, rhsProof) :=
+      M.simplifyBinaryProof
+        (ReifiedBVExpr.mkBVRefl lhs.width)
+        lhsEval lhsProof?
+        rhsEval rhsProof? | return none
     return mkApp6 congrThm lhsExpr rhsExpr lhsEval rhsEval lhsProof rhsProof
 
   unaryReflection (innerExpr : Expr) (op : BVUnOp) (congrThm : Name) :
@@ -316,9 +329,9 @@ where
     let proof := unaryCongrProof inner innerExpr (mkConst congrThm)
     return some ⟨inner.width, bvExpr, proof, expr⟩
 
-  unaryCongrProof (inner : ReifiedBVExpr) (innerExpr : Expr) (congrProof : Expr) : M Expr := do
+  unaryCongrProof (inner : ReifiedBVExpr) (innerExpr : Expr) (congrProof : Expr) : M (Option Expr) := do
     let innerEval ← ReifiedBVExpr.mkEvalExpr inner.width inner.expr
-    let innerProof ← inner.evalsAtAtoms
+    let some innerProof ← inner.evalsAtAtoms | return none
     return mkApp4 congrProof (toExpr inner.width) innerExpr innerEval innerProof
 
   goBvLit (x : Expr) : M (Option ReifiedBVExpr) := do

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

@@ -37,7 +37,7 @@ partial def of (h : Expr) : LemmaM (Option SatAtBVLogical) := do
     let proof := do
       let evalLogic ← ReifiedBVLogical.mkEvalExpr bvLogical.expr
       -- this is evalLogic = lhsExpr
-      let evalProof ← bvLogical.evalsAtAtoms
+      let evalProof := (← bvLogical.evalsAtAtoms).getD (ReifiedBVLogical.mkRefl evalLogic)
       -- h is lhsExpr = true
       -- we prove evalLogic = true by evalLogic = lhsExpr = true
       return ReifiedBVLogical.mkTrans evalLogic lhsExpr (mkConst ``Bool.true) evalProof h