faster?

ImportGraph/Imports.lean

@@ -176,15 +176,13 @@ end Lean.NameMap
 Returns a `List (Name × List Name)` with a key for each module `n` in `amongst`,
 whose corresponding value is the list of modules `m` in `amongst` which are transitively imported by `n`,
 but no declaration in `n` makes use of a declaration in `m`.
-
-The current implementation is too slow to run on the entirety of Mathlib,
-although it should be fine for any sequential chain of imports in Mathlib.
 -/
 def unusedTransitiveImports (amongst : List Name) : CoreM (List (Name × List Name)) := do
   let env ← getEnv
   let transitiveImports := env.importGraph.transitiveClosure
+  let transitivelyRequired ← env.transitivelyRequiredModules' amongst
   amongst.mapM fun n => do return (n,
-    let unused := (transitiveImports.find? n).getD {} \ (← env.transitivelyRequiredModules n)
+    let unused := (transitiveImports.find? n).getD {} \ (transitivelyRequired.find? n |>.getD {})
     amongst.filter (fun m => unused.contains m))
 
 /--

ImportGraph/RequiredModules.lean

@@ -115,6 +115,32 @@ def Environment.transitivelyRequiredModules (env : Environment) (module : Name)
     |>.filter (env.getModuleFor? · = some module)
   (NameSet.ofList constants).transitivelyRequiredModules env
 
+/--
+Computes all the modules transitively required by the specified modules.
+Should be equivalent to calling `transitivelyRequiredModules` on each module, but shares more of the work.
+-/
+partial def Environment.transitivelyRequiredModules' (env : Environment) (modules : List Name) :
+    CoreM (NameMap NameSet) := do
+  let mut c2m : NameMap NameSet := {}
+  let mut result : NameMap NameSet := {}
+  for m in modules do
+    let mut r : NameSet := {}
+    for n in env.header.moduleData[(env.header.moduleNames.getIdx? m).getD 0]!.constNames do
+      c2m ← process c2m n
+      r := r.union ((c2m.find? n).getD {})
+    result := result.insert m r
+  return result
+where process (constantsToModules : NameMap NameSet) (const : Name) : CoreM (NameMap NameSet) := do
+  if constantsToModules.contains const then
+    return constantsToModules
+  let mut c2m := constantsToModules
+  let ci ← getConstInfo const
+  let mut r : NameSet := {}
+  for n in ci.getUsedConstantsAsSet do
+    c2m ← process c2m n
+    r := r.union ((c2m.find? n).getD {})
+  return c2m.insert const r
+
 /--
 Return the names of the modules in which constants used in the current file were defined.