feat: add stage1 extension for storing LCNF declarations

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

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

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

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

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)