From 91a0ac8a12ccf3484b795083d9b71b7a3464493f Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 12 Jan 2022 16:05:18 -0800
Subject: [PATCH] feat: elaborate new `termination_by` syntax

---
 src/Lean/Elab/PreDefinition/WF/Main.lean |  2 +-
 src/Lean/Elab/PreDefinition/WF/Rel.lean  | 59 +++++++++++++++++++++---
 tests/lean/run/mutwf2.lean               | 15 ++++++
 3 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/src/Lean/Elab/PreDefinition/WF/Main.lean b/src/Lean/Elab/PreDefinition/WF/Main.lean
index b422f70e3a..e8690b30e5 100644
--- a/src/Lean/Elab/PreDefinition/WF/Main.lean
+++ b/src/Lean/Elab/PreDefinition/WF/Main.lean
@@ -53,7 +53,7 @@ def wfRecursion (preDefs : Array PreDefinition) (wf? : Option TerminationWF) (de
       addAsAxiom preDef
     let unaryPreDefs ← packDomain preDefs
     packMutual unaryPreDefs
-  let wfRel ← elabWFRel unaryPreDef wf?
+  let wfRel ← elabWFRel preDefs unaryPreDef wf?
   let preDefNonRec ← withoutModifyingEnv do
     addAsAxiom unaryPreDef
     mkFix unaryPreDef wfRel decrTactic?
diff --git a/src/Lean/Elab/PreDefinition/WF/Rel.lean b/src/Lean/Elab/PreDefinition/WF/Rel.lean
index 1415e62129..8d9cf66bd3 100644
--- a/src/Lean/Elab/PreDefinition/WF/Rel.lean
+++ b/src/Lean/Elab/PreDefinition/WF/Rel.lean
@@ -3,6 +3,9 @@ Copyright (c) 2021 Microsoft Corporation. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
+import Lean.Meta.Tactic.Apply
+import Lean.Meta.Tactic.Rename
+import Lean.Meta.Tactic.Intro
 import Lean.Elab.SyntheticMVars
 import Lean.Elab.PreDefinition.Basic
 import Lean.Elab.PreDefinition.WF.TerminationHint
@@ -11,18 +14,62 @@ namespace Lean.Elab.WF
 open Meta
 open Term
 
-def elabWFRel (unaryPreDef : PreDefinition) (wf? : Option TerminationWF) : TermElabM Expr := do
+private def getRefFromElems (elems : Array TerminationByElement) : Syntax := Id.run do
+  for elem in elems do
+    if !elem.implicit then
+      return elem.ref
+  return elems[0].ref
+
+private partial def unpackMutual (preDefs : Array PreDefinition) (mvarId : MVarId) (fvarId : FVarId) : TermElabM (Array (FVarId × MVarId)) := do
+  let rec go (i : Nat) (mvarId : MVarId) (fvarId : FVarId) (result : Array (FVarId × MVarId)) : TermElabM (Array (FVarId × MVarId)) := do
+    if i < preDefs.size - 1 then
+      let #[s₁, s₂] ← cases mvarId fvarId | unreachable!
+      go (i + 1) s₂.mvarId s₂.fields[0].fvarId! (result.push (s₁.fields[0].fvarId!, s₁.mvarId))
+    else
+      return result.push (fvarId, mvarId)
+  go 0 mvarId fvarId #[]
+
+private partial def unpackUnary (preDef : PreDefinition) (mvarId : MVarId) (fvarId : FVarId) (element : TerminationByElement) : TermElabM MVarId := do
+  let varNames ← lambdaTelescope preDef.value fun xs body => do
+    let mut varNames ← xs.mapM fun x => return (← getLocalDecl x.fvarId!).userName
+    if element.vars.size > varNames.size then
+      throwErrorAt element.vars[varNames.size] "too many variable names"
+    for i in [:element.vars.size] do
+      let varStx := element.vars[i]
+      if varStx.isIdent then
+        varNames := varNames.set! (varNames.size - element.vars.size + i) varStx.getId
+    return varNames
+  let rec go (i : Nat) (mvarId : MVarId) (fvarId : FVarId) : TermElabM MVarId := do
+    if i < varNames.size - 1 then
+      let #[s] ← cases mvarId fvarId #[{ varNames := [varNames[i]] }] | unreachable!
+      go (i+1) s.mvarId s.fields[1].fvarId!
+    else
+      rename mvarId fvarId varNames.back
+  go 0 mvarId fvarId
+
+def elabWFRel (preDefs : Array PreDefinition) (unaryPreDef : PreDefinition) (wf? : Option TerminationWF) : TermElabM Expr := do
+  let α := unaryPreDef.type.bindingDomain!
+  let u ← getLevel α
+  let expectedType := mkApp (mkConst ``WellFoundedRelation [u]) α
   match wf? with
   | some (TerminationWF.core wfStx) => withDeclName unaryPreDef.declName do
-      let α := unaryPreDef.type.bindingDomain!
-      let u ← getLevel α
-      let expectedType := mkApp (mkConst ``WellFoundedRelation [u]) α
       let wfRel ← instantiateMVars (← withSynthesize <| elabTermEnsuringType wfStx expectedType)
       let pendingMVarIds ← getMVars wfRel
       discard <| logUnassignedUsingErrorInfos pendingMVarIds
       return wfRel
-  | some (TerminationWF.ext clique) =>
-    throwError "`termination_by` syntax is being modified/simplified. To use the old syntax, please use `termination_by'` instead"
+  | some (TerminationWF.ext elements) => withRef (getRefFromElems elements) do
+    let mainMVarId := (← mkFreshExprSyntheticOpaqueMVar expectedType).mvarId!
+    let [fMVarId, wfRelMVarId, _] ← apply mainMVarId (← mkConstWithFreshMVarLevels ``invImage) | throwError "failed to apply 'invImage'"
+    let (d, fMVarId) ← intro1 fMVarId
+    let subgoals ← unpackMutual preDefs fMVarId d
+    for (d, mvarId) in subgoals, element in elements, preDef in preDefs do
+      let mvarId ← unpackUnary preDef mvarId d element
+      withMVarContext mvarId do
+        let value ← Term.withSynthesize <| elabTermEnsuringType element.body (← getMVarType mvarId)
+        assignExprMVar mvarId value
+    let wfRelVal ← synthInstance (← inferType (mkMVar wfRelMVarId))
+    assignExprMVar wfRelMVarId wfRelVal
+    instantiateMVars (mkMVar mainMVarId)
   | none =>
     -- TODO: try to synthesize some default relation
     throwError "'termination_by' modifier missing"
diff --git a/tests/lean/run/mutwf2.lean b/tests/lean/run/mutwf2.lean
index d4695b017f..2488b6adfc 100644
--- a/tests/lean/run/mutwf2.lean
+++ b/tests/lean/run/mutwf2.lean
@@ -1,3 +1,4 @@
+namespace Ex1
 mutual
   def isEven : Nat → Bool
     | 0 => true
@@ -13,3 +14,17 @@ termination_by' measure fun
 #print isEven
 #print isOdd
 #print isEven._mutual
+end Ex1
+
+
+namespace Ex2
+mutual
+  def isEven : Nat → Bool
+    | 0 => true
+    | n+1 => isOdd n
+  def isOdd : Nat → Bool
+    | 0 => false
+    | n+1 => isEven n
+end
+termination_by _ n => n
+end Ex2