From 981cd44eaa72edf57f0ec60139ef5ae8e81b291c Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 9 Nov 2019 20:20:18 -0800
Subject: [PATCH] feat: add `lambdaTelescope`

---
 library/Init/Lean/Meta.lean | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/library/Init/Lean/Meta.lean b/library/Init/Lean/Meta.lean
index dd511be13b..9c34497736 100644
--- a/library/Init/Lean/Meta.lean
+++ b/library/Init/Lean/Meta.lean
@@ -388,6 +388,39 @@ do c? ← isClassQuick type;
    | LOption.some c => pure (some c)
    | LOption.undef  => isClassExpensive whnf type
 
+/-- Similar to `forallTelescopeAuxAux` but for lambda expressions. -/
+@[specialize] private partial def lambdaTelescopeAux {α}
+    (whnf             : Expr → MetaM Expr)
+    (k                : Array Expr → Expr → MetaM α)
+    : LocalContext → Array Expr → Nat → Expr → MetaM α
+| lctx, fvars, j, Expr.lam n bi d b => do
+  let d := d.instantiateRevRange j fvars.size fvars;
+  fvarId ← mkFreshId;
+  let lctx := lctx.mkLocalDecl fvarId n d bi;
+  let fvar := Expr.fvar fvarId;
+  lambdaTelescopeAux lctx (fvars.push fvar) j b
+| lctx, fvars, j, type =>
+  let type := type.instantiateRevRange j fvars.size fvars;
+  adaptReader (fun (ctx : Context) => { lctx := lctx, .. ctx }) $
+    withNewLocalInstances (isClassExpensive whnf) fvars j $ do
+      newType ← whnf type;
+      if newType.isForall then
+        lambdaTelescopeAux lctx fvars fvars.size newType
+      else
+        k fvars type
+
+/-- Similar to `forallTelescope` but for lambda expressions. -/
+@[specialize] private def lambdaTelescope {α}
+    (whnf             : Expr → MetaM Expr)
+    (type : Expr) (k : Array Expr → Expr → MetaM α) : MetaM α :=
+do newType ← whnf type;
+   if newType.isLambda then
+     savingCache $ do
+       lctx ← getLCtx;
+       lambdaTelescopeAux whnf k lctx #[] 0 newType
+   else do
+     k #[] type
+
 @[specialize] private def getForallResultType
     (whnf      : Expr → MetaM Expr)
     (fType : Expr) (args : Array Expr) : MetaM Expr :=