perf: cache imported environment

Lake/Config/Load.lean

@@ -41,12 +41,47 @@ unsafe def evalScriptDecl
 
 namespace Package
 
+deriving instance BEq, Hashable for Import
+
+/- Cache for the imported header environment of Lake configuration files. -/
+initialize importEnvCache : IO.Ref (Std.HashMap (List Import) Environment) ← IO.mkRef {}
+
+/-- Like `Lean.Elab.processHeader`, but using `importEnvCache`. -/
+def processHeader (header : Syntax) (messages : MessageLog) (inputCtx : Parser.InputContext) (trustLevel : UInt32)
+: IO (Environment × MessageLog) := do
+  try
+    let imports := Elab.headerToImports header
+    if let some env := (← importEnvCache.get).find? imports then return (env, messages)
+    let env ← importModules imports {} trustLevel
+    importEnvCache.modify (·.insert imports env)
+    return (env, messages)
+  catch e =>
+    let env ← mkEmptyEnvironment
+    let spos := header.getPos?.getD 0
+    let pos  := inputCtx.fileMap.toPosition spos
+    return (env, messages.add { fileName := inputCtx.fileName, data := toString e, pos })
+
+/-- Like `Lean.Elab.runFrontend`, but using `importEnvCache`. -/
+def runFrontend (input : String) (opts : Options) (fileName : String) (mainModuleName : Name) (trustLevel : UInt32 := 1024)
+: IO (Environment × Bool) := do
+  let inputCtx := Parser.mkInputContext input fileName
+  let (header, parserState, messages) ← Parser.parseHeader inputCtx
+  let (env, messages) ← processHeader header messages inputCtx trustLevel
+  let env := env.setMainModule mainModuleName
+
+  let mut commandState := Elab.Command.mkState env {} opts
+  let s ← Elab.IO.processCommands inputCtx parserState commandState
+  for msg in s.commandState.messages.toList do
+    IO.eprint (← msg.toString)
+
+  pure (s.commandState.env, !s.commandState.messages.hasErrors)
+
 /-- Unsafe implementation of `load`. -/
 unsafe def loadUnsafe (dir : FilePath) (args : List String := [])
 (configFile := dir / defaultConfigFile) (leanOpts := Options.empty)
 : IO Package := do
   let input ← IO.FS.readFile configFile
-  let (env, ok) ← Elab.runFrontend input leanOpts configFile.toString configModName
+  let (env, ok) ← runFrontend input leanOpts configFile.toString configModName
   if ok then
     match packageAttr.ext.getState env |>.toList with
     | [] => error s!"configuration file is missing a `package` declaration"