feat: proof construction for grind order (#10590)

This PR implements proof term construction for `grind order`.

Author: leodemoura

src/Lean/Meta/Tactic/Grind/Order.lean

@@ -9,10 +9,14 @@ public import Lean.Meta.Tactic.Grind.Order.Types
 public import Lean.Meta.Tactic.Grind.Order.Internalize
 public import Lean.Meta.Tactic.Grind.Order.StructId
 public import Lean.Meta.Tactic.Grind.Order.OrderM
+public import Lean.Meta.Tactic.Grind.Order.Assert
+public import Lean.Meta.Tactic.Grind.Order.Util
 public section
 namespace Lean.Meta.Grind.Order
 builtin_initialize registerTraceClass `grind.order
+builtin_initialize registerTraceClass `grind.order.assert
 builtin_initialize registerTraceClass `grind.order.internalize
+builtin_initialize registerTraceClass `grind.order.internalize.term
 
 builtin_initialize
   orderExt.setMethods

src/Lean/Meta/Tactic/Grind/Order/Assert.lean

@@ -0,0 +1,82 @@
+/-
+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
+-/
+module
+prelude
+public import Lean.Meta.Tactic.Grind.Order.OrderM
+import Init.Grind.Propagator
+import Lean.Meta.Tactic.Grind.PropagatorAttr
+import Lean.Meta.Tactic.Grind.Order.Util
+import Lean.Meta.Tactic.Grind.Order.Proof
+namespace Lean.Meta.Grind.Order
+/--
+Returns a proof for `u + k ≤ v` (or `u ≤ v + k`) where `k` is the
+shortest path between `u` and `v`.
+-/
+partial def mkProofForPath (u v : NodeId) : OrderM Expr := do
+  go (← getProof u v)
+where
+  go (p : ProofInfo) : OrderM Expr := do
+    if u == p.w then
+      return p.proof
+    else
+      let p' ← getProof u p.w
+      go (← mkTrans p' p v)
+
+/-- Adds `p` to the list of things to be propagated. -/
+def pushToPropagate (p : ToPropagate) : OrderM Unit :=
+  modifyStruct fun s => { s with propagate := p :: s.propagate }
+
+/-
+def propagateEqTrue (e : Expr) (u v : NodeId) (k k' : Int) : OrderM Unit := do
+  let kuv ← mkProofForPath u v
+  let u ← getExpr u
+  let v ← getExpr v
+  pushEqTrue e <| mkPropagateEqTrueProof u v k kuv k'
+
+private def propagateEqFalse (e : Expr) (u v : NodeId) (k k' : Int) : OrderM Unit := do
+  let kuv ← mkProofForPath u v
+  let u ← getExpr u
+  let v ← getExpr v
+  pushEqFalse e <| mkPropagateEqFalseProof u v k kuv k'
+-/
+
+/--
+Auxiliary function for implementing theory propagation.
+Traverses the constraints `c` (representing an expression `e`) s.t.
+`c.u = u` and `c.v = v`, it removes `c` from the list of constraints
+associated with `(u, v)` IF
+- `e` is already assigned, or
+- `f c e` returns true
+-/
+@[inline]
+private def updateCnstrsOf (u v : NodeId) (f : Cnstr NodeId → Expr → OrderM Bool) : OrderM Unit := do
+  if let some cs := (← getStruct).cnstrsOf.find? (u, v) then
+    let cs' ← cs.filterM fun (c, e) => do
+      if (← isEqTrue e <||> isEqFalse e) then
+        return false -- constraint was already assigned
+      else
+        return !(← f c e)
+    modifyStruct fun s => { s with cnstrsOf := s.cnstrsOf.insert (u, v) cs' }
+
+def assertTrue (c : Cnstr NodeId) (p : Expr) : OrderM Unit := do
+  trace[grind.order.assert] "{p} = True: {← c.pp}"
+
+def assertFalse (c : Cnstr NodeId) (p : Expr) : OrderM Unit := do
+  trace[grind.order.assert] "{p} = False: {← c.pp}"
+
+def getStructIdOf? (e : Expr) : GoalM (Option Nat) := do
+  return (← get').exprToStructId.find? { expr := e }
+
+builtin_grind_propagator propagateLE ↓LE.le := fun e => do
+  let some structId ← getStructIdOf? e | return ()
+  OrderM.run structId do
+  let some c ← getCnstr? e | return ()
+  if (← isEqTrue e) then
+    assertTrue c e
+  else if (← isEqFalse e) then
+    assertFalse c e
+
+end Lean.Meta.Grind.Order

src/Lean/Meta/Tactic/Grind/Order/Internalize.lean

@@ -140,12 +140,60 @@ def mkCnstr? (e : Expr) (kind : CnstrKind) (lhs rhs : Expr) : OrderM (Option (Cn
   else
     return some { kind, u := lhs, v := rhs }
 
+def setStructId (e : Expr) : OrderM Unit := do
+  let structId ← getStructId
+  modify' fun s => { s with
+    exprToStructId := s.exprToStructId.insert { expr := e } structId
+  }
+
+def mkNode (e : Expr) : OrderM NodeId := do
+  if let some nodeId := (← getStruct).nodeMap.find? { expr := e } then
+    return nodeId
+  let nodeId : NodeId := (← getStruct).nodes.size
+  trace[grind.order.internalize.term] "{e} ↦ #{nodeId}"
+  modifyStruct fun s => { s with
+    nodes   := s.nodes.push e
+    nodeMap := s.nodeMap.insert { expr := e } nodeId
+    sources := s.sources.push {}
+    targets := s.targets.push {}
+    proofs  := s.proofs.push {}
+  }
+  setStructId e
+  /-
+  **Note**: not all node expressions have been internalized because this solver
+  creates auxiliary terms.
+  -/
+  if (← alreadyInternalized e) then
+    orderExt.markTerm e
+  return nodeId
+
 def internalizeCnstr (e : Expr) (kind : CnstrKind) (lhs rhs : Expr) : OrderM Unit := do
-  let some cnstr ← mkCnstr? e kind lhs rhs | return ()
-  trace[grind.order.internalize] "{cnstr.u}, {cnstr.v}, {cnstr.k}"
+  let some c ← mkCnstr? e kind lhs rhs | return ()
+  trace[grind.order.internalize] "{c.u}, {c.v}, {c.k}"
   if grind.debug.get (← getOptions) then
-    if let some h := cnstr.h? then check h
-  -- **TODO**: update data-structures
+    if let some h := c.h? then check h
+  let u ← mkNode c.u
+  let v ← mkNode c.v
+  let c := { c with u, v }
+/-
+  -- **TODO**: propagate
+  if let some k ← getDist? u v then
+    if k ≤ c.k then
+      propagateEqTrue e u v k c.k
+      return ()
+  if let some k ← getDist? v u then
+    if k + c.k < 0 then
+      propagateEqFalse e v u k c.k
+      return ()
+  trace[grind.offset.internalize] "{e} ↦ {c}"
+-/
+  setStructId e
+  modifyStruct fun s => { s with
+    cnstrs   := s.cnstrs.insert { expr := e } c
+    cnstrsOf :=
+      let cs := if let some cs := s.cnstrsOf.find? (u, v) then (c, e) :: cs else [(c, e)]
+      s.cnstrsOf.insert (u, v) cs
+  }
 
 def hasLt : OrderM Bool :=
   return (← getStruct).ltFn?.isSome
@@ -183,19 +231,23 @@ def adapt (α : Expr) (e : Expr) : GoalM (Expr × Expr) := do
   else
     return (α, e)
 
+def alreadyInternalized (e : Expr) : OrderM Bool := do
+  let s ← getStruct
+  return s.cnstrs.contains { expr := e } || s.nodeMap.contains { expr := e }
+
 public def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
   unless (← getConfig).order do return ()
   let some α := getType? e | return ()
   let (α, e) ← adapt α e
   if isForbiddenParent parent? then return ()
-  let some structId ← getStructId? α | return ()
-  OrderM.run structId do
-  match_expr e with
-  | LE.le _ _ lhs rhs => internalizeCnstr e .le lhs rhs
-  | LT.lt _ _ lhs rhs => if (← hasLt) then internalizeCnstr e .lt lhs rhs
-  | Eq _ lhs rhs => internalizeCnstr e .eq lhs rhs
-  | _ =>
-    -- TODO
-    return ()
+  if let some structId ← getStructId? α then OrderM.run structId do
+    if (← alreadyInternalized e) then return ()
+    match_expr e with
+    | LE.le _ _ lhs rhs => internalizeCnstr e .le lhs rhs
+    | LT.lt _ _ lhs rhs => if (← hasLt) then internalizeCnstr e .lt lhs rhs
+    | Eq _ lhs rhs => internalizeCnstr e .eq lhs rhs
+    | _ =>
+      -- TODO
+      return ()
 
 end Lean.Meta.Grind.Order

src/Lean/Meta/Tactic/Grind/Order/OrderM.lean

@@ -32,4 +32,26 @@ def modifyStruct (f : Struct → Struct) : OrderM Unit := do
   let structId ← getStructId
   modify' fun s => { s with structs := s.structs.modify structId f }
 
+def getExpr (u : NodeId) : OrderM Expr := do
+  return (← getStruct).nodes[u]!
+
+def getDist? (u v : NodeId) : OrderM (Option Weight) := do
+  return (← getStruct).targets[u]!.find? v
+
+def getProof? (u v : NodeId) : OrderM (Option ProofInfo) := do
+  return (← getStruct).proofs[u]!.find? v
+
+def getNodeId (e : Expr) : OrderM NodeId := do
+  let some nodeId := (← getStruct).nodeMap.find? { expr := e }
+    | throwError "internal `grind` error, term has not been internalized by order module{indentExpr e}"
+  return nodeId
+
+def getProof (u v : NodeId) : OrderM ProofInfo := do
+  let some p ← getProof? u v
+    | throwError "internal `grind` error, failed to construct proof for{indentExpr (← getExpr u)}\nand{indentExpr (← getExpr v)}"
+  return p
+
+def getCnstr? (e : Expr) : OrderM (Option (Cnstr NodeId)) :=
+  return (← getStruct).cnstrs.find? { expr := e }
+
 end Lean.Meta.Grind.Order

src/Lean/Meta/Tactic/Grind/Order/Proof.lean

@@ -0,0 +1,88 @@
+/-
+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
+-/
+module
+prelude
+public import Lean.Meta.Tactic.Grind.Order.OrderM
+public import Lean.Meta.Tactic.Grind.Arith.CommRing.RingM
+import Init.Grind.Order
+namespace Lean.Meta.Grind.Order
+/--
+Returns `declName α leInst isPreorderInst`
+-/
+def mkLePreorderPrefix (declName : Name) : OrderM Expr := do
+  let s ← getStruct
+  return mkApp3 (mkConst declName [s.u]) s.type s.leInst s.isPreorderInst
+
+/--
+Returns `declName α leInst ltInst lawfulOrderLtInst isPreorderInst`
+-/
+def mkLeLtPrefix (declName : Name) : OrderM Expr := do
+  let s ← getStruct
+  return mkApp5 (mkConst declName [s.u]) s.type s.leInst s.ltInst?.get! s.lawfulOrderLTInst?.get! s.isPreorderInst
+
+/--
+Returns `declName α leInst ltInst lawfulOrderLtInst isPreorderInst ringInst ordRingInst`
+-/
+def mkOrdRingPrefix (declName : Name) : OrderM Expr := do
+  let s ← getStruct
+  let h ← mkLeLtPrefix declName
+  return mkApp2 h s.ringInst?.get! s.orderedRingInst?.get!
+
+/--
+Assume `p₁` is `{ w := u, k := k₁, proof := p₁ }` and `p₂` is `{ w := w, k := k₂, proof := p₂ }`
+`p₁` is the proof for edge `u → w` and `p₂` the proof for edge `w -> v`.
+Then, this function returns a proof for edge `u -> v`.
+
+Remark: for orders that do not support offsets.
+-/
+def mkTransCore (p₁ : ProofInfo) (p₂ : ProofInfo) (v : NodeId) : OrderM ProofInfo := do
+  let { w := u, k.strict := s₁, proof := h₁, .. } := p₁
+  let { w, k.strict := s₂, proof := h₂, .. } := p₂
+  let h ← match s₁, s₂ with
+    | false, false => mkLePreorderPrefix ``Grind.Order.le_trans
+    | false, true  => mkLeLtPrefix ``Grind.Order.le_lt_trans
+    | true,  false => mkLeLtPrefix ``Grind.Order.lt_le_trans
+    | true,  true  => mkLeLtPrefix ``Grind.Order.lt_trans
+  let ns := (← getStruct).nodes
+  let h := mkApp5 h ns[u]! ns[w]! ns[v]! h₁ h₂
+  return { w := p₁.w, k.strict := s₁ || s₂, proof := h }
+
+/--
+Assume `p₁` is `{ w := u, k := k₁, proof := p₁ }` and `p₂` is `{ w := w, k := k₂, proof := p₂ }`
+`p₁` is the proof for edge `u -(k₁)→ w` and `p₂` the proof for edge `w -(k₂)-> v`.
+Then, this function returns a proof for edge `u -(k₁+k₂) -> v`.
+
+Remark: for orders that support offsets.
+-/
+def mkTransOffset (p₁ : ProofInfo) (p₂ : ProofInfo) (v : NodeId) : OrderM ProofInfo := do
+  let { w := u, k.k := k₁, k.strict := s₁, proof := h₁, .. } := p₁
+  let { w, k.k := k₂, k.strict := s₂, proof := h₂, .. } := p₂
+  let h ← match s₁, s₂ with
+    | false, false => mkOrdRingPrefix ``Grind.Order.le_trans_k
+    | false, true  => mkOrdRingPrefix ``Grind.Order.le_lt_trans_k
+    | true,  false => mkOrdRingPrefix ``Grind.Order.lt_le_trans_k
+    | true,  true  => mkOrdRingPrefix ``Grind.Order.lt_trans_k
+  let k  := k₁ + k₂
+  let ns := (← getStruct).nodes
+  let h := mkApp6 h ns[u]! ns[w]! ns[v]! (toExpr k₁) (toExpr k₂) (toExpr k)
+  let h := mkApp3 h h₁ h₂ eagerReflBoolTrue
+  return { w := p₁.w, k.k := k, k.strict := s₁ || s₂, proof := h }
+
+/--
+Assume `p₁` is `{ w := u, k := k₁, proof := p₁ }` and `p₂` is `{ w := w, k := k₂, proof := p₂ }`
+`p₁` is the proof for edge `u -(k₁)→ w` and `p₂` the proof for edge `w -(k₂)-> v`.
+Then, this function returns a proof for edge `u -(k₁+k₂) -> v`.
+
+Remark: if the order does not support offset `k₁` and `k₂` are zero.
+-/
+public def mkTrans (p₁ : ProofInfo) (p₂ : ProofInfo) (v : NodeId) : OrderM ProofInfo := do
+  let s ← getStruct
+  if s.orderedRingInst?.isSome then
+    mkTransOffset p₁ p₂ v
+  else
+    mkTransCore p₁ p₂ v
+
+end Lean.Meta.Grind.Order

src/Lean/Meta/Tactic/Grind/Order/StructId.lean

@@ -55,23 +55,25 @@ where
         pure (inst?, some ltFn)
     else
       pure (none, none)
-    let (ringId?, orderedRingInst?, isCommRing) ← if lawfulOrderLTInst?.isNone then
-      pure (none, none, false)
+    let (ringId?, ringInst?, orderedRingInst?, isCommRing) ← if lawfulOrderLTInst?.isNone then
+      pure (none, none, none, false)
     else if let some ringId ← getCommRingId? type then
+      let ringInst ← RingM.run ringId do return (← getRing).ringInst
       let semiringInst ← RingM.run ringId do return (← getRing).semiringInst
       let some ordRingInst ← mkOrderedRingInst? u type semiringInst leInst ltInst?.get! isPreorderInst
-        | pure (none, none, true)
-      pure (some ringId, some ordRingInst, true)
+        | pure (none, none, none, true)
+      pure (some ringId, some ringInst, some ordRingInst, true)
     else if let some ringId ← getNonCommRingId? type then
       let semiringInst ← NonCommRingM.run ringId do return (← getRing).semiringInst
+      let ringInst ← NonCommRingM.run ringId do return (← getRing).ringInst
       let some ordRingInst ← mkOrderedRingInst? u type semiringInst leInst ltInst?.get! isPreorderInst
-        | pure (none, none, true)
-      pure (some ringId, some ordRingInst, false)
+        | pure (none, none, none, true)
+      pure (some ringId, some ringInst, some ordRingInst, false)
     else
-      pure (none, none, false)
+      pure (none, none, none, false)
     let id := (← get').structs.size
     let struct := {
-      id,  type, u, leInst, isPreorderInst, ltInst?, leFn, isPartialInst?, orderedRingInst?
+      id,  type, u, leInst, isPreorderInst, ltInst?, leFn, isPartialInst?, ringInst?, orderedRingInst?
       isLinearPreInst?, ltFn?, lawfulOrderLTInst?, ringId?, isCommRing
      }
     modify' fun s => { s with structs := s.structs.push struct }