perf: parallelize core docs and disable the kernel (#245)

DocGen4/Load.lean

@@ -30,7 +30,8 @@ def load (task : Process.AnalyzeTask) : IO (Process.AnalyzerResult × Hierarchy)
     maxHeartbeats := 100000000,
     options := ⟨[
       (`pp.tagAppFns, true),
-      (`pp.funBinderTypes, true)
+      (`pp.funBinderTypes, true),
+      (`debug.skipKernelTC, true)
     ]⟩,
     -- TODO: Figure out whether this could cause some bugs
     fileName := default,
@@ -39,7 +40,4 @@ def load (task : Process.AnalyzeTask) : IO (Process.AnalyzerResult × Hierarchy)
 
   Prod.fst <$> Meta.MetaM.toIO (Process.process task) config { env := env } {} {}
 
-def loadCore : IO (Process.AnalyzerResult × Hierarchy) := do
-  load <| .loadAll #[`Init, `Lean, `Lake, `Std]
-
 end DocGen4

DocGen4/Process/Analyze.lean

@@ -84,12 +84,13 @@ def shouldRender : ModuleMember → Bool
 end ModuleMember
 
 inductive AnalyzeTask where
-| loadAll (load : Array Name) : AnalyzeTask
-| loadAllLimitAnalysis (analyze : Array Name) : AnalyzeTask
+| analyzePrefixModules (topLevel : Name) : AnalyzeTask
+| analyzeConcreteModules (modules : Array Name) : AnalyzeTask
 
-def AnalyzeTask.getLoad : AnalyzeTask → Array Name
-| loadAll load => load
-| loadAllLimitAnalysis load => load
+def AnalyzeTask.getLoad (task : AnalyzeTask) : Array Name :=
+  match task with
+  | .analyzePrefixModules topLevel => #[topLevel]
+  | .analyzeConcreteModules modules => modules
 
 def getAllModuleDocs (relevantModules : Array Name) : MetaM (Std.HashMap Name Module) := do
   let env ← getEnv
@@ -108,9 +109,13 @@ of this `MetaM` run and mentioned by the `AnalyzeTask`.
 -/
 def process (task : AnalyzeTask) : MetaM (AnalyzerResult × Hierarchy) := do
   let env ← getEnv
-  let relevantModules := match task with
-    | .loadAll _ => Std.HashSet.insertMany {} env.header.moduleNames
-    | .loadAllLimitAnalysis analysis => Std.HashSet.insertMany {} analysis
+  let relevantModules :=
+    match task with
+    | .analyzePrefixModules topLevel =>
+      let modules := env.header.moduleNames.filter (topLevel.isPrefixOf ·)
+      Std.HashSet.insertMany (Std.HashSet.empty modules.size) modules
+    | .analyzeConcreteModules modules =>
+      Std.HashSet.insertMany (Std.HashSet.empty modules.size) modules
   let allModules := env.header.moduleNames
 
   let mut res ← getAllModuleDocs relevantModules.toArray

Main.lean

@@ -23,7 +23,7 @@ def runSingleCmd (p : Parsed) : IO UInt32 := do
     | none => ".lake/build"
   let relevantModules := #[p.positionalArg! "module" |>.as! String |> String.toName]
   let sourceUri := p.positionalArg! "sourceUri" |>.as! String
-  let (doc, hierarchy) ← load <| .loadAllLimitAnalysis relevantModules
+  let (doc, hierarchy) ← load <| .analyzeConcreteModules relevantModules
   let baseConfig ← getSimpleBaseContext buildDir hierarchy
   htmlOutputResults baseConfig doc (some sourceUri)
   return 0
@@ -41,7 +41,8 @@ def runGenCoreCmd (p : Parsed) : IO UInt32 := do
   let buildDir := match p.flag? "build" with
     | some dir => dir.as! String
     | none => ".lake/build"
-  let (doc, hierarchy) ← loadCore
+  let module := p.positionalArg! "module" |>.as! String |> String.toName
+  let (doc, hierarchy) ← load <| .analyzePrefixModules module
   let baseConfig ← getSimpleBaseContext buildDir hierarchy
   htmlOutputResults baseConfig doc none
   return 0
@@ -92,10 +93,13 @@ def indexCmd := `[Cli|
 
 def genCoreCmd := `[Cli|
   genCore VIA runGenCoreCmd;
-  "Generate documentation for the core Lean modules: Init, Lean, Lake and Std since they are not Lake projects"
+  "Generate documentation for the specified Lean core module as they are not lake projects."
 
   FLAGS:
     b, build : String; "Build directory."
+
+  ARGS:
+    module : String; "The module to generate the HTML for."
 ]
 
 def bibPrepassCmd := `[Cli|

lakefile.lean

@@ -173,30 +173,34 @@ module_facet docs (mod) : FilePath := do
           return (docFile, trace)
 
 -- TODO: technically speaking this target does not show all file dependencies
-target coreDocs : FilePath := do
+def coreTarget (component : Lean.Name) : FetchM (BuildJob FilePath) := do
   let exeJob ← «doc-gen4».fetch
   let bibPrepassJob ← bibPrepass.fetch
   let dataPath := (← getWorkspace).root.buildDir / "doc-data"
-  let dataFile := dataPath / "declaration-data-Lean.bmp"
-  let buildDir := (←getWorkspace).root.buildDir
+  let dataFile := dataPath / s!"declaration-data-{component}.bmp"
+  let buildDir := (← getWorkspace).root.buildDir
   bibPrepassJob.bindAsync fun _ bibPrepassTrace => do
     exeJob.bindSync fun exeFile exeTrace => do
       let depTrace := mixTraceArray #[exeTrace, bibPrepassTrace]
       let trace ← buildFileUnlessUpToDate dataFile depTrace do
         proc {
           cmd := exeFile.toString
-          args := #["genCore", "--build", buildDir.toString]
+          args := #["genCore", component.toString, "--build", buildDir.toString]
           env := ← getAugmentedEnv
         }
       return (dataFile, trace)
 
+target coreDocs : Unit := do
+  let coreComponents := #[`Init, `Std, `Lake, `Lean]
+  BuildJob.mixArray <| ← coreComponents.mapM coreTarget
+
 library_facet docs (lib) : FilePath := do
   let mods ← lib.modules.fetch
   let moduleJobs ← BuildJob.mixArray <| ← mods.mapM (fetch <| ·.facet `docs)
