Skip to content

Commit

Permalink
perf: parallelize core docs and disable the kernel
Browse files Browse the repository at this point in the history
  • Loading branch information
hargoniX committed Dec 12, 2024
1 parent 65417de commit 3edfca9
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 26 deletions.
6 changes: 2 additions & 4 deletions DocGen4/Load.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
21 changes: 13 additions & 8 deletions DocGen4/Process/Analyze.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
10 changes: 7 additions & 3 deletions Main.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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|
Expand Down
26 changes: 15 additions & 11 deletions lakefile.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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 := #[
Expand All @@ -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
Expand All @@ -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
Expand Down

0 comments on commit 3edfca9

Please sign in to comment.