From 011521013d85265d037849ef5ee14036415b2103 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Fri, 23 Sep 2022 16:45:04 -0700
Subject: [PATCH] feat: use phase at `inferConstType`, save specialization

---
 src/Lean/Compiler/LCNF/InferType.lean  |  7 +++++--
 src/Lean/Compiler/LCNF/Specialize.lean |  4 +++-
 tests/lean/run/lcnfCheckIssue.lean     | 14 ++++++++++++++
 3 files changed, 22 insertions(+), 3 deletions(-)
 create mode 100644 tests/lean/run/lcnfCheckIssue.lean

diff --git a/src/Lean/Compiler/LCNF/InferType.lean b/src/Lean/Compiler/LCNF/InferType.lean
index e3c64673ba..b6175c770e 100644
--- a/src/Lean/Compiler/LCNF/InferType.lean
+++ b/src/Lean/Compiler/LCNF/InferType.lean
@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 import Lean.Compiler.LCNF.CompilerM
 import Lean.Compiler.LCNF.Types
+import Lean.Compiler.LCNF.PhaseExt
 
 namespace Lean.Compiler.LCNF
 /-! # Type inference for LCNF -/
@@ -168,11 +169,13 @@ def mkForallParams (params : Array Param) (type : Expr) : InferTypeM Expr :=
   withReader (fun lctx => lctx.mkLocalDecl fvarId binderName type binderInfo) do
     k (.fvar fvarId)
 
-def inferConstType (declName : Name) (us : List Level) : CoreM Expr :=
+def inferConstType (declName : Name) (us : List Level) : CompilerM Expr := do
   if declName == ``lcAny || declName == ``lcErased then
     return anyTypeExpr
+  else if let some decl ← getDecl? (← getPhase) declName then
+    return decl.instantiateTypeLevelParams us
   else
-    instantiateLCNFTypeLevelParams declName us
+    instantiateLCNFTypeLevelParams declName us -- TODO: delete after we compile decls at definition time
 
 mutual
 
diff --git a/src/Lean/Compiler/LCNF/Specialize.lean b/src/Lean/Compiler/LCNF/Specialize.lean
index 1260882170..a7756b105e 100644
--- a/src/Lean/Compiler/LCNF/Specialize.lean
+++ b/src/Lean/Compiler/LCNF/Specialize.lean
@@ -10,6 +10,7 @@ import Lean.Compiler.LCNF.SpecInfo
 import Lean.Compiler.LCNF.PrettyPrinter
 import Lean.Compiler.LCNF.ToExpr
 import Lean.Compiler.LCNF.Level
+import Lean.Compiler.LCNF.PhaseExt
 
 namespace Lean.Compiler.LCNF
 namespace Specialize
@@ -360,7 +361,8 @@ mutual
       trace[Compiler.specialize.step] "new: {specDecl.name}"
       cacheSpec key specDecl.name
       let specDecl ← specDecl.etaExpand
-      saveLCNFType specDecl.name specDecl.type
+      specDecl.saveBase
+      saveLCNFType specDecl.name specDecl.type -- TODO: delete after we start compiling at decl time
       let specDecl ← specDecl.simp {} -- TODO: `simp` config
       let specDecl ← specDecl.simp { etaPoly := true, inlinePartial := true, implementedBy := true }
       let value ← withReader (fun _ => { declName := specDecl.name }) do
diff --git a/tests/lean/run/lcnfCheckIssue.lean b/tests/lean/run/lcnfCheckIssue.lean
new file mode 100644
index 0000000000..d963d7096e
--- /dev/null
+++ b/tests/lean/run/lcnfCheckIssue.lean
@@ -0,0 +1,14 @@
+import Lean
+
+abbrev Sequence (α : Type) := List α
+
+def bigop (init : β) (seq : Sequence α) (op : β → β → β) (f : α → Bool × β) : β := Id.run do
+  let mut result := init
+  for a in seq do
+    let (ok, b) := f a
+    if ok then
+      result := op result b
+  return result
+
+set_option trace.Compiler.result true
+#eval Lean.Compiler.compile #[``bigop]