feat: typify the index API

DocGen4/Output.lean

@@ -123,42 +123,15 @@ def getSimpleBaseContext (hierarchy : Hierarchy) : SiteBaseContext :=
 def htmlOutputIndex (baseConfig : SiteBaseContext) : IO Unit := do
   htmlOutputSetup baseConfig
 
-  let mut allDecls : List (String × Json) := []
-  let mut allInstances : HashMap String (Array String) := .empty
-  let mut importedBy : HashMap String (Array String) := .empty
-  let mut allModules : List (String × Json) := []
-  let mut instancesFor : HashMap String (Array String) := .empty
+  let mut index : JsonIndex := { }
   for entry in ←System.FilePath.readDir declarationsBasePath do
     if entry.fileName.startsWith "declaration-data-" && entry.fileName.endsWith ".bmp" then
       let fileContent ← FS.readFile entry.path
       let .ok jsonContent := Json.parse fileContent | unreachable!
       let .ok (module : JsonModule) := fromJson? jsonContent | unreachable!
-      allModules := (module.name, Json.str <| moduleNameToLink (String.toName module.name) |>.run baseConfig) :: allModules
-      allDecls := (module.declarations.map (λ d => (d.name, toJson d))) ++ allDecls
-      for inst in module.instances do
-        let mut insts := allInstances.findD inst.className #[]
-        insts := insts.push inst.name
-        allInstances := allInstances.insert inst.className insts
-        for typeName in inst.typeNames do
-          let mut instsFor := instancesFor.findD typeName #[]
-          instsFor := instsFor.push inst.name
-          instancesFor := instancesFor.insert typeName instsFor
-      for imp in module.imports do
-        let mut impBy := importedBy.findD imp #[]
-        impBy := impBy.push module.name
-        importedBy := importedBy.insert imp impBy
+      index := index.addModule module |>.run baseConfig
 
-  let postProcessInstances := allInstances.toList.map (λ(k, v) => (k, toJson v))
-  let postProcessImportedBy := importedBy.toList.map (λ(k, v) => (k, toJson v))
-  let postProcessInstancesFor := instancesFor.toList.map (λ(k, v) => (k, toJson v))
-
-  let finalJson := Json.mkObj [
-    ("declarations", Json.mkObj allDecls),
-    ("instances", Json.mkObj postProcessInstances),
-    ("importedBy", Json.mkObj postProcessImportedBy),
-    ("modules", Json.mkObj allModules),
-    ("instancesFor", Json.mkObj postProcessInstancesFor)
-  ]
+  let finalJson := toJson index
   -- The root JSON for find
   FS.writeFile (declarationsBasePath / "declaration-data.bmp") finalJson.compress
 

DocGen4/Output/ToJson.lean

@@ -27,6 +27,53 @@ structure JsonModule where
   imports : Array String
   deriving FromJson, ToJson
 
+structure JsonIndex where
+  declarations : List (String × JsonDeclaration) := []
+  instances : HashMap String (Array String) := .empty
+  importedBy : HashMap String (Array String) := .empty
+  modules : List (String × String) := []
+  instancesFor : HashMap String (Array String) := .empty
+
+instance : ToJson JsonIndex where
+  toJson idx := Id.run do
+    let jsonDecls := Json.mkObj <| idx.declarations.map (λ (k, v) => (k, toJson v))
+    let jsonInstances := Json.mkObj <| idx.instances.toList.map (λ (k, v) => (k, toJson v))
+    let jsonImportedBy := Json.mkObj <| idx.importedBy.toList.map (λ (k, v) => (k, toJson v))
+    let jsonModules := Json.mkObj <| idx.modules.map (λ (k, v) => (k, toJson v))
+    let jsonInstancesFor := Json.mkObj <| idx.instancesFor.toList.map (λ (k, v) => (k, toJson v))
+    let finalJson := Json.mkObj [
+      ("declarations", jsonDecls),
+      ("instances", jsonInstances),
+      ("importedBy", jsonImportedBy),
+      ("modules", jsonModules),
+      ("instancesFor", jsonInstancesFor)
+    ]
+    pure finalJson
+
+def JsonIndex.addModule (index : JsonIndex) (module : JsonModule) : BaseHtmlM JsonIndex := do
+  let mut index := index
+  let newModule := (module.name, ←moduleNameToLink (String.toName module.name))
+  let newDecls := module.declarations.map (λ d => (d.name, d))
+  index := { index with
+    modules := newModule :: index.modules
+    declarations := newDecls ++ index.declarations
+  }
+  -- TODO: In theory one could sort instances and imports by name and batch the writes
+  for inst in module.instances do
+    let mut insts := index.instances.findD inst.className #[]
+    insts := insts.push inst.name
+    index := { index with instances := index.instances.insert inst.className insts }
+    for typeName in inst.typeNames do
+      let mut instsFor := index.instancesFor.findD typeName #[]
+      instsFor := instsFor.push inst.name
+      index := { index with instancesFor := index.instancesFor.insert typeName instsFor }
+
+  for imp in module.imports do
+    let mut impBy := index.importedBy.findD imp #[]
+    impBy := impBy.push module.name
+    index := { index with importedBy := index.importedBy.insert imp impBy }
+  pure index
+
 def DocInfo.toJson (module : Name) (info : Process.DocInfo) : HtmlM JsonDeclaration := do
   let name := info.getName.toString
   let doc := info.getDocString.getD ""