From c9c794ee8a32467c9f2bd6c85ff1873bb9544921 Mon Sep 17 00:00:00 2001 From: Leonardo de Moura Date: Sat, 7 Jun 2025 22:53:28 -0400 Subject: [PATCH] feat: reification and denotation for linarith module in `grind` (#8680) This PR adds the `reify?` and `denoteExpr` for the new linarith module in `grind`. --- src/Lean/Meta/Tactic/Grind/Arith/Linear.lean | 2 + .../Tactic/Grind/Arith/Linear/DenoteExpr.lean | 66 +++++++++++ .../Tactic/Grind/Arith/Linear/IneqCnstr.lean | 37 +++++-- .../Meta/Tactic/Grind/Arith/Linear/Reify.lean | 104 ++++++++++++++++++ .../Tactic/Grind/Arith/Linear/StructId.lean | 26 +++-- src/Lean/Meta/Tactic/Grind/Arith/Main.lean | 8 +- 6 files changed, 221 insertions(+), 22 deletions(-) create mode 100644 src/Lean/Meta/Tactic/Grind/Arith/Linear/DenoteExpr.lean create mode 100644 src/Lean/Meta/Tactic/Grind/Arith/Linear/Reify.lean diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Linear.lean b/src/Lean/Meta/Tactic/Grind/Arith/Linear.lean index f2dc4d7499..bc93d7ddec 100644 --- a/src/Lean/Meta/Tactic/Grind/Arith/Linear.lean +++ b/src/Lean/Meta/Tactic/Grind/Arith/Linear.lean @@ -9,6 +9,8 @@ import Lean.Meta.Tactic.Grind.Arith.Linear.Util import Lean.Meta.Tactic.Grind.Arith.Linear.Var import Lean.Meta.Tactic.Grind.Arith.Linear.StructId import Lean.Meta.Tactic.Grind.Arith.Linear.IneqCnstr +import Lean.Meta.Tactic.Grind.Arith.Linear.Reify +import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr namespace Lean diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Linear/DenoteExpr.lean b/src/Lean/Meta/Tactic/Grind/Arith/Linear/DenoteExpr.lean new file mode 100644 index 0000000000..d6f0206e8c --- /dev/null +++ b/src/Lean/Meta/Tactic/Grind/Arith/Linear/DenoteExpr.lean @@ -0,0 +1,66 @@ +/- +Copyright (c) 2025 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 Lean.Meta.Tactic.Grind.Arith.Linear.Util +import Lean.Meta.Tactic.Grind.Arith.Linear.Var + +namespace Lean.Meta.Grind.Arith.Linear +/-! +Helper functions for converting reified terms back into their denotations. +-/ + +variable [Monad M] [MonadGetStruct M] + +def _root_.Lean.Grind.Linarith.Poly.denoteExpr (p : Poly) : M Expr := do + match p with + | .nil => return (← getStruct).zero + | .add k x p => go p (← denoteTerm k x) +where + denoteTerm (k : Int) (x : Var) : M Expr := do + if k == 1 then + return (← getStruct).vars[x]! + else + return mkApp2 (← getStruct).hmulFn (mkIntLit k) (← getStruct).vars[x]! + + go (p : Poly) (acc : Expr) : M Expr := do + match p with + | .nil => return acc + | .add k m p => go p (mkApp2 (← getStruct).addFn acc (← denoteTerm k m)) + +def _root_.Lean.Grind.Linarith.Expr.denoteExpr (e : LinExpr) : M Expr := do + go e +where + go : LinExpr → M Expr + | .zero => return (← getStruct).zero + | .var x => return (← getStruct).vars[x]! + | .add a b => return mkApp2 (← getStruct).addFn (← go a) (← go b) + | .sub a b => return mkApp2 (← getStruct).subFn (← go a) (← go b) + | .mul k a => return mkApp2 (← getStruct).hmulFn (mkIntLit k) (← go a) + | .neg a => return mkApp (← getStruct).negFn (← go a) + +private def mkEq (a b : Expr) : M Expr := do + let s ← getStruct + return mkApp3 (mkConst ``Eq [s.u.succ]) s.type a b + +def EqCnstr.denoteExpr (c : EqCnstr) : M Expr := do + mkEq (← c.p.denoteExpr) (← getStruct).zero + +def DiseqCnstr.denoteExpr (c : DiseqCnstr) : M Expr := do + return mkNot (← mkEq (← c.p.denoteExpr) (← getStruct).zero) + +private def denoteIneq (p : Poly) (strict : Bool) : M Expr := do + if strict then + return mkApp2 (← getStruct).ltFn (← p.denoteExpr) (← getStruct).zero + else + return mkApp2 (← getStruct).leFn (← p.denoteExpr) (← getStruct).zero + +def IneqCnstr.denoteExpr (c : IneqCnstr) : M Expr := do + denoteIneq c.p c.strict + +def NotIneqCnstr.denoteExpr (c : NotIneqCnstr) : M Expr := do + return mkNot (← denoteIneq c.p c.strict) + +end Lean.Meta.Grind.Arith.Linear diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Linear/IneqCnstr.lean b/src/Lean/Meta/Tactic/Grind/Arith/Linear/IneqCnstr.lean index 63361a15ae..e06c110d33 100644 --- a/src/Lean/Meta/Tactic/Grind/Arith/Linear/IneqCnstr.lean +++ b/src/Lean/Meta/Tactic/Grind/Arith/Linear/IneqCnstr.lean @@ -6,17 +6,38 @@ Authors: Leonardo de Moura prelude import Lean.Meta.Tactic.Grind.Arith.Linear.Var import Lean.Meta.Tactic.Grind.Arith.Linear.StructId +import Lean.Meta.Tactic.Grind.Arith.Linear.Reify +import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr namespace Lean.Meta.Grind.Arith.Linear -def propagateIneq (e : Expr) (eqTrue : Bool) (strict : Bool) : GoalM Unit := do - trace[grind.linarith] "{e}, {eqTrue}, {strict}" - let args := e.getAppArgs - unless args.size == 4 do return () - let α := args[0]! +def isLeInst (struct : Struct) (inst : Expr) : Bool := + isSameExpr struct.leFn.appArg! inst +def isLtInst (struct : Struct) (inst : Expr) : Bool := + isSameExpr struct.ltFn.appArg! inst + +def propagateIneq (e : Expr) (eqTrue : Bool) : GoalM Unit := do + let numArgs := e.getAppNumArgs + unless numArgs == 4 do return () + let α := e.getArg! 0 numArgs let some structId ← getStructId? α | return () - trace[grind.linarith] "structId: {structId}" - -- TODO - return () + LinearM.run structId do + let inst := e.getArg! 1 numArgs + let struct ← getStruct + let strict ← if isLeInst struct inst then + pure false + else if isLtInst struct inst then + pure true + else + trace[grind.linarith] "invalid {e}, {(← getStruct).leFn}, {(← getStruct).ltFn}" + return () + let some lhs ← reify? (e.getArg! 2 numArgs) (skipVar := false) | trace[grind.linarith] "lhs failed {e}"; return () + let some rhs ← reify? (e.getArg! 3 numArgs) (skipVar := false) | trace[grind.linarith] "rhs failed {e}"; return () + let p := (lhs.sub rhs).norm + -- TODO + trace[grind.linarith] "{e}, {eqTrue}, strict: {strict}, structId: {structId}" + trace[grind.linarith] "{← p.denoteExpr}" + trace[grind.linarith] "structId: {structId}" + return () end Lean.Meta.Grind.Arith.Linear diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Linear/Reify.lean b/src/Lean/Meta/Tactic/Grind/Arith/Linear/Reify.lean new file mode 100644 index 0000000000..1def6a828c --- /dev/null +++ b/src/Lean/Meta/Tactic/Grind/Arith/Linear/Reify.lean @@ -0,0 +1,104 @@ +/- +Copyright (c) 2025 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 Lean.Meta.Tactic.Grind.Simp +import Lean.Meta.Tactic.Grind.Arith.Linear.Var + +namespace Lean.Meta.Grind.Arith.Linear + +def isAddInst (struct : Struct) (inst : Expr) : Bool := + isSameExpr struct.addFn.appArg! inst +def isZeroInst (struct : Struct) (inst : Expr) : Bool := + isSameExpr struct.zero.appArg! inst +def isHMulInst (struct : Struct) (inst : Expr) : Bool := + isSameExpr struct.hmulFn.appArg! inst +def isSMulInst (struct : Struct) (inst : Expr) : Bool := + if let some smulFn := struct.smulFn? then + isSameExpr smulFn.appArg! inst + else + false +def isSubInst (struct : Struct) (inst : Expr) : Bool := + isSameExpr struct.subFn.appArg! inst +def isNegInst (struct : Struct) (inst : Expr) : Bool := + isSameExpr struct.negFn.appArg! inst + +def reportInstIssue (e : Expr) : GoalM Unit := do + reportIssue! "`grind linarith` term with unexpected instance{indentExpr e}" + +/-- +Converts a Lean `IntModule` expression `e` into a `LinExpr` + +If `skipVar` is `true`, then the result is `none` if `e` is not an interpreted `IntModule` term. +We use `skipVar := false` when processing inequalities, and `skipVar := true` for equalities and disequalities +-/ +partial def reify? (e : Expr) (skipVar : Bool) : LinearM (Option LinExpr) := do + let toVar (e : Expr) : LinearM LinExpr := do + return .var (← mkVar e) + let asVar (e : Expr) : LinearM LinExpr := do + reportInstIssue e + return .var (← mkVar e) + let isOfNatZero (e : Expr) : LinearM Bool := do + let_expr OfNat.ofNat _ n _ := e | return false + let some k ← getNatValue? n | return false + unless k == 0 do return false + withDefault <| isDefEq e (← getStruct).zero + let rec go (e : Expr) : LinearM LinExpr := do + match_expr e with + | HAdd.hAdd _ _ _ i a b => + if isAddInst (← getStruct) i then return .add (← go a) (← go b) else asVar e + | HSub.hSub _ _ _ i a b => + if isSubInst (← getStruct) i then return .sub (← go a) (← go b) else asVar e + | HMul.hMul _ _ _ i a b => + if isHMulInst (← getStruct) i then + let some k ← getIntValue? a | pure () + return .mul k (← go b) + asVar e + | HSMul.hSMul _ _ _ i a b => + if isSMulInst (← getStruct) i then + let some k ← getIntValue? a | pure () + return .mul k (← go b) + asVar e + | Neg.neg _ i a => + if isNegInst (← getStruct) i then return .neg (← go a) else asVar e + | Zero.zero _ i => + if isZeroInst (← getStruct) i then return .zero else asVar e + | OfNat.ofNat _ _ _ => + if (← isOfNatZero e) then return .zero else toVar e + | _ => toVar e + let asTopVar (e : Expr) : LinearM (Option LinExpr) := do + reportInstIssue e + if skipVar then + return none + else + return some (← asVar e) + match_expr e with + | HAdd.hAdd _ _ _ i a b => + if isAddInst (← getStruct ) i then return some (.add (← go a) (← go b)) else asTopVar e + | HSub.hSub _ _ _ i a b => + if isSubInst (← getStruct ) i then return some (.sub (← go a) (← go b)) else asTopVar e + | HMul.hMul _ _ _ i a b => + if isHMulInst (← getStruct) i then + let some k ← getIntValue? a | pure () + return some (.mul k (← go b)) + asTopVar e + | HSMul.hSMul _ _ _ i a b => + if isSMulInst (← getStruct) i then + let some k ← getIntValue? a | pure () + return some (.mul k (← go b)) + asTopVar e + | Neg.neg _ i a => + if isNegInst (← getStruct ) i then return some (.neg (← go a)) else asTopVar e + | Zero.zero _ i => + if isZeroInst (← getStruct) i then return some .zero else asTopVar e + | OfNat.ofNat _ _ _ => + if (← isOfNatZero e) then return some .zero else asTopVar e + | _ => + if skipVar then + return none + else + return some (← asVar e) + +end Lean.Meta.Grind.Arith.Linear diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Linear/StructId.lean b/src/Lean/Meta/Tactic/Grind/Arith/Linear/StructId.lean index dc3c14f742..ea726de7fe 100644 --- a/src/Lean/Meta/Tactic/Grind/Arith/Linear/StructId.lean +++ b/src/Lean/Meta/Tactic/Grind/Arith/Linear/StructId.lean @@ -6,6 +6,7 @@ Authors: Leonardo de Moura prelude import Init.Grind.Ordered.Module import Lean.Meta.Tactic.Grind.Simp +import Lean.Meta.Tactic.Grind.Internalize import Lean.Meta.Tactic.Grind.Arith.Linear.Util import Lean.Meta.Tactic.Grind.Arith.Linear.Var @@ -14,6 +15,11 @@ namespace Lean.Meta.Grind.Arith.Linear private def internalizeFn (fn : Expr) : GoalM Expr := do shareCommon (← canon fn) +private def internalizeConst (c : Expr) : GoalM Expr := do + let c ← shareCommon (← canon c) + internalize c none + return c + open Grind.Linarith (Poly) def getStructId? (type : Expr) : GoalM (Option Nat) := do @@ -63,15 +69,15 @@ where throwError "`grind linarith` expected{indentExpr parentInst}\nto be definitionally equal to{indentExpr heteroToField}" let some intModuleInst ← getInst? ``Grind.IntModule | return none let zeroInst ← getInst ``Zero - let zero := mkApp2 (mkConst ``Zero.zero [u]) type zeroInst + let zero ← internalizeConst <| mkApp2 (mkConst ``Zero.zero [u]) type zeroInst let addInst ← getBinHomoInst ``HAdd - let addFn := mkApp4 (mkConst ``HAdd.hAdd [u, u, u]) type type type addInst + let addFn ← internalizeFn <| mkApp4 (mkConst ``HAdd.hAdd [u, u, u]) type type type addInst let subInst ← getBinHomoInst ``HSub - let subFn := mkApp4 (mkConst ``HSub.hSub [u, u, u]) type type type subInst + let subFn ← internalizeFn <| mkApp4 (mkConst ``HSub.hSub [u, u, u]) type type type subInst let negInst ← getInst ``Neg - let negFn := mkApp2 (mkConst ``Neg.neg [u]) type negInst + let negFn ← internalizeFn <| mkApp2 (mkConst ``Neg.neg [u]) type negInst let hmulInst ← getHMulInst - let hmulFn := mkApp4 (mkConst ``HMul.hMul [0, u, u]) Int.mkType type type hmulInst + let hmulFn ← internalizeFn <| mkApp4 (mkConst ``HMul.hMul [0, u, u]) Int.mkType type type hmulInst ensureToFieldDefEq zeroInst intModuleInst ``Grind.IntModule.toZero ensureToHomoFieldDefEq addInst intModuleInst ``Grind.IntModule.toAdd ``instHAdd ensureToHomoFieldDefEq subInst intModuleInst ``Grind.IntModule.toSub ``instHSub @@ -80,8 +86,8 @@ where let some preorderInst ← getInst? ``Grind.Preorder | return none let leInst ← getInst ``LE let ltInst ← getInst ``LT - let leFn := mkApp2 (mkConst ``LE.le [u]) type leInst - let ltFn := mkApp2 (mkConst ``LT.lt [u]) type ltInst + let leFn ← internalizeFn <| mkApp2 (mkConst ``LE.le [u]) type leInst + let ltFn ← internalizeFn <| mkApp2 (mkConst ``LT.lt [u]) type ltInst ensureToFieldDefEq leInst preorderInst ``Grind.Preorder.toLE ensureToFieldDefEq ltInst preorderInst ``Grind.Preorder.toLT let partialInst? ← checkToFieldDefEq? (some preorderInst) (← getInst? ``Grind.PartialOrder) ``Grind.PartialOrder.toPreorder @@ -89,9 +95,9 @@ where let isOrderedType := mkApp3 (mkConst ``Grind.IntModule.IsOrdered [u]) type preorderInst intModuleInst let .some isOrdInst ← trySynthInstance isOrderedType | return none let getSMulFn? : GoalM (Option Expr) := do - let smulType := mkApp2 (mkConst ``SMul [0, u]) Int.mkType type + let smulType := mkApp3 (mkConst ``HSMul [0, u, u]) Int.mkType type type let .some smulInst ← trySynthInstance smulType | return none - let smulFn := mkApp3 (mkConst ``SMul.smul [0, u]) Int.mkType type smulInst + let smulFn ← internalizeFn <| mkApp4 (mkConst ``HSMul.hSMul [0, u, u]) Int.mkType type smulInst smulInst if (← withDefault <| isDefEq hmulFn smulFn) then return smulFn reportIssue! "`grind linarith` expected{indentExpr hmulFn}\nto be definitionally equal to{indentExpr smulFn}" @@ -100,7 +106,7 @@ where let ringInst? ← getInst? ``Grind.Ring let getOne? : GoalM (Option Expr) := do let some oneInst ← getInst? ``One | return none - let one := mkApp2 (mkConst ``One.one [u]) type oneInst + let one ← internalizeConst <| mkApp2 (mkConst ``One.one [u]) type oneInst let one' ← mkNumeral type 1 unless (← withDefault <| isDefEq one one') do reportIssue! "`grind linarith` expected{indentExpr one}\nto be definitionally equal to{indentExpr one'}" return some one diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Main.lean b/src/Lean/Meta/Tactic/Grind/Arith/Main.lean index 4fc0607b5a..c813c1ce2e 100644 --- a/src/Lean/Meta/Tactic/Grind/Arith/Main.lean +++ b/src/Lean/Meta/Tactic/Grind/Arith/Main.lean @@ -32,18 +32,18 @@ builtin_grind_propagator propagateLE ↓LE.le := fun e => do if let some c ← Offset.isCnstr? e then Offset.assertTrue c (← mkEqTrueProof e) Cutsat.propagateIfSupportedLe e (eqTrue := true) - Linear.propagateIneq e (eqTrue := true) (strict := false) + Linear.propagateIneq e (eqTrue := true) else if (← isEqFalse e) then if let some c ← Offset.isCnstr? e then Offset.assertFalse c (← mkEqFalseProof e) Cutsat.propagateIfSupportedLe e (eqTrue := false) - Linear.propagateIneq e (eqTrue := false) (strict := false) + Linear.propagateIneq e (eqTrue := false) builtin_grind_propagator propagateLT ↓LT.lt := fun e => do if (← isEqTrue e) then - Linear.propagateIneq e (eqTrue := true) (strict := true) + Linear.propagateIneq e (eqTrue := true) else if (← isEqFalse e) then - Linear.propagateIneq e (eqTrue := false) (strict := true) + Linear.propagateIneq e (eqTrue := false) def check : GoalM Bool := do let c₁ ← Cutsat.check