feat(linarith): add linarith? mode which suggests a linarith only call (#28533)

Adds `linarith?`, which traces the internal to watch which inequalities are being used, and suggests a `linarith only` call via the "try this:" mechanism.

Also adds a `+minimize` flag, on by default, which then greedily tries to drop hypotheses from the used set, to see if the problem is still possible with a smaller set. Currently I don't have test case showing this is worthwhile, however. Contributions welcome.

Author: kim-em

Mathlib/Tactic/Linarith/Frontend.lean

@@ -164,6 +164,9 @@ structure LinarithConfig : Type where
   splitHypotheses : Bool := true
   /-- Split `≠` in hypotheses, by branching in cases `<` and `>`. -/
   splitNe : Bool := false
+  /-- If true, `linarith?` attempts to greedily remove unused hypotheses from its
+  suggestion. -/
+  minimize : Bool := true
   /-- Override the list of preprocessors. -/
   preprocessors : List GlobalBranchingPreprocessor := defaultPreprocessors
   /-- Specify an oracle for identifying candidate contradictions.
@@ -238,66 +241,84 @@ def ExprMultiMap.insert {α : Type} (self : ExprMultiMap α) (k : Expr) (v : α)
   return self.push (k, [v])
 
 /--
-`partitionByType l` takes a list `l` of proofs of comparisons. It sorts these proofs by
-the type of the variables in the comparison, e.g. `(a : ℚ) < 1` and `(b : ℤ) > c` will be separated.
-Returns a map from a type to a list of comparisons over that type.
+`partitionByTypeIdx l` takes a list `l` of pairs `(h, i)` where `h` is a proof of a
+comparison and `i` records the original position of `h`. The proofs are grouped by the
+type of the variables appearing in the comparison, e.g. `(a : ℚ) < 1` and
+`(b : ℤ) > c` will be separated. The resulting map associates each type with the
+list of `(h, i)` pairs over that type.
 -/
-def partitionByType (l : List Expr) : MetaM (ExprMultiMap Expr) :=
-  l.foldlM (fun m h => do m.insert (← typeOfIneqProof h) h) #[]
+def partitionByTypeIdx (l : List (Expr × Nat)) : MetaM (ExprMultiMap (Expr × Nat)) :=
+  l.foldlM (fun m ⟨h, i⟩ => do m.insert (← typeOfIneqProof h) (h, i)) #[]
 
 /--
-Given a list `ls` of lists of proofs of comparisons, `findLinarithContradiction cfg ls` will try to
-prove `False` by calling `linarith` on each list in succession. It will stop at the first proof of
-`False`, and fail if no contradiction is found with any list.
+Given a list `ls` of pairs `(α, L)` where each `L` is a list of indexed proofs of
+comparisons over the type `α`, `findLinarithContradiction cfg g ls` tries each list in
+succession, invoking `linarith` until one produces a contradiction. It returns the
+resulting proof of `False` together with the indices of the hypotheses that had
+nonzero coefficients in the final certificate.
 -/
-def findLinarithContradiction (cfg : LinarithConfig) (g : MVarId) (ls : List (Expr × List Expr)) :
-    MetaM Expr :=
+def findLinarithContradiction (cfg : LinarithConfig) (g : MVarId)
+    (ls : List (Expr × List (Expr × Nat))) : MetaM (Expr × List Nat) :=
   try
     ls.firstM (fun ⟨α, L⟩ =>
-      withTraceNode `linarith (return m!"{exceptEmoji ·} running on type {α}") <|
-        proveFalseByLinarith cfg.transparency cfg.oracle cfg.discharger g L)
+      withTraceNode `linarith (return m!"{exceptEmoji ·} running on type {α}") do
+        let (pf, idxs) ←
+          proveFalseByLinarith cfg.transparency cfg.oracle cfg.discharger g (L.map Prod.fst)
+        let idxs := idxs.map fun i => L[i]!.2
+        return (pf, idxs))
   catch e => throwError "linarith failed to find a contradiction\n{g}\n{e.toMessageData}"
 
 /--
-Given a list `hyps` of proofs of comparisons, `runLinarith cfg hyps prefType`
-preprocesses `hyps` according to the list of preprocessors in `cfg`.
-This results in a list of branches (typically only one),
-each of which must succeed in order to close the goal.
-
-In each branch, we partition the list of hypotheses by type, and run `linarith` on each class
-in the partition; one of these must succeed in order for `linarith` to succeed on this branch.
-If `prefType` is given, it will first use the class of proofs of comparisons over that type.
+Given a list `hyps` of proofs of comparisons, `runLinarith cfg prefType g hyps` preprocesses
+`hyps` according to the list of preprocessors in `cfg`. This results in a list of branches
+(typically only one), each of which must succeed in order to close the goal.
+
+In each branch, the hypotheses are partitioned by type and `linarith` is run on each class in
+turn; one of these must succeed in order for `linarith` to succeed on the branch. If `prefType`
+is provided, the corresponding class is tried first.
+
+On success, the metavariable `g` is assigned and the function returns the indices of the
+original hypotheses that were used with nonzero coefficient in the final proof.
 -/
 -- If it succeeds, the passed metavariable should have been assigned.
 def runLinarith (cfg : LinarithConfig) (prefType : Option Expr) (g : MVarId)
-    (hyps : List Expr) : MetaM Unit := do
-  let singleProcess (g : MVarId) (hyps : List Expr) : MetaM Expr := g.withContext do
-    linarithTraceProofs s!"after preprocessing, linarith has {hyps.length} facts:" hyps
-    let mut hyp_set ← partitionByType hyps
-    trace[linarith] "hypotheses appear in {hyp_set.size} different types"
-    -- If we have a preferred type, strip it from `hyp_set` and prepare a handler with a custom
-    -- trace message
-    let pref : MetaM _ ← do
-      if let some t := prefType then
-        let (i, vs) ← hyp_set.find t
-        hyp_set := hyp_set.eraseIdxIfInBounds i
-        pure <|
-          withTraceNode `linarith (return m!"{exceptEmoji ·} running on preferred type {t}") <|
-            proveFalseByLinarith cfg.transparency cfg.oracle cfg.discharger g vs
-      else
-        pure failure
-    pref <|> findLinarithContradiction cfg g hyp_set.toList
+    (hyps : List Expr) : MetaM (List Nat) := do
+  let singleProcess (g : MVarId) (hyps : List (Expr × Nat)) : MetaM (Expr × List Nat) :=
+    g.withContext do
+      linarithTraceProofs
+        s!"after preprocessing, linarith has {hyps.length} facts:" (hyps.map Prod.fst)
+      let mut hyp_set ← partitionByTypeIdx hyps
+      trace[linarith] "hypotheses appear in {hyp_set.size} different types"
+      -- If we have a preferred type, strip it from `hyp_set` and prepare a handler with a custom
+      -- trace message
+      let pref : MetaM _ ← do
+        if let some t := prefType then
+          let (i, vs) ← hyp_set.find t
+          hyp_set := hyp_set.eraseIdxIfInBounds i
+          pure <|
+            withTraceNode `linarith (return m!"{exceptEmoji ·} running on preferred type {t}") do
+              let (pf, idxs) ←
+                proveFalseByLinarith cfg.transparency cfg.oracle cfg.discharger g (vs.map Prod.fst)
+              let idxs := idxs.map fun j => vs[j]!.2
+              return (pf, idxs)
+        else
+          pure failure
+      pref <|> findLinarithContradiction cfg g hyp_set.toList
   let mut preprocessors := cfg.preprocessors
   if cfg.splitNe then
     preprocessors := Linarith.removeNe :: preprocessors
   if cfg.splitHypotheses then
     preprocessors := Linarith.splitConjunctions.globalize.branching :: preprocessors
   let branches ← preprocess preprocessors g hyps
+  let mut used : List Nat := []
   for (g, es) in branches do
-    let r ← singleProcess g es
+    let esIdx := es.zipIdx
+    let (r, idxs) ← singleProcess g esIdx
     g.assign r
+    used := idxs ++ used
   -- Verify that we closed the goal. Failure here should only result from a bad `Preprocessor`.
   (Expr.mvar g).ensureHasNoMVars
+  return used.eraseDups
 
 -- /--
 -- `filterHyps restr_type hyps` takes a list of proofs of comparisons `hyps`, and filters it
@@ -311,24 +332,27 @@ def runLinarith (cfg : LinarithConfig) (prefType : Option Expr) (g : MVarId)
 --     | none => return false)
 
 /--
