refactor: grind offset as solver extension (#10311)

This PR uses the new solver extension framework to implement `grind
offset`.

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

@@ -5,17 +5,14 @@ Authors: Leonardo de Moura
 -/
 module
 prelude
-public import Lean.Meta.Tactic.Grind.Arith.Offset
+public import Lean.Meta.Tactic.Grind.Types
 import Lean.Meta.Tactic.Grind.Arith.Cutsat.EqCnstr
 import Lean.Meta.Tactic.Grind.Arith.Linear.Internalize
-
 public section
-
 namespace Lean.Meta.Grind.Arith
 
 @[export lean_grind_arith_internalize]
 def internalizeImpl (e : Expr) (parent? : Option Expr) : GoalM Unit := do
-  Offset.internalize e parent?
   Cutsat.internalize e parent?
 
 end Lean.Meta.Grind.Arith

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

@@ -4,9 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Tactic.Grind.Arith.Offset
 public import Lean.Meta.Tactic.Grind.Arith.Cutsat.Inv
 
 public section
@@ -14,7 +12,6 @@ public section
 namespace Lean.Meta.Grind.Arith
 
 def checkInvariants : GoalM Unit := do
-  Offset.checkInvariants
   Cutsat.checkInvariants
 
 end Lean.Meta.Grind.Arith

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

@@ -4,33 +4,30 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Tactic.Grind.PropagatorAttr
-public import Lean.Meta.Tactic.Grind.Arith.Offset
-public import Lean.Meta.Tactic.Grind.Arith.Cutsat.LeCnstr
-public import Lean.Meta.Tactic.Grind.Arith.Cutsat.Search
-public import Lean.Meta.Tactic.Grind.Arith.Linear.IneqCnstr
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Search
-
+public import Lean.Meta.Tactic.Grind.Types
+import Lean.Meta.Tactic.Grind.PropagatorAttr
+import Lean.Meta.Tactic.Grind.Arith.Offset.Types
+import Lean.Meta.Tactic.Grind.Arith.Offset.Main
+import Lean.Meta.Tactic.Grind.Arith.Offset.Proof
+import Lean.Meta.Tactic.Grind.Arith.Cutsat.LeCnstr
+import Lean.Meta.Tactic.Grind.Arith.Cutsat.Search
+import Lean.Meta.Tactic.Grind.Arith.Linear.IneqCnstr
+import Lean.Meta.Tactic.Grind.Arith.Linear.Search
 public section
-
 namespace Lean.Meta.Grind.Arith
 
-namespace Offset
-def isCnstr? (e : Expr) : GoalM (Option (Cnstr NodeId)) :=
-  return (← get).arith.offset.cnstrs.find? { expr := e }
+private def Offset.isCnstr? (e : Expr) : GoalM (Option (Cnstr NodeId)) :=
+  return (← get').cnstrs.find? { expr := e }
 
-def assertTrue (c : Cnstr NodeId) (p : Expr) : GoalM Unit := do
+private def Offset.assertTrue (c : Cnstr NodeId) (p : Expr) : GoalM Unit := do
   addEdge c.u c.v c.k (← mkOfEqTrue p)
 
-def assertFalse (c : Cnstr NodeId) (p : Expr) : GoalM Unit := do
+private def Offset.assertFalse (c : Cnstr NodeId) (p : Expr) : GoalM Unit := do
   let p := mkOfNegEqFalse (← get').nodes c p
   let c := c.neg
   addEdge c.u c.v c.k p
 
-end Offset
-
 builtin_grind_propagator propagateLE ↓LE.le := fun e => do
   if (← isEqTrue e) then
     if let some c ← Offset.isCnstr? e then

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

@@ -4,16 +4,13 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.Tactic.Grind.Arith.Offset.Main
 public import Lean.Meta.Tactic.Grind.Arith.Offset.Proof
 public import Lean.Meta.Tactic.Grind.Arith.Offset.Util
 public import Lean.Meta.Tactic.Grind.Arith.Offset.Types
-
 public section
-
-namespace Lean
+namespace Lean.Meta.Grind.Arith.Offset
 
 builtin_initialize registerTraceClass `grind.offset
 builtin_initialize registerTraceClass `grind.offset.dist
@@ -28,4 +25,10 @@ builtin_initialize registerTraceClass `grind.offset.model
 builtin_initialize registerTraceClass `grind.debug.offset
 builtin_initialize registerTraceClass `grind.debug.offset.proof
 
-end Lean
+builtin_initialize
+  offsetExt.setMethods
+    (internalize := Offset.internalize)
+    (newEq       := Offset.processNewEq)
+    (checkInv    := Offset.checkInvariants)
+
+end Lean.Meta.Grind.Arith.Offset

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

@@ -4,15 +4,12 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Init.Grind.Offset
-public import Lean.Meta.Tactic.Grind.Types
-public import Lean.Meta.Tactic.Grind.Arith.Offset.Proof
-public import Lean.Meta.Tactic.Grind.Arith.Offset.Util
-
+public import Lean.Meta.Tactic.Grind.Arith.Offset.Types
+import Lean.Meta.Tactic.Grind.Arith.Offset.Proof
+import Lean.Meta.Tactic.Grind.Arith.Offset.Util
 public section
-
 namespace Lean.Meta.Grind.Arith.Offset
 /-!
 This module implements a decision procedure for offset constraints of the form:
@@ -35,11 +32,11 @@ The main advantage of this module over a full linear integer arithmetic procedur
 its ability to efficiently detect all implied equalities and inequalities.
 -/
 
-def get' : GoalM State := do
-  return (← get).arith.offset
+def get' : GoalM State :=
+  offsetExt.getState
 
 @[inline] def modify' (f : State → State) : GoalM Unit := do
-  modify fun s => { s with arith.offset := f s.arith.offset }
+  offsetExt.modifyState f
 
 def mkNode (expr : Expr) : GoalM NodeId := do
   if let some nodeId := (← get').nodeMap.find? { expr } then
@@ -53,7 +50,7 @@ def mkNode (expr : Expr) : GoalM NodeId := do
     targets := s.targets.push {}
     proofs  := s.proofs.push {}
   }
-  markAsOffsetTerm expr
+  offsetExt.markTerm expr
   return nodeId
 
 private def getExpr (u : NodeId) : GoalM Expr := do
@@ -147,7 +144,8 @@ private def propagateEqFalse (e : Expr) (u v : NodeId) (k k' : Int) : GoalM Unit
 
 /-- Propagates all pending constraints and equalities and resets to "to do" list. -/
 private def propagatePending : GoalM Unit := do
-  let todo ← modifyGet fun s => (s.arith.offset.propagate, { s with arith.offset.propagate := [] })
+  let todo := (← get').propagate
+  modify' fun s => { s with propagate := [] }
   for p in todo do
     match p with
     | .eqTrue e u v k k' => propagateEqTrue e u v k k'
@@ -345,8 +343,7 @@ def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
     else if let some k := isNatNum? e then
       internalizeTerm e z k
 
-@[export lean_process_new_offset_eq]
-def processNewEqImpl (a b : Expr) : GoalM Unit := do
+def processNewEq (a b : Expr) : GoalM Unit := do
   unless isSameExpr a b do
     trace[grind.offset.eq.to] "{a}, {b}"
     let u ← getNodeId a

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

@@ -4,19 +4,15 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Basic
-public import Lean.Meta.Tactic.Grind.Types
-public import Lean.Meta.Tactic.Grind.Util
-
+public import Lean.Meta.Tactic.Grind.Arith.Offset.Types
+import Lean.Meta.Tactic.Grind.Util
 public section
-
 namespace Lean.Meta.Grind.Arith.Offset
 
 /-- Construct a model that satisfies all offset constraints -/
 def mkModel (goal : Goal) : MetaM (Array (Expr × Nat)) := do
-  let s := goal.arith.offset
+  let s ← offsetExt.getStateCore goal
   let dbg := grind.debug.get (← getOptions)
   let nodes := s.nodes
   let isInterpreted (u : Nat) : Bool := isNatNum s.nodes[u]!

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

@@ -4,14 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Init.Grind.Offset
-public import Init.Grind.Lemmas
-public import Lean.Meta.Tactic.Grind.Types
-
+public import Lean.Meta.Tactic.Grind.Arith.Offset.Types
+import Init.Grind.Offset
+import Init.Grind.Lemmas
 public section
-
 namespace Lean.Meta.Grind.Arith.Offset
 /-!
 Helper functions for constructing proof terms in the offset constraint procedure.

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

@@ -4,16 +4,12 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Data.AssocList
 public import Lean.Data.PersistentArray
-public import Lean.Meta.Tactic.Grind.ExprPtr
-public import Lean.Meta.Tactic.Grind.Arith.Util
+public import Lean.Meta.Tactic.Grind.Types
 public import Lean.Meta.Tactic.Grind.Arith.Offset.Util
-
 public section
-
 namespace Lean.Meta.Grind.Arith.Offset
 
 abbrev NodeId := Nat
@@ -48,7 +44,7 @@ structure State where
   nodes    : PArray Expr := {}
   /-- Mapping from `Expr` to a node representing it. -/
   nodeMap  : PHashMap ExprPtr NodeId := {}
-  /-- Mapping from `Expr` representing inequalites to constraints. -/
+  /-- Mapping from `Expr` representing inequalities to constraints. -/
   cnstrs   : PHashMap ExprPtr (Cnstr NodeId) := {}
   /--
   Mapping from pairs `(u, v)` to a list of offset constraints on `u` and `v`.
@@ -76,4 +72,6 @@ structure State where
   propagate : List ToPropagate := []
   deriving Inhabited
 
+builtin_initialize offsetExt : SolverExtension State ← registerSolverExtension (return {})
+
 end Lean.Meta.Grind.Arith.Offset

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

@@ -5,14 +5,12 @@ Authors: Leonardo de Moura
 -/
 module
 prelude
-public import Lean.Meta.Tactic.Grind.Arith.Offset.Types
 public import Lean.Meta.Tactic.Grind.Arith.Cutsat.Types
 public section
 namespace Lean.Meta.Grind.Arith
 
 /-- State for the arithmetic procedures. -/
 structure State where
-  offset : Offset.State := {}
   cutsat : Cutsat.State := {}
   deriving Inhabited
 

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

@@ -120,26 +120,6 @@ inductive PendingTheoryPropagation where
   | /-- Propagate the disequalities in `ps`. -/
     diseqs (ps : ParentSet)
 
-/--
-Helper function for combining `ENode.offset?` fields and detecting what needs
-to be propagated to the offset constraint module.
--/
-private def checkOffsetEq (rhsRoot lhsRoot : ENode) : GoalM PendingTheoryPropagation := do
-  match lhsRoot.offset? with
-  | some lhsOffset =>
-    if let some rhsOffset := rhsRoot.offset? then
-      return .eq lhsOffset rhsOffset
-    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 }
-      return .none
-  | none => return .none
-
-def propagateOffset : PendingTheoryPropagation → GoalM Unit
-  | .eq lhs rhs => Arith.Offset.processNewEq lhs rhs
-  | _ => return ()
-
 /--
 Helper function for combining `ENode.cutsat?` fields and detecting what needs
 to be propagated to the cutsat module.
@@ -295,7 +275,6 @@ where
     }
     propagateBeta lams₁ fns₁
     propagateBeta lams₂ fns₂
-    let offsetTodo ← checkOffsetEq rhsRoot lhsRoot
     let cutsatTodo ← checkCutsatEq rhsRoot lhsRoot
     let todo ← Solvers.mergeTerms rhsRoot lhsRoot
     resetParentsOf lhsRoot.self
@@ -308,7 +287,6 @@ where
       for e in toPropagateDown do
         propagateDown e
       propagateUnitConstFuns lams₁ lams₂
-      propagateOffset offsetTodo
       propagateCutsat cutsatTodo
       todo.propagate
   updateRoots (lhs : Expr) (rootNew : Expr) : GoalM Unit := do

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

@@ -9,6 +9,7 @@ public import Init.Grind.Util
 public import Init.Grind.PP
 public import Lean.Meta.Tactic.Grind.Types
 public import Lean.Meta.Tactic.Grind.Arith.Model
+public import Lean.Meta.Tactic.Grind.Arith.Offset.Types
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.PP
 public import Lean.Meta.Tactic.Grind.Arith.Linear.PP
 public import Lean.Meta.Tactic.Grind.AC.PP
@@ -125,7 +126,7 @@ private def ppOffset : M Unit := do
   unless grind.debug.get (← getOptions) do
     return ()
   let goal ← read
-  let s := goal.arith.offset
+  let s ← Arith.Offset.offsetExt.getStateCore goal
   let nodes := s.nodes
   if nodes.isEmpty then return ()
   let model ← Arith.Offset.mkModel goal

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

@@ -440,15 +440,8 @@ structure ENode where
   /-- Modification time -/
   mt : Nat := 0
   /--
-  The `offset?` field is used to propagate equalities from the `grind` congruence closure module
-  to the offset constraints module. When `grind` merges two equivalence classes, and both have
-  an associated `offset?` set to `some e`, the equality is propagated. This field is
-  assigned during the internalization of offset terms.
-  -/
-  offset? : Option Expr := none
-  /--
   The `cutsat?` field is used to propagate equalities from the `grind` congruence closure module
-  to the cutsat module. Its implementation is similar to the `offset?` field.
+  to the cutsat module.
   -/
   cutsat? : Option Expr := none
   /-- Solver terms attached to this E-node. -/
@@ -1089,31 +1082,12 @@ def setENode (e : Expr) (n : ENode) : GoalM Unit :=
     congrTable := unsafe unsafeCast s.congrTable
   }
 
-/--
-Notifies 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 Arith.Offset.processNewEq (a b : 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 root ← getRootENode e
-  if let some e' := root.offset? then
-    Arith.Offset.processNewEq e e'
-  else
-    setENode root.self { root with offset? := some e }
-
 /--
 Notifies the cutsat module that `a = b` where
 `a` and `b` are terms that have been internalized by this module.