feat: implementation of separate staged builds

DocGen4/Load.lean

@@ -34,12 +34,19 @@ def lakeSetup (imports : List String) : IO (Except UInt32 Lake.Workspace) := do
   | .error err =>
     throw $ IO.userError err.toString
 
+def envOfImports (imports : List Name) : IO Environment := do
+ importModules (imports.map (Import.mk · false)) Options.empty
+
+def loadInit (imports : List Name) : IO Hierarchy := do
+ let env ← envOfImports imports
+ pure $ Hierarchy.fromArray env.header.moduleNames
+
 /--
 Load a list of modules from the current Lean search path into an `Environment`
 to process for documentation.
 -/
-def load (imports : List Name) (transitiveModules : Bool) : IO (Process.AnalyzerResult × Hierarchy) := do
+def load (task : Process.AnalyzeTask) : IO (Process.AnalyzerResult × Hierarchy) := do
-  let env ← importModules (List.map (Import.mk · false) imports) Options.empty
+  let env ← envOfImports task.getLoad
   IO.println "Processing modules"
   let config := {
     -- TODO: parameterize maxHeartbeats
@@ -49,6 +56,7 @@ def load (imports : List Name) (transitiveModules : Bool) : IO (Process.Analyzer
     fileName := default,
     fileMap := default,
   }
-  Prod.fst <$> Meta.MetaM.toIO (Process.process imports transitiveModules) config { env := env } {} {}
+
+  Prod.fst <$> Meta.MetaM.toIO (Process.process task) config { env := env } {} {}
 
 end DocGen4

DocGen4/Output.lean

@@ -54,10 +54,6 @@ def htmlOutputSetup (config : SiteBaseContext) : IO Unit := do
   for (fileName, content) in findStatic do
     FS.writeFile (findBasePath / fileName) content
 
-  -- The root JSON for find
-  let topLevelModules := config.hierarchy.getChildren.toArray.map (Json.str ∘ toString ∘ Prod.fst)
-  FS.writeFile (basePath / "declaration-data.bmp") (Json.arr topLevelModules).compress
-
   let alectryonStatic := #[
     ("alectryon.css", alectryonCss),
     ("alectryon.js", alectryonJs),
@@ -132,6 +128,16 @@ def getSimpleBaseContext (hierarchy : Hierarchy) : SiteBaseContext :=
     hierarchy
   }
 
+def finalizeDeclarationData : IO Unit := do
+  let mut topLevelModules := #[]
+  for entry in ←System.FilePath.readDir basePath do
+    if entry.fileName.startsWith "declaration-data-" && entry.fileName.endsWith ".bmp" then
+      let module := entry.fileName.drop "declaration-data-".length |>.dropRight ".bmp".length
+      topLevelModules := topLevelModules.push (Json.str module)
+
+  -- The root JSON for find
+  FS.writeFile (basePath / "declaration-data.bmp") (Json.arr topLevelModules).compress
+
 /--
 The main entrypoint for outputting the documentation HTML based on an
 `AnalyzerResult`.
