feat: trim grind linarith proof context (#9947)

This PR optimizes the proof terms produced by `grind linarith`. It is
similar to #9945, but for the `linarith` module in `grind`.
It removes unused entries from the context objects when generating the
final proof, significantly reducing the amount of junk in the resulting
terms.

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

@@ -22,6 +22,7 @@ public import Lean.Meta.Tactic.Grind.Arith.Linear.Internalize
 public import Lean.Meta.Tactic.Grind.Arith.Linear.Model
 public import Lean.Meta.Tactic.Grind.Arith.Linear.PP
 public import Lean.Meta.Tactic.Grind.Arith.Linear.MBTC
+public import Lean.Meta.Tactic.Grind.Arith.Linear.VarRename
 
 public section
 

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

@@ -11,6 +11,9 @@ public import Lean.Meta.Tactic.Grind.Arith.CommRing.Proof
 public import Lean.Meta.Tactic.Grind.Arith.Linear.Util
 public import Lean.Meta.Tactic.Grind.Arith.Linear.ToExpr
 public import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr
+public import Lean.Meta.Tactic.Grind.Arith.VarRename
+import Lean.Meta.Tactic.Grind.Arith.Linear.VarRename
+import Lean.Meta.Tactic.Grind.Arith.CommRing.VarRename
 
 public section
 
@@ -22,35 +25,20 @@ open CommRing (RingExpr)
 Returns a context of type `RArray α` containing the variables of the given structure, where
 `α` is `(← getStruct).type`.
 -/
-def toContextExpr : LinearM Expr := do
+def toContextExpr (vars : Array Expr) : LinearM Expr := do
   let struct ← getStruct
-  if h : 0 < struct.vars.size then
-    RArray.toExpr struct.type id (RArray.ofFn (struct.vars[·]) h)
+  if h : 0 < vars.size then
+    RArray.toExpr struct.type id (RArray.ofFn (vars[·]) h)
   else
     RArray.toExpr struct.type id (RArray.leaf struct.zero)
 
-/--
-Similar to `toContextExpr`, but for the `CommRing` module.
-Recall that this module interfaces with the `CommRing` module and their variable contexts
-may not be necessarily identical. For example, for this module, the term `x*y` does not have an interpretation
-and we have a "variable" representing it, but it is interpreted in the `CommRing` module, and such
-variable does not exist there. On the other direction, suppose we have the inequality `z < 0`, and
-`z` does not occur in any equality or disequality. Then, the `CommRing` does not even "see" `z`, and
-`z` does not occur in its context, but it occurs in the variable context created by this module.
--/
-def toRingContextExpr : LinearM Expr := do
-  if (← isCommRing) then
-    withRingM do CommRing.toContextExpr
-  else
-    let struct ← getStruct
-    RArray.toExpr struct.type id (RArray.leaf struct.zero)
-
 structure ProofM.State where
-  cache       : Std.HashMap UInt64 Expr := {}
-  polyMap     : Std.HashMap Poly Expr := {}
-  exprMap     : Std.HashMap LinExpr Expr := {}
-  ringPolyMap : Std.HashMap CommRing.Poly Expr := {}
-  ringExprMap : Std.HashMap RingExpr Expr := {}
+  cache         : Std.HashMap UInt64 Expr := {}
+  varDecls      : Std.HashMap Var Expr := {}
+  polyDecls     : Std.HashMap Poly Expr := {}
+  exprDecls     : Std.HashMap LinExpr Expr := {}
+  ringPolyDecls : Std.HashMap CommRing.Poly Expr := {}
+  ringExprDecls : Std.HashMap RingExpr Expr := {}
 
 structure ProofM.Context where
   ctx     : Expr
@@ -79,47 +67,82 @@ private abbrev caching (c : α) (k : ProofM Expr) : ProofM Expr := do
     modify fun s => { s with cache := s.cache.insert addr h }
     return h
 
+local macro "declare! " decls:ident a:ident : term =>
+  `(do if let some x := (← get).$decls[$a]? then
+         return x
+       let x := mkFVar (← mkFreshFVarId);
+       modify fun s => { s with $decls:ident := (s.$decls).insert $a x };
+       return x)
+
+def mkVarDecl (x : Var) : ProofM Expr := do
+  declare! varDecls x
+
 def mkPolyDecl (p : Poly) : ProofM Expr := do
-  if let some x := (← get).polyMap[p]? then
-    return x
-  let x := mkFVar (← mkFreshFVarId)
-  modify fun s => { s with polyMap := s.polyMap.insert p x }
-  return x
+  declare! polyDecls p
 
 def mkExprDecl (e : LinExpr) : ProofM Expr := do
-  if let some x := (← get).exprMap[e]? then
-    return x
-  let x := mkFVar (← mkFreshFVarId)
-  modify fun s => { s with exprMap := s.exprMap.insert e x }
-  return x
+  declare! exprDecls e
 
 def mkRingPolyDecl (p : CommRing.Poly) : ProofM Expr := do
-  if let some x := (← get).ringPolyMap[p]? then
-    return x
-  let x := mkFVar (← mkFreshFVarId)
-  modify fun s => { s with ringPolyMap := s.ringPolyMap.insert p x }
-  return x
+  declare! ringPolyDecls p
 
 def mkRingExprDecl (e : RingExpr) : ProofM Expr := do
-  if let some x := (← get).ringExprMap[e]? then
-    return x
-  let x := mkFVar (← mkFreshFVarId)
-  modify fun s => { s with ringExprMap := s.ringExprMap.insert e x }
-  return x
+  declare! ringExprDecls e
+
+private def mkContext (h : Expr) : ProofM Expr := do
+  let varDecls     := (← get).varDecls
+  let polyDecls    := (← get).polyDecls
+  let exprDecls    := (← get).exprDecls
+  let usedVars     :=
+    collectMapVars varDecls collectVar >>
+    collectMapVars polyDecls (·.collectVars) >>
+    collectMapVars exprDecls (·.collectVars) <| {}
+  let vars'        := usedVars.toArray
+  let varRename    := mkVarRename vars'
+  let vars         := (← getStruct).vars
+  let vars         := vars'.map fun x => vars[x]!
+  let varFVars     := vars'.map fun x => varDecls[x]?.getD default
+  let varIdsAsExpr := List.range vars'.size |>.toArray |>.map toExpr
+  let h := h.replaceFVars varFVars varIdsAsExpr
+  let h := mkLetOfMap exprDecls h `e (mkConst ``Grind.Linarith.Expr) fun e => toExpr <| e.renameVars varRename
+  let h := mkLetOfMap polyDecls h `p (mkConst ``Grind.Linarith.Poly) fun p => toExpr <| p.renameVars varRename
+  let h := h.abstract #[(← read).ctx]
+  if h.hasLooseBVars then
+    let struct ← getStruct
+    let ctxType := mkApp (mkConst ``RArray [struct.u]) struct.type
+    let ctxVal ← toContextExpr vars
+    return .letE `ctx ctxType ctxVal h (nondep := false)
+  else
+    return h
+
+private def mkRingContext (h : Expr) : ProofM Expr := do
+  unless (← isCommRing) do return h
+  let ring ← withRingM do CommRing.getRing
+  let vars := ring.vars
+  let usedVars     := collectMapVars (← get).ringPolyDecls (·.collectVars) >> collectMapVars (← get).ringExprDecls (·.collectVars) <| {}
+  let vars'        := usedVars.toArray
+  let varRename    := mkVarRename vars'
+  let vars         := vars'.map fun x => vars[x]!
+  let h := mkLetOfMap (← get).ringExprDecls h `re (mkConst ``Grind.CommRing.Expr) fun e => toExpr <| e.renameVars varRename
+  let h := mkLetOfMap (← get).ringPolyDecls h `rp (mkConst ``Grind.CommRing.Poly) fun p => toExpr <| p.renameVars varRename
+  let h := h.abstract #[(← read).ringCtx]
+  if h.hasLooseBVars then
+    let struct ← getStruct
+    let ctxType := mkApp (mkConst ``RArray [struct.u]) struct.type
+    let ctxVal ← toContextExpr vars
+    return .letE `rctx ctxType ctxVal h (nondep := false)
+  else
+    return h
 
 private abbrev withProofContext (x : ProofM Expr) : LinearM Expr := do
-  let struct ← getStruct
-  withLetDecl `ctx  (mkApp (mkConst ``RArray [struct.u]) struct.type) (← toContextExpr) fun ctx => do
-  withLetDecl `rctx (mkApp (mkConst ``RArray [struct.u]) struct.type) (← toRingContextExpr) fun ringCtx =>
+  let ctx := mkFVar (← mkFreshFVarId)
+  let ringCtx := mkFVar (← mkFreshFVarId)
   go { ctx, ringCtx } |>.run' {}
 where
   go : ProofM Expr := do
     let h ← x
-    let h := mkLetOfMap (← get).polyMap h `p (mkConst ``Grind.Linarith.Poly) toExpr
-    let h := mkLetOfMap (← get).exprMap h `e (mkConst ``Grind.Linarith.Expr) toExpr
-    let h := mkLetOfMap (← get).ringPolyMap h `rp (mkConst ``Grind.CommRing.Poly) toExpr
-    let h := mkLetOfMap (← get).ringExprMap h `re (mkConst ``Grind.CommRing.Expr) toExpr
-    mkLetFVars #[(← getContext), (← getRingContext)] h
+    let h ← mkRingContext h
+    mkContext h
 
 /--
 Returns the prefix of a theorem with name `declName` where the first three arguments are
@@ -273,7 +296,7 @@ partial def EqCnstr.toExprProof (c' : EqCnstr) : ProofM Expr := caching c' do
     return mkApp5 h (← mkPolyDecl c.p) (← mkPolyDecl c'.p) (toExpr k) eagerReflBoolTrue (← c.toExprProof)
   | .subst x c₁ c₂ =>
     let h ← mkIntModThmPrefix ``Grind.Linarith.eq_eq_subst
-    return mkApp7 h (toExpr x) (← mkPolyDecl c₁.p) (← mkPolyDecl c₂.p) (← mkPolyDecl c'.p) eagerReflBoolTrue
+    return mkApp7 h (← mkVarDecl x) (← mkPolyDecl c₁.p) (← mkPolyDecl c₂.p) (← mkPolyDecl c'.p) eagerReflBoolTrue
       (← c₁.toExprProof) (← c₂.toExprProof)
 
 partial def UnsatProof.toExprProofCore (h : UnsatProof) : ProofM Expr := do

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

@@ -0,0 +1,42 @@
+/-
+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 Init.Grind.Ordered.Linarith
+public import Lean.Meta.Tactic.Grind.Arith.VarRename
+
+namespace Lean.Grind.Linarith
+open Lean.Meta.Grind.Arith
+
+public def Poly.renameVars (p : Poly) (f : VarRename) : Poly :=
+  match p with
+  | .nil => p
+  | .add k x p => .add k (f x) (renameVars p f)
+
+public def Expr.renameVars (e : Expr) (f : VarRename) : Expr :=
+  match e with
+  | .zero => e
+  | .var x => .var (f x)
+  | .neg a => .neg (renameVars a f)
+  | .add a b => .add (renameVars a f) (renameVars b f)
+  | .sub a b => .sub (renameVars a f) (renameVars b f)
+  | .natMul k a => .natMul k (renameVars a f)
+  | .intMul k a => .intMul k (renameVars a f)
+
+public def Poly.collectVars (p : Poly) : VarCollector :=
+  match p with
+  | .nil => id
+  | .add _ x p => collectVar x >> p.collectVars
+
+public def Expr.collectVars (e : Expr) : VarCollector :=
+  match e with
+  | .zero => id
+  | .var x => collectVar x
+  | .neg a | .natMul _ a | .intMul _ a => a.collectVars
+  | .add a b | .sub a b => a.collectVars >> b.collectVars
+
+end Lean.Grind.Linarith

tests/lean/run/grind_linarith_2.lean

@@ -7,14 +7,14 @@ example [IntModule α] [LE α] [LT α] [Preorder α] [OrderedAdd α] (a b : α)
 
 /--
 trace: [grind.debug.proof] Classical.byContradiction fun h =>
-      let ctx := RArray.branch 1 (RArray.leaf One.one) (RArray.branch 2 (RArray.leaf a) (RArray.leaf b));
-      let rctx := RArray.branch 1 (RArray.leaf a) (RArray.leaf b);
-      let re_1 := (CommRing.Expr.var 0).add (CommRing.Expr.var 1);
-      let re_2 := (CommRing.Expr.var 1).add (CommRing.Expr.var 0);
-      let rp_1 := CommRing.Poly.num 0;
+      let ctx := RArray.leaf One.one;
+      let p_1 := Poly.nil;
       let e_1 := Expr.zero;
       let e_2 := Expr.intMul 0 (Expr.var 0);
-      let p_1 := Poly.nil;
+      let rctx := RArray.branch 1 (RArray.leaf a) (RArray.leaf b);
+      let rp_1 := CommRing.Poly.num 0;
+      let re_1 := (CommRing.Expr.var 0).add (CommRing.Expr.var 1);
+      let re_2 := (CommRing.Expr.var 1).add (CommRing.Expr.var 0);
       diseq_unsat ctx
         (diseq_norm ctx e_2 e_1 p_1 (eagerReduce (Eq.refl true))
           (CommRing.diseq_norm rctx re_1 re_2 rp_1 (eagerReduce (Eq.refl true)) h))

tests/lean/run/grind_linarith_trim_context.lean

@@ -0,0 +1,32 @@
+/--
+trace: [grind.debug.proof] fun h h_1 h_2 h_3 h_4 h_5 h_6 h_7 h_8 =>
+      let ctx := RArray.branch 1 (RArray.leaf One.one) (RArray.leaf (f 2));
+      let p_1 := Poly.nil;
+      let p_2 := Poly.add 1 1 Poly.nil;
+      let p_3 := Poly.add 1 0 Poly.nil;
+      let p_4 := Poly.add (-1) 1 (Poly.add 1 0 Poly.nil);
+      let p_5 := Poly.add (-1) 0 Poly.nil;
+      let e_1 := (Expr.intMul 1 (Expr.var 1)).add (Expr.intMul 0 (Expr.var 0));
+      let e_2 := Expr.zero;
+      let e_3 := (Expr.intMul (-1) (Expr.var 1)).add (Expr.intMul 1 (Expr.var 0));
+      let rctx := RArray.leaf (f 2);
+      let rp_1 := CommRing.Poly.add 1 (CommRing.Mon.mult { x := 0, k := 1 } CommRing.Mon.unit) (CommRing.Poly.num 0);
+      let rp_2 := CommRing.Poly.add (-1) (CommRing.Mon.mult { x := 0, k := 1 } CommRing.Mon.unit) (CommRing.Poly.num 1);
+      let re_1 := CommRing.Expr.var 0;
+      let re_2 := CommRing.Expr.num 0;
+      let re_3 := ((CommRing.Expr.num 1).neg.mul (CommRing.Expr.var 0)).add (CommRing.Expr.num 1);
+      lt_unsat ctx
+        (le_lt_combine ctx p_3 p_5 p_1 (eagerReduce (Eq.refl true))
+          (le_le_combine ctx p_4 p_2 p_3 (eagerReduce (Eq.refl true))
+            (le_norm ctx e_3 e_2 p_4 (eagerReduce (Eq.refl true))
+              (CommRing.le_norm rctx re_3 re_2 rp_2 (eagerReduce (Eq.refl true)) h_8))
+            (le_norm ctx e_1 e_2 p_2 (eagerReduce (Eq.refl true))
+              (CommRing.le_norm rctx re_1 re_2 rp_1 (eagerReduce (Eq.refl true)) h_1)))
+          (zero_lt_one ctx p_5 (eagerReduce (Eq.refl true)) (Eq.refl One.one)))
+-/
+#guard_msgs in
+open Lean Grind Linarith in
+set_option trace.grind.debug.proof true in -- Context should contain only `f 2` and `One`
+example [CommRing α] [LE α] [LT α] [LinearOrder α] [OrderedRing α] (f : Nat → α) :
+    f 1 <= 0 → f 2 <= 0 → f 3 <= 0 → f 4 <= 0 → f 5 <= 0 → f 6 <= 0 → f 7 <= 0 → f 8 <= 0 → -1 * f 2 + 1 <= 0 → False := by
+  grind