feat: basic skeleton for termination_by' vs termination_by

src/Lean/Elab/PreDefinition/Main.lean

@@ -67,8 +67,8 @@ def addPreDefinitions (preDefs : Array PreDefinition) (hints : TerminationHints)
     trace[Elab.definition.body] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
   let preDefs ← preDefs.mapM ensureNoUnassignedMVarsAtPreDef
   let cliques ← partitionPreDefs preDefs
-  let mut terminationBy ← liftMacroM <| WF.expandTerminationHint hints.terminationBy? (cliques.map fun ds => ds.map (·.declName))
-  let mut decreasingBy ← liftMacroM <| WF.expandTerminationHint hints.decreasingBy? (cliques.map fun ds => ds.map (·.declName))
+  let mut terminationBy ← liftMacroM <| WF.expandTerminationBy hints.terminationBy? (cliques.map fun ds => ds.map (·.declName))
+  let mut decreasingBy  ← liftMacroM <| WF.expandTerminationHint hints.decreasingBy? (cliques.map fun ds => ds.map (·.declName))
   for preDefs in cliques do
     trace[Elab.definition.scc] "{preDefs.map (·.declName)}"
     if preDefs.size == 1 && isNonRecursive preDefs[0] then
@@ -85,16 +85,16 @@ def addPreDefinitions (preDefs : Array PreDefinition) (hints : TerminationHints)
           withRef preDef.ref <| throwError "invalid use of 'partial', '{preDef.declName}' is not a function{indentExpr preDef.type}"
       addAndCompilePartial preDefs
     else
-      let mut wfStx? := none
+      let mut wf? := none
       let mut decrTactic? := none
-      if let some { value := wfStx, .. } := terminationBy.find? (preDefs.map (·.declName)) then
-        wfStx? := some wfStx
+      if let some wf := terminationBy.find? (preDefs.map (·.declName)) then
+        wf? := some wf
         terminationBy := terminationBy.erase (preDefs.map (·.declName))
       if let some { ref, value := decrTactic } := decreasingBy.find? (preDefs.map (·.declName)) then
         decrTactic? := some (← withRef ref `(by $decrTactic))
         decreasingBy := decreasingBy.erase (preDefs.map (·.declName))
-      if wfStx?.isSome || decrTactic?.isSome then
-        wfRecursion preDefs wfStx? decrTactic?
+      if wf?.isSome || decrTactic?.isSome then
+        wfRecursion preDefs wf? decrTactic?
       else
         withRef (preDefs[0].ref) <| mapError
           (orelseMergeErrors

src/Lean/Elab/PreDefinition/WF/Main.lean

@@ -47,13 +47,13 @@ private partial def addNonRecPreDefs (preDefs : Array PreDefinition) (preDefNonR
     trace[Elab.definition.wf] "{preDef.declName} := {value}"
     addNonRec { preDef with value } (applyAttrAfterCompilation := false)
 
-def wfRecursion (preDefs : Array PreDefinition) (wfStx? : Option Syntax) (decrTactic? : Option Syntax) : TermElabM Unit := do
+def wfRecursion (preDefs : Array PreDefinition) (wf? : Option TerminationWF) (decrTactic? : Option Syntax) : TermElabM Unit := do
   let unaryPreDef ← withoutModifyingEnv do
     for preDef in preDefs do
       addAsAxiom preDef
     let unaryPreDefs ← packDomain preDefs
     packMutual unaryPreDefs
-  let wfRel ← elabWFRel unaryPreDef wfStx?
+  let wfRel ← elabWFRel unaryPreDef wf?
   let preDefNonRec ← withoutModifyingEnv do
     addAsAxiom unaryPreDef
     mkFix unaryPreDef wfRel decrTactic?

src/Lean/Elab/PreDefinition/WF/Rel.lean

@@ -5,14 +5,15 @@ Authors: Leonardo de Moura
 -/
 import Lean.Elab.SyntheticMVars
 import Lean.Elab.PreDefinition.Basic
+import Lean.Elab.PreDefinition.WF.TerminationHint
 
 namespace Lean.Elab.WF
 open Meta
 open Term
 
-def elabWFRel (unaryPreDef : PreDefinition) (wfStx? : Option Syntax) : TermElabM Expr := do
-  if let some wfStx := wfStx? then
-    withDeclName unaryPreDef.declName do
+def elabWFRel (unaryPreDef : PreDefinition) (wf? : Option TerminationWF) : TermElabM Expr := do
+  match wf? with
+  | some (TerminationWF.core wfStx) => withDeclName unaryPreDef.declName do
       let α := unaryPreDef.type.bindingDomain!
       let u ← getLevel α
       let expectedType := mkApp (mkConst ``WellFoundedRelation [u]) α
@@ -20,7 +21,7 @@ def elabWFRel (unaryPreDef : PreDefinition) (wfStx? : Option Syntax) : TermElabM
       let pendingMVarIds ← getMVars wfRel
       discard <| logUnassignedUsingErrorInfos pendingMVarIds
       return wfRel
-  else
+  | none =>
     -- TODO: try to synthesize some default relation
     throwError "'termination_by' modifier missing"
 

src/Lean/Elab/PreDefinition/WF/TerminationHint.lean

@@ -71,4 +71,31 @@ def TerminationHint.ensureIsEmpty (t : TerminationHint) : MacroM Unit := do
   | TerminationHint.many m  => m.forM fun _ v => Macro.throwErrorAt v.ref "unused termination hint element"
   | _ => pure ()
 
+inductive TerminationBy where
+  | core (hint : TerminationHint)
+
+inductive TerminationWF where
+  | core (stx : Syntax)
+
+def expandTerminationBy (hint? : Option Syntax) (cliques : Array (Array Name)) : MacroM TerminationBy :=
+  if let some hint := hint? then
+    if hint.isOfKind ``Parser.Command.terminationByCore then
+      return TerminationBy.core (← expandTerminationHint hint? cliques)
+    else
+      withRef hint <| Macro.throwError "`termination_by` syntax is being modified/simplified. To use the old syntax, please use `termination_by'` instead"
+  else
+      return TerminationBy.core TerminationHint.none
+
+def TerminationBy.erase (t : TerminationBy) (clique : Array Name) : TerminationBy :=
+  match t with
+  | core hint => core (hint.erase clique)
+
+def TerminationBy.find? (t : TerminationBy) (clique : Array Name) : Option TerminationWF :=
+  match t with
+  | core hint => hint.find? clique |>.map fun v => TerminationWF.core v.value
+
+def TerminationBy.ensureIsEmpty (t : TerminationBy) : MacroM Unit :=
+  match t with
+  | core hint => hint.ensureIsEmpty
+
 end Lean.Elab.WF