From 9290b491bb16ec897a5cb9ffd98c8416af30bfb2 Mon Sep 17 00:00:00 2001
From: Joachim Breitner <mail@joachim-breitner.de>
Date: Mon, 4 Dec 2023 22:42:24 +0100
Subject: [PATCH] refactor: WF.Fix: gather subgoals (#3017)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This is pure refactoring: Instead of solving each subgoal as we
encounter it while traversing the syntax tree, we leave the `MVar`
there, at the end collect them all using `getMVarsNoDelayed`, and then
solve them.

This is a refactoring preparing for two upcoming changes:

 * removing unexpected duplicate goals that can arise from term
   duplication
 * running interactive tactics on all, not each goal (#2921)

In order to not regress with error locations, we have to associated the
`TermElabM`’s syntax refernce with the `MVar` somehow. I do this using
the existing `mkRecAppWithSyntax` expression annotation, on the `MVar`’s
type. Alternatives would be stack another `StateT` on the traversal
and accumulate `Array (MVarId, Syntax)` explicitly, but that did not
seem to be more appealing.
---
 src/Lean/Elab/PreDefinition/WF/Fix.lean | 54 ++++++++++++++++---------
 1 file changed, 35 insertions(+), 19 deletions(-)

diff --git a/src/Lean/Elab/PreDefinition/WF/Fix.lean b/src/Lean/Elab/PreDefinition/WF/Fix.lean
index e553ba3987..10a9d9d336 100644
--- a/src/Lean/Elab/PreDefinition/WF/Fix.lean
+++ b/src/Lean/Elab/PreDefinition/WF/Fix.lean
@@ -17,25 +17,20 @@ import Lean.Elab.PreDefinition.Structural.BRecOn
 namespace Lean.Elab.WF
 open Meta
 
-private def applyDefaultDecrTactic (mvarId : MVarId) : TermElabM Unit := do
-  let remainingGoals ← Tactic.run mvarId do
-    Tactic.evalTactic (← `(tactic| decreasing_tactic))
-  unless remainingGoals.isEmpty do
-    Term.reportUnsolvedGoals remainingGoals
-
-private def mkDecreasingProof (decreasingProp : Expr) (decrTactic? : Option Syntax) : TermElabM Expr := do
-  let mvar ← mkFreshExprSyntheticOpaqueMVar decreasingProp
+/-
+Creates a subgoal for a recursive call, as an unsolved `MVar`. The goal is cleaned up, and
+the current syntax reference is stored in the `MVar`’s type as a `RecApp` marker, for
+use by `solveDecreasingGoals` below.
+-/
+private def mkDecreasingProof (decreasingProp : Expr) : TermElabM Expr := do
+  -- We store the current Ref in the MVar as a RecApp annotation around the type
+  let ref ← getRef
+  let mvar ← mkFreshExprSyntheticOpaqueMVar (mkRecAppWithSyntax decreasingProp ref)
   let mvarId := mvar.mvarId!
-  let mvarId ← mvarId.cleanup
-  match decrTactic? with
-  | none => applyDefaultDecrTactic mvarId
-  | some decrTactic =>
-    -- make info from `runTactic` available
-    pushInfoTree (.hole mvarId)
-    Term.runTactic mvarId decrTactic
-  instantiateMVars mvar
+  let _mvarId ← mvarId.cleanup
+  return mvar
 
-private partial def replaceRecApps (recFnName : Name) (fixedPrefixSize : Nat) (decrTactic? : Option Syntax) (F : Expr) (e : Expr) : TermElabM Expr := do
+private partial def replaceRecApps (recFnName : Name) (fixedPrefixSize : Nat) (F : Expr) (e : Expr) : TermElabM Expr := do
   trace[Elab.definition.wf] "replaceRecApps:{indentExpr e}"
   trace[Elab.definition.wf] "{F} : {← inferType F}"
   loop F e |>.run' {}
@@ -47,7 +42,7 @@ where
       let args := e.getAppArgs
       let r := mkApp F (← loop F args[fixedPrefixSize]!)
       let decreasingProp := (← whnf (← inferType r)).bindingDomain!
-      let r := mkApp r (← mkDecreasingProof decreasingProp decrTactic?)
+      let r := mkApp r (← mkDecreasingProof decreasingProp)
       return mkAppN r (← args[fixedPrefixSize+1:].toArray.mapM (loop F))
 
   processApp (F : Expr) (e : Expr) : StateRefT (HasConstCache recFnName) TermElabM Expr := do
@@ -164,6 +159,26 @@ private partial def processPSigmaCasesOn (x F val : Expr) (k : (F : Expr) → (v
   else
     k F val
 
+private def applyDefaultDecrTactic (mvarId : MVarId) : TermElabM Unit := do
+  let remainingGoals ← Tactic.run mvarId do
+    Tactic.evalTactic (← `(tactic| decreasing_tactic))
+  unless remainingGoals.isEmpty do
+    Term.reportUnsolvedGoals remainingGoals
+
+def solveDecreasingGoals (decrTactic? : Option Syntax) (value : Expr) : MetaM Expr := do
+  let goals ← getMVarsNoDelayed value
+  goals.forM fun goal => Lean.Elab.Term.TermElabM.run' <|
+    match decrTactic? with
+    | none => do
+      let some ref := getRecAppSyntax? (← goal.getType)
+        | throwError "MVar does not look like like a recursive call"
+        withRef ref <| applyDefaultDecrTactic goal
+    | some decrTactic => do
+      -- make info from `runTactic` available
+      pushInfoTree (.hole goal)
+      Term.runTactic goal decrTactic
+  instantiateMVars value
+
 def mkFix (preDef : PreDefinition) (prefixArgs : Array Expr) (wfRel : Expr) (decrTactic? : Option Syntax) : TermElabM Expr := do
   let type ← instantiateForall preDef.type prefixArgs
   let (wfFix, varName) ← forallBoundedTelescope type (some 1) fun x type => do
@@ -186,7 +201,8 @@ def mkFix (preDef : PreDefinition) (prefixArgs : Array Expr) (wfRel : Expr) (dec
       let F   := xs[1]!
       let val := preDef.value.beta (prefixArgs.push x)
       let val ← processSumCasesOn x F val fun x F val => do
-        processPSigmaCasesOn x F val (replaceRecApps preDef.declName prefixArgs.size decrTactic?)
+        processPSigmaCasesOn x F val (replaceRecApps preDef.declName prefixArgs.size)
+      let val ← solveDecreasingGoals decrTactic? val
       mkLambdaFVars prefixArgs (mkApp wfFix (← mkLambdaFVars #[x, F] val))
 
 end Lean.Elab.WF