feat: track occurrences in linarith (leanprover#8801)

This PR implements the infrastructure for variable elimination in the
`grind linarith` procedure.

Author: leodemoura

src/Lean/Meta/Tactic/Grind/Arith/Cutsat/LeCnstr.lean

@@ -114,11 +114,10 @@ def LeCnstr.assertImpl (c : LeCnstr) : GoalM Unit := do
     return ()
   let c ← refineWithDiseq c
   trace[grind.cutsat.assert.store] "{← c.pp}"
+  c.p.updateOccs
   if a < 0 then
-    c.p.updateOccs
     modify' fun s => { s with lowers := s.lowers.modify x (·.push c) }
   else
-    c.p.updateOccs
     modify' fun s => { s with uppers := s.uppers.modify x (·.push c) }
   if (← c.satisfied) == .false then
     resetAssignmentFrom x

src/Lean/Meta/Tactic/Grind/Arith/Linear/IneqCnstr.lean

@@ -38,6 +38,7 @@ def IneqCnstr.assert (c : IneqCnstr) : LinearM Unit := do
       trace[grind.linarith.trivial] "{← c.denoteExpr}"
   | .add a x _ =>
     trace[grind.linarith.assert.store] "{← c.denoteExpr}"
+    c.p.updateOccs
     if a < 0 then
       modifyStruct fun s => { s with lowers := s.lowers.modify x (·.push c) }
     else

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

@@ -25,12 +25,25 @@ def _root_.Lean.Grind.Linarith.Poly.checkCoeffs : Poly → Bool
   | .nil => true
   | .add k _ p => k != 0 && checkCoeffs p
 
+def _root_.Lean.Grind.Linarith.Poly.checkOccs (p : Poly) : LinearM Unit := do
+  let .add _ y p := p | return ()
+  let rec go (p : Poly) : LinearM Unit := do
+    let .add _ x p := p | return ()
+    assert! (← getOccursOf x).contains y
+    go p
+  go p
+
+def _root_.Lean.Grind.Linarith.Poly.checkNoElimVars (p : Poly) : LinearM Unit := do
+  let .add _ x p := p | return ()
+  assert! !(← eliminated x)
+  checkNoElimVars p
+
 def _root_.Lean.Grind.Linarith.Poly.checkCnstrOf (p : Poly) (x : Var) : LinearM Unit := do
   assert! p.isSorted
   assert! p.checkCoeffs
   unless (← inconsistent) do
-    -- p.checkNoElimVars
-    -- p.checkOccs
+    p.checkNoElimVars
+    p.checkOccs
     pure ()
   let .add _ y _ := p | unreachable!
   assert! x == y

src/Lean/Meta/Tactic/Grind/Arith/Linear/PropagateEq.lean

@@ -72,6 +72,7 @@ def DiseqCnstr.assert (c : DiseqCnstr) : LinearM Unit := do
     setInconsistent (.diseq c)
   | .add _ x _ =>
     trace[grind.linarith.assert.store] "{← c.denoteExpr}"
+    c.p.updateOccs
     modifyStruct fun s => { s with diseqs := s.diseqs.modify x (·.push c) }
     if (← c.satisfied) == .false then
       resetAssignmentFrom x

src/Lean/Meta/Tactic/Grind/Arith/Linear/StructId.lean

@@ -179,6 +179,7 @@ where
         -- Create `1` variable, and assert strict lower bound `0 < 1`
         let x ← mkVar one (mark := false)
         let p := Poly.add (-1) x .nil
+        p.updateOccs
         modifyStruct fun s => { s with
           lowers := s.lowers.modify x fun cs => cs.push { p, h := .oneGtZero, strict := true }
         }

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

@@ -19,6 +19,14 @@ deriving instance Hashable for Poly
 deriving instance Hashable for Grind.Linarith.Expr
 
 mutual
+/-- An equality constraint and its justification/proof. -/
+structure EqCnstr where
+  p      : Poly
+  h      : EqCnstrProof
+
+inductive EqCnstrProof where
+  | rfl -- TODO
+
 /-- An inequality constraint and its justification/proof. -/
 structure IneqCnstr where
   p      : Poly
@@ -58,6 +66,8 @@ end
 instance : Inhabited DiseqCnstr where
   default := { p := .nil, h := .core default default .zero .zero }
 
+abbrev VarSet := RBTree Var compare
+
 /--
 State for each algebraic structure by this module.
 Each type must at least implement the instance `IntModule`.
@@ -142,6 +152,24 @@ structure Struct where
   -/
   diseqSplits : PHashMap Poly FVarId := {}
   /--
+  Mapping from variable to equation constraint used to eliminate it. `solved` variables should not occur in
+  `diseqs`, `lowers`, or `uppers`.
+  -/
+  elimEqs : PArray (Option EqCnstr) := {}
+  /--
+  Elimination stack. For every variable in `elimStack`. If `x` in `elimStack`, then `elimEqs[x]` is not `none`.
+  -/
+  elimStack : List Var := []
+  /--
+  Mapping from variable to occurrences.
+  For example, an entry `x ↦ {y, z}` means that `x` may occur in `lowers`, or `uppers`, or `diseqs` of
+  variables `y` and `z`.
+  If `x` occurs in `diseqs[y]`, `lowers[y]`, or `uppers[y]`, then `y` is in `occurs[x]`,
+  but the reverse is not true.
+  If `x` is in `elimStack`, then `occurs[x]` is the empty set.
+  -/
+  occurs : PArray VarSet := {}
+  /--
   Linear constraints that are not supported.
   We use this information for diagnostics.
   TODO: store constraints instead.

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

@@ -194,4 +194,31 @@ def inconsistent : LinearM Bool := do
   if (← isInconsistent) then return true
   return (← getStruct).conflict?.isSome
 
+/-- Returns `true` if `x` has been eliminated using an equality constraint. -/
+def eliminated (x : Var) : LinearM Bool :=
+  return (← getStruct).elimEqs[x]!.isSome
+
+/-- Returns occurrences of `x`. -/
+def getOccursOf (x : Var) : LinearM VarSet :=
+  return (← getStruct).occurs[x]!
+
+/--
+Adds `y` as an occurrence of `x`.
+That is, `x` occurs in `lowers[y]`, `uppers[y]`, or `diseqs[y]`.
+-/
+def addOcc (x : Var) (y : Var) : LinearM Unit := do
+  unless (← getOccursOf x).contains y do
+    modifyStruct fun s => { s with occurs := s.occurs.modify x fun ys => ys.insert y }
+
+/--
+Given `p` a polynomial being inserted into `lowers`, `uppers`, or `diseqs`,
+get its leading variable `y`, and adds `y` as an occurrence for the remaining variables in `p`.
+-/
+partial def _root_.Lean.Grind.Linarith.Poly.updateOccs (p : Poly) : LinearM Unit := do
+  let .add _ y p := p | throwError "`grind linarith` internal error, unexpected constant polynomial"
+  let rec go (p : Poly) : LinearM Unit := do
+    let .add _ x p := p | return ()
+    addOcc x y; go p
+  go p
+
 end Lean.Meta.Grind.Arith.Linear

src/Lean/Meta/Tactic/Grind/Arith/Linear/Var.lean

@@ -19,6 +19,8 @@ def mkVar (e : Expr) (mark := true) : LinearM Var := do
     lowers     := s.lowers.push {}
     uppers     := s.uppers.push {}
     diseqs     := s.diseqs.push {}
+    occurs     := s.occurs.push {}
+    elimEqs    := s.elimEqs.push none
   }
   setTermStructId e
   if mark then