improve mkLetFun, add instantiateForallWithParamInfos

src/Lean/Meta/AppBuilder.lean

@@ -27,13 +27,17 @@ def mkExpectedTypeHint (e : Expr) (expectedType : Expr) : MetaM Expr := do
 /--
 `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`.
-NB: `x` must not be a let-bound free variable.
 -/
 def mkLetFun (x : Expr) (v : Expr) (e : Expr) : MetaM Expr := do
-  let f ← mkLambdaFVars #[x] e
+  -- If `x` is an `ldecl`, then the result of `mkLambdaFVars` is a let expression.
+  let ensureLambda : Expr → Expr
+    | .letE n t _ b _ => .lam n t b .default
+    | e@(.lam ..)     => e
+    | _               => unreachable!
+  let f ← ensureLambda <$> mkLambdaFVars (usedLetOnly := false) #[x] e
   let ety ← inferType e
   let α ← inferType x
-  let β ← mkLambdaFVars #[x] ety
+  let β ← ensureLambda <$> mkLambdaFVars (usedLetOnly := false) #[x] ety
   let u1 ← getLevel α
   let u2 ← getLevel ety
   return mkAppN (.const ``letFun [u1, u2]) #[α, β, v, f]

src/Lean/Meta/Basic.lean

@@ -1896,6 +1896,76 @@ private partial def instantiateLambdaAux (ps : Array Expr) (i : Nat) (e : Expr)
 def instantiateLambda (e : Expr) (ps : Array Expr) : MetaM Expr :=
   instantiateLambdaAux ps 0 e
 
+/--
+A simpler version of `ParamInfo` for information about the parameter of a forall or lambda.
+-/
+structure ExprParamInfo where
+  /-- The name of the parameter. -/
+  name : Name
+  /-- The type of the parameter. -/
+  type : Expr
+  /-- The binder annotation for the parameter. -/
+  binderInfo     : BinderInfo := BinderInfo.default
+  deriving Inhabited
+
+/--
+Given `e` of the form `∀ (p₁ : P₁) … (p₁ : P₁), B[p_1,…,p_n]` and `arg₁ : P₁, …, argₙ : Pₙ`, returns
+* the names `p₁, …, pₙ`,
+* the binder infos,
+* the binder types `P₁, P₂[arg₁], …, P[arg₁,…,argₙ₋₁]`, and
+* the type `B[arg₁,…,argₙ]`.
+
+It uses `whnf` to reduce `e` if it is not a forall.
+
+See also `Lean.Meta.instantiateForall`.
+-/
+def instantiateForallWithParamInfos (e : Expr) (args : Array Expr) (cleanupAnnotations : Bool := false) :
+    MetaM (Array ExprParamInfo × Expr) := do
+  let mut e := e
+  let mut res := Array.mkEmpty args.size
+  let mut j := 0
+  for i in [0:args.size] do
+    unless e.isForall do
+      e ← whnf (e.instantiateRevRange j i args)
+      j := i
+    match e with
+    | .forallE name type e' binderInfo =>
+      let type := type.instantiateRevRange j i args
+      let type := if cleanupAnnotations then type.cleanupAnnotations else type
+      res := res.push { name, type, binderInfo }
+      e := e'
+    | _ => throwError "invalid `instantiateForallWithParams`, too many parameters{indentExpr e}"
+  return (res, e)
+
+/--
+Given `e` of the form `fun (p₁ : P₁) … (p₁ : P₁) => t[p_1,…,p_n]` and `arg₁ : P₁, …, argₙ : Pₙ`, returns
+* the names `p₁, …, pₙ`,
+* the binder infos,
+* the binder types `P₁, P₂[arg₁], …, P[arg₁,…,argₙ₋₁]`, and
+* the term `t[arg₁,…,argₙ]`.
+
+It uses `whnf` to reduce `e` if it is not a lambda.
+
+See also `Lean.Meta.instantiateLambda`.
+-/
+def instantiateLambdaWithParamInfos (e : Expr) (args : Array Expr) (cleanupAnnotations : Bool := false) :
+    MetaM (Array ExprParamInfo × Expr) := do
+  let mut e := e
+  let mut res := Array.mkEmpty args.size
+  let mut j := 0
+  for i in [0:args.size] do
+    unless e.isLambda do
+      e ← whnf (e.instantiateRevRange j i args)
+      j := i
+    match e with
+    | .forallE name type e' binderInfo =>
+      let type := type.instantiateRevRange j i args
+      let type := if cleanupAnnotations then type.cleanupAnnotations else type
+      res := res.push { name, type, binderInfo }
+      e := e'
+    | _ => throwError "invalid `instantiateForallWithParams`, too many parameters{indentExpr e}"
+  return (res, e)
+
 /-- Pretty-print the given expression. -/
 def ppExprWithInfos (e : Expr) : MetaM FormatWithInfos := do
   let ctxCore  ← readThe Core.Context

src/Lean/Meta/Tactic/Lets.lean

@@ -19,6 +19,7 @@ namespace Lean.Meta
 
 Extracting `let`s means to locate `let`/`letFun`s in a term and to extract them
 from the term, extending the local context with new declarations in the process.
