feat: lake: save elaborated config as an olean

src/Lean/Elab/Import.lean

@@ -8,9 +8,9 @@ import Lean.Data.Json
 
 namespace Lean.Elab
 
-def headerToImports (header : Syntax) : List Import :=
-  let imports := if header[0].isNone then [{ module := `Init : Import }] else []
-  imports ++ header[1].getArgs.toList.map fun stx =>
+def headerToImports (header : Syntax) : Array Import :=
+  let imports := if header[0].isNone then #[{ module := `Init : Import }] else #[]
+  imports ++ header[1].getArgs.map fun stx =>
     -- `stx` is of the form `(Module.import "import" "runtime"? id)
     let runtime := !stx[1].isNone
     let id      := stx[2].getId
@@ -27,7 +27,7 @@ def processHeader (header : Syntax) (opts : Options) (messages : MessageLog) (in
     let pos  := inputCtx.fileMap.toPosition spos
     pure (env, messages.add { fileName := inputCtx.fileName, data := toString e, pos := pos })
 
-def parseImports (input : String) (fileName : Option String := none) : IO (List Import × Position × MessageLog) := do
+def parseImports (input : String) (fileName : Option String := none) : IO (Array Import × Position × MessageLog) := do
   let fileName := fileName.getD "<input>"
   let inputCtx := Parser.mkInputContext input fileName
   let (header, parserState, messages) ← Parser.parseHeader inputCtx

src/Lean/Environment.lean

@@ -33,6 +33,8 @@ structure Import where
   runtimeOnly : Bool := false
   deriving Repr, Inhabited
 
+instance : Coe Name Import := ⟨({module := ·})⟩
+
 instance : ToString Import := ⟨fun imp => toString imp.module ++ if imp.runtimeOnly then " (runtime)" else ""⟩
 
 /--
@@ -731,73 +733,77 @@ def throwAlreadyImported (s : ImportState) (const2ModIdx : HashMap Name ModuleId
   let constModName := s.moduleNames[const2ModIdx[cname].get!.toNat]!
   throw <| IO.userError s!"import {modName} failed, environment already contains '{cname}' from {constModName}"
 
+abbrev ImportStateM := StateRefT ImportState IO
+
+@[inline] nonrec def ImportStateM.run (x : ImportStateM α) (s : ImportState := {}) : IO (α × ImportState) :=
+  x.run s
+
+partial def importModulesCore (imports : Array Import) : ImportStateM Unit := do
+  for i in imports do
+    if i.runtimeOnly || (← get).moduleNameSet.contains i.module then
+      continue
+    modify fun s => { s with moduleNameSet := s.moduleNameSet.insert i.module }
+    let mFile ← findOLean i.module
+    unless (← mFile.pathExists) do
+      throw <| IO.userError s!"object file '{mFile}' of module {i.module} does not exist"
+    let (mod, region) ← readModuleData mFile
+    importModulesCore mod.imports
+    modify fun s => { s with
+      moduleData  := s.moduleData.push mod
+      regions     := s.regions.push region
+      moduleNames := s.moduleNames.push i.module
+    }
+
+def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0) : IO Environment := do
+  let numConsts := s.moduleData.foldl (init := 0) fun numConsts mod =>
+    numConsts + mod.constants.size + mod.extraConstNames.size
+  let mut const2ModIdx : HashMap Name ModuleIdx := mkHashMap (capacity := numConsts)
+  let mut constantMap : HashMap Name ConstantInfo := mkHashMap (capacity := numConsts)
+  for h:modIdx in [0:s.moduleData.size] do
+    let mod := s.moduleData[modIdx]'h.upper
+    for cname in mod.constNames, cinfo in mod.constants do
+      match constantMap.insert' cname cinfo with
+      | (constantMap', replaced) =>
+        constantMap := constantMap'
+        if replaced then
+          throwAlreadyImported s const2ModIdx modIdx cname
+      const2ModIdx := const2ModIdx.insert cname modIdx
+    for cname in mod.extraConstNames do
+      const2ModIdx := const2ModIdx.insert cname modIdx
+  let constants : ConstMap := SMap.fromHashMap constantMap false
+  let exts ← mkInitialExtensionStates
+  let env : Environment := {
+    const2ModIdx    := const2ModIdx
+    constants       := constants
+    extraConstNames := {}
+    extensions      := exts
+    header          := {
+      quotInit     := !imports.isEmpty -- We assume `core.lean` initializes quotient module
+      trustLevel   := trustLevel
+      imports      := imports
+      regions      := s.regions
+      moduleNames  := s.moduleNames
+      moduleData   := s.moduleData
+    }
+  }
+  let env ← setImportedEntries env s.moduleData
+  let env ← finalizePersistentExtensions env s.moduleData opts
+  pure env
+
 @[export lean_import_modules]
-partial def importModules (imports : List Import) (opts : Options) (trustLevel : UInt32 := 0) : IO Environment := profileitIO "import" opts do
+def importModules (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0) : IO Environment := profileitIO "import" opts do
+  let imports := imports
   for imp in imports do
     if imp.module matches .anonymous then
       throw <| IO.userError "import failed, trying to import module with anonymous name"
   withImporting do
-    let (_, s) ← importMods imports |>.run {}
-    let mut numConsts := 0
-    for mod in s.moduleData do
-      numConsts := numConsts + mod.constants.size + mod.extraConstNames.size
-    let mut modIdx : Nat := 0
-    let mut const2ModIdx : HashMap Name ModuleIdx := mkHashMap (capacity := numConsts)
-    let mut constantMap : HashMap Name ConstantInfo := mkHashMap (capacity := numConsts)
-    for mod in s.moduleData do
-      for cname in mod.constNames, cinfo in mod.constants do
-        match constantMap.insert' cname cinfo with
-        | (constantMap', replaced) =>
-          constantMap := constantMap'
-          if replaced then
-            throwAlreadyImported s const2ModIdx modIdx cname
-        const2ModIdx := const2ModIdx.insert cname modIdx
-      for cname in mod.extraConstNames do
-        const2ModIdx := const2ModIdx.insert cname modIdx
-      modIdx := modIdx + 1
-    let constants : ConstMap := SMap.fromHashMap constantMap false
-    let exts ← mkInitialExtensionStates
-    let env : Environment := {
-      const2ModIdx    := const2ModIdx
-      constants       := constants
-      extraConstNames := {}
-      extensions      := exts
-      header          := {
-        quotInit     := !imports.isEmpty -- We assume `core.lean` initializes quotient module
-        trustLevel   := trustLevel
-        imports      := imports.toArray
-        regions      := s.regions
-        moduleNames  := s.moduleNames
-        moduleData   := s.moduleData
-      }
-    }
-    let env ← setImportedEntries env s.moduleData
-    let env ← finalizePersistentExtensions env s.moduleData opts
-    pure env
-where
-  importMods : List Import → StateRefT ImportState IO Unit
-  | []    => pure ()
-  | i::is => do
-    if i.runtimeOnly || (← get).moduleNameSet.contains i.module then
-      importMods is
-    else do
-      modify fun s => { s with moduleNameSet := s.moduleNameSet.insert i.module }
-      let mFile ← findOLean i.module
-      unless (← mFile.pathExists) do
-        throw <| IO.userError s!"object file '{mFile}' of module {i.module} does not exist"
-      let (mod, region) ← readModuleData mFile
-      importMods mod.imports.toList
-      modify fun s => { s with
-        moduleData  := s.moduleData.push mod
-        regions     := s.regions.push region
-        moduleNames := s.moduleNames.push i.module
-      }
-      importMods is
+    let (_, s) ← importModulesCore imports |>.run
+    finalizeImport s imports opts trustLevel
 
 /--
   Create environment object from imports and free compacted regions after calling `act`. No live references to the
   environment object or imported objects may exist after `act` finishes. -/
-unsafe def withImportModules {α : Type} (imports : List Import) (opts : Options) (trustLevel : UInt32 := 0) (x : Environment → IO α) : IO α := do
+unsafe def withImportModules {α : Type} (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0) (x : Environment → IO α) : IO α := do
   let env ← importModules imports opts trustLevel
   try x env finally env.freeRegions
 

src/Lean/Server/FileWorker.lean

@@ -207,7 +207,7 @@ section Initialization
           -- NOTE: we assume for now that `lakefile.lean` does not have any non-stdlib deps
           -- NOTE: lake does not exist in stage 0 (yet?)
           if path.fileName != "lakefile.lean" && (← System.FilePath.pathExists lakePath) then
-            let pkgSearchPath ← lakeSetupSearchPath lakePath m (Lean.Elab.headerToImports headerStx).toArray hOut
+            let pkgSearchPath ← lakeSetupSearchPath lakePath m (Lean.Elab.headerToImports headerStx) hOut
             srcSearchPath ← initSrcSearchPath (← getBuildDir) pkgSearchPath
         Elab.processHeader headerStx opts msgLog m.mkInputContext
       catch e =>  -- should be from `lake print-paths`

src/lake/.gitignore

@@ -1,3 +1,4 @@
 /build
+/lakefile.olean
 produced.out
 result*

src/lake/Lake/CLI/Help.lean

@@ -25,6 +25,7 @@ OPTIONS:
   --old                 only rebuild modified modules (ignore transitive deps)
   --rehash, -H          hash all files for traces (do not trust `.hash` files)
   --update, -U          update manifest before building
+  --reconfigure, -R     elaborate configuration files instead of using OLeans
 
 COMMANDS:
   new <name> <temp>     create a Lean package in a new directory

src/lake/Lake/CLI/Init.lean

@@ -22,6 +22,7 @@ def toolchainFileName : FilePath :=
 
 def gitignoreContents :=
 s!"/{defaultBuildDir}
+/{defaultConfigFile.withExtension "olean"}
 /{defaultPackagesDir}/*
 "
 

src/lake/Lake/CLI/Main.lean

@@ -34,6 +34,7 @@ structure LakeOptions where
   wantsHelp : Bool := false
   verbosity : Verbosity := .normal
   updateDeps : Bool := false
+  reconfigure : Bool := false
   oldMode : Bool := false
   trustHash : Bool := true
 
@@ -65,6 +66,7 @@ def LakeOptions.mkLoadConfig (opts : LakeOptions) : EIO CliError LoadConfig :=
     configFile := opts.rootDir / opts.configFile
     configOpts := opts.configOpts
     leanOpts := Lean.Options.empty
+    reconfigure := opts.reconfigure
   }
 
 /-- Make a `BuildConfig` from a `LakeOptions`. -/
@@ -139,22 +141,24 @@ def lakeShortOption : (opt : Char) → CliM PUnit
 | 'f' => do let configFile ← takeOptArg "-f" "path"; modifyThe LakeOptions ({· with configFile})
 | 'K' => do setConfigOpt <| ← takeOptArg "-K" "key-value pair"
 | 'U' => modifyThe LakeOptions ({· with updateDeps := true})
+| 'R' => modifyThe LakeOptions ({· with reconfigure := true})
 | 'h' => modifyThe LakeOptions ({· with wantsHelp := true})
 | 'H' => modifyThe LakeOptions ({· with trustHash := false})
 | opt => throw <| CliError.unknownShortOption opt
 
 def lakeLongOption : (opt : String) → CliM PUnit
-| "--quiet"   => modifyThe LakeOptions ({· with verbosity := .quiet})
-| "--verbose" => modifyThe LakeOptions ({· with verbosity := .verbose})
-| "--update"  => modifyThe LakeOptions ({· with updateDeps := true})
-| "--old"     => modifyThe LakeOptions ({· with oldMode := true})
-| "--rehash"  => modifyThe LakeOptions ({· with trustHash := false})
-| "--dir"     => do let rootDir ← takeOptArg "--dir" "path"; modifyThe LakeOptions ({· with rootDir})
-| "--file"    => do let configFile ← takeOptArg "--file" "path"; modifyThe LakeOptions ({· with configFile})
-| "--lean"    => do setLean <| ← takeOptArg "--lean" "path or command"
-| "--help"    => modifyThe LakeOptions ({· with wantsHelp := true})
-| "--"        => do let subArgs ← takeArgs; modifyThe LakeOptions ({· with subArgs})
-| opt         => throw <| CliError.unknownLongOption opt
+| "--quiet"       => modifyThe LakeOptions ({· with verbosity := .quiet})
+| "--verbose"     => modifyThe LakeOptions ({· with verbosity := .verbose})
+| "--update"      => modifyThe LakeOptions ({· with updateDeps := true})
+| "--reconfigure" => modifyThe LakeOptions ({· with reconfigure := true})
+| "--old"         => modifyThe LakeOptions ({· with oldMode := true})
+| "--rehash"      => modifyThe LakeOptions ({· with trustHash := false})
+| "--dir"         => do let rootDir ← takeOptArg "--dir" "path"; modifyThe LakeOptions ({· with rootDir})
+| "--file"        => do let configFile ← takeOptArg "--file" "path"; modifyThe LakeOptions ({· with configFile})
+| "--lean"        => do setLean <| ← takeOptArg "--lean" "path or command"
+| "--help"        => modifyThe LakeOptions ({· with wantsHelp := true})
+| "--"            => do let subArgs ← takeArgs; modifyThe LakeOptions ({· with subArgs})
+| opt             =>  throw <| CliError.unknownLongOption opt
 
 def lakeOption :=
   option {

src/lake/Lake/Config/Package.lean

@@ -248,19 +248,16 @@ namespace Package
 @[inline] def deps (self : Package) : Array Package  :=
   self.opaqueDeps.map (·.get)
 
-/--
-The path for storing the package's remote dependencies relative to `dir`.
-Either its `packagesDir` configuration or `defaultPackagesDir`.
--/
-def relPkgsDir (self : Package) : FilePath :=
-  self.config.packagesDir.getD defaultPackagesDir
+/-- The path for storing the package's remote dependencies relative to `dir` (i.e., `packagesDir`). -/
+@[inline] def relPkgsDir (self : Package) : FilePath :=
+  self.config.packagesDir
 
 /-- The package's `dir` joined with its `relPkgsDir` -/
-def pkgsDir (self : Package) : FilePath :=
+@[inline] def pkgsDir (self : Package) : FilePath :=
   self.dir / self.relPkgsDir
 
 /-- The package's JSON manifest of remote dependencies. -/
-def manifestFile (self : Package) : FilePath :=
+@[inline] def manifestFile (self : Package) : FilePath :=
   self.dir / self.config.manifestFile
 
 /-- The package's `dir` joined with its `buildDir` configuration. -/

src/lake/Lake/Config/WorkspaceConfig.lean

@@ -16,5 +16,5 @@ structure WorkspaceConfig where
   The directory to which Lake should download remote dependencies.
   Defaults to `defaultPackagesDir` (i.e., `lake-packages`).
   -/
-  packagesDir : Option FilePath := none
+  packagesDir : FilePath := defaultPackagesDir
   deriving Inhabited, Repr

src/lake/Lake/Load/Config.lean

@@ -26,3 +26,5 @@ structure LoadConfig where
   configOpts : NameMap String := {}
   /-- The Lean options with which to elaborate the configuration file. -/
   leanOpts : Options := {}
+  /-- If true, Lake will elaborate configuration files instead of using OLeans. -/
+  reconfigure : Bool := false

src/lake/Lake/Load/Elab.lean

@@ -6,26 +6,37 @@ Authors: Mac Malone
 import Lean.Elab.Frontend
 import Lake.DSL.Extensions
 import Lake.Load.Config
+import Lake.Build.Trace
 import Lake.Util.Log
 
 namespace Lake
 open Lean System
 
 deriving instance BEq, Hashable for Import
 
-/- Cache for the imported header environment of Lake configuration files. -/
-initialize importEnvCache : IO.Ref (HashMap (List Import) Environment) ← IO.mkRef {}
+/- Cache for the import state of Lake configuration files. -/
+initialize importStateCache : IO.Ref (HashMap (Array Import) ImportState) ← IO.mkRef {}
+
+/--
+Like `importModulesCore`, but fetch the
+resulting import state from the cache if possible. -/
+def importModulesUsingCache (imports : Array Import) : IO ImportState := do
+  match (← importStateCache.get).find? imports with
+  | none =>
+    let (_, s) ← importModulesCore imports |>.run
+    importStateCache.modify (·.insert imports s)
+    return s
+  | some s =>
+    return s
 
 /-- Like `Lean.Elab.processHeader`, but using `importEnvCache`. -/
-def processHeader (header : Syntax) (opts : Options) (trustLevel : UInt32)
+def processHeader (header : Syntax) (opts : Options)
 (inputCtx : Parser.InputContext) : StateT MessageLog IO Environment := do
   try
     let imports := Elab.headerToImports header
-    if let some env := (← importEnvCache.get).find? imports then
-      return env
-    let env ← importModules imports opts trustLevel
-    importEnvCache.modify (·.insert imports env)
-    return env
+    withImporting do
+      let s ← importModulesUsingCache imports
+      finalizeImport s imports opts 1024
   catch e =>
     let pos := inputCtx.fileMap.toPosition <| header.getPos?.getD 0
     modify (·.add { fileName := inputCtx.fileName, data := toString e, pos })
@@ -35,19 +46,19 @@ def processHeader (header : Syntax) (opts : Options) (trustLevel : UInt32)
 def configModuleName : Name := `lakefile
 
 /-- Elaborate `configFile` with the given package directory and options. -/
