feat: literals, lower and upper bounds in the offset constraint modul…

…e within `grind` (#6648)

This PR adds support for numerals, lower & upper bounds to the offset
constraint module in the `grind` tactic. `grind` can now solve examples
such as:
```
example (f : Nat → Nat) :
        f 2 = a →
        b ≤ 1 → b ≥ 1 →
        c = b + 1 →
        f c = a := by
  grind
```
In the example above, the literal `2` and the lower&upper bounds, `b ≤
1` and `b ≥ 1`, are now processed by offset constraint module.

src/Init/Grind/Offset.lean

@@ -87,5 +87,6 @@ Helper theorems for equality propagation
 theorem Nat.le_of_eq_1 (u v : Nat) : u = v → u ≤ v := by omega
 theorem Nat.le_of_eq_2 (u v : Nat) : u = v → v ≤ u := by omega
 theorem Nat.eq_of_le_of_le (u v : Nat) : u ≤ v → v ≤ u → u = v := by omega
+theorem Nat.le_offset (a k : Nat) : k ≤ a + k := by omega
 
 end Lean.Grind

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

@@ -8,7 +8,7 @@ import Lean.Meta.Tactic.Grind.Arith.Offset
 
 namespace Lean.Meta.Grind.Arith
 
-def internalize (e : Expr) (parent : Expr) : GoalM Unit := do
-  Offset.internalize e parent
+def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
+  Offset.internalize e parent?
 
 end Lean.Meta.Grind.Arith

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

@@ -39,7 +39,7 @@ def mkModel (goal : Goal) : MetaM (Array (Expr × Nat)) := do
     We should not include the assignment for auxiliary offset terms since
     they do not provide any additional information.
     -/
-    unless isNatOffset? e |>.isSome do
+    if (isNatOffset? e).isNone && isNatNum? e != some 0 then
       r := r.push (e, val)
   return r
 

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

@@ -239,18 +239,48 @@ private def internalizeCnstr (e : Expr) (c : Cnstr Expr) : GoalM Unit := do
       s.cnstrsOf.insert (u, v) cs
   }
 
-def internalize (e : Expr) (parent : Expr) : GoalM Unit := do
-  if let some c := isNatOffsetCnstr? e then
+private def getZeroNode : GoalM NodeId := do
+  mkNode (← getNatZeroExpr)
+
+/-- Internalize `e` of the form `b + k` -/
+private def internalizeTerm (e : Expr) (b : Expr) (k : Nat) : GoalM Unit := do
+  -- `e` is of the form `b + k`
+  let u ← mkNode e
+  let v ← mkNode b
+  -- `u = v + k`. So, we add edges for `u ≤ v + k` and `v + k ≤ u`.
+  let h := mkApp (mkConst ``Nat.le_refl) e
+  addEdge u v k h
+  addEdge v u (-k) h
+  -- `0 + k ≤ u`
+  let z ← getZeroNode
+  addEdge z u (-k) <| mkApp2 (mkConst ``Grind.Nat.le_offset) b (toExpr k)
+
+/--
+Returns `true`, if `parent?` is relevant for internalization.
+For example, we do not want to internalize an offset term that
+is the child of an addition. This kind of term will be processed by the
+more general linear arithmetic module.
+-/
+private def isRelevantParent (parent? : Option Expr) : GoalM Bool := do
+  let some parent := parent? | return false
+  let z ← getNatZeroExpr
+  return !isNatAdd parent && (isNatOffsetCnstr? parent z).isNone
+
+private def isEqParent (parent? : Option Expr) : Bool := Id.run do
+  let some parent := parent? | return false
+  return parent.isEq
+
+def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
+  let z ← getNatZeroExpr
+  if let some c := isNatOffsetCnstr? e z then
     internalizeCnstr e c
