From b37178d54706323839f313acebe6a3c0c8b3f55d Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Fri, 19 Aug 2022 18:40:44 -0700
Subject: [PATCH] perf: improve `toLCNFType`

---
 src/Lean/Compiler/LCNFTypes.lean | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/Lean/Compiler/LCNFTypes.lean b/src/Lean/Compiler/LCNFTypes.lean
index 3067fcbb62..57f05ff588 100644
--- a/src/Lean/Compiler/LCNFTypes.lean
+++ b/src/Lean/Compiler/LCNFTypes.lean
@@ -93,18 +93,25 @@ partial def toLCNFType (type : Expr) : MetaM Expr := do
         return b
       else
         return (Expr.lam n d (b.abstract #[x]) bi).eta
-  | .forallE n d b bi =>
-    withLocalDecl n bi d fun x => do
-      let borrowed := isMarkedBorrowed d
-      let mut d ← toLCNFType d
-      if borrowed then
-        d := markBorrowed d
-      let b ← toLCNFType (b.instantiate1 x)
-      return .forallE n d (b.abstract #[x]) bi
+  | .forallE .. => visitForall type #[]
   | .app ..  => type.withApp visitApp
   | .fvar .. => visitApp type #[]
   | _        => return anyTypeExpr
 where
+  visitForall (e : Expr) (xs : Array Expr) : MetaM Expr := do
+    match e with
+    | .forallE n d b bi =>
+      let d := d.instantiateRev xs
+      withLocalDecl n bi d fun x => do
+        let borrowed := isMarkedBorrowed d
+        let mut d := (← toLCNFType d).abstract xs
+        if borrowed then
+          d := markBorrowed d
+        return .forallE n d (← visitForall b (xs.push x)) bi
+    | _ =>
+      let e ← toLCNFType (e.instantiateRev xs)
+      return e.abstract xs
+
   visitApp (f : Expr) (args : Array Expr) := do
     let fNew ← match f with
       | .const declName us =>