Skip to content

Commit 3e0168d

Browse files
authored
feat: proof term construction infrastructure for linarith in grind (#8687)
This PR implements the infrastructure for constructing proof terms in the linarith procedure in `grind`. It also adds the `ToExpr` instances for the reified objects.
1 parent fcaae1d commit 3e0168d

File tree

11 files changed

+313
-14
lines changed

11 files changed

+313
-14
lines changed

src/Init/Grind/Tactics.lean

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,7 @@ structure Config where
175175
-/
176176
zeta := true
177177
/--
178-
When `true` (default: `false`), uses procedure for handling equalities over commutative rings.
178+
When `true` (default: `true`), uses procedure for handling equalities over commutative rings.
179179
-/
180180
ring := true
181181
ringSteps := 10000
@@ -184,6 +184,11 @@ structure Config where
184184
proof terms, instead of a single-step Nullstellensatz certificate
185185
-/
186186
ringNull := false
187+
/--
188+
When `true` (default: `true`), uses procedure for handling linear arithmetic for `IntModule`, and
189+
`CommRing`.
190+
-/
191+
linarith := true
187192
deriving Inhabited, BEq
188193

189194
end Lean.Grind

src/Lean/Elab/Tactic/Grind.lean

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ def grind
144144
let result ← Grind.main mvar'.mvarId! params fallback
145145
if result.hasFailed then
146146
throwError "`grind` failed\n{← result.toMessageData}"
147+
trace[grind.debug.proof] "{← instantiateMVars mvar'}"
147148
-- `grind` proofs are often big
148149
let e ← if (← isProp type) then
149150
mkAuxTheorem type (← instantiateMVarsProfiling mvar') (zetaDelta := true)

src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Util.lean

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
55
-/
66
prelude
77
import Lean.Meta.Tactic.Grind.Types
8+
import Lean.Meta.Tactic.Grind.Arith.Util
89

910
namespace Int.Linear
1011
def Poly.isZero : Poly → Bool
@@ -59,16 +60,6 @@ def eliminated (x : Var) : GoalM Bool :=
5960
@[extern "lean_grind_cutsat_assert_eq"] -- forward definition
6061
opaque EqCnstr.assert (c : EqCnstr) : GoalM Unit
6162

62-
-- TODO: PArray.shrink and PArray.resize
63-
partial def shrink (a : PArray Rat) (sz : Nat) : PArray Rat :=
64-
if a.size > sz then shrink a.pop sz else a
65-
66-
partial def resize (a : PArray Rat) (sz : Nat) : PArray Rat :=
67-
if a.size > sz then shrink a sz else go a
68-
where
69-
go (a : PArray Rat) : PArray Rat :=
70-
if a.size < sz then go (a.push 0) else a
71-
7263
/-- Resets the assignment of any variable bigger or equal to `x`. -/
7364
def resetAssignmentFrom (x : Var) : GoalM Unit := do
7465
modify' fun s => { s with assignment := shrink s.assignment x }

src/Lean/Meta/Tactic/Grind/Arith/Linear.lean

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import Lean.Meta.Tactic.Grind.Arith.Linear.StructId
1111
import Lean.Meta.Tactic.Grind.Arith.Linear.IneqCnstr
1212
import Lean.Meta.Tactic.Grind.Arith.Linear.Reify
1313
import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr
14+
import Lean.Meta.Tactic.Grind.Arith.Linear.ToExpr
15+
import Lean.Meta.Tactic.Grind.Arith.Linear.Proof
1416

1517
namespace Lean
1618

src/Lean/Meta/Tactic/Grind/Arith/Linear/IneqCnstr.lean

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import Lean.Meta.Tactic.Grind.Arith.Linear.Var
1111
import Lean.Meta.Tactic.Grind.Arith.Linear.StructId
1212
import Lean.Meta.Tactic.Grind.Arith.Linear.Reify
1313
import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr
14+
import Lean.Meta.Tactic.Grind.Arith.Linear.Proof
1415

1516
namespace Lean.Meta.Grind.Arith.Linear
1617

@@ -21,7 +22,21 @@ def isLtInst (struct : Struct) (inst : Expr) : Bool :=
2122

2223
def IneqCnstr.assert (c : IneqCnstr) : LinearM Unit := do
2324
trace[grind.linarith.assert] "{← c.denoteExpr}"
24-
-- TODO
25+
match c.p with
26+
| .nil =>
27+
if c.strict then
28+
trace[grind.linarith.unsat] "{← c.denoteExpr}"
29+
setInconsistent (.lt c)
30+
else
31+
trace[grind.linarith.trivial] "{← c.denoteExpr}"
32+
| .add a x _ =>
33+
trace[grind.cutsat.assert.store] "{← c.denoteExpr}"
34+
if a < 0 then
35+
modifyStruct fun s => { s with lowers := s.lowers.modify x (·.push c) }
36+
else
37+
modifyStruct fun s => { s with uppers := s.uppers.modify x (·.push c) }
38+
if (← c.satisfied) == .false then
39+
resetAssignmentFrom x
2540

2641
def NotIneqCnstr.assert (c : NotIneqCnstr) : LinearM Unit := do
2742
trace[grind.linarith.assert] "{← c.denoteExpr}"
@@ -71,6 +86,7 @@ def propagateIntModuleIneq (e : Expr) (lhs rhs : Expr) (strict : Bool) (eqTrue :
7186
c.assert
7287

7388
def propagateIneq (e : Expr) (eqTrue : Bool) : GoalM Unit := do
89+
unless (← getConfig).linarith do return ()
7490
let numArgs := e.getAppNumArgs
7591
unless numArgs == 4 do return ()
7692
let α := e.getArg! 0 numArgs
Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
/-
2+
Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
Released under Apache 2.0 license as described in the file LICENSE.
4+
Authors: Leonardo de Moura
5+
-/
6+
prelude
7+
import Lean.Meta.Tactic.Grind.Arith.CommRing.Proof
8+
import Lean.Meta.Tactic.Grind.Arith.Linear.Util
9+
import Lean.Meta.Tactic.Grind.Arith.Linear.ToExpr
10+
11+
namespace Lean.Meta.Grind.Arith.Linear
12+
13+
open CommRing (RingExpr)
14+
15+
/--
16+
Returns a context of type `RArray α` containing the variables of the given structure, where
17+
`α` is `(← getStruct).type`.
18+
-/
19+
def toContextExpr : LinearM Expr := do
20+
let struct ← getStruct
21+
if h : 0 < struct.vars.size then
22+
RArray.toExpr struct.type id (RArray.ofFn (struct.vars[·]) h)
23+
else
24+
RArray.toExpr struct.type id (RArray.leaf struct.zero)
25+
26+
/--
27+
Similar to `toContextExpr`, but for the `CommRing` module.
28+
Recall that this module interfaces with the `CommRing` module and their variable contexts
29+
may not be necessarily identical. For example, for this module, the term `x*y` does not have an interpretation
30+
and we have a "variable" representing it, but it is interpreted in the `CommRing` module, and such
31+
variable does not exist there. On the other direction, suppose we have the inequality `z < 0`, and
32+
`z` does not occur in any equality or disequality. Then, the `CommRing` does not even "see" `z`, and
33+
`z` does not occur in its context, but it occurs in the variable context created by this module.
34+
-/
35+
def toRingContextExpr : LinearM Expr := do
36+
if (← isCommRing) then
37+
withRingM do CommRing.toContextExpr
38+
else
39+
let struct ← getStruct
40+
RArray.toExpr struct.type id (RArray.leaf struct.zero)
41+
42+
structure ProofM.State where
43+
cache : Std.HashMap UInt64 Expr := {}
44+
polyMap : Std.HashMap Poly Expr := {}
45+
exprMap : Std.HashMap LinExpr Expr := {}
46+
ringExprMap : Std.HashMap RingExpr Expr := {}
47+
48+
structure ProofM.Context where
49+
ctx : Expr
50+
ringCtx : Expr
51+
52+
/-- Auxiliary monad for constructing linarith proofs. -/
53+
abbrev ProofM := ReaderT ProofM.Context (StateRefT ProofM.State LinearM)
54+
55+
/-- Returns a Lean expression representing the variable context used to construct linarith proofs. -/
56+
private abbrev getContext : ProofM Expr := do
57+
return (← read).ctx
58+
59+
/--
60+
Returns a Lean expression representing the auxiliary `CommRing` variable context
61+
used to construct linarith proofs.
62+
-/
63+
private abbrev getRingContext : ProofM Expr := do
64+
return (← read).ringCtx
65+
66+
private abbrev caching (c : α) (k : ProofM Expr) : ProofM Expr := do
67+
let addr := unsafe (ptrAddrUnsafe c).toUInt64 >>> 2
68+
if let some h := (← get).cache[addr]? then
69+
return h
70+
else
71+
let h ← k
72+
modify fun s => { s with cache := s.cache.insert addr h }
73+
return h
74+
75+
def mkPolyDecl (p : Poly) : ProofM Expr := do
76+
if let some x := (← get).polyMap[p]? then
77+
return x
78+
let x := mkFVar (← mkFreshFVarId)
79+
modify fun s => { s with polyMap := s.polyMap.insert p x }
80+
return x
81+
82+
def mkExprDecl (e : LinExpr) : ProofM Expr := do
83+
if let some x := (← get).exprMap[e]? then
84+
return x
85+
let x := mkFVar (← mkFreshFVarId)
86+
modify fun s => { s with exprMap := s.exprMap.insert e x }
87+
return x
88+
89+
def mkRingExprDecl (e : RingExpr) : ProofM Expr := do
90+
if let some x := (← get).ringExprMap[e]? then
91+
return x
92+
let x := mkFVar (← mkFreshFVarId)
93+
modify fun s => { s with ringExprMap := s.ringExprMap.insert e x }
94+
return x
95+
96+
private abbrev withProofContext (x : ProofM Expr) : LinearM Expr := do
97+
let struct ← getStruct
98+
withLetDecl `ctx (mkApp (mkConst ``RArray [struct.u]) struct.type) (← toContextExpr) fun ctx => do
99+
withLetDecl `rctx (mkApp (mkConst ``RArray [struct.u]) struct.type) (← toRingContextExpr) fun ringCtx =>
100+
go { ctx, ringCtx } |>.run' {}
101+
where
102+
go : ProofM Expr := do
103+
let h ← x
104+
let h ← mkLetOfMap (← get).polyMap h `p (mkConst ``Grind.Linarith.Poly) toExpr
105+
let h ← mkLetOfMap (← get).exprMap h `e (mkConst ``Grind.Linarith.Expr) toExpr
106+
let h ← mkLetOfMap (← get).ringExprMap h `r (mkConst ``Grind.CommRing.Expr) toExpr
107+
mkLetFVars #[(← getContext), (← getRingContext)] h
108+
109+
/--
110+
Returns the prefix of a theorem with name `declName` where the first three arguments are
111+
`{α} [IntModule α] (ctx : Context α)`
112+
-/
113+
private def mkIntModThmPrefix (declName : Name) : ProofM Expr := do
114+
let s ← getStruct
115+
return mkApp3 (mkConst declName [s.u]) s.type s.intModuleInst (← getContext)
116+
117+
/--
118+
Returns the prefix of a theorem with name `declName` where the first four arguments are
119+
`{α} [IntModule α] [Preorder α] (ctx : Context α)`
120+
-/
121+
private def mkIntModPreThmPrefix (declName : Name) : ProofM Expr := do
122+
let s ← getStruct
123+
return mkApp4 (mkConst declName [s.u]) s.type s.intModuleInst s.preorderInst (← getContext)
124+
125+
/--
126+
Returns the prefix of a theorem with name `declName` where the first five arguments are
127+
`{α} [IntModule α] [Preorder α] [IntModule.IsOrdered α] (ctx : Context α)`
128+
This is the most common theorem prefix at `Linarith.lean`
129+
-/
130+
private def mkIntModPreOrdThmPrefix (declName : Name) : ProofM Expr := do
131+
let s ← getStruct
132+
return mkApp5 (mkConst declName [s.u]) s.type s.intModuleInst s.preorderInst s.isOrdInst (← getContext)
133+
134+
mutual
135+
partial def EqCnstr.toExprProof (c' : EqCnstr) : ProofM Expr := caching c' do
136+
throwError "NIY"
137+
138+
partial def IneqCnstr.toExprProof (c' : IneqCnstr) : ProofM Expr := caching c' do
139+
match c'.h with
140+
| .core e lhs rhs =>
141+
let h ← mkIntModPreOrdThmPrefix (if c'.strict then ``Grind.Linarith.lt_norm else ``Grind.Linarith.le_norm)
142+
return mkApp5 h (← mkExprDecl lhs) (← mkExprDecl rhs) (← mkPolyDecl c'.p) reflBoolTrue (mkOfEqTrueCore e (← mkEqTrueProof e))
143+
| _ => throwError "NIY"
144+
145+
partial def DiseqCnstr.toExprProof (c' : DiseqCnstr) : ProofM Expr := caching c' do
146+
throwError "NIY"
147+
148+
partial def NotIneqCnstr.toExprProof (c' : NotIneqCnstr) : ProofM Expr := caching c' do
149+
throwError "NIY"
150+
151+
partial def UnsatProof.toExprProofCore (h : UnsatProof) : ProofM Expr := do
152+
match h with
153+
| .lt c => return mkApp (← mkIntModPreThmPrefix ``Grind.Linarith.lt_unsat) (← c.toExprProof)
154+
| .diseq _ => throwError "NIY"
155+
156+
end
157+
158+
def UnsatProof.toExprProof (h : UnsatProof) : LinearM Expr := do
159+
withProofContext do h.toExprProofCore
160+
161+
def setInconsistent (h : UnsatProof) : LinearM Unit := do
162+
if (← getStruct).caseSplits then
163+
-- Let the search procedure in `SearchM` resolve the conflict.
164+
modifyStruct fun s => { s with conflict? := some h }
165+
else
166+
let h ← h.toExprProof
167+
closeGoal h
168+
169+
end Lean.Meta.Grind.Arith.Linear
Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
/-
2+
Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
Released under Apache 2.0 license as described in the file LICENSE.
4+
Authors: Leonardo de Moura
5+
-/
6+
prelude
7+
import Init.Grind.Ordered.Linarith
8+
import Lean.ToExpr
9+
10+
namespace Lean.Meta.Grind.Arith.Linear
11+
open Grind.Linarith
12+
13+
/-!
14+
`ToExpr` instances for `Linarith.Poly` types.
15+
-/
16+
17+
def ofPoly (p : Poly) : Expr :=
18+
match p with
19+
| .nil => mkConst ``Poly.nil
20+
| .add k x p => mkApp3 (mkConst ``Poly.add) (toExpr k) (toExpr x) (ofPoly p)
21+
22+
instance : ToExpr Poly where
23+
toExpr := ofPoly
24+
toTypeExpr := mkConst ``Poly
25+
26+
open Lean.Grind
27+
28+
def ofLinExpr (e : Linarith.Expr) : Expr :=
29+
match e with
30+
| .zero => mkConst ``Linarith.Expr.zero
31+
| .var x => mkApp (mkConst ``Linarith.Expr.var) (toExpr x)
32+
| .add a b => mkApp2 (mkConst ``Linarith.Expr.add) (ofLinExpr a) (ofLinExpr b)
33+
| .sub a b => mkApp2 (mkConst ``Linarith.Expr.sub) (ofLinExpr a) (ofLinExpr b)
34+
| .neg a => mkApp (mkConst ``Linarith.Expr.neg) (ofLinExpr a)
35+
| .mul k a => mkApp2 (mkConst ``Linarith.Expr.mul) (toExpr k) (ofLinExpr a)
36+
37+
instance : ToExpr Linarith.Expr where
38+
toExpr := ofLinExpr
39+
toTypeExpr := mkConst ``Linarith.Expr
40+
41+
end Lean.Meta.Grind.Arith.Linear

src/Lean/Meta/Tactic/Grind/Arith/Linear/Types.lean

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ export Std.Internal (Rat)
1616
abbrev LinExpr := Lean.Grind.Linarith.Expr
1717

1818
deriving instance Hashable for Poly
19+
deriving instance Hashable for Grind.Linarith.Expr
1920

2021
mutual
2122
/-- A equality constraint and its justification/proof. -/
@@ -69,8 +70,8 @@ inductive NotIneqCnstrProof where
6970
-- TODO: norm, and combineEq
7071

7172
inductive UnsatProof where
72-
| diseqUnsat (c : DiseqCnstr)
73-
| ltUnsat (c : IneqCnstr)
73+
| diseq (c : DiseqCnstr)
74+
| lt (c : IneqCnstr)
7475
-- TODO: IneqCnstr + NotIneqCnstr
7576

7677
end

src/Lean/Meta/Tactic/Grind/Arith/Linear/Util.lean

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,4 +104,48 @@ def setTermStructId (e : Expr) : LinearM Unit := do
104104
return ()
105105
modify' fun s => { s with exprToStructId := s.exprToStructId.insert { expr := e } structId }
106106

107+
/--
108+
Tries to evaluate the polynomial `p` using the partial model/assignment built so far.
109+
The result is `none` if the polynomial contains variables that have not been assigned.
110+
-/
111+
def _root_.Lean.Grind.Linarith.Poly.eval? (p : Poly) : LinearM (Option Rat) := do
112+
let a := (← getStruct).assignment
113+
let rec go (v : Rat) : Poly → Option Rat
114+
| .nil => some v
115+
| .add k x p =>
116+
if _ : x < a.size then
117+
go (v + k*a[x]) p
118+
else
119+
none
120+
return go 0 p
121+
/--
122+
Returns `.true` if `c` is satisfied by the current partial model,
123+
`.undef` if `c` contains unassigned variables, and `.false` otherwise.
124+
-/
125+
def IneqCnstr.satisfied (c : IneqCnstr) : LinearM LBool := do
126+
let some v ← c.p.eval? | return .undef
127+
if c.strict then
128+
return decide (v < 0) |>.toLBool
129+
else
130+
return decide (v <= 0) |>.toLBool
131+
132+
def EqCnstr.satisfied (c : EqCnstr) : LinearM LBool := do
133+
let some v ← c.p.eval? | return .undef
134+
return decide (v == 0) |>.toLBool
135+
136+
def DiseqCnstr.satisfied (c : DiseqCnstr) : LinearM LBool := do
137+
let some v ← c.p.eval? | return .undef
138+
return decide (v != 0) |>.toLBool
139+
140+
def NotEqCnstr.satisfied (c : NotIneqCnstr) : LinearM LBool := do
141+
let some v ← c.p.eval? | return .undef
142+
if c.strict then
143+
return decide (v >= 0) |>.toLBool
144+
else
145+
return decide (v > 0) |>.toLBool
146+
147+
/-- Resets the assignment of any variable bigger or equal to `x`. -/
148+
def resetAssignmentFrom (x : Var) : LinearM Unit := do
149+
modifyStruct fun s => { s with assignment := shrink s.assignment x }
150+
107151
end Lean.Meta.Grind.Arith.Linear

0 commit comments

Comments
 (0)