-  else if let some (b, k) := isNatOffset? e then
-    if (isNatOffsetCnstr? parent).isSome then return ()
-    -- `e` is of the form `b + k`
-    let u ← mkNode e
-    let v ← mkNode b
-    -- `u = v + k`. So, we add edges for `u ≤ v + k` and `v + k ≤ u`.
-    let h := mkApp (mkConst ``Nat.le_refl) e
-    addEdge u v k h
-    addEdge v u (-k) h
+  else if (← isRelevantParent parent?) then
+    if let some (b, k) := isNatOffset? e then
+      internalizeTerm e b k
+    else if let some k := isNatNum? e then
+      -- core module has support for detecting equality between literals
+      unless isEqParent parent? do
+        internalizeTerm e z k
 
 @[export lean_process_new_offset_eq]
 def processNewOffsetEqImpl (a b : Expr) : GoalM Unit := do
@@ -262,6 +292,17 @@ def processNewOffsetEqImpl (a b : Expr) : GoalM Unit := do
     addEdge u v 0 <| mkApp3 (mkConst ``Grind.Nat.le_of_eq_1) a b h
     addEdge v u 0 <| mkApp3 (mkConst ``Grind.Nat.le_of_eq_2) a b h
 
+@[export lean_process_new_offset_eq_lit]
+def processNewOffsetEqLitImpl (a b : Expr) : GoalM Unit := do
+  unless isSameExpr a b do
+    trace[grind.offset.eq.to] "{a}, {b}"
+    let some k := isNatNum? b | unreachable!
+    let u ← getNodeId a
+    let z ← mkNode (← getNatZeroExpr)
+    let h ← mkEqProof a b
+    addEdge u z k <| mkApp3 (mkConst ``Grind.Nat.le_of_eq_1) a b h
+    addEdge z u (-k) <| mkApp3 (mkConst ``Grind.Nat.le_of_eq_2) a b h
+
 def traceDists : GoalM Unit := do
   let s ← get'
   for u in [:s.targets.size], es in s.targets.toArray do

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

@@ -32,6 +32,16 @@ def isNatAdd? (e : Expr) : Option (Expr × Expr) :=
   let_expr HAdd.hAdd _ _ _ i a b := e | none
   if isInstAddNat i then some (a, b) else none
 
+/--
+Returns `true` if `e` is of the form
+```
+@HAdd.hAdd Nat Nat Nat (instHAdd Nat instAddNat) _ _
+```
+-/
+def isNatAdd (e : Expr) : Bool :=
+  let_expr HAdd.hAdd _ _ _ i _ _ := e | false
+  isInstAddNat i
+
 /-- Returns `some k` if `e` `@OfNat.ofNat Nat _ (instOfNatNat k)` -/
 def isNatNum? (e : Expr) : Option Nat := Id.run do
   let_expr OfNat.ofNat _ _ inst := e | none
@@ -67,8 +77,12 @@ def Offset.toMessageData [inst : ToMessageData α] (c : Offset.Cnstr α) : Messa
 instance : ToMessageData (Offset.Cnstr Expr) where
   toMessageData c := Offset.toMessageData c
 
-/-- Returns `some cnstr` if `e` is offset constraint. -/
-def isNatOffsetCnstr? (e : Expr) : Option (Offset.Cnstr Expr) :=
+/--
+Returns `some cnstr` if `e` is offset constraint.
+Remark: `z` is `0` numeral. It is an extra argument because we
+want to be able to provide the one that has already been internalized.
+-/
+def isNatOffsetCnstr? (e : Expr) (z : Expr) : Option (Offset.Cnstr Expr) :=
   match_expr e with
   | LE.le _ inst a b => if isInstLENat inst then go a b else none
   | _ => none
@@ -77,7 +91,11 @@ where
     if let some (u, k) := isNatOffset? u then
       some { u, k := - k, v }
     else if let some (v, k) := isNatOffset? v then
-      some { u, v, k := k }
+      some { u, v, k }
+    else if let some k := isNatNum? u then
+      some { u := z, v, k := - k }
+    else if let some k := isNatNum? v then
+      some { u, v := z, k }
     else
       some { u, v }
 

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

@@ -10,6 +10,7 @@ import Lean.Meta.Tactic.Grind.Types
 import Lean.Meta.Tactic.Grind.Inv
 import Lean.Meta.Tactic.Grind.PP
 import Lean.Meta.Tactic.Grind.Ctor