-def elabConfigFile (pkgDir : FilePath) (configOpts : NameMap String)
+def elabConfigFile (pkgDir : FilePath) (lakeOpts : NameMap String)
 (leanOpts := Options.empty) (configFile := pkgDir / defaultConfigFile) : LogIO Environment := do
 
   -- Read file and initialize environment
   let input ← IO.FS.readFile configFile
   let inputCtx := Parser.mkInputContext input configFile.toString
   let (header, parserState, messages) ← Parser.parseHeader inputCtx
-  let (env, messages) ← processHeader header leanOpts 1024 inputCtx messages
+  let (env, messages) ← processHeader header leanOpts inputCtx messages
   let env := env.setMainModule configModuleName
 
   -- Configure extensions
   let env := dirExt.setState env pkgDir
-  let env := optsExt.setState env configOpts
+  let env := optsExt.setState env lakeOpts
 
   -- Elaborate File
   let commandState := Elab.Command.mkState env messages leanOpts
@@ -65,3 +76,30 @@ def elabConfigFile (pkgDir : FilePath) (configOpts : NameMap String)
     error s!"{configFile}: package configuration has errors"
   else
     return s.commandState.env
+
+/--
+If `reconfigure` is not set and up-to-date OLean for the configuration file exists,
+import it. Otherwise, elaborate the configuration and store save it to the OLean.
+-/
+def importConfigFile (pkgDir : FilePath) (lakeOpts : NameMap String)
+(leanOpts := Options.empty) (configFile := pkgDir / defaultConfigFile) (reconfigure := true) : LogIO Environment := do
+  let olean := configFile.withExtension "olean"
+  let useOLean ← id do
+    if reconfigure then return false
+    unless (← olean.pathExists) do return false
+    unless (← getMTime olean) > (← getMTime configFile) do return false
+    return true
+  if useOLean then
+    withImporting do
+    let (mod, region) ← readModuleData olean
+    let s ← importModulesUsingCache mod.imports
+    let s := {s with
+      moduleData  := s.moduleData.push mod
+      regions     := s.regions.push region
+      moduleNames := s.moduleNames.push configModuleName
+    }
+    finalizeImport s #[configModuleName] leanOpts 1024
+  else
+    let env ← elabConfigFile pkgDir lakeOpts leanOpts configFile
+    Lean.writeModule env olean
+    return env