From 9ac4eb5062639ca5ffb46aa4168c66675ff33555 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20B=C3=B6ving?= <hargonix@gmail.com>
Date: Sun, 16 Jan 2022 14:22:53 +0100
Subject: [PATCH] feat: imports, imported by

---
 DocGen4/Output/Module.lean | 52 +++++++++++++++++++++++++++++++++++---
 DocGen4/Process.lean       | 27 ++++++++++++--------
 lean-toolchain             |  2 +-
 3 files changed, 67 insertions(+), 14 deletions(-)

diff --git a/DocGen4/Output/Module.lean b/DocGen4/Output/Module.lean
index 58d75e1..914e26e 100644
--- a/DocGen4/Output/Module.lean
+++ b/DocGen4/Output/Module.lean
@@ -89,22 +89,68 @@ def docInfoToHtml (module : Name) (doc : DocInfo) : HtmlM Html := do
     </div>
   </div>
 
-def moduleToNavLink (module : Name) : Html :=
+def declarationToNavLink (module : Name) : Html :=
   <div «class»="nav_link">
     <a «class»="break_within" href={s!"#{module.toString}"}>{module.toString}</a>
   </div>
 
+-- TODO: Similar functions are used all over the place, we should dedup them
+def moduleToNavLink (module : Name) : HtmlM Html := do
+  <a href={←moduleNameToLink module}>{module.toString}</a>
+
+def getImports (module : Name) : HtmlM (Array Name) := do
+  let res ← getResult
+  let some idx ← res.moduleNames.findIdx? (. == module) | unreachable!
+  let adj := res.importAdj.get! idx
+  let mut imports := #[]
+  for i in [:adj.size] do
+    if adj.get! i then
+      imports := imports.push (res.moduleNames.get! i)
+  imports
+
+def getImportedBy (module : Name) : HtmlM (Array Name) := do
+  let res ← getResult
+  let some idx ← res.moduleNames.findIdx? (. == module) | unreachable!
+  let adj := res.importAdj
+  let mut impBy := #[]
+  for i in [:adj.size] do
+    if adj.get! i |>.get! idx then
+      impBy := impBy.push (res.moduleNames.get! i)
+  impBy
+
+def importedByHtml (moduleName : Name) : HtmlM (Array Html) := do
+  let imports := (←getImportedBy moduleName) |>.qsort Name.lt
+  imports.mapM (λ i => do <li>{←moduleToNavLink i}</li>)
+
+
+def importsHtml (moduleName : Name) : HtmlM (Array Html) := do
+  let imports := (←getImports moduleName) |>.qsort Name.lt
+  imports.mapM (λ i => do <li>{←moduleToNavLink i}</li>)
+
 def internalNav (members : Array Name) (moduleName : Name) : HtmlM Html := do
   <nav «class»="internal_nav">
     <h3><a «class»="break_within" href="#top">{moduleName.toString}</a></h3>
     -- TODO: Proper source links
     <p «class»="gh_nav_link"><a href={←getSourceUrl moduleName none}>source</a></p>
-    [members.map moduleToNavLink]
+    <div «class»="imports">
+      <details>
+        <summary>Imports</summary>
+        <ul>
+          [←importsHtml moduleName]
+        </ul>
+      </details>
+      <details>
+        <summary>Imported by</summary>
+        <ul>
+          [←importedByHtml moduleName]
+        </ul>
+      </details>
+    </div>
+    [members.map declarationToNavLink]
   </nav>
 
 def moduleToHtml (module : Module) : HtmlM Html := withReader (setCurrentName module.name) do
   let docInfos ← module.members.mapM (λ i => docInfoToHtml module.name i)
-  -- TODO: This is missing imports, imported by
   templateExtends (baseHtmlArray module.name.toString) $ #[
     ←internalNav (module.members.map DocInfo.getName) module.name,
     Html.element "main" false #[] docInfos
diff --git a/DocGen4/Process.lean b/DocGen4/Process.lean
index eb58f72..bc9eebd 100644
--- a/DocGen4/Process.lean
+++ b/DocGen4/Process.lean
@@ -129,7 +129,7 @@ def prettyPrintTerm (expr : Expr) : MetaM CodeWithInfos := do
     options := ← getOptions
     currNamespace := ← getCurrNamespace
     openDecls := ← getOpenDecls
-    fileMap := arbitrary
+    fileMap := default
   }
   tagExprInfos ctx infos tt
 
@@ -325,6 +325,8 @@ structure AnalyzerResult where
   moduleNames : Array Name
   moduleInfo : HashMap Name Module
   hierarchy : Hierarchy
+  -- Indexed by ModIdx
+  importAdj : Array (Array Bool)
   deriving Inhabited
 
 def process : MetaM AnalyzerResult := do
@@ -343,24 +345,29 @@ def process : MetaM AnalyzerResult := do
     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"
+      let some modidx ← env.getModuleIdxFor? cinfo.fst | unreachable!
+      let moduleName := env.allImportedModuleNames.get! modidx
+      let module := res.find! moduleName
+      res := res.insert moduleName {module with members := module.members.push dinfo}
     | none => ()
 
-  -- This could probably be faster if we did an insertion sort above instead
+  -- TODO: This is definitely not the most efficient way to store this data
+  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
   for (moduleName, module) in res.toArray do
     res := res.insert moduleName {module with members := module.members.qsort DocInfo.lineOrder}
+    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)
 
   return {
     name2ModIdx := env.const2ModIdx,
     moduleNames := env.header.moduleNames,
     moduleInfo := res,
-    hierarchy := Hierarchy.fromArray env.header.moduleNames
+    hierarchy := Hierarchy.fromArray env.header.moduleNames,
+    importAdj := adj
   }
 
 end DocGen4
diff --git a/lean-toolchain b/lean-toolchain
index ff3b224..19c9eec 100644
--- a/lean-toolchain
+++ b/lean-toolchain
@@ -1 +1 @@
-leanprover/lean4:nightly-2022-01-15
+leanprover/lean4:nightly-2022-01-16