+import Lean.Meta.Tactic.Grind.Util
 import Lean.Meta.Tactic.Grind.Internalize
 
 namespace Lean.Meta.Grind
@@ -90,13 +91,21 @@ private partial def updateMT (root : Expr) : GoalM Unit := do
 Helper function for combining `ENode.offset?` fields and propagating an equality
 to the offset constraint module.
 -/
-private def propagateOffsetEq (root : Expr) (roofOffset? otherOffset? : Option Expr) : GoalM Unit := do
-  let some otherOffset := otherOffset? | return ()
-  if let some rootOffset := roofOffset? then
-    processNewOffsetEq rootOffset otherOffset
-  else
-    let n ← getENode root
-    setENode root { n with offset? := otherOffset? }
+private def propagateOffsetEq (rhsRoot lhsRoot : ENode) : GoalM Unit := do
+  match lhsRoot.offset? with
+  | some lhsOffset =>
+    if let some rhsOffset := rhsRoot.offset? then
+      Arith.processNewOffsetEq lhsOffset rhsOffset
+    else if isNatNum rhsRoot.self then
+      Arith.processNewOffsetEqLit lhsOffset rhsRoot.self
+    else
+      -- We have to retrieve the node because other fields have been updated
+      let rhsRoot ← getENode rhsRoot.self
+      setENode rhsRoot.self { rhsRoot with offset? := lhsOffset }
+  | none =>
+    if isNatNum lhsRoot.self then
+    if let some rhsOffset := rhsRoot.offset? then
+      Arith.processNewOffsetEqLit rhsOffset lhsRoot.self
 
 private partial def addEqStep (lhs rhs proof : Expr) (isHEq : Bool) : GoalM Unit := do
   let lhsNode ← getENode lhs
@@ -166,7 +175,7 @@ where
     copyParentsTo parents rhsNode.root
     unless (← isInconsistent) do
       updateMT rhsRoot.self
-    propagateOffsetEq rhsNode.root rhsRoot.offset? lhsRoot.offset?
+    propagateOffsetEq rhsRoot lhsRoot
     unless (← isInconsistent) do
       for parent in parents do
         propagateUp parent
@@ -218,9 +227,10 @@ private def storeFact (fact : Expr) : GoalM Unit := do
 
 /-- Internalizes `lhs` and `rhs`, and then adds equality `lhs = rhs`. -/
 def addNewEq (lhs rhs proof : Expr) (generation : Nat) : GoalM Unit := do
-  storeFact (← mkEq lhs rhs)
-  internalize lhs generation
-  internalize rhs generation
+  let eq ← mkEq lhs rhs
+  storeFact eq
+  internalize lhs generation eq
+  internalize rhs generation eq
   addEq lhs rhs proof
 
 /-- Adds a new `fact` justified by the given proof and using the given generation. -/
@@ -257,8 +267,8 @@ where
       internalize p generation
       addEq p (← getFalseExpr) (← mkEqFalse proof)
     else
-      internalize lhs generation
-      internalize rhs generation
+      internalize lhs generation p
+      internalize rhs generation p
       addEqCore lhs rhs proof isHEq
 
 /-- Adds a new hypothesis. -/

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

@@ -98,16 +98,14 @@ private def pushCastHEqs (e : Expr) : GoalM Unit := do
   | [email protected] α a motive b h v => pushHEq e v (mkApp6 (mkConst ``Grind.eqRecOn_heq f.constLevels!) α a motive b h v)
   | _ => return ()
 
-def noParent := mkBVar 0
-
 mutual
 /-- Internalizes the nested ground terms in the given pattern. -/
 private partial def internalizePattern (pattern : Expr) (generation : Nat) : GoalM Expr := do
   if pattern.isBVar || isPatternDontCare pattern then
     return pattern
   else if let some e := groundPattern? pattern then
     let e ← shareCommon (← canon (← normalizeLevels (← unfoldReducible e)))
