feat: encode let_fun using a letFun function (#2973)

Switches from encoding `let_fun` using an annotated `(fun x : t => b) v`
expression to a function application `letFun v (fun x : t => b)`.

---------

Co-authored-by: Sebastian Ullrich <sebasti@nullri.ch>

src/Init/Prelude.lean

@@ -66,6 +66,19 @@ example (b : Bool) : Function.const Bool 10 b = 10 :=
 @[inline] def Function.const {α : Sort u} (β : Sort v) (a : α) : β → α :=
   fun _ => a
 
+/--
+The encoding of `let_fun x := v; b` is `letFun v (fun x => b)`.
+This is equal to `(fun x => b) v`, so the value of `x` is not accessible to `b`.
+This is in contrast to `let x := v; b`, where the value of `x` is accessible to `b`.
+
+There is special support for `letFun`.
+Both WHNF and `simp` are aware of `letFun` and can reduce it when zeta reduction is enabled,
+despite the fact it is marked `irreducible`.
+For metaprogramming, the function `Lean.Expr.letFun?` can be used to recognize a `let_fun` expression
+to extract its parts as if it were a `let` expression.
+-/
+@[irreducible] def letFun {α : Sort u} {β : α → Sort v} (v : α) (f : (x : α) → β x) : β v := f v
+
 set_option checkBinderAnnotations false in
 /--
 `inferInstance` synthesizes a value of any target type by typeclass

src/Lean/Compiler/LCNF/ToLCNF.lean

@@ -658,7 +658,9 @@ where
       visit (f.beta e.getAppArgs)
 
   visitApp (e : Expr) : M Arg := do
-    if let .const declName _ := e.getAppFn then
+    if let some (args, n, t, v, b) := e.letFunAppArgs? then
+      visitCore <| mkAppN (.letE n t v b (nonDep := true)) args
+    else if let .const declName _ := e.getAppFn then
       if declName == ``Quot.lift then
         visitQuotLift e
       else if declName == ``Quot.mk then
@@ -725,11 +727,8 @@ where
       pushElement (.fun funDecl)
       return .fvar funDecl.fvarId
 
-  visitMData (mdata : MData) (e : Expr) : M Arg := do
-    if let some (.app (.lam n t b ..) v) := letFunAnnotation? (.mdata mdata e) then
-      visitLet (.letE n t v b (nonDep := true)) #[]
-    else
-      visit e
+  visitMData (_mdata : MData) (e : Expr) : M Arg := do
+    visit e
 
   visitProj (s : Name) (i : Nat) (e : Expr) : M Arg := do
     match (← visit e) with

src/Lean/Elab/Binders.lean

@@ -668,12 +668,11 @@ def elabLetDeclAux (id : Syntax) (binders : Array Syntax) (typeStx : Syntax) (va
       let body ← instantiateMVars body
       mkLetFVars #[x] body (usedLetOnly := usedLetOnly)
   else
-    let f ← withLocalDecl id.getId (kind := kind) .default type fun x => do
+    withLocalDecl id.getId (kind := kind) .default type fun x => do
       addLocalVarInfo id x
       let body ← elabTermEnsuringType body expectedType?
       let body ← instantiateMVars body
-      mkLambdaFVars #[x] body (usedLetOnly := false)
-    pure <| mkLetFunAnnotation (mkApp f val)
+      mkLetFun x val body
   if elabBodyFirst then
     forallBoundedTelescope type binders.size fun xs type => do
       -- the original `fvars` from above are gone, so add back info manually

src/Lean/Elab/PreDefinition/Eqns.lean

@@ -21,16 +21,25 @@ structure EqnInfoCore where
   value       : Expr
   deriving Inhabited
 
-partial def expand : Expr → Expr
-  | Expr.letE _ _ v b _ => expand (b.instantiate1 v)
-  | Expr.mdata _ b      => expand b
-  | e => e
+/--
+Zeta reduces `let` and `let_fun` while consuming metadata.
+Returns true if progress is made.
+-/
+partial def expand (progress : Bool) (e : Expr) : Bool × Expr :=
+  match e with
+  | Expr.letE _ _ v b _ => expand true (b.instantiate1 v)
+  | Expr.mdata _ b      => expand true b
+  | e =>
+    if let some (_, _, v, b) := e.letFun? then
+      expand true (b.instantiate1 v)
+    else
+      (progress, e)
 
 def expandRHS? (mvarId : MVarId) : MetaM (Option MVarId) := do
   let target ← mvarId.getType'
   let some (_, lhs, rhs) := target.eq? | return none
-  unless rhs.isLet || rhs.isMData do return none
-  return some (← mvarId.replaceTargetDefEq (← mkEq lhs (expand rhs)))
+  let (true, rhs') := expand false rhs | return none
+  return some (← mvarId.replaceTargetDefEq (← mkEq lhs rhs'))
 
 def funext? (mvarId : MVarId) : MetaM (Option MVarId) := do
   let target ← mvarId.getType'

src/Lean/Expr.lean

@@ -1653,6 +1653,47 @@ def setAppPPExplicitForExposingMVars (e : Expr) : Expr :=
     mkAppN f args |>.setPPExplicit true
   | _      => e
 
+/--
+Returns true if `e` is a `let_fun` expression, which is an expression of the form `letFun v f`.
+Ideally `f` is a lambda, but we do not require that here.
+Warning: if the `let_fun` is applied to additional arguments (such as in `(let_fun f := id; id) 1`), this function returns `false`.
+-/
+def isLetFun (e : Expr) : Bool := e.isAppOfArity ``letFun 4
+
+/--
+Recognizes a `let_fun` expression.
+For `let_fun n : t := v; b`, returns `some (n, t, v, b)`, which are the first four arguments to `Lean.Expr.letE`.
+Warning: if the `let_fun` is applied to additional arguments (such as in `(let_fun f := id; id) 1`), this function returns `none`.
+
+`let_fun` expressions are encoded as `letFun v (fun (n : t) => b)`.
+They can be created using `Lean.Meta.mkLetFun`.
+
+If in the encoding of `let_fun` the last argument to `letFun` is eta reduced, this returns `Name.anonymous` for the binder name.
+-/
+def letFun? (e : Expr) : Option (Name × Expr × Expr × Expr) :=
+  match e with
+  | .app (.app (.app (.app (.const ``letFun _) t) _β) v) f =>
+    match f with
+    | .lam n _ b _ => some (n, t, v, b)
+    | _ => some (.anonymous, t, v, .app f (.bvar 0))
+  | _ => none
+
+/--
+Like `Lean.Expr.letFun?`, but handles the case when the `let_fun` expression is possibly applied to additional arguments.
+Returns those arguments in addition to the values returned by `letFun?`.
+-/
+def letFunAppArgs? (e : Expr) : Option (Array Expr × Name × Expr × Expr × Expr) := do
+  guard <| 4 ≤ e.getAppNumArgs
+  guard <| e.isAppOf ``letFun
+  let args := e.getAppArgs
+  let t := args[0]!
+  let v := args[2]!
+  let f := args[3]!
+  let rest := args.extract 4 args.size
+  match f with
+  | .lam n _ b _ => some (rest, n, t, v, b)
+  | _ => some (rest, .anonymous, t, v, .app f (.bvar 0))
+
 end Expr
 
 /--
@@ -1670,28 +1711,6 @@ def annotation? (kind : Name) (e : Expr) : Option Expr :=
   | .mdata d b => if d.size == 1 && d.getBool kind false then some b else none
   | _          => none
 
-/--
-Annotate `e` with the `let_fun` annotation. This annotation is used as hint for the delaborator.
-If `e` is of the form `(fun x : t => b) v`, then `mkLetFunAnnotation e` is delaborated at
-`let_fun x : t := v; b`
--/
-def mkLetFunAnnotation (e : Expr) : Expr :=
-  mkAnnotation `let_fun e
-
-/--
-Return `some e'` if `e = mkLetFunAnnotation e'`
--/
-def letFunAnnotation? (e : Expr) : Option Expr :=
-  annotation? `let_fun e
-
-/--
-Return true if `e = mkLetFunAnnotation e'`, and `e'` is of the form `(fun x : t => b) v`
--/
-def isLetFun (e : Expr) : Bool :=
-  match letFunAnnotation? e with
-  | none   => false
-  | some e => e.isApp && e.appFn!.isLambda
-
 /--
 Auxiliary annotation used to mark terms marked with the "inaccessible" annotation `.(t)` and
 `_` in patterns.

src/Lean/Meta/AppBuilder.lean

@@ -24,6 +24,19 @@ def mkExpectedTypeHint (e : Expr) (expectedType : Expr) : MetaM Expr := do
   let u ← getLevel expectedType
   return mkApp2 (mkConst ``id [u]) expectedType e
 
+/--
+`mkLetFun x v e` creates the encoding for the `let_fun x := v; e` expression.
+The expression `x` can either be a free variable or a metavariable, and the function suitably abstracts `x` in `e`.
+-/
+def mkLetFun (x : Expr) (v : Expr) (e : Expr) : MetaM Expr := do
+  let f ← mkLambdaFVars #[x] e
+  let ety ← inferType e
+  let α ← inferType x
+  let β ← mkLambdaFVars #[x] ety
+  let u1 ← getLevel α
+  let u2 ← getLevel ety
+  return mkAppN (.const ``letFun [u1, u2]) #[α, β, v, f]
+
 /-- Return `a = b`. -/
 def mkEq (a b : Expr) : MetaM Expr := do
   let aType ← inferType a

src/Lean/Meta/ExprDefEq.lean

@@ -1491,10 +1491,16 @@ private def isDefEqMVarSelf (mvar : Expr) (args₁ args₂ : Array Expr) : MetaM
     else
       pure false
 
-/-- Remove unnecessary let-decls -/
-private def consumeLet : Expr → Expr
+/--
+Removes unnecessary let-decls (both true `let`s and `let_fun`s).
+-/
+private partial def consumeLet : Expr → Expr
   | e@(Expr.letE _ _ _ b _) => if b.hasLooseBVars then e else consumeLet b
-  | e                       => e
+  | e =>
+    if let some (_, _, _, b) := e.letFun? then
+      if b.hasLooseBVars then e else consumeLet b
+    else
+      e
 
 mutual
 

src/Lean/Meta/Tactic/Simp/Main.lean

@@ -250,6 +250,9 @@ private partial def reduce (e : Expr) : SimpM Expr := withIncRecDepth do
     match (← reduceRecMatcher? e) with
     | some e => return (← reduce e)
     | none   => pure ()
+  if cfg.zeta then
+    if let some (args, _, _, v, b) := e.letFunAppArgs? then
+      return (← reduce <| mkAppN (b.instantiate1 v) args)
   match (← unfold? e) with
   | some e' =>
     trace[Meta.Tactic.simp.rewrite] "unfold {mkConst e.getAppFn.constName!}, {e} ==> {e'}"

src/Lean/Meta/WHNF.lean

@@ -560,6 +560,10 @@ where
       | .const ..  => pure e
       | .letE _ _ v b _ => if config.zeta then go <| b.instantiate1 v else return e
       | .app f ..       =>
+        if config.zeta then
+          if let some (args, _, _, v, b) := e.letFunAppArgs? then
+            -- When zeta reducing enabled, always reduce `letFun` no matter the current reducibility level
+            return (← go <| mkAppN (b.instantiate1 v) args)
         let f := f.getAppFn
         let f' ← go f
         if config.beta && f'.isLambda then

src/Lean/PrettyPrinter/Delaborator/Basic.lean

@@ -277,17 +277,6 @@ end Delaborator
 open SubExpr (Pos PosMap)
 open Delaborator (OptionsPerPos topDownAnalyze)
 
-/-- Custom version of `Lean.Core.betaReduce` to beta reduce expressions for the `pp.beta` option.
-We do not want to beta reduce the application in `let_fun` annotations. -/
-private partial def betaReduce' (e : Expr) : CoreM Expr :=
-  Core.transform e (pre := fun e => do
-    if isLetFun e then
-      return .done <| e.updateMData! (.app (← betaReduce' e.mdataExpr!.appFn!) (← betaReduce' e.mdataExpr!.appArg!))
-    else if e.isHeadBetaTarget then
-      return .visit e.headBeta
-    else
-      return .continue)
-
 def delabCore (e : Expr) (optionsPerPos : OptionsPerPos := {}) (delab := Delaborator.delab) : MetaM (Term × PosMap Elab.Info) := do
   /- Using `erasePatternAnnotations` here is a bit hackish, but we do it
      `Expr.mdata` affects the delaborator. TODO: should we fix that? -/
@@ -302,7 +291,7 @@ def delabCore (e : Expr) (optionsPerPos : OptionsPerPos := {}) (delab := Delabor
     catch _ => pure ()
   withOptions (fun _ => opts) do
     let e ← if getPPInstantiateMVars opts then instantiateMVars e else pure e
-    let e ← if getPPBeta opts then betaReduce' e else pure e
+    let e ← if getPPBeta opts then Core.betaReduce e else pure e
     let optionsPerPos ←
       if !getPPAll opts && getPPAnalyze opts && optionsPerPos.isEmpty then
         topDownAnalyze e

src/Lean/PrettyPrinter/Delaborator/Builtins.lean

@@ -395,19 +395,34 @@ def delabAppMatch : Delab := whenPPOption getPPNotation <| whenPPOption getPPMat
     return Syntax.mkApp stx st.moreArgs
 
 /--
-  Delaborate applications of the form `(fun x => b) v` as `let_fun x := v; b`
+Annotation key to use in hack for overapplied `let_fun` notation.
 -/
-def delabLetFun : Delab := do
-  let stxV ← withAppArg delab
-  withAppFn do
-    let Expr.lam n _ b _ ← getExpr | unreachable!
-    let n ← getUnusedName n b
-    let stxB ← withBindingBody n delab
-    if ← getPPOption getPPLetVarTypes <||> getPPOption getPPAnalysisLetVarType then
-      let stxT ← withBindingDomain delab
-      `(let_fun $(mkIdent n) : $stxT := $stxV; $stxB)
-    else
-      `(let_fun $(mkIdent n) := $stxV; $stxB)
+def delabLetFunKey : Name := `_delabLetFun
+
+/--
+Delaborates applications of the form `letFun v (fun x => b)` as `let_fun x := v; b`.
+-/
+@[builtin_delab app.letFun]
+def delabLetFun : Delab := whenPPOption getPPNotation do
+  let e ← getExpr
+  let nargs := e.getAppNumArgs
+  if 4 < nargs then
+    -- It's overapplied. Hack: insert metadata around the first four arguments and delaborate again.
+    -- This will cause the app delaborator to delaborate `(letFun v f) x1 ... xn` with `letFun v f` as the function.
+    let args := e.getAppArgs
+    let f := mkAppN e.getAppFn (args.extract 0 4)
+    let e' := mkAppN (mkAnnotation delabLetFunKey f) (args.extract 4 args.size)
+    return (← withTheReader SubExpr ({· with expr := e'}) delab)
+  guard <| e.getAppNumArgs == 4
+  let Expr.lam n _ b _ := e.appArg! | failure
+  let n ← getUnusedName n b
+  let stxV ← withAppFn <| withAppArg delab
+  let stxB ← withAppArg <| withBindingBody n delab
+  if ← getPPOption getPPLetVarTypes <||> getPPOption getPPAnalysisLetVarType then
+    let stxT ← SubExpr.withNaryArg 0 delab
+    `(let_fun $(mkIdent n) : $stxT := $stxV; $stxB)
+  else
+    `(let_fun $(mkIdent n) := $stxV; $stxB)
 
 @[builtin_delab mdata]
 def delabMData : Delab := do
@@ -417,8 +432,6 @@ def delabMData : Delab := do
       `(.($s)) -- We only include the inaccessible annotation when we are delaborating patterns
     else
       return s
-  else if isLetFun (← getExpr) && getPPNotation (← getOptions) then
-    withMDataExpr <| delabLetFun
   else if let some _ := isLHSGoal? (← getExpr) then
     withMDataExpr <| withAppFn <| withAppArg <| delab
   else