feat: model construction for offset constraints (#6636)

This PR implements model construction for offset constraints in the
`grind` tactic.

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

@@ -0,0 +1,39 @@
+/-
+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.Basic
+import Lean.Meta.Tactic.Grind.Types
+
+namespace Lean.Meta.Grind.Arith.Offset
+/-- Construct a model that statisfies all offset constraints -/
+def mkModel (goal : Goal) : MetaM (Array (Expr × Nat)) := do
+  let s := goal.arith.offset
+  let nodes := s.nodes
+  let mut pre : Array (Option Int) := mkArray nodes.size none
+  for u in [:nodes.size] do
+    let val? := s.sources[u]!.foldl (init := @none Int) fun val? v k => Id.run do
+      let some va := pre[v]! | return val?
+      let val' := va - k
+      let some val := val? | return val'
+      if val' > val then return val' else val?
+    let val? := s.targets[u]!.foldl (init := val?) fun val? v k => Id.run do
+      let some va := pre[v]! | return val?
+      let val' := va + k
+      let some val := val? | return val'
+      if val' < val then return val' else val?
+    let val := val?.getD 0
+    pre := pre.set! u (some val)
+  let min := pre.foldl (init := 0) fun min val? => Id.run do
+    let some val := val? | return min
+    if val < min then val else min
+  let mut r := {}
+  for u in [:nodes.size] do
+    let some val := pre[u]! | unreachable!
+    let val := (val - min).toNat
+    r := r.push (nodes[u]!, val)
+  return r
+
+end Lean.Meta.Grind.Arith.Offset

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

@@ -7,6 +7,7 @@ prelude
 import Init.Grind.Util
 import Init.Grind.PP
 import Lean.Meta.Tactic.Grind.Types
+import Lean.Meta.Tactic.Grind.Arith.Model
 
 namespace Lean.Meta.Grind
 
@@ -111,11 +112,22 @@ private def ppActiveTheorems (goal : Goal) : MetaM MessageData := do
   else
     return .trace { cls := `ematch } "E-matching" m
 
+def ppOffset (goal : Goal) : MetaM MessageData := do
+  let s := goal.arith.offset
+  let nodes := s.nodes
+  if nodes.isEmpty then return ""
+  let model ← Arith.Offset.mkModel goal
+  let mut ms := #[]
+  for (e, val) in model do
+    ms := ms.push <| .trace { cls := `assign } m!"{e} := {val}" #[]
+  return .trace { cls := `offset } "Assignment satisfying offset contraints" ms
+
 def goalToMessageData (goal : Goal) : MetaM MessageData := goal.mvarId.withContext do
   let mut m : Array MessageData := #[.ofGoal goal.mvarId]
   m := m.push <| ppExprArray `facts "Asserted facts" goal.facts.toArray `prop
   m := m ++ (← ppEqcs goal)
   m := m.push (← ppActiveTheorems goal)
+  m := m.push (← ppOffset goal)
   addMessageContextFull <| MessageData.joinSep m.toList ""
 
 def goalsToMessageData (goals : List Goal) : MetaM MessageData :=

tests/lean/run/grind_pre.lean

@@ -74,7 +74,11 @@ x✝ : ¬g (i + 1) j ⋯ = i + j + 1
   [prop] j + 1 ≤ i
   [prop] ¬g (i + 1) j ⋯ = i + j + 1[eqc] True propositions
   [prop] j + 1 ≤ i[eqc] False propositions
-  [prop] g (i + 1) j ⋯ = i + j + 1
+  [prop] g (i + 1) j ⋯ = i + j + 1[offset] Assignment satisfying offset contraints
+  [assign] j := 0
+  [assign] i := 1
+  [assign] g (i + 1) j ⋯ := 0
+  [assign] i + j := 0
 -/
 #guard_msgs (error) in
 example (i j : Nat) (h : i + 1 > j + 1) : g (i+1) j = f ((fun x => x) i) + f j + 1 := by