-    internalize e generation
+    internalize e generation none
     return mkGroundPattern e
   else pattern.withApp fun f args => do
     return mkAppN f (← args.mapM (internalizePattern · generation))
@@ -148,7 +146,7 @@ private partial def activateTheoremPatterns (fName : Name) (generation : Nat) :
           trace_goal[grind.ematch] "reinsert `{thm.origin.key}`"
           modify fun s => { s with thmMap := s.thmMap.insert thm }
 
-partial def internalize (e : Expr) (generation : Nat) (parent : Expr := noParent) : GoalM Unit := do
+partial def internalize (e : Expr) (generation : Nat) (parent? : Option Expr := none) : GoalM Unit := do
   if (← alreadyInternalized e) then return ()
   trace_goal[grind.internalize] "{e}"
   match e with
@@ -176,6 +174,7 @@ partial def internalize (e : Expr) (generation : Nat) (parent : Expr := noParent
     if (← isLitValue e) then
       -- We do not want to internalize the components of a literal value.
       mkENode e generation
+      Arith.internalize e parent?
     else e.withApp fun f args => do
       checkAndAddSplitCandidate e
       pushCastHEqs e
@@ -199,7 +198,7 @@ partial def internalize (e : Expr) (generation : Nat) (parent : Expr := noParent
       mkENode e generation
       addCongrTable e
       updateAppMap e
-      Arith.internalize e parent
+      Arith.internalize e parent?
       propagateUp e
 end
 

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

@@ -40,17 +40,20 @@ def GrindM.run (x : GrindM α) (mainDeclName : Name) (config : Grind.Config) (fa
   let scState := ShareCommon.State.mk _
   let (falseExpr, scState) := ShareCommon.State.shareCommon scState (mkConst ``False)
   let (trueExpr, scState)  := ShareCommon.State.shareCommon scState (mkConst ``True)
+  let (natZExpr, scState)  := ShareCommon.State.shareCommon scState (mkNatLit 0)
   let simprocs ← Grind.getSimprocs
   let simp ← Grind.getSimpContext
-  x (← mkMethods fallback).toMethodsRef { mainDeclName, config, simprocs, simp } |>.run' { scState, trueExpr, falseExpr }
+  x (← mkMethods fallback).toMethodsRef { mainDeclName, config, simprocs, simp } |>.run' { scState, trueExpr, falseExpr, natZExpr }
 
 private def mkGoal (mvarId : MVarId) : GrindM Goal := do
   let trueExpr ← getTrueExpr
   let falseExpr ← getFalseExpr
+  let natZeroExpr ← getNatZeroExpr
   let thmMap ← getEMatchTheorems
   GoalM.run' { mvarId, thmMap } do
     mkENodeCore falseExpr (interpreted := true) (ctor := false) (generation := 0)
     mkENodeCore trueExpr (interpreted := true) (ctor := false) (generation := 0)
+    mkENodeCore natZeroExpr (interpreted := true) (ctor := false) (generation := 0)
 
 private def initCore (mvarId : MVarId) : GrindM (List Goal) := do
   mvarId.ensureProp

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

@@ -80,6 +80,7 @@ structure State where
   simpStats  : Simp.Stats := {}
   trueExpr   : Expr
   falseExpr  : Expr
+  natZExpr   : Expr
   /--
   Used to generate trace messages of the for `[grind] working on <tag>`,
   and implement the macro `trace_goal`.
@@ -104,6 +105,10 @@ def getTrueExpr : GrindM Expr := do
 def getFalseExpr : GrindM Expr := do
   return (← get).falseExpr
 
+/-- Returns the internalized `0 : Nat` numeral.  -/
+def getNatZeroExpr : GrindM Expr := do
+  return (← get).natZExpr
+
 def getMainDeclName : GrindM Name :=
   return (← readThe Context).mainDeclName
 
@@ -128,9 +133,9 @@ Applies hash-consing to `e`. Recall that all expressions in a `grind` goal have
 been hash-consed. We perform this step before we internalize expressions.
 -/
 def shareCommon (e : Expr) : GrindM Expr := do
-  modifyGet fun { canon, scState, nextThmIdx, congrThms, trueExpr, falseExpr, simpStats, lastTag } =>
+  modifyGet fun { canon, scState, nextThmIdx, congrThms, trueExpr, falseExpr, natZExpr, simpStats, lastTag } =>
     let (e, scState) := ShareCommon.State.shareCommon scState e
-    (e, { canon, scState, nextThmIdx, congrThms, trueExpr, falseExpr, simpStats, lastTag })
+    (e, { canon, scState, nextThmIdx, congrThms, trueExpr, falseExpr, natZExpr, simpStats, lastTag })
 
 /--
 Canonicalizes nested types, type formers, and instances in `e`.
@@ -649,20 +654,40 @@ def mkENode (e : Expr) (generation : Nat) : GoalM Unit := do
   let interpreted ← isInterpreted e
   mkENodeCore e interpreted ctor generation
 
+/--
+Notify the offset constraint module that `a = b` where
+`a` and `b` are terms that have been internalized by this module.
+-/
 @[extern "lean_process_new_offset_eq"] -- forward definition
-opaque processNewOffsetEq (a b : Expr) : GoalM Unit
+opaque Arith.processNewOffsetEq (a b : Expr) : GoalM Unit
+
+/--
+Notify the offset constraint module that `a = k` where
+`a` is term that has been internalized by this module,
+and `k` is a numeral.
+-/
+@[extern "lean_process_new_offset_eq_lit"] -- forward definition
+opaque Arith.processNewOffsetEqLit (a k : Expr) : GoalM Unit
+
+/-- Returns `true` if `e` is a numeral and has type `Nat`. -/
+def isNatNum (e : Expr) : Bool := Id.run do
+  let_expr OfNat.ofNat _ _ inst := e | false
+  let_expr instOfNatNat _ := inst | false
+  true
 
 /--
 Marks `e` as a term of interest to the offset constraint module.
 If the root of `e`s equivalence class has already a term of interest,
 a new equality is propagated to the offset module.
 -/
 def markAsOffsetTerm (e : Expr) : GoalM Unit := do
-  let n ← getRootENode e
-  if let some e' := n.offset? then
-    processNewOffsetEq e e'
+  let root ← getRootENode e
+  if let some e' := root.offset? then
+    Arith.processNewOffsetEq e e'
+  else if isNatNum root.self && !isSameExpr e root.self then
+    Arith.processNewOffsetEqLit e root.self
   else
-    setENode n.self { n with offset? := some e }
+    setENode root.self { root with offset? := some e }
 
 /-- Returns `true` is `e` is the root of its congruence class. -/
 def isCongrRoot (e : Expr) : GoalM Bool := do

tests/lean/run/grind_offset_cnstr.lean

@@ -373,3 +373,32 @@ example (f : Nat → Nat) (a b c d e : Nat) :
         e < c →
         b = d := by
   grind
+
+example (a : Nat) : a < 2 → a < 5 := by
+  grind
+
+example (a b : Nat) : 2 < a → a ≤ b → 2 < b := by
+  grind
+
+example (a b : Nat) : 2 < a → a ≤ b → 0 < b := by
+  grind
+
+example (f : Nat → Nat) : f 1 = a → b ≤ 1 → b ≥ 1 → f b = a := by
+  grind
+
+example (f : Nat → Nat) : f 2 = a → b ≤ 1 → b ≥ 1 → c = b + 1 → f c = a := by
+  grind
+
+example (a : Nat) : a < 2 → a = 5 → False := by
+  grind
+
+example (a : Nat) : a < 2 → a = b → b = c → c = 5 → False := by
+  grind
+
+#guard_msgs (info) in -- none of the numerals should be internalized by the offset module
+set_option trace.grind.offset.internalize true in
+example (a b c d e : Nat) : a = 1 → b = 2 → c = 3 → d = 4 → e = 5 → a ≠ e := by
+  grind
+
+example (a b : Nat) : a + 1 = b → b = 0 → False := by
+  grind