@@ -140,6 +146,7 @@ def htmlOutput (result : AnalyzerResult) (hierarchy : Hierarchy) (ws : Lake.Work
   let baseConfig := getSimpleBaseContext hierarchy
   htmlOutputSetup baseConfig
   htmlOutputResults baseConfig result ws inkPath
+  finalizeDeclarationData
 
 end DocGen4
 

DocGen4/Process/Analyze.lean

@@ -96,23 +96,55 @@ def getRelevantModules (imports : List Name) : MetaM (HashSet Name) := do
       relevant := relevant.insert module
   pure relevant
 
-/--
+inductive AnalyzeTask where
-Run the doc-gen analysis on all modules that are loaded into the `Environment`
+| loadAll (load : List Name) : AnalyzeTask
-of this `MetaM` run.
+| loadAllLimitAnalysis (load : List Name) (analyze : List Name) : AnalyzeTask
--/
+
-def process (imports : List Name) (transitiveModules : Bool) : MetaM (AnalyzerResult × Hierarchy) := do
+def AnalyzeTask.getLoad : AnalyzeTask → List Name
+| loadAll load => load
+| loadAllLimitAnalysis load _ => load
+
+def AnalyzeTask.getAnalyze : AnalyzeTask → List Name
+| loadAll load => load
+| loadAllLimitAnalysis _ analysis => analysis
+
+def getAllModuleDocs (relevantModules : Array Name) : MetaM (HashMap Name Module) := do
   let env ← getEnv
-  let relevantNames := if transitiveModules then (HashSet.fromArray env.header.moduleNames) else (←getRelevantModules imports)
+  let mut res := mkHashMap relevantModules.size
-  let mut res := mkHashMap relevantNames.size
+  for module in relevantModules do
-  for module in relevantNames.toArray do
     let modDocs := getModuleDoc? env module |>.getD #[] |>.map .modDoc
     res := res.insert module (Module.mk module modDocs)
+  pure res
+
+-- TODO: This is definitely not the most efficient way to store this data
+def buildImportAdjMatrix (allModules : Array Name) : MetaM (Array (Array Bool)) := do
+  let env ← getEnv
+  let mut adj := Array.mkArray allModules.size (Array.mkArray allModules.size false)
+  for moduleName in allModules do
+    let some modIdx := env.getModuleIdx? moduleName | unreachable!
+    let moduleData := env.header.moduleData.get! modIdx
+    for imp in moduleData.imports do
+      let some importIdx := env.getModuleIdx? imp.module | unreachable!
+      adj := adj.set! modIdx (adj.get! modIdx |>.set! importIdx true)
+  pure adj
+
+/--
+Run the doc-gen analysis on all modules that are loaded into the `Environment`
+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 _ => pure $ HashSet.fromArray env.header.moduleNames
+    | .loadAllLimitAnalysis _ analysis => getRelevantModules analysis
+  let allModules := env.header.moduleNames
+
+  let mut res ← getAllModuleDocs relevantModules.toArray
 
   for (name, cinfo) in env.constants.toList do
     let some modidx := env.getModuleIdxFor? name | unreachable!
     let moduleName := env.allImportedModuleNames.get! modidx
-    -- skip irrelevant names
+    if !relevantModules.contains moduleName then
-    if !relevantNames.contains moduleName.getRoot then
       continue
 
     try
@@ -122,7 +154,7 @@ def process (imports : List Name) (transitiveModules : Bool) : MetaM (AnalyzerRe
         fileName := ←getFileName,
         fileMap := ←getFileMap
       }
-      let analysis := Prod.fst <$> Meta.MetaM.toIO (DocInfo.ofConstant (name, cinfo)) config { env := env} {} {}
+      let analysis := Prod.fst <$> Meta.MetaM.toIO (DocInfo.ofConstant (name, cinfo)) config { env := env } {} {}
       if let some dinfo ← analysis then
         let moduleName := env.allImportedModuleNames.get! modidx
         let module := res.find! moduleName
@@ -130,21 +162,16 @@ def process (imports : List Name) (transitiveModules : Bool) : MetaM (AnalyzerRe
     catch e =>
       IO.println s!"WARNING: Failed to obtain information for: {name}: {←e.toMessageData.toString}"
 
-  -- TODO: This is definitely not the most efficient way to store this data
+  let adj ← buildImportAdjMatrix allModules
-  let mut adj := Array.mkArray res.size (Array.mkArray res.size false)
+
-  -- TODO: This could probably be faster if we did an insertion sort above instead
+  -- TODO: This could probably be faster if we did sorted insert above instead
   for (moduleName, module) in res.toArray do
     res := res.insert moduleName {module with members := module.members.qsort ModuleMember.order}
-    let some modIdx := env.getModuleIdx? moduleName | unreachable!
-    let moduleData := env.header.moduleData.get! modIdx
-    for imp in moduleData.imports do
-      let some importIdx := env.getModuleIdx? imp.module | unreachable!
-      adj := adj.set! modIdx (adj.get! modIdx |>.set! importIdx true)
 
-  let hierarchy := Hierarchy.fromArray relevantNames.toArray
+  let hierarchy := Hierarchy.fromArray allModules
   let analysis := {
     name2ModIdx := env.const2ModIdx,
-    moduleNames := env.header.moduleNames,
+    moduleNames := allModules,
     moduleInfo := res,
     importAdj := adj
   }

Main.lean

@@ -14,54 +14,97 @@ def findLeanInk? (p : Parsed) : IO (Option System.FilePath) := do
       throw $ IO.userError "Invalid path to LeanInk binary provided"
   | none => pure none
 
+def getTopLevelModules (p : Parsed) : IO (List String) :=  do
+  let topLevelModules := p.variableArgsAs! String |>.toList
+  if topLevelModules.length == 0 then
+    throw $ IO.userError "No topLevelModules provided."
+  pure topLevelModules
+
+def runInitCmd (p : Parsed) : IO UInt32 := do
+  let topLevelModules ← getTopLevelModules p
+  let res ← lakeSetup topLevelModules
+  match res with
+  | Except.ok _ =>
+    let modules := topLevelModules.map Name.mkSimple
+    let hierarchy ← loadInit modules
+    let baseConfig := getSimpleBaseContext hierarchy
+    htmlOutputSetup baseConfig
+    pure 0
+  | Except.error rc => pure rc
+
+def runSingleCmd (p : Parsed) : IO UInt32 := do
+    let topLevelModules ← getTopLevelModules p
+    let relevantModules := [p.positionalArg! "module" |>.as! String]
+    let res ← lakeSetup (relevantModules ++ topLevelModules)
+    match res with
+    | Except.ok ws =>
+      let relevantModules := relevantModules.map Name.mkSimple
+      let topLevelModules := topLevelModules.map Name.mkSimple
+      let (doc, hierarchy) ← load (.loadAllLimitAnalysis topLevelModules relevantModules)
+      IO.println "Outputting HTML"
+      let baseConfig := getSimpleBaseContext hierarchy
+      htmlOutputResults baseConfig doc ws (←findLeanInk? p)
+      pure 0
+    | Except.error rc => pure rc
+
+def runFinalizeCmd (p : Parsed) : IO UInt32 := do
+  finalizeDeclarationData
+  pure 0
+
 def runDocGenCmd (p : Parsed) : IO UInt32 := do
   let modules : List String := p.variableArgsAs! String |>.toList
-  if p.hasFlag "single" && p.hasFlag "setup" then
+  if modules.length == 0 then
-    throw $ IO.userError "Can't have single and setup at the same time"
+    throw $ IO.userError "No modules provided."
-  else
+
   let res ← lakeSetup modules
-    let modules := modules.map Name.mkSimple
   match res with
   | Except.ok ws =>
     IO.println s!"Loading modules from: {←searchPathRef.get}"
-      --if p.hasFlag "single" then
+    let modules := modules.map Name.mkSimple
-      --  if modules.length ≠ 1 then
+    let (doc, hierarchy) ← load (.loadAll modules)
-      --    throw $ IO.userError "Called single with more than a single module"
-      --  else
-      --    let (doc, hierarchy) ← load modules false
-      --    IO.println "Outputting HTML"
-      --    let baseConfig := getSimpleBaseContext hierarchy
-      --    htmlOutputResults baseConfig doc ws leanHash (←findLeanInk? p)
-      --    pure 0
-      --else if p.hasFlag "setup" then
-      --  let config := {
-      --    fileName := default,
-      --    fileMap := default,
-      --  }
-      --  let env ← importModules (List.map (Import.mk · false) modules) Options.empty
-      --  let relevantModules ← Prod.fst <$> Meta.MetaM.toIO (Process.getRelevantModules modules) config { env := env } {}
-      --  let hierarchy := Hierarchy.fromArray relevantModules.toArray
-      --  let baseConfig := getSimpleBaseContext hierarchy
-      --  htmlOutputSetup baseConfig
-      --  pure 0
-      --else
-        let (doc, hierarchy) ← load modules true
     IO.println "Outputting HTML"
     htmlOutput doc hierarchy ws (←findLeanInk? p)
     pure 0
   | Except.error rc => pure rc
 
+def initCmd := `[Cli|
+  init VIA runInitCmd;
+  "Generate the initial directory structure and static files."
+  ARGS:
+    ...topLevelModule : String; "The top level modules this documentation will be for."
+]
+
+def singleCmd := `[Cli|
+  single VIA runSingleCmd;
+  "Only generate the documentation for the module it was given, might contain broken links unless all documentation is generated."
+
+  FLAGS:
+    ink : String; "Path to a LeanInk binary to use for rendering the Lean sources."
+
+  ARGS:
+    module : String; "The module to generate the HTML for. Does not have to be part of topLevelModules."
+    ...topLevelModules : String; "The top level modules this documentation will be for."
+]
+
+def finalizeCmd := `[Cli|
+  finalize VIA runFinalizeCmd;
+  "Finalize the documentation that has been generated by single."
+]
+
 def docGenCmd : Cmd := `[Cli|
   "doc-gen4" VIA runDocGenCmd; ["0.0.1"]
   "A documentation generator for Lean 4."
 
   FLAGS:
     ink : String; "Path to a LeanInk binary to use for rendering the Lean sources."
-    --single; "Generate documentation only for a single module, will cause broken links if there are others"
-    --setup; "Only output the files that are always required"
 
   ARGS:
-    ...modules : String; "The modules to generate the HTML for"
+    ...modules : String; "The modules to generate the HTML for."
+
+  SUBCOMMANDS:
+    initCmd;
+    singleCmd;
+    finalizeCmd
 ]
 
 def main (args : List String) : IO UInt32 :=