feat: Sort everything into modules instead of just declaration lists

DocGen4/Load.lean

@@ -1,9 +1,10 @@
 import Lean
 import DocGen4.Process
+import Std.Data.HashMap
 
 namespace DocGen4
 
-open Lean System
+open Lean System Std
 
 def printSearchPath : IO PUnit := do
   IO.println s!"{←searchPathRef.get}"
@@ -11,11 +12,11 @@ def printSearchPath : IO PUnit := do
 def setSearchPath (path : List FilePath) : IO PUnit := do
     searchPathRef.set path
 
-def load (imports : List Name) : IO (List DocInfo) := do
+def load (imports : List Name) : IO (HashMap Name Module) := do
   let env ← importModules (List.map (Import.mk · false) imports) Options.empty
-  let doc ← Prod.fst <$> (Meta.MetaM.toIO (process env) {} { env := env} {} {})
-  for i in doc do
-    let s ← Core.CoreM.toIO i.prettyPrint {} { env := env }
+  let doc ← Prod.fst <$> (Meta.MetaM.toIO (process) {} { env := env} {} {})
+  for (_, module) in doc.toList do
+    let s ← Core.CoreM.toIO module.prettyPrint {} { env := env }
     IO.println s.fst
   return doc
 

DocGen4/Process.lean

@@ -1,56 +1,30 @@
 import Lean
 import Lean.Meta.Match.MatcherInfo
 import Lean.PrettyPrinter
+import Std.Data.HashMap
 
 namespace DocGen4
 
-open Lean Meta PrettyPrinter
-
-def prettyPrintTerm (expr : Expr) : MetaM Syntax :=
-  delab Name.anonymous [] expr
+open Lean Meta PrettyPrinter Std
 
 structure Info where
   name : Name
   type : Syntax
   doc : Option String
-  module : Name
   deriving Repr
 
-def Info.ofConstantVal (v : ConstantVal) : MetaM Info := do
-  let env ← getEnv
-  let type := (←prettyPrintTerm v.type )
-  let doc := findDocString? env v.name
-  match (env.getModuleIdxFor? v.name) with
-  | some modidx =>
-    let module := env.allImportedModuleNames.get! modidx
-    return Info.mk v.name type doc module
-  | none => panic! "impossible"
-
 structure AxiomInfo extends Info where
   isUnsafe : Bool
   deriving Repr
 
-def AxiomInfo.ofAxiomVal (v : AxiomVal) : MetaM AxiomInfo := do
-  let info ← Info.ofConstantVal v.toConstantVal
-  return AxiomInfo.mk info v.isUnsafe
-
 structure TheoremInfo extends Info where
   deriving Repr
 
-def TheoremInfo.ofTheoremVal (v : TheoremVal) : MetaM TheoremInfo := do
-  let info ← Info.ofConstantVal v.toConstantVal
-  return TheoremInfo.mk info
-
 structure OpaqueInfo extends Info where
   value : Syntax
   isUnsafe : Bool
   deriving Repr
 
-def OpaqueInfo.ofOpaqueVal (v : OpaqueVal) : MetaM OpaqueInfo := do
-  let info ← Info.ofConstantVal v.toConstantVal
-  let value ← prettyPrintTerm v.value
-  return OpaqueInfo.mk info value v.isUnsafe
-  
 inductive DocInfo where
 | axiomInfo (info : AxiomInfo) : DocInfo
 | theoremInfo (info : TheoremInfo) : DocInfo
@@ -63,6 +37,34 @@ inductive DocInfo where
 | classInductiveInfo : DocInfo
 deriving Repr
 
+structure Module where
+  name : Name
+  doc : Option String
+  members : Array DocInfo
+  deriving Inhabited, Repr
+
+def prettyPrintTerm (expr : Expr) : MetaM Syntax :=
+  delab Name.anonymous [] expr
+
+def Info.ofConstantVal (v : ConstantVal) : MetaM Info := do
+  let env ← getEnv
+  let type := (←prettyPrintTerm v.type )
+  let doc := findDocString? env v.name
+  return Info.mk v.name type doc
+
+def AxiomInfo.ofAxiomVal (v : AxiomVal) : MetaM AxiomInfo := do
+  let info ← Info.ofConstantVal v.toConstantVal
+  return AxiomInfo.mk info v.isUnsafe
+
+def TheoremInfo.ofTheoremVal (v : TheoremVal) : MetaM TheoremInfo := do
+  let info ← Info.ofConstantVal v.toConstantVal
+  return TheoremInfo.mk info
+
+def OpaqueInfo.ofOpaqueVal (v : OpaqueVal) : MetaM OpaqueInfo := do
+  let info ← Info.ofConstantVal v.toConstantVal
+  let value ← prettyPrintTerm v.value
+  return OpaqueInfo.mk info value v.isUnsafe
+
 namespace DocInfo
 
 private def isBlackListed (declName : Name) : MetaM Bool := do
@@ -93,19 +95,44 @@ def ofConstant : (Name × ConstantInfo) → MetaM (Option DocInfo) := λ (name,
 
 def prettyPrint (i : DocInfo) : CoreM String := do
   match i with
-  | axiomInfo i => s!"axiom {i.name} : {←PrettyPrinter.formatTerm i.type} from {i.module}, doc string: {i.doc}"
-  | theoremInfo i => s!"theorem {i.name} : {←PrettyPrinter.formatTerm i.type} from {i.module}, doc string: {i.doc}"
-  | opaqueInfo i => s!"constant {i.name} : {←PrettyPrinter.formatTerm i.type} from {i.module}, doc string: {i.doc}"
+  | axiomInfo i => s!"axiom {i.name} : {←PrettyPrinter.formatTerm i.type}, doc string: {i.doc}"
+  | theoremInfo i => s!"theorem {i.name} : {←PrettyPrinter.formatTerm i.type}, doc string: {i.doc}"
+  | opaqueInfo i => s!"constant {i.name} : {←PrettyPrinter.formatTerm i.type}, doc string: {i.doc}"
   | _ => ""
 
 end DocInfo
 
-def process : Environment → MetaM (List DocInfo) := λ env => do
-  let mut res := []
+namespace Module
+
+def prettyPrint (m : Module) : CoreM String := do
+  let pretty := s!"Module {m.name}, doc string: {m.doc} with members:\n"
+  Array.foldlM (λ p mem => return p ++ "  " ++ (←mem.prettyPrint) ++ "\n") pretty m.members
+
+end Module
+
+def process : MetaM (HashMap Name Module) := do
+  let env ← getEnv
+  let mut res := mkHashMap env.header.moduleNames.size
+  for module in env.header.moduleNames do
+    -- TODO: Check why modules can have multiple doc strings and add that later on
+    let moduleDoc := match getModuleDoc? env module with
+    | none => none
+    | some #[] => none
+    | some doc => doc.get! 0
+
+    res := res.insert module (Module.mk module moduleDoc #[])
+
   for cinfo in env.constants.toList do
-    let dinfo := ←DocInfo.ofConstant cinfo
-    match dinfo with
-    | some d => res := d :: res
+    let d := ←DocInfo.ofConstant cinfo
+    match d with
+    | some dinfo =>
+      match (env.getModuleIdxFor? cinfo.fst) with
+      | some modidx =>
+        -- TODO: Check whether this is still efficient
+        let moduleName := env.allImportedModuleNames.get! modidx
+        let module := res.find! moduleName
+        res := res.insert moduleName {module with members := module.members.push dinfo}
+      | none => panic! "impossible"
     | none => ()
   return res