-  let coreJob ← coreDocs.fetch
+  let coreJobs  ← coreDocs.fetch
   let exeJob ← «doc-gen4».fetch
   -- Shared with DocGen4.Output
-  let buildDir := (←getWorkspace).root.buildDir
+  let buildDir := (← getWorkspace).root.buildDir
   let basePath := buildDir / "doc"
   let dataFile := basePath / "declarations" / "declaration-data.bmp"
   let staticFiles := #[
@@ -219,8 +223,8 @@ library_facet docs (lib) : FilePath := do
     basePath / "find" / "index.html",
     basePath / "find" / "find.js"
   ]
-  coreJob.bindAsync fun _ coreInputTrace => do
-    exeJob.bindAsync fun exeFile exeTrace => do
+  exeJob.bindAsync fun exeFile exeTrace => do
+    coreJobs.bindAsync fun _ coreInputTrace => do
       moduleJobs.bindSync fun _ inputTrace => do
         let depTrace := mixTraceArray #[inputTrace, exeTrace, coreInputTrace]
         let trace ← buildFileUnlessUpToDate dataFile depTrace do
@@ -237,14 +241,14 @@ library_facet docs (lib) : FilePath := do
 library_facet docsHeader (lib) : FilePath := do
   let mods ← lib.modules.fetch
   let moduleJobs ← BuildJob.mixArray <| ← mods.mapM (fetch <| ·.facet `docs)
-  let coreJob ← coreDocs.fetch
   let exeJob ← «doc-gen4».fetch
+  let coreJobs ← coreDocs.fetch
   -- Shared with DocGen4.Output
   let buildDir := (←getWorkspace).root.buildDir
   let basePath := buildDir / "doc"
   let dataFile := basePath / "declarations" / "header-data.bmp"
-  coreJob.bindAsync fun _ coreInputTrace => do
-    exeJob.bindAsync fun exeFile exeTrace => do
+  exeJob.bindAsync fun exeFile exeTrace => do
+   coreJobs.bindAsync fun _ coreInputTrace => do
       moduleJobs.bindSync fun _ inputTrace => do
         let depTrace := mixTraceArray #[inputTrace, exeTrace, coreInputTrace]
         let trace ← buildFileUnlessUpToDate dataFile depTrace do