feat: add WF.mkFix

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

@@ -0,0 +1,80 @@
+/-
+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.Match.Match
+import Lean.Elab.PreDefinition.Basic
+import Lean.Elab.PreDefinition.Structural.Basic
+
+namespace Lean.Elab.WF
+open Meta
+
+private def mkDecreasingProof (decreasingProp : Expr) : TermElabM Expr := do
+  let mvar ← mkFreshExprSyntheticOpaqueMVar decreasingProp
+  let mvarId := mvar.mvarId!
+  trace[Elab.definition.wf] "{MessageData.ofGoal mvarId}"
+  -- TODO: invoke tactic to close the goal
+  admit mvarId
+  instantiateMVars mvar
+
+private partial def replaceRecApps (recFnName : Name) (F : Expr) (e : Expr) : TermElabM Expr :=
+  let rec loop (F : Expr) (e : Expr) : TermElabM Expr := do
+    match e with
+    | Expr.lam n d b c =>
+      withLocalDecl n c.binderInfo (← loop F d) fun x => do
+        mkLambdaFVars #[x] (← loop F (b.instantiate1 x))
+    | Expr.forallE n d b c =>
+      withLocalDecl n c.binderInfo (← loop F d) fun x => do
+        mkForallFVars #[x] (← loop F (b.instantiate1 x))
+    | Expr.letE n type val body _ =>
+      withLetDecl n (← loop F type) (← loop F val) fun x => do
+        mkLetFVars #[x] (← loop F (body.instantiate1 x))
+    | Expr.mdata d e _   => return mkMData d (← loop F e)
+    | Expr.proj n i e _  => return mkProj n i (← loop F e)
+    | Expr.app _ _ _ =>
+      let processApp (e : Expr) : TermElabM Expr :=
+        e.withApp fun f args => do
+          if f.isConstOf recFnName && args.size == 1 then
+            let r := mkApp F args[0]
+            let decreasingProp := (← whnf (← inferType r)).bindingDomain!
+            return mkApp r (← mkDecreasingProof decreasingProp)
+          else
+            return mkAppN (← loop F f) (← args.mapM (loop F))
+      let matcherApp? ← matchMatcherApp? e
+      match matcherApp? with
+      | some matcherApp =>
+        if !Structural.recArgHasLooseBVarsAt recFnName 0 e then
+          processApp e
+        else
+          trace[Elab.definition.structural] "below before matcherApp.addArg: {F} : {← inferType F}"
+          let matcherApp ← mapError (matcherApp.addArg F) (fun msg => "failed to add functional argument to 'matcher' application" ++ indentD msg)
+          let altsNew ← (Array.zip matcherApp.alts matcherApp.altNumParams).mapM fun (alt, numParams) =>
+            lambdaTelescope alt fun xs altBody => do
+              unless xs.size >= numParams do
+                throwError "unexpected matcher application alternative{indentExpr alt}\nat application{indentExpr e}"
+              let FAlt := xs[numParams - 1]
+              mkLambdaFVars xs (← loop FAlt altBody)
+          pure { matcherApp with alts := altsNew }.toExpr
+      | none => processApp e
+    | e => Structural.ensureNoRecFn recFnName e
+  loop F e
+
+def mkFix (preDef : PreDefinition) (wfRel : Expr) : TermElabM PreDefinition := do
+  let wfFix ← forallBoundedTelescope preDef.type (some 1) fun x type => do
+    let x := x[0]
+    let α ← inferType x
+    let u ← getLevel α
+    let v ← getLevel type
+    let motive ← mkLambdaFVars #[x] type
+    let rel := mkProj ``WellFoundedRelation 0 wfRel
+    let wf  := mkProj ``WellFoundedRelation 1 wfRel
+    return mkApp4 (mkConst ``WellFounded.fix [u, v]) α motive rel wf
+  forallBoundedTelescope (← whnf (← inferType wfFix)).bindingDomain! (some 2) fun xs _ => do
+    let x   := xs[0]
+    let F   := xs[1]
+    let value ← replaceRecApps preDef.declName F (preDef.value.betaRev #[x])
+    let value := mkApp wfFix (← mkLambdaFVars #[x, F] value)
+    return { preDef with value }
+
+end Lean.Elab.WF

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

@@ -8,13 +8,14 @@ import Lean.Elab.PreDefinition.WF.TerminationBy
 import Lean.Elab.PreDefinition.WF.PackDomain
 import Lean.Elab.PreDefinition.WF.PackMutual
 import Lean.Elab.PreDefinition.WF.Rel
+import Lean.Elab.PreDefinition.WF.Fix
 
 namespace Lean.Elab
 open WF
 open Meta
 
 def wfRecursion (preDefs : Array PreDefinition) (wfStx? : Option Syntax) : TermElabM Unit := do
-  withoutModifyingEnv do
+  let unaryPreDef ← withoutModifyingEnv do
     for preDef in preDefs do
       addAsAxiom preDef
     let unaryPreDefs ← packDomain preDefs
@@ -23,11 +24,15 @@ def wfRecursion (preDefs : Array PreDefinition) (wfStx? : Option Syntax) : TermE
       trace[Elab.definition.wf] "{preDef.declName}, {preDef.levelParams}, {preDef.value}"
     let unaryPreDef ← packMutual unaryPreDefs
     trace[Elab.definition.wf] "{unaryPreDef.declName} := {unaryPreDef.value}"
-    check unaryPreDef.value -- TODO: remove
-    let wfRel ← elabWFRel unaryPreDef wfStx?
-    trace[Elab.definition.wf] "{wfRel}"
-  -- TODO
-  throwError "well-founded recursion has not been implemented yet"
+    return unaryPreDef
+  let wfRel ← elabWFRel unaryPreDef wfStx?
+  trace[Elab.definition.wf] "{wfRel}"
+  let preDefNonRec ← withoutModifyingEnv do
+    addAsAxiom unaryPreDef
+    mkFix unaryPreDef wfRel
+  addNonRec preDefNonRec
+  -- TODO: define preDefs
+  -- addAndCompilePartialRec preDefs
 
 builtin_initialize registerTraceClass `Elab.definition.wf