From c52203ff57647a321edb2c4f284e7a969ca040bf Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 21 Sep 2022 17:47:50 -0700
Subject: [PATCH] feat: add `baseExt` environment extension for storing code
 generator results

---
 src/Lean/Compiler/LCNF/Main.lean        |  7 +++-
 src/Lean/Compiler/LCNF/PassManager.lean |  6 +--
 src/Lean/Compiler/LCNF/Passes.lean      |  4 +-
 src/Lean/Compiler/LCNF/PhaseExt.lean    | 56 +++++++++++++++++++++++++
 src/Lean/Compiler/LCNF/Stage1.lean      |  2 +
 5 files changed, 70 insertions(+), 5 deletions(-)
 create mode 100644 src/Lean/Compiler/LCNF/PhaseExt.lean

diff --git a/src/Lean/Compiler/LCNF/Main.lean b/src/Lean/Compiler/LCNF/Main.lean
index 2c4b02ef78..c98050b3e3 100644
--- a/src/Lean/Compiler/LCNF/Main.lean
+++ b/src/Lean/Compiler/LCNF/Main.lean
@@ -11,6 +11,7 @@ import Lean.Compiler.LCNF.ToDecl
 import Lean.Compiler.LCNF.Check
 import Lean.Compiler.LCNF.Stage1
 import Lean.Compiler.LCNF.PullLetDecls
+import Lean.Compiler.LCNF.PhaseExt
 import Lean.Compiler.LCNF.CSE
 
 namespace Lean.Compiler.LCNF
@@ -67,7 +68,7 @@ def run (declNames : Array Name) : CompilerM (Array Decl) := withAtLeastMaxRecDe
   if declNames.isEmpty then return #[]
   let mut decls ← declNames.mapM toDecl
   let mut manager := { passes := #[{ name := `init, run := pure, phase := .base }] }
-  let installers := PassInstaller.passInstallerExt.getState (←getEnv)
+  let installers := PassInstaller.passInstallerExt.getState (← getEnv)
   manager ← installers.foldlM (init := manager) PassInstaller.runFromDecl
   for pass in manager.passes do
     trace[Compiler] s!"Running pass: {pass.name}"
@@ -86,6 +87,10 @@ def compileStage1Impl (declNames : Array Name) : CoreM (Array Decl) :=
   CompilerM.run do
     PassManager.run declNames
 
+def showDecl (phase : Phase) (declName : Name) : CoreM Format := do
+  let some decl ← getDecl? phase declName | return "<not-available>"
+  ppDecl' decl
+
 builtin_initialize
   registerTraceClass `Compiler.init (inherited := true)
   registerTraceClass `Compiler.test (inherited := true)
diff --git a/src/Lean/Compiler/LCNF/PassManager.lean b/src/Lean/Compiler/LCNF/PassManager.lean
index d5a0505b3d..661dd565cb 100644
--- a/src/Lean/Compiler/LCNF/PassManager.lean
+++ b/src/Lean/Compiler/LCNF/PassManager.lean
@@ -73,9 +73,9 @@ structure PassManager where
 namespace Phase
 
 def toNat : Phase → Nat
-| .base => 0
-| .mono => 1
-| .impure => 2
+  | .base => 0
+  | .mono => 1
+  | .impure => 2
 
 instance : LT Phase where
   lt l r := l.toNat < r.toNat
diff --git a/src/Lean/Compiler/LCNF/Passes.lean b/src/Lean/Compiler/LCNF/Passes.lean
index cf3d6a2649..f3bc159e80 100644
--- a/src/Lean/Compiler/LCNF/Passes.lean
+++ b/src/Lean/Compiler/LCNF/Passes.lean
@@ -11,6 +11,7 @@ import Lean.Compiler.LCNF.PullFunDecls
 import Lean.Compiler.LCNF.ReduceJpArity
 import Lean.Compiler.LCNF.JoinPoints
 import Lean.Compiler.LCNF.Specialize
+import Lean.Compiler.LCNF.PhaseExt
 
 namespace Lean.Compiler.LCNF
 
@@ -24,7 +25,8 @@ namespace Lean.Compiler.LCNF
     reduceJpArity,
     simp { etaPoly := true, inlinePartial := true, implementedBy := true } (occurrence := 1),
     specialize,
-    simp (occurrence := 2)
+    simp (occurrence := 2),
+    saveBase -- End of base phase
   ]
 
 end Lean.Compiler.LCNF
diff --git a/src/Lean/Compiler/LCNF/PhaseExt.lean b/src/Lean/Compiler/LCNF/PhaseExt.lean
new file mode 100644
index 0000000000..29109d2aaf
--- /dev/null
+++ b/src/Lean/Compiler/LCNF/PhaseExt.lean
@@ -0,0 +1,56 @@
+/-
+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.LCNF.PassManager
+
+namespace Lean.Compiler.LCNF
+
+abbrev BaseExtState := Std.PHashMap Name Decl
+
+private abbrev declLt (a b : Decl) :=
+  Name.quickLt a.name b.name
+
+private abbrev sortDecls (decls : Array Decl) : Array Decl :=
+  decls.qsort declLt
+
+private abbrev findAtSorted? (decls : Array Decl) (declName : Name) : Option Decl :=
+  let tmpDecl : Decl := default
+  let tmpDecl := { tmpDecl with name := declName }
+  decls.binSearch tmpDecl declLt
+
+builtin_initialize baseExt : SimplePersistentEnvExtension Decl BaseExtState ← do
+  registerSimplePersistentEnvExtension {
+    name          := `baseDecls
+    addImportedFn := fun _ => {}
+    addEntryFn    := fun decls decl => decls.insert decl.name decl
+    toArrayFn     := fun es => sortDecls es.toArray
+  }
+
+def getBaseDeclCore? (env : Environment) (declName : Name) : Option Decl :=
+  match env.getModuleIdxFor? declName with
+  | some modIdx => findAtSorted? (baseExt.getModuleEntries env modIdx) declName
+  | none        => dbg_trace "getBaseDeclCore?"; baseExt.getState env |>.find? declName
+
+def getBaseDecl? (declName : Name) : CoreM (Option Decl) := do
+  IO.println s!">> getBaseDecl? {declName}"
+  return getBaseDeclCore? (← getEnv) declName
+
+def saveBaseDeclCore (env : Environment) (decl : Decl) : Environment :=
+  baseExt.addEntry env decl
+
+def Decl.saveBase (decl : Decl) : CoreM Unit := do
+  IO.println "saving"
+  IO.println decl.name
+  modifyEnv (saveBaseDeclCore · decl)
+
+def getDecl? (phase : Phase) (declName : Name) : CoreM (Option Decl) :=
+  match phase with
+  | .base => getBaseDecl? declName
+  | _  => return none -- TODO
+
+def saveBase : Pass :=
+  .mkPerDeclaration `saveBase (fun decl => do decl.saveBase; return decl) .base
+
+end Lean.Compiler.LCNF
\ No newline at end of file
diff --git a/src/Lean/Compiler/LCNF/Stage1.lean b/src/Lean/Compiler/LCNF/Stage1.lean
index de1ed126d4..6e4ee8d35c 100644
--- a/src/Lean/Compiler/LCNF/Stage1.lean
+++ b/src/Lean/Compiler/LCNF/Stage1.lean
@@ -7,6 +7,8 @@ import Lean.CoreM
 import Lean.MonadEnv
 import Lean.Compiler.LCNF.Basic
 
+-- TODO: delete this file
+
 namespace Lean.Compiler.LCNF
 
 structure Stage1ExtState where