perf: skip canonicalization of Decidable instances and add congruence-closure support (#9267)

This PR optimizes support for `Decidable` instances in `grind`. Because
`Decidable` is a subsingleton, the canonicalizer no longer wastes time
normalizing such instances, a significant performance bottleneck in
benchmarks like `grind_bitvec2.lean`. In addition, the
congruence-closure module now handles `Decidable` instances, and can
solve examples such as:
```lean
example (p q : Prop) (h₁ : Decidable p) (h₂ : Decidable (p ∧ q)) : (p ↔ q) → h₁ ≍ h₂ := by
  grind
```
This commit is contained in:
Leonardo de Moura 2025-07-08 14:55:40 -07:00 committed by GitHub
parent cec0c82f1c
commit 192c0c8e67
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 200 additions and 126 deletions

View file

@ -138,6 +138,9 @@ theorem exists_and_left {α : Sort u} {p : α → Prop} {b : Prop} : (∃ x, b
theorem exists_and_right {α : Sort u} {p : α → Prop} {b : Prop} : (∃ x, p x ∧ b) = ((∃ x, p x) ∧ b) := by
apply propext; apply _root_.exists_and_right
theorem zero_sub (a : Nat) : 0 - a = 0 := by
simp
init_grind_norm
/- Pre theorems -/
not_and not_or not_ite not_forall not_exists
@ -181,7 +184,7 @@ init_grind_norm
Nat.add_eq Nat.sub_eq Nat.mul_eq Nat.zero_eq Nat.le_eq
Nat.div_zero Nat.mod_zero Nat.div_one Nat.mod_one
Nat.sub_sub Nat.pow_zero Nat.pow_one Nat.sub_self
Nat.one_pow
Nat.one_pow Nat.zero_sub
-- Int
Int.lt_eq
Int.emod_neg Int.ediv_neg

View file

@ -16,6 +16,11 @@ namespace Lean.Grind
/-- A helper gadget for annotating nested proofs in goals. -/
def nestedProof (p : Prop) {h : p} : p := h
/-- A helper gadget for annotating nested decidable instances in goals. -/
-- Remark: we currently have special gadgets for the two most common subsingletons in Lean, and are the only
-- currently supported in `grind`. We may add a generic `nestedSubsingleton` inn the future.
def nestedDecidable {p : Prop} (h : Decidable p) : Decidable p := h
/--
Gadget for marking `match`-expressions that should not be reduced by the `grind` simplifier, but the discriminants should be normalized.
We use it when adding instances of `match`-equations to prevent them from being simplified to true.
@ -58,6 +63,9 @@ def PreMatchCond (p : Prop) : Prop := p
theorem nestedProof_congr (p q : Prop) (h : p = q) (hp : p) (hq : q) : @nestedProof p hp ≍ @nestedProof q hq := by
subst h; apply HEq.refl
theorem nestedDecidable_congr (p q : Prop) (h : p = q) (hp : Decidable p) (hq : Decidable q) : @nestedDecidable p hp ≍ @nestedDecidable q hq := by
subst h; cases hp <;> cases hq <;> simp <;> contradiction
@[app_unexpander nestedProof]
meta def nestedProofUnexpander : PrettyPrinter.Unexpander := fun stx => do
match stx with

View file

@ -12,7 +12,7 @@ import Lean.Meta.Tactic.Grind.Cases
import Lean.Meta.Tactic.Grind.Injection
import Lean.Meta.Tactic.Grind.Core
import Lean.Meta.Tactic.Grind.Canon
import Lean.Meta.Tactic.Grind.MarkNestedProofs
import Lean.Meta.Tactic.Grind.MarkNestedSubsingletons
import Lean.Meta.Tactic.Grind.Inv
import Lean.Meta.Tactic.Grind.Proof
import Lean.Meta.Tactic.Grind.Propagate

View file

@ -94,13 +94,13 @@ def canonElemCore (parent : Expr) (f : Expr) (i : Nat) (e : Expr) (useIsDefEqBou
-- Moreover, we store the canonicalizer state in the `Goal` because we case-split
-- and different locals are added in different branches.
modify' fun s => { s with canon := s.canon.insert e c }
trace_goal[grind.debugn.canon] "found {e} ===> {c}"
trace_goal[grind.debug.canon] "found {e} ===> {c}"
return c
if useIsDefEqBounded then
-- If `e` and `c` are not types, we use `isDefEqBounded`
if (← isDefEqBounded e c parent) then
modify' fun s => { s with canon := s.canon.insert e c }
trace_goal[grind.debugn.canon] "found using `isDefEqBounded`: {e} ===> {c}"
trace_goal[grind.debug.canon] "found using `isDefEqBounded`: {e} ===> {c}"
return c
trace_goal[grind.debug.canon] "({f}, {i}) ↦ {e}"
modify' fun s => { s with canon := s.canon.insert e e, argMap := s.argMap.insert key ((e, eType)::cs) }
@ -122,7 +122,7 @@ private inductive ShouldCanonResult where
canonImplicit
| /-
Term is not a proof, type (former), nor an instance.
Thus, it must be recursively visited by the canonizer.
Thus, it must be recursively visited by the canonicalizer.
-/
visit
deriving Inhabited
@ -156,25 +156,32 @@ def shouldCanon (pinfos : Array ParamInfo) (i : Nat) (arg : Expr) : MetaM Should
else
return .visit
unsafe def canonImpl (e : Expr) : GoalM Expr := do
visit e |>.run' mkPtrMap
/-- Canonicalizes nested types, type formers, and instances in `e`. -/
partial def canon (e : Expr) : GoalM Expr := do profileitM Exception "grind canon" (← getOptions) do
trace_goal[grind.debug.canon] "{e}"
visit e |>.run' {}
where
visit (e : Expr) : StateRefT (PtrMap Expr Expr) GoalM Expr := do
visit (e : Expr) : StateRefT (Std.HashMap ExprPtr Expr) GoalM Expr := do
unless e.isApp || e.isForall do return e
-- Check whether it is cached
if let some r := (← get).find? e then
if let some r := (← get).get? { expr := e } then
return r
let e' ← match e with
| .app .. => e.withApp fun f args => do
if f.isConstOf ``Lean.Grind.nestedProof && args.size == 2 then
if f.isConstOf ``Grind.nestedProof && args.size == 2 then
let prop := args[0]!
let prop' ← visit prop
if let some r := (← get').proofCanon.find? prop' then
pure r
else
let e' := if ptrEq prop prop' then e else mkAppN f (args.set! 0 prop')
let e' := if isSameExpr prop prop' then e else mkAppN f (args.set! 0 prop')
modify' fun s => { s with proofCanon := s.proofCanon.insert prop' e' }
pure e'
else if f.isConstOf ``Grind.nestedDecidable && args.size == 2 then
let prop := args[0]!
let prop' ← visit prop
let e' := if isSameExpr prop prop' then e else mkAppN f (args.set! 0 prop')
pure e'
else
let pinfos := (← getFunInfo f).paramInfo
let mut modified := false
@ -183,11 +190,17 @@ where
let arg := args[i]
trace_goal[grind.debug.canon] "[{repr (← shouldCanon pinfos i arg)}]: {arg} : {← inferType arg}"
let arg' ← match (← shouldCanon pinfos i arg) with
| .canonType => canonType e f i arg
| .canonInst => canonInst e f i arg
| .canonImplicit => canonImplicit e f i (← visit arg)
| .visit => visit arg
unless ptrEq arg arg' do
| .canonType => canonType e f i arg
| .canonImplicit => canonImplicit e f i (← visit arg)
| .visit => visit arg
| .canonInst =>
if arg.isAppOfArity ``Grind.nestedDecidable 2 then
let prop := arg.appFn!.appArg!
let prop' ← visit prop
if isSameExpr prop prop' then pure arg else pure (mkApp2 arg.appFn!.appFn! prop' arg.appArg!)
else
canonInst e f i arg
unless isSameExpr arg arg' do
args := args.set i arg'
modified := true
pure <| if modified then mkAppN f args.toArray else e
@ -198,14 +211,11 @@ where
let b' ← if b.hasLooseBVars then pure b else visit b
pure <| e.updateForallE! d' b'
| _ => unreachable!
modify fun s => s.insert e e'
modify fun s => s.insert { expr := e } e'
return e'
end Canon
/-- Canonicalizes nested types, type formers, and instances in `e`. -/
def canon (e : Expr) : GoalM Expr := do profileitM Exception "grind canon" (← getOptions) do
trace_goal[grind.debug.canon] "{e}"
unsafe Canon.canonImpl e
export Canon (canon)
end Lean.Meta.Grind

View file

@ -393,13 +393,19 @@ where
checkAndAddSplitCandidate e
pushCastHEqs e
addMatchEqns f generation
if f.isConstOf ``Lean.Grind.nestedProof && args.size == 2 then
if args.size == 2 && f.isConstOf ``Grind.nestedProof then
-- We only internalize the proposition. We can skip the proof because of
-- proof irrelevance
let c := args[0]!
internalizeImpl c generation e
registerParent e c
pushEqTrue c <| mkApp2 (mkConst ``eq_true) c args[1]!
else if args.size == 2 && f.isConstOf ``Grind.nestedDecidable then
-- We only internalize the proposition. We can skip the instance because it is
-- a subsingleton
let c := args[0]!
internalizeImpl c generation e
registerParent e c
else if f.isConstOf ``ite && args.size == 5 then
let c := args[1]!
internalizeImpl c generation e

View file

@ -1,98 +0,0 @@
/-
Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
prelude
import Init.Grind.Util
import Lean.Util.PtrSet
import Lean.Meta.Transform
import Lean.Meta.Basic
import Lean.Meta.InferType
import Lean.Meta.Tactic.Grind.Util
namespace Lean.Meta.Grind
private unsafe def markNestedProofImpl (e : Expr) (visit : Expr → StateRefT (PtrMap Expr Expr) MetaM Expr)
: StateRefT (PtrMap Expr Expr) MetaM Expr := do
let prop ← inferType e
/-
We must unfold reducible constants occurring in `prop` because the congruence closure
module in `grind` assumes they have been expanded.
See `grind_mark_nested_proofs_bug.lean` for an example.
TODO: We may have to normalize `prop` too.
-/
/- We must also apply beta-reduction to improve the effectiveness of the congruence closure procedure. -/
let prop ← Core.betaReduce prop
let prop ← unfoldReducible prop
/- We must mask proofs occurring in `prop` too. -/
let prop ← visit prop
let prop ← eraseIrrelevantMData prop
/- We must fold kernel projections like it is done in the preprocessor. -/
let prop ← foldProjs prop
let prop ← normalizeLevels prop
return mkApp2 (mkConst ``Lean.Grind.nestedProof) prop e
unsafe def markNestedProofsImpl (e : Expr) : MetaM Expr := do
visit e |>.run' mkPtrMap
where
visit (e : Expr) : StateRefT (PtrMap Expr Expr) MetaM Expr := do
if (← isProof e) then
if e.isAppOf ``Lean.Grind.nestedProof then
return e -- `e` is already marked
if let some r := (← get).find? e then
return r
let e' ← markNestedProofImpl e visit
modify fun s => s.insert e e'
return e'
-- Remark: we have to process `Expr.proj` since we only
-- fold projections later during term internalization
unless e.isApp || e.isForall || e.isProj || e.isMData do
return e
-- Check whether it is cached
if let some r := (← get).find? e then
return r
let e' ← match e with
| .app .. => e.withApp fun f args => do
let mut modified := false
let mut args := args
for i in *...args.size do
let arg := args[i]!
let arg' ← visit arg
unless ptrEq arg arg' do
args := args.set! i arg'
modified := true
if modified then
pure <| mkAppN f args
else
pure e
| .proj _ _ b =>
pure <| e.updateProj! (← visit b)
| .mdata _ b =>
pure <| e.updateMData! (← visit b)
| .forallE _ d b _ =>
-- Recall that we have `ForallProp.lean`.
let d' ← visit d
let b' ← if b.hasLooseBVars then pure b else visit b
pure <| e.updateForallE! d' b'
| _ => unreachable!
modify fun s => s.insert e e'
return e'
/--
Wrap nested proofs `e` with `Lean.Grind.nestedProof`-applications.
Recall that the congruence closure module has special support for `Lean.Grind.nestedProof`.
-/
def markNestedProofs (e : Expr) : MetaM Expr :=
unsafe markNestedProofsImpl e
/--
Given a proof `e`, mark it with `Lean.Grind.nestedProof`
-/
def markProof (e : Expr) : MetaM Expr := do
if e.isAppOf ``Lean.Grind.nestedProof then
return e -- `e` is already marked
else
unsafe markNestedProofImpl e markNestedProofsImpl.visit |>.run' mkPtrMap
end Lean.Meta.Grind

View file

@ -0,0 +1,119 @@
/-
Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
prelude
import Init.Grind.Util
import Lean.Util.PtrSet
import Lean.Meta.Transform
import Lean.Meta.Basic
import Lean.Meta.InferType
import Lean.Meta.Tactic.Grind.ExprPtr
import Lean.Meta.Tactic.Grind.Util
namespace Lean.Meta.Grind
private abbrev M := StateRefT (Std.HashMap ExprPtr Expr) MetaM
def isMarkedSubsingletonConst (e : Expr) : Bool := Id.run do
let .const declName _ := e | false
return declName == ``Grind.nestedProof || declName == ``Grind.nestedDecidable
def isMarkedSubsingletonApp (e : Expr) : Bool :=
isMarkedSubsingletonConst e.getAppFn
/-- Returns `some p` if `e` is of the form `Decidable p` -/
private def isDecidable (e : Expr) : MetaM (Option Expr) := do
match_expr (← whnfCore e) with
| Decidable p => return some p
| _ => return none
/--
Wrap nested proofs and decidable instances in `e` with `Lean.Grind.nestedProof` and `Lean.Grind.nestedDecidable`-applications.
Recall that the congruence closure module has special support for them.
-/
-- TODO: consider other subsingletons in the future? We decided to not support them to avoid the overhead of
-- synthesizing `Subsingleton` instances.
partial def markNestedSubsingletons (e : Expr) : MetaM Expr := do
visit e |>.run' {}
where
visit (e : Expr) : M Expr := do
if isMarkedSubsingletonApp e then
return e -- `e` is already marked
-- check whether result is cached
if let some r := (← get).get? { expr := e } then
return r
-- check whether `e` is a proposition or `Decidable`
let type ← inferType e
if (← isProp type) then
let e' := mkApp2 (mkConst ``Grind.nestedProof) (← preprocess type) e
modify fun s => s.insert { expr := e } e'
return e'
if let some p ← isDecidable type then
let e' := mkApp2 (mkConst ``Grind.nestedDecidable) (← preprocess p) e
modify fun s => s.insert { expr := e } e'
return e'
-- Remark: we have to process `Expr.proj` since we only
-- fold projections later during term internalization
unless e.isApp || e.isForall || e.isProj || e.isMData do
return e
let e' ← match e with
| .app .. => e.withApp fun f args => do
let mut modified := false
let mut args := args.toVector
for h : i in *...args.size do
let arg := args[i]
let arg' ← visit arg
unless isSameExpr arg arg' do
args := args.set i arg'
modified := true
if modified then
pure <| mkAppN f args.toArray
else
pure e
| .proj _ _ b =>
pure <| e.updateProj! (← visit b)
| .mdata _ b =>
pure <| e.updateMData! (← visit b)
| .forallE _ d b _ =>
-- Recall that we have `ForallProp.lean`.
let d' ← visit d
let b' ← if b.hasLooseBVars then pure b else visit b
pure <| e.updateForallE! d' b'
| _ => unreachable!
modify fun s => s.insert { expr := e } e'
return e'
preprocess (e : Expr) : M Expr := do
/-
We must unfold reducible constants occurring in `prop` because the congruence closure
module in `grind` assumes they have been expanded.
See `grind_mark_nested_proofs_bug.lean` for an example.
TODO: We may have to normalize `prop` too.
-/
/- We must also apply beta-reduction to improve the effectiveness of the congruence closure procedure. -/
let e ← Core.betaReduce e
let e ← unfoldReducible e
/- We must mask proofs occurring in `prop` too. -/
let e ← visit e
let e ← eraseIrrelevantMData e
/- We must fold kernel projections like it is done in the preprocessor. -/
let e ← foldProjs e
normalizeLevels e
def markNestedProof (e : Expr) : M Expr := do
let prop ← inferType e
let prop ← markNestedSubsingletons.preprocess prop
return mkApp2 (mkConst ``Grind.nestedProof) prop e
/--
Given a proof `e`, mark it with `Lean.Grind.nestedProof`
-/
def markProof (e : Expr) : MetaM Expr := do
if e.isAppOf ``Grind.nestedProof then
return e -- `e` is already marked
else
markNestedProof e |>.run' {}
end Lean.Meta.Grind

View file

@ -124,6 +124,19 @@ mutual
let h := mkApp5 (mkConst ``Lean.Grind.nestedProof_congr) p q (← mkEqProofCore p q false) hp hq
mkEqOfHEqIfNeeded h heq
/--
Given `lhs` and `rhs` proof terms of the form `nestedDecidable p hp` and `nestedDecidable q hq`,
constructs a congruence proof for `nestedDecidable p hp ≍ nestedDecidable q hq`.
`p` and `q` are in the same equivalence class.
-/
private partial def mkNestedDecidableCongr (lhs rhs : Expr) (heq : Bool) : GoalM Expr := do
let p := lhs.appFn!.appArg!
let hp := lhs.appArg!
let q := rhs.appFn!.appArg!
let hq := rhs.appArg!
let h := mkApp5 (mkConst ``Lean.Grind.nestedDecidable_congr) p q (← mkEqProofCore p q false) hp hq
mkEqOfHEqIfNeeded h heq
/--
Constructs a congruence proof for `lhs` and `rhs` using `congr`, `congrFun`, and `congrArg`.
This function assumes `isCongrDefaultProofTarget` returned `true`.
@ -217,9 +230,11 @@ mutual
let g := rhs.getAppFn
let numArgs := lhs.getAppNumArgs
assert! rhs.getAppNumArgs == numArgs
if f.isConstOf ``Lean.Grind.nestedProof && g.isConstOf ``Lean.Grind.nestedProof && numArgs == 2 then
if numArgs == 2 && f.isConstOf ``Grind.nestedProof && g.isConstOf ``Grind.nestedProof then
mkNestedProofCongr lhs rhs heq
else if f.isConstOf ``Eq && g.isConstOf ``Eq && numArgs == 3 then
else if numArgs == 2 && f.isConstOf ``Grind.nestedDecidable && g.isConstOf ``Grind.nestedDecidable then
mkNestedDecidableCongr lhs rhs heq
else if numArgs == 3 && f.isConstOf ``Eq && g.isConstOf ``Eq then
mkEqCongrProof lhs rhs heq
else if (← isCongrDefaultProofTarget lhs rhs f g numArgs) then
mkCongrDefaultProof lhs rhs heq

View file

@ -13,13 +13,13 @@ A lighter version of `preprocess` which produces a definitionally equal term,
but ensures assumptions made by `grind` are satisfied.
-/
def preprocessLight (e : Expr) : GoalM Expr := do
shareCommon (← canon (← normalizeLevels (← foldProjs (← eraseIrrelevantMData (← markNestedProofs (← unfoldReducible e))))))
shareCommon (← canon (← normalizeLevels (← foldProjs (← eraseIrrelevantMData (← markNestedSubsingletons (← unfoldReducible e))))))
/--
If `e` has not been internalized yet, instantiate metavariables, unfold reducible, canonicalize,
and internalize the result.
This is an auxliary function used at `proveEq?` and `proveHEq?`.
This is an auxiliary function used at `proveEq?` and `proveHEq?`.
-/
private def ensureInternalized (e : Expr) : GoalM Expr := do
if (← alreadyInternalized e) then

View file

@ -10,7 +10,7 @@ import Lean.Meta.Tactic.Simp.Main
import Lean.Meta.Tactic.Grind.Util
import Lean.Meta.Tactic.Grind.Types
import Lean.Meta.Tactic.Grind.MatchDiscrOnly
import Lean.Meta.Tactic.Grind.MarkNestedProofs
import Lean.Meta.Tactic.Grind.MarkNestedSubsingletons
import Lean.Meta.Tactic.Grind.Canon
namespace Lean.Meta.Grind
@ -41,7 +41,7 @@ def preprocess (e : Expr) : GoalM Simp.Result := do
let e' := r.expr
let e' ← unfoldReducible e'
let e' ← abstractNestedProofs e'
let e' ← markNestedProofs e'
let e' ← markNestedSubsingletons e'
let e' ← eraseIrrelevantMData e'
let e' ← foldProjs e'
let e' ← normalizeLevels e'

View file

@ -477,6 +477,7 @@ private def congrHash (enodes : ENodeMap) (e : Expr) : UInt64 :=
mixHash (hashRoot enodes d) (hashRoot enodes b)
else match_expr e with
| Grind.nestedProof p _ => hashRoot enodes p
| Grind.nestedDecidable p _ => mixHash 13 (hashRoot enodes p)
| Eq _ lhs rhs => goEq lhs rhs
| _ => go e 17
where
@ -500,6 +501,9 @@ private partial def isCongruent (enodes : ENodeMap) (a b : Expr) : Bool :=
| Grind.nestedProof p₁ _ =>
let_expr Grind.nestedProof p₂ _ := b | false
hasSameRoot enodes p₁ p₂
| Grind.nestedDecidable p₁ _ =>
let_expr Grind.nestedDecidable p₂ _ := b | false
hasSameRoot enodes p₁ p₂
| Eq α₁ lhs₁ rhs₁ =>
let_expr Eq α₂ lhs₂ rhs₂ := b | false
if isSameExpr α₁ α₂ then

View file

@ -85,3 +85,10 @@ info: appV_assoc': [@appV #6 #5 (@HAdd.hAdd `[Nat] `[Nat] `[Nat] `[instHAdd] #4
@[grind? =]
theorem appV_assoc' (a : Vector α n) (b : Vector α m) (c : Vector α n') :
appV a (appV b c) ≍ appV (appV a b) c := sorry
example (p : Prop) (h₁ h₂ : Decidable p) : h₁ = h₂ := by
grind
example (p q : Prop) (h₁ : Decidable p) (h₂ : Decidable (p ∧ q)) : (p ↔ q) → h₁ ≍ h₂ := by
grind