+A related process is lifting `lets`, which means to move `let`/`letFun`s toward the root of a term.
 -/
 
 namespace ExtractLets
@@ -40,7 +41,7 @@ structure State where
   valueMap : ExprStructMap FVarId := {}
   deriving Inhabited
 
--- The `Bool` is whether this is a top-level let.
+-- The `Bool` in the cache key is whether we are looking at a "top-level" expression.
 abbrev M := ReaderT ExtractLetsConfig <| MonadCacheT (Bool × ExprStructEq) Expr <| StateRefT State MetaM
 
 /-- Returns `true` if `nextName?` would return a name. -/
@@ -88,8 +89,7 @@ def isExtractableLet (fvars : List Expr) (n : Name) (t v : Expr) : M (Bool × Na
   if (← hasNextName) && extractable fvars t && extractable fvars v then
     if let some n ← nextNameForBinderName? n then
       return (true, n)
-  -- In lift mode, we temporarily extract non-extractable lets, but we do not make use of `givenNames` for them,
-  -- which would count against the number of extracted lets.
+  -- In lift mode, we temporarily extract non-extractable lets, but we do not make use of `givenNames` for them.
   -- These will be flushed as let/letFun expressions, and we wish to preserve the original binder name.
   if (← read).lift then
     return (true, n)
@@ -136,7 +136,7 @@ Does not require that any of the declarations are in context.
 Assumes that `e` contains no metavariables with local contexts that contain any of these metavariables
 (the extraction procedure creates no new metavariables, so this is the case).
 
-This should *not* be used when closing lets for new goals, since
+This should *not* be used when closing lets for new goal metavariables, since
 1. The goal contains the decls in its local context, violating the assumption.
 2. We need to use true `let`s in that case, since tactics may zeta-delta reduce these declarations.
 -/
@@ -146,16 +146,7 @@ def mkLetDecls (decls : Array LocalDecl') (e : Expr) : MetaM Expr := do
       if isLet then
         return .letE decl.userName decl.type decl.value (e.abstract #[decl.toExpr]) false
       else
-        let n := decl.userName
-        let x := decl.toExpr
-        let α := decl.type
-        let ety ← inferType e
-        let β := Expr.lam n α (ety.abstract #[x]) .default
-        let v := decl.value
-        let f := Expr.lam n α (e.abstract #[x]) .default
-        let u1 ← getLevel α
-        let u2 ← getLevel ety
-        return mkAppN (.const ``letFun [u1, u2]) #[α, β, v, f]
+        mkLetFun decl.toExpr decl.value e
 
 /--
 Makes sure the declaration for `fvarId` is marked with `isLet := true`.
@@ -177,7 +168,7 @@ Used for `merge` feature.
 def withDeclInContext (fvarId : FVarId) (k : M α) : M α := do
   let decls := (← get).decls
   if (← getLCtx).contains fvarId then
-    -- Either pre-existing or already added.
+    -- Is either pre-existing or already added.
     k
   else if let some idx := decls.findIdx? (·.decl.fvarId == fvarId) then
     withEnsuringDeclsInContext decls[0:idx+1] k
@@ -186,6 +177,7 @@ def withDeclInContext (fvarId : FVarId) (k : M α) : M α := do
 
 /--
 Initializes the `valueMap` with all the local definitions that aren't implementation details.
+Used for `merge` feature when `useContext` is enabled.
 -/
 def initializeValueMap : M Unit := do
   let lctx ← getLCtx
@@ -298,17 +290,10 @@ where
       if cfg.implicits then
         return mkAppN f' (← args.mapM (extractCore fvars))
       else
+        let (paramInfos, _) ← instantiateForallWithParamInfos (← inferType f) args
         let mut args := args
-        let mut fty ← inferType f
-        let mut j := 0
         for i in [0:args.size] do
-          unless fty.isForall do
-            fty ← withTransparency .all <| whnf (fty.instantiateRevRange j i args)
-            j := i
-          let .forallE _ _ fty' bi := fty
-            | throwError "Lean.Meta.ExtractLets.extractCore: expecting function, type is{indentD fty}"
-          fty := fty'
-          if bi.isExplicit then
+          if paramInfos[i]!.binderInfo.isExplicit then
             args := args.set! i (← extractCore fvars args[i]!)
         return mkAppN f' args
 

tests/lean/run/extract_lets.lean

@@ -718,6 +718,23 @@ example : ∀ n : Nat, (let x := n; x) = (let x' := n; x') := by
   intros
   rfl
 
+/-!
+Without merging
+-/
+/--
+info: ⊢ ∀ (n : Nat),
+    let_fun x := n;
+    let_fun x' := n;
+    x = x'
+-/
+#guard_msgs in
+example : ∀ n : Nat, (have x := n; x) = (have x' := n; x') := by
+  fail_if_success extract_lets
+  extract_lets +lift -merge
+  trace_state
+  intros
+  rfl
+
 /-!
 Make sure `+lift` doesn't lift things that transitively depend on a binder.
 -/