feat: add splitNext grind action (#10801)

This PR implements the `splitNext` action for `grind`.

Author: leodemoura

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

@@ -151,33 +151,6 @@ Executes `x`, but behaves like a `skip` if it is not applicable.
 def skipIfNA (x : Action) : Action := fun goal _ kp =>
   x goal kp kp
 
-private def isTargetFalse (mvarId : MVarId) : MetaM Bool := do
-  return (← mvarId.getType).isFalse
-
-private def getFalseProof? (mvarId : MVarId) : MetaM (Option Expr) := mvarId.withContext do
-  let proof ← instantiateMVars (mkMVar mvarId)
-  if (← isTargetFalse mvarId) then
-    return some proof
-  else if proof.isAppOfArity ``False.elim 2 || proof.isAppOfArity ``False.casesOn 2 then
-    return some proof.appArg!
-  else
-    return none
-
-/--
-Returns the maximum free variable id occurring in `e`
--/
-private def findMaxFVarIdx? (e : Expr) : MetaM (Option Nat) := do
-  let go (e : Expr) : StateT (Option Nat) MetaM Bool := do
-    unless e.hasFVar do return false
-    let .fvar fvarId := e | return true
-    let localDecl ← fvarId.getDecl
-    modify fun
-      | none => some localDecl.index
-      | some index => some (max index localDecl.index)
-    return false
-  let (_, s?) ← e.forEach' go |>.run none
-  return s?
-
 private def mkGrindSeq (s : List (TSyntax `grind)) : TSyntax ``Parser.Tactic.Grind.grindSeq :=
   let s := s.map (·.raw)
   let s := s.intersperse (mkNullNode #[])
@@ -236,73 +209,6 @@ def ungroup : Action := fun goal _ kp => do
   else
     return r
 
-/--
-Returns `some falseProof` if we can use non-chronological backtracking with `subgoal`.
-That is, `subgoal` was closed using `falseProof`, but its proof does not use any of the
-new hypotheses. A hypothesis is new if its `index >= oldNumIndices`.
--/
-def useNCB? (oldNumIndices : Nat) (subgoal : Goal) : MetaM (Option Expr) := do
-  let some falseProof ← getFalseProof? subgoal.mvarId
-    | return none
-  let some max ← subgoal.mvarId.withContext <| findMaxFVarIdx? falseProof
-    | return some falseProof -- Proof actually closes any pending split
-  if max < oldNumIndices then
-    return some falseProof
-  else
-    return none
-
-/--
-Helper function for implementing tactics that perform case-splits
-**Note**: We will probably delete this function.
--/
-def splitCore
-    (goal : Goal)
-    (anchor? : Option (Nat × UInt64))
-    (s : MVarId → GrindM (List MVarId))
-    (kp : ActionCont) : GrindM ActionResult := do
-  let mvarDecl ← goal.mvarId.getDecl
-  let numIndices := mvarDecl.lctx.numIndices
-  let mvarId ← goal.mkAuxMVar
-  let mvarIds ← s mvarId
-  let subgoals := mvarIds.map fun mvarId => { goal with mvarId }
-  let traceEnabled := (← getConfig).trace
-  let mut seqNew : Array (List (TSyntax `grind)) := #[]
-  let mut stuckNew : Array Goal := #[]
-  for subgoal in subgoals do
-    match (← kp subgoal) with
-    | .stuck gs =>
-      -- *TODO*: Add support for saving multiple failures
-      return .stuck gs
-    | .closed seq =>
-      if let some falseProof ← useNCB? numIndices subgoal then
-        goal.mvarId.assignFalseProof falseProof
-        return .closed seq
-      else
-        seqNew := seqNew.push seq
-  if stuckNew.isEmpty then
-    goal.mvarId.assign (← instantiateMVars (mkMVar mvarId))
-    if traceEnabled then
-      let seqListNew ← if h : seqNew.size = 1 then
-        pure seqNew[0]
-      else
-        seqNew.toList.mapM fun s => mkGrindNext s
-      let mut seqListNew := seqListNew
-      if let some anchor := anchor? then
-        let hexnum := mkNode `hexnum #[mkAtom (anchorToString anchor.1 anchor.2)]
-        /-
-        *TODO*: We need to distinguish between user-facing `cases` which `intros` new hypotheses
-        automatically, and auto-generated `cases` produced by `grind?` and `finish?` which does not
-        `intros` automatically. Each branch provides includes its own `intros`.
-        *Current strategy*: Use only one `cases` (`intros`) automatically and add `rename_i`.
-        -/
-        let cases ← `(grind| cases #$hexnum)
-        seqListNew := cases :: seqListNew
-      return .closed seqListNew
-    else
-      return .closed []
-  else
-    return .stuck stuckNew.toList
-
 section
 /-!
 Some sanity check properties.

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

@@ -7,11 +7,13 @@ module
 prelude
 public import Lean.Meta.Tactic.Grind.Types
 public import Lean.Meta.Tactic.Grind.SearchM
+public import Lean.Meta.Tactic.Grind.Action
 import Lean.Meta.Tactic.Grind.Intro
 import Lean.Meta.Tactic.Grind.Cases
 import Lean.Meta.Tactic.Grind.Util
 import Lean.Meta.Tactic.Grind.CasesMatch
 import Lean.Meta.Tactic.Grind.Internalize
+import Lean.Meta.Tactic.Grind.Anchor
 public section
 namespace Lean.Meta.Grind
 
@@ -242,6 +244,131 @@ private def casesWithTrace (mvarId : MVarId) (major : Expr) : GoalM (List MVarId
       saveCases declName false
   cases mvarId major
 
+namespace Action
+
+/--
+Given a `mvarId` associated with a subgoal created by `splitCore`, inspects the
+proof term assigned to `mvarId` and tries to extract the proof of `False` that does not
+depend on hypotheses introduced in the subgoal.
+For example: suppose the subgoal is of the form `p → q → False` where `p` and `q` are new
+hypotheses introduced during case analysis. If the proof is of the form `fun _ _ => h`, returns
+`some h`.
+-/
+private def getFalseProof? (mvarId : MVarId) : MetaM (Option Expr) := mvarId.withContext do
+  let proof ← instantiateMVars (mkMVar mvarId)
+  go proof
+where
+  go (proof : Expr) : MetaM (Option Expr) := do
+    match_expr proof with
+    | False.elim _ p => return some p
+    | False.casesOn _ p => return some p
+    | id α p => if α.isFalse then return some p else return none
+    | _ =>
+      /-
+      **Note**: `intros` tactics may hide the `False` proof behind a `casesOn`
+      For example: suppose the subgoal has a type of the form `p₁ → q₁ ∧ q₂ → p₂ → False`
+      The proof will be of the form `fun _ h => h.casesOn (fun _ _ => hf)` where `hf` is the proof
+      of `False` we are looking for.
+      Non-chronological backtracking currently fails in this kind of example.
+      -/
+      let .lam _ _ b _ := proof | return none
+      if b.hasLooseBVars then return none
+      go b
+
+/--
+Performs a case-split using `c`.
+Remark: `numCases` and `isRec` are computed using `checkSplitStatus`.
+-/
+private def splitCore (c : SplitInfo) (numCases : Nat) (isRec : Bool) (stopAtFirstFailure : Bool) : Action := fun goal _ kp => do
+  let mvarDecl ← goal.mvarId.getDecl
+  let numIndices := mvarDecl.lctx.numIndices
+  let mvarId ← goal.mkAuxMVar
+  let cExpr := c.getExpr
+  let (mvarIds, goal) ← GoalM.run goal do
+    let gen ← getGeneration cExpr
+    let genNew := if numCases > 1 || isRec then gen+1 else gen
+    saveSplitDiagInfo cExpr genNew numCases c.source
+    markCaseSplitAsResolved cExpr
+    trace_goal[grind.split] "{cExpr}, generation: {gen}"
+    let mvarIds ← if let .imp e h _ := c then
+      casesWithTrace mvarId (mkGrindEM (e.forallDomain h))
+    else if (← isMatcherApp cExpr) then
+      casesMatch mvarId cExpr
+    else
+      casesWithTrace mvarId (← mkCasesMajor cExpr)
+  let subgoals := mvarIds.map fun mvarId => { goal with mvarId }
+  let traceEnabled := (← getConfig).trace
+  let mut seqNew : Array (List (TSyntax `grind)) := #[]
+  let mut stuckNew : Array Goal := #[]
+  for subgoal in subgoals do
+    match (← kp subgoal) with
+    | .stuck gs =>
+      if stopAtFirstFailure then
+        /-
+        **Note**: We don't need to assign `goal.mvarId` when `stopAtFirstFailure = true`
+        because the caller will not be able to process the all failure/stuck goals anyway.
+        -/
+        return .stuck gs
+      else
+        stuckNew := stuckNew ++ gs
+    | .closed seq =>
+      if let some falseProof ← getFalseProof? subgoal.mvarId then
+        goal.mvarId.assignFalseProof falseProof
+        return .closed seq
+      else if !seq.isEmpty then
+        /- **Note**: if the sequence is empty, it means the user will never see this goal. -/
+        seqNew := seqNew.push seq
+  if (← goal.mvarId.getType).isFalse then
+    /- **Note**: We add the marker to assist `getFalseExpr?` -/
+    goal.mvarId.assign (mkExpectedPropHint (← instantiateMVars (mkMVar mvarId)) (mkConst ``False))
+  else
+    goal.mvarId.assign (← instantiateMVars (mkMVar mvarId))
+  if stuckNew.isEmpty then
+    if traceEnabled then
+      let seqListNew ← if h : seqNew.size = 1 then
+        pure seqNew[0]
+      else
+        seqNew.toList.mapM fun s => mkGrindNext s
+      let mut seqListNew := seqListNew
+      let anchor ← goal.withContext <| getAnchor cExpr
+      -- **TODO**: compute the exact number of digits
+      let numDigits := 4
+      let anchorPrefix := anchor >>> (64 - 16)
+      let hexnum := mkNode `hexnum #[mkAtom (anchorToString numDigits anchorPrefix)]
+      let cases ← `(grind| cases #$hexnum)
+      seqListNew := cases :: seqListNew
+      return .closed seqListNew
+    else
+      return .closed []
+  else
+    return .stuck stuckNew.toList
+
+/--
+Selects a case-split from the list of candidates, performs the split and applies
+continuation to all subgoals.
+If a subgoal is solved without using new hypotheses, closes the original goal using this proof. That is,
+it performs non-chronological backtracking.
+If `stopsAtFirstFailure = true`, it stops the search as soon as the given continuation cannot solve a subgoal.
+-/
+def splitNext (stopAtFirstFailure := true) : Action := fun goal kna kp => do
+  let (r, goal) ← GoalM.run goal selectNextSplit?
+  let .some c numCases isRec _ := r
+    | kna goal
+  let cExpr := c.getExpr
+  let gen := goal.getGeneration cExpr
+  let x : Action := splitCore c numCases isRec stopAtFirstFailure >> intros gen
+  x goal kna kp
+
+end Action
+
+/-!
+**------------------------------------------**
+**------------------------------------------**
+**TODO** Delete rest of the file
+**------------------------------------------**
+**------------------------------------------**
+-/
+
 /--
 Performs a case-split using `c`.
 Remarks:

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

@@ -923,10 +923,15 @@ def Goal.getENode (goal : Goal) (e : Expr) : CoreM ENode := do
 def getENode (e : Expr) : GoalM ENode := do
   (← get).getENode e
 
+def Goal.getGeneration (goal : Goal) (e : Expr) : Nat :=
+  if let some n := goal.getENode? e then
+    n.generation
+  else
+    0
+
 /-- Returns the generation of the given term. Is assumes it has been internalized -/
-def getGeneration (e : Expr) : GoalM Nat := do
-  let some n ← getENode? e | return 0
-  return n.generation
+def getGeneration (e : Expr) : GoalM Nat :=
+  return (← get).getGeneration e
 
 /-- Returns `true` if `e` is in the equivalence class of `True`. -/
 def isEqTrue (e : Expr) : GoalM Bool := do
@@ -1244,7 +1249,10 @@ If type of `mvarId` is not `False`, then use `False.elim`.
 def _root_.Lean.MVarId.assignFalseProof (mvarId : MVarId) (falseProof : Expr) : MetaM Unit := mvarId.withContext do
   let target ← mvarId.getType
   if target.isFalse then
-    mvarId.assign falseProof
+    /-
+    **Note**: We add the marker to assist `getFalseExpr?` used to implement
+    non-chronological backtracking. -/
+    mvarId.assign (mkExpectedPropHint falseProof (mkConst ``False))
   else
     mvarId.assign (← mkFalseElim target falseProof)
 
@@ -1341,12 +1349,6 @@ def propagateUp (e : Expr) : GoalM Unit := do
 def propagateDown (e : Expr) : GoalM Unit := do
   (← getMethods).propagateDown e
 
-def Goal.getGeneration (goal : Goal) (e : Expr) : Nat :=
-  if let some n := goal.getENode? e then
-    n.generation
-  else
-    0
-
 /-- Returns expressions in the given expression equivalence class. -/
 partial def Goal.getEqc (goal : Goal) (e : Expr) (sort := false) : List Expr :=
   let eqc := go e e #[]

tests/lean/run/grind_cutsat_trim_context.lean

@@ -1,11 +1,12 @@
 module
 /--
 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.leaf (f 2);
-      let p_1 := Poly.add 1 0 (Poly.num 0);
-      let p_2 := Poly.add (-1) 0 (Poly.num 1);
-      let p_3 := Poly.num 1;
-      le_unsat ctx p_3 (eagerReduce (Eq.refl true)) (le_combine ctx p_2 p_1 p_3 (eagerReduce (Eq.refl true)) h_8 h_1)
+      id
+        (let ctx := RArray.leaf (f 2);
+        let p_1 := Poly.add 1 0 (Poly.num 0);
+        let p_2 := Poly.add (-1) 0 (Poly.num 1);
+        let p_3 := Poly.num 1;
+        le_unsat ctx p_3 (eagerReduce (Eq.refl true)) (le_combine ctx p_2 p_1 p_3 (eagerReduce (Eq.refl true)) h_8 h_1))
 -/
 #guard_msgs in -- `cutsat` context should contain only `f 2`
 open Lean Int Linear in

tests/lean/run/grind_linarith_2.lean

@@ -8,17 +8,18 @@ example [IntModule α] [LE α] [LT α] [IsPreorder α] [OrderedAdd α] (a b : α
 
 /--
 trace: [grind.debug.proof] Classical.byContradiction fun h =>
-      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 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))
+      id
+        (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 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)))
 -/
 #guard_msgs in
 open Linarith in