-`linarith only_on hyps cfg` tries to close the goal using linear arithmetic. It fails
-if it does not succeed at doing this.
+`linarithUsedHyps only_on hyps cfg g` runs `linarith` with the supplied hypotheses. It
+fails if the goal cannot be closed. When successful, it returns the subset of `hyps` that
+were actually used (i.e. had a nonzero coefficient) in the final certificate.
 
 * `hyps` is a list of proofs of comparisons to include in the search.
 * If `only_on` is true, the search will be restricted to `hyps`. Otherwise it will use all
   comparisons in the local context.
 * If `cfg.transparency := semireducible`,
   it will unfold semireducible definitions when trying to match atomic expressions.
 -/
-partial def linarith (only_on : Bool) (hyps : List Expr) (cfg : LinarithConfig := {})
-    (g : MVarId) : MetaM Unit := g.withContext do
+partial def linarithUsedHyps (only_on : Bool) (hyps : List Expr)
+    (cfg : LinarithConfig := {}) (g : MVarId) : MetaM (List Expr) := g.withContext do
   -- if the target is an equality, we run `linarith` twice, to prove ≤ and ≥.
   if (← whnfR (← instantiateMVars (← g.getType))).isEq then
     trace[linarith] "target is an equality: splitting"
     if let some [g₁, g₂] ← try? (g.apply (← mkConst' ``eq_of_not_lt_of_not_gt)) then
-      withTraceNode `linarith (return m!"{exceptEmoji ·} proving ≥") <| linarith only_on hyps cfg g₁
-      withTraceNode `linarith (return m!"{exceptEmoji ·} proving ≤") <| linarith only_on hyps cfg g₂
-      return
+      let h₁ ← withTraceNode `linarith (return m!"{exceptEmoji ·} proving ≥") <|
+        linarithUsedHyps only_on hyps cfg g₁
+      let h₂ ← withTraceNode `linarith (return m!"{exceptEmoji ·} proving ≤") <|
+        linarithUsedHyps only_on hyps cfg g₂
+      return h₁ ++ h₂
 
   /- If we are proving a comparison goal (and not just `False`), we consider the type of the
     elements in the comparison to be the "preferred" type. That is, if we find comparison
@@ -347,17 +371,33 @@ partial def linarith (only_on : Bool) (hyps : List Expr) (cfg : LinarithConfig :
   | (some (t, v), g) => pure (g, some t, some v)
 
   g.withContext do
-  -- set up the list of hypotheses, considering the `only_on` and `restrict_type` options
-    let hyps ← (if only_on then return new_var.toList ++ hyps
-      else return (← getLocalHyps).toList ++ hyps)
+    -- set up the list of hypotheses, considering the `only_on` and `restrict_type` options
+    let hyps ←
+      (if only_on then return new_var.toList ++ hyps
+        else return (← getLocalHyps).toList ++ hyps)
 
     -- TODO in mathlib3 we could specify a restriction to a single type.
     -- I haven't done that here because I don't know how to store a `Type` in `LinarithConfig`.
     -- There's only one use of the `restrict_type` configuration option in mathlib3,
     -- and it can be avoided just by using `linarith only`.
 
     linarithTraceProofs "linarith is running on the following hypotheses:" hyps
-    runLinarith cfg target_type g hyps
+    let usedIdxs ← runLinarith cfg target_type g hyps
+    let used := usedIdxs.filterMap (hyps[·]?)
+    let used := match new_var with
+      | some nv => used.filter (fun h => !(h == nv))
+      | none => used
+    return used
+
+/--
+Run the core `linarith` procedure on the goal `g` using the hypotheses `hyps`.
+If `only_on` is true, the search is restricted to `hyps`; otherwise all suitable
+local hypotheses are considered. This is the workhorse behind the user-facing
+`linarith` tactic.
+-/
+partial def linarith (only_on : Bool) (hyps : List Expr) (cfg : LinarithConfig := {})
+    (g : MVarId) : MetaM Unit := do
+  discard <| linarithUsedHyps only_on hyps cfg g
 
 end Linarith
 
@@ -416,6 +456,8 @@ optional arguments:
   disequality hypotheses. (`false` by default.)
 * If `exfalso` is `false`, `linarith` will fail when the goal is neither an inequality nor `False`.
   (`true` by default.)
+* If `minimize` is `false`, `linarith?` will report all hypotheses appearing in its initial
+  proof without attempting to drop redundancies. (`true` by default.)
 * `restrict_type` (not yet implemented in mathlib4)
   will only use hypotheses that are inequalities over `tp`. This is useful
   if you have e.g. both integer- and rational-valued inequalities in the local context, which can
@@ -428,8 +470,19 @@ routine.
 -/
 syntax (name := linarith) "linarith" "!"? linarithArgsRest : tactic
 
+/--
+`linarith?` behaves like `linarith` but, on success, it prints a suggestion of
+the form `linarith only [...]` listing a minimized set of hypotheses used in the
+final proof.  Use `linarith?!` for the higher-reducibility variant and set the
+`minimize` flag in the configuration to control whether greedy minimization is
+performed.
+-/
+syntax (name := linarith?) "linarith?" "!"? linarithArgsRest : tactic
+
 @[inherit_doc linarith] macro "linarith!" rest:linarithArgsRest : tactic =>
   `(tactic| linarith ! $rest:linarithArgsRest)
+@[inherit_doc linarith?] macro "linarith?!" rest:linarithArgsRest : tactic =>
+  `(tactic| linarith? ! $rest:linarithArgsRest)
 
 /--
 An extension of `linarith` with some preprocessing to allow it to solve some nonlinear arithmetic
@@ -458,6 +511,46 @@ elab_rules : tactic
     let cfg := (← elabLinarithConfig cfg).updateReducibility bang.isSome
     commitIfNoEx do liftMetaFinishingTactic <| Linarith.linarith o.isSome args.toList cfg
 
+elab_rules : tactic
+  | `(tactic| linarith?%$tk $[!%$bang]? $cfg:optConfig $[only%$o]? $[[$args,*]]?) =>
+      withMainContext do
+        let args ←
+          ((args.map (TSepArray.getElems)).getD {}).mapM (elabTermWithoutNewMVars `linarith)
+        let cfg := (← elabLinarithConfig cfg).updateReducibility bang.isSome
+        let g ← getMainGoal
+        let st ← saveState
+        try
+          let used₀ ← Linarith.linarithUsedHyps o.isSome args.toList cfg g
+          -- Check that all used hypotheses are fvars (not arbitrary terms)
+          if used₀.any (fun e => e.fvarId?.isNone) then
+            throwError "linarith? currently only supports named hypothesis, not terms"
+          let used ←
+            if cfg.minimize then
+              let rec minimize (hs : List Expr) (i : Nat) : TacticM (List Expr) := do
+                if _h : i < hs.length then
+                  let rest := hs.eraseIdx i
+                  st.restore
+                  try
+                    let _ ← Linarith.linarith true rest cfg g
+                    minimize rest i
+                  catch _ => minimize hs (i+1)
+                else
+                  return hs
+              minimize used₀ 0
+            else
+              pure used₀
+          st.restore
+          discard <| Linarith.linarith true used cfg g
+          replaceMainGoal []
+          -- TODO: we should check for, and deal with, shadowed names here.
+          let idsList ← used.mapM fun e => do
+            pure (Lean.mkIdent (← e.fvarId!.getUserName))
+          let sugg ← `(tactic| linarith only [$(idsList.toArray),*])
+          Lean.Meta.Tactic.TryThis.addSuggestion tk sugg
+        catch e =>
+          discard <| st.restore
+          throw e
+
 -- TODO restore this when `add_tactic_doc` is ported
 -- add_tactic_doc
 -- { name       := "linarith",

Mathlib/Tactic/Linarith/Verification.lean

@@ -140,19 +140,22 @@ def mkNegOneLtZeroProof (tp : Expr) : MetaM Expr := do
   mkAppM `neg_neg_of_pos #[zero_lt_one]
 
 /--
-`addNegEqProofs l` inspects the list of proofs `l` for proofs of the form `t = 0`. For each such
-proof, it adds a proof of `-t = 0` to the list.
+`addNegEqProofsIdx l` inspects a list `l` of pairs `(h, i)` where `h` proves
+`tᵢ Rᵢ 0` and `i` records the original index of the hypothesis. For each
+equality proof `t = 0` in the list, it appends a proof of `-t = 0` with the
+same index `i`. All other entries are preserved.
 -/
-def addNegEqProofs : List Expr → MetaM (List Expr)
+def addNegEqProofsIdx : List (Expr × Nat) → MetaM (List (Expr × Nat))
   | [] => return []
-  | (h::tl) => do
+  | (⟨h, i⟩::tl) => do
     let (iq, t) ← parseCompAndExpr (← inferType h)
     match iq with
     | Ineq.eq => do
-      let nep := mkAppN (← mkAppM `Iff.mpr #[← mkAppOptM ``neg_eq_zero #[none, none, t]]) #[h]
-      let tl ← addNegEqProofs tl
-      return h::nep::tl
-    | _ => return h :: (← addNegEqProofs tl)
+      let nep :=
+        mkAppN (← mkAppM `Iff.mpr #[← mkAppOptM ``neg_eq_zero #[none, none, t]]) #[h]
+      let tl ← addNegEqProofsIdx tl
+      return (h, i)::(nep, i)::tl
+    | _ => return (h, i) :: (← addNegEqProofsIdx tl)
 
 /--
 `proveEqZeroUsing tac e` tries to use `tac` to construct a proof of `e = 0`.
@@ -188,15 +191,19 @@ tactic, which is typically `ring`. We prove (2) by folding over the set of hypot
 `transparency : TransparencyMode` controls the transparency level with which atoms are identified.
 -/
 def proveFalseByLinarith (transparency : TransparencyMode) (oracle : CertificateOracle)
-    (discharger : TacticM Unit) : MVarId → List Expr → MetaM Expr
+    (discharger : TacticM Unit) : MVarId → List Expr → MetaM (Expr × List Nat)
   | _, [] => throwError "no args to linarith"
   | g, l@(h::_) => do
       Lean.Core.checkSystem decl_name%.toString
       -- for the elimination to work properly, we must add a proof of `-1 < 0` to the list,
       -- along with negated equality proofs.
-      let l' ← detailTrace "addNegEqProofs" <| addNegEqProofs l
-      let inputs ← detailTrace "mkNegOneLtZeroProof" <|
-        return (← mkNegOneLtZeroProof (← typeOfIneqProof h))::l'.reverse
+      let lidx := l.zipIdx
+      let l' ← detailTrace "addNegEqProofs" <| addNegEqProofsIdx lidx
+      let inputsTagged : List (Expr × Option Nat) ←
+        detailTrace "mkNegOneLtZeroProof" <|
+          return ((← mkNegOneLtZeroProof (← typeOfIneqProof h)), none) ::
+            (l'.reverse.map fun ⟨e, i⟩ => (e, some i))
+      let inputs := inputsTagged.map Prod.fst
       trace[linarith.detail] "inputs:{indentD <| toMessageData (← inputs.mapM inferType)}"
       let (comps, max_var) ← detailTrace "linearFormsAndMaxVar" <|
         linearFormsAndMaxVar transparency inputs
@@ -212,28 +219,36 @@ def proveFalseByLinarith (transparency : TransparencyMode) (oracle : Certificate
               throwError "linarith failed to find a contradiction"
           trace[linarith] "found a contradiction: {certificate.toList}"
           return certificate
-      let (sm, zip) ←
+      let (sm, zip, idxs) ←
         withTraceNode `linarith (return m!"{exceptEmoji ·} Building final expression") do
-          let enum_inputs := inputs.zipIdx
+          let enum_inputs := inputsTagged.zipIdx
           -- construct a list pairing nonzero coeffs with the proof of their corresponding
-          -- comparison
-          let zip := enum_inputs.filterMap fun ⟨e, n⟩ => (certificate[n]?).map (e, ·)
-          let mls ← zip.mapM fun ⟨e, n⟩ => do mulExpr n (← leftOfIneqProof e)
+          -- comparison and track the original index
+          let used := enum_inputs.filterMap fun ⟨⟨e, orig?⟩, n⟩ =>
+            (certificate[n]?).map fun c => (e, c, orig?)
+          let zip := used.map fun ⟨e, c, _⟩ => (e, c)
+          let mls ← used.mapM fun ⟨e, c, _⟩ => do mulExpr c (← leftOfIneqProof e)
           -- `sm` is the sum of input terms, scaled to cancel out all variables.
           let sm ← addExprs mls
           -- let sm ← instantiateMVars sm
           trace[linarith] "{indentD sm}\nshould be both 0 and negative"
-          return (sm, zip)
+          let idxs :=
+            (used.foldl (fun acc (_, _, orig?) =>
+                match orig? with
+                | some i => i :: acc
+                | none => acc) []).eraseDups
+          return (sm, zip, idxs)
       -- we prove that `sm = 0`, typically with `ring`.
       let sm_eq_zero ← detailTrace "proveEqZeroUsing" <| proveEqZeroUsing discharger sm
       -- we also prove that `sm < 0`
       let sm_lt_zero ← detailTrace "mkLTZeroProof" <| mkLTZeroProof zip
-      detailTrace "Linarith.lt_irrefl" do
+      let pf ← detailTrace "Linarith.lt_irrefl" do
         -- this is a contradiction.
         let pftp ← inferType sm_lt_zero
         let ⟨_, nep, _⟩ ← g.rewrite pftp sm_eq_zero
         let pf' ← mkAppM ``Eq.mp #[nep, sm_lt_zero]
         mkAppM ``Linarith.lt_irrefl #[pf']
+      return (pf, idxs)
 where
   /-- Log `f` under `linarith.detail`, with exception emojis and the provided name. -/
   detailTrace {α} (s : String) (f : MetaM α) : MetaM α :=