From 126ad49401c75b80b18da682355331de28c8bc08 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sun, 14 Aug 2022 08:30:24 -0700
Subject: [PATCH] feat: add stage1 extension for storing LCNF declarations

---
 src/Lean/Compiler.lean        |  1 +
 src/Lean/Compiler/Decl.lean   |  4 ++++
 src/Lean/Compiler/Main.lean   | 17 +++++++++-----
 src/Lean/Compiler/Stage1.lean | 42 +++++++++++++++++++++++++++++++++++
 tests/lean/run/lcnf1.lean     | 10 ++++++---
 5 files changed, 66 insertions(+), 8 deletions(-)
 create mode 100644 src/Lean/Compiler/Stage1.lean

diff --git a/src/Lean/Compiler.lean b/src/Lean/Compiler.lean
index 3a0af52cb0..53ae98ac60 100644
--- a/src/Lean/Compiler.lean
+++ b/src/Lean/Compiler.lean
@@ -18,6 +18,7 @@ import Lean.Compiler.CompilerM
 import Lean.Compiler.LCNFTypes
 import Lean.Compiler.InferType
 import Lean.Compiler.LCNF
+import Lean.Compiler.Stage1
 import Lean.Compiler.Decl
 import Lean.Compiler.Main
 import Lean.Compiler.AtMostOnce -- TODO: delete after we port code generator to Lean
diff --git a/src/Lean/Compiler/Decl.lean b/src/Lean/Compiler/Decl.lean
index 0bac380214..87d186d61b 100644
--- a/src/Lean/Compiler/Decl.lean
+++ b/src/Lean/Compiler/Decl.lean
@@ -111,4 +111,8 @@ def Decl.check (decl : Decl) (cfg : Check.Config := {}): CoreM Unit := do
 def Decl.mapValue (decl : Decl) (f : Expr → CompilerM Expr) : CoreM Decl := do
   return { decl with value := (← f decl.value |>.run' { nextIdx := (← getMaxLetVarIdx decl.value) + 1 }) }
 
+def Decl.toString (decl : Decl) : CoreM String := do
+  Meta.MetaM.run' do
+    return s!"{decl.name} : {← Meta.ppExpr decl.type} :=\n{← Meta.ppExpr decl.value}"
+
 end Lean.Compiler
diff --git a/src/Lean/Compiler/Main.lean b/src/Lean/Compiler/Main.lean
index 4590ccaf40..ba9c0c7d02 100644
--- a/src/Lean/Compiler/Main.lean
+++ b/src/Lean/Compiler/Main.lean
@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 import Lean.Compiler.Decl
 import Lean.Compiler.TerminalCases
 import Lean.Compiler.CSE
+import Lean.Compiler.Stage1
 
 namespace Lean.Compiler
 
@@ -41,11 +42,8 @@ def checkpoint (step : Name) (decls : Array Decl) (cfg : Check.Config := {}): Co
       trace[Compiler.step] "{decl.name} : {decl.type} :=\n{decl.value}"
       decl.check cfg
 
-/--
-Run the code generation pipeline for all declarations in `declNames`
-that fulfill the requirements of `shouldGenerateCode`.
--/
-def compile (declNames : Array Name) : CoreM Unit := do profileitM Exception "compiler new" (← getOptions) do
+@[export lean_compile_stage1]
+def compileStage1Impl (declNames : Array Name) : CoreM (Array Decl) := do
   let declNames ← declNames.filterM shouldGenerateCode
   let decls ← declNames.mapM toDecl
   checkpoint `init decls { terminalCasesOnly := false }
@@ -54,6 +52,15 @@ def compile (declNames : Array Name) : CoreM Unit := do profileitM Exception "co
   -- Remark: add simplification step here, `cse` is useful after simplification
   let decls ← decls.mapM (·.cse)
   checkpoint `cse decls
+  saveStage1Decls decls
+  return decls
+
+/--
+Run the code generation pipeline for all declarations in `declNames`
+that fulfill the requirements of `shouldGenerateCode`.
+-/
+def compile (declNames : Array Name) : CoreM Unit := do profileitM Exception "compiler new" (← getOptions) do
+  discard <| compileStage1Impl declNames
 
 builtin_initialize
   registerTraceClass `Compiler
diff --git a/src/Lean/Compiler/Stage1.lean b/src/Lean/Compiler/Stage1.lean
new file mode 100644
index 0000000000..86b7e265ca
--- /dev/null
+++ b/src/Lean/Compiler/Stage1.lean
@@ -0,0 +1,42 @@
+/-
+Copyright (c) 2022 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Lean.Compiler.Decl
+
+namespace Lean.Compiler
+
+structure Stage1ExtState where
+  decls : Std.PHashMap Name Decl := {}
+  deriving Inhabited
+
+builtin_initialize stage1Ext : EnvExtension Stage1ExtState ←
+  registerEnvExtension (pure {})
+
+/- "Forward declaration" -/
+@[extern "lean_compile_stage1"]
+opaque compileStage1 : Array Name → CoreM (Array Decl)
+
+/--
+Save/Cache the given LCNF declarations in the stage1 environment extension.
+-/
+def saveStage1Decls (decls : Array Decl) : CoreM Unit :=
+  modifyEnv fun env => stage1Ext.modifyState env fun s => decls.foldl (init := s) fun s decl =>
+    { s with decls := s.decls.insert decl.name decl }
+
+/--
+Retrieve stage1 LCNF declaration for the given declaration.
+
+We currently do not store stage1 declarations on .olean files,
+and we generate
+-/
+def getStage1Decl? (declName : Name) : CoreM (Option Decl) := do
+  match stage1Ext.getState (← getEnv) |>.decls.find? declName with
+  | some decl => return decl
+  | none =>
+    let info ← getConstInfo declName
+    let decls ← compileStage1 info.all.toArray
+    return decls.find? fun decl => declName == decl.name
+
+end Lean.Compiler
\ No newline at end of file
diff --git a/tests/lean/run/lcnf1.lean b/tests/lean/run/lcnf1.lean
index ab5a4796b8..034da3f23d 100644
--- a/tests/lean/run/lcnf1.lean
+++ b/tests/lean/run/lcnf1.lean
@@ -95,15 +95,19 @@ def Tuple.example (a b : Nat) :=
 #eval Compiler.compile #[``Tuple.example]
 
 def gebner1 (x : UInt64) : UInt64 := assert! x > 0; x - 1
-set_option pp.letVarTypes true in
+-- set_option pp.letVarTypes true in
 #eval Compiler.compile #[``gebner1]
 
 def gebner2 (x : UInt64) := x &&& ((1 : UInt64) <<< 5 : UInt64)
-set_option pp.letVarTypes true in
+-- set_option pp.letVarTypes true in
 #eval Compiler.compile #[``gebner2]
 
 #eval Compiler.compile #[``Lean.Meta.instMonadMetaM]
 #eval Compiler.compile #[``Lean.Core.instMonadCoreM]
 #eval Compiler.compile #[``instMonadEIO]
-set_option pp.explicit true in
+-- set_option pp.explicit true in
 #eval Compiler.compile #[``EStateM.instMonadEStateM]
+
+#eval do
+  let some decl ← Compiler.getStage1Decl? ``List.length | throwError "not found"
+  IO.println (← decl.toString)