hmm

ImportGraph/Cli.lean

@@ -14,10 +14,18 @@ open Cli
 open Lean
 open ImportGraph
 
-/-- Write an import graph, represented as a `NameMap (Array Name)` to the ".dot" graph format. -/
-def asDotGraph (graph : NameMap (Array Name)) (header := "import_graph") : String := Id.run do
+/-- Write an import graph, represented as a `NameMap (Array Name)` to the ".dot" graph format.
+    Nodes in the `unused` set will be shaded light gray. -/
+def asDotGraph
+    (graph : NameMap (Array Name)) (unused : NameSet := {}) (header := "import_graph") :
+    String := Id.run do
   let mut lines := #[s!"digraph \"{header}\" " ++ "{"]
+  -- Add node styling for unused nodes
+  lines := lines.push "  node [style=filled];"
+  -- Add nodes and edges
   for (n, is) in graph do
+    if unused.contains n then
+      lines := lines.push s!"  \"{n}\" [fillcolor=lightgray];"
     for i in is do
       lines := lines.push s!"  \"{i}\" -> \"{n}\";"
   lines := lines.push "}"
@@ -37,6 +45,10 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
   searchPathRef.set compile_time_search_path%
   let dotFile ← try unsafe withImportModules #[{module := to}] {} (trustLevel := 1024) fun env => do
     let mut graph := env.importGraph
+    let ctx := { options := {}, fileName := "<input>", fileMap := default }
+    let state := { env }
+    let used ← Prod.fst <$> (CoreM.toIO (env.transitivelyRequiredModules to) ctx state)
+    let unused := graph.fold (fun acc n _ => if used.contains n then acc else acc.insert n) NameSet.empty
     if let Option.some f := from? then
       graph := graph.downstreamOf (NameSet.empty.insert f)
     if ¬(args.hasFlag "include-deps") then
@@ -53,7 +65,7 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
       graph := graph.filterGraph filterMathlibMeta (replacement := `«Mathlib.Tactics»)
     if args.hasFlag "reduce" then
       graph := graph.transitiveReduction
-    return asDotGraph graph
+    return asDotGraph graph (unused := unused)
   catch err =>
     -- TODO: try to build `to` first, so this doesn't happen
     throw <| IO.userError <| s!"{err}\nIf the error above says `unknown package`, " ++

ImportGraph/RequiredModules.lean

@@ -41,18 +41,14 @@ and the constants they use transitively.
 -/
 def Name.transitivelyUsedConstants (n : Name) : CoreM NameSet := do
   let mut usedConstants : NameSet := {}
-  let mut thisBatch : NameSet := NameSet.empty.insert n
-  let mut nextBatch : NameSet := {}
-  let mut done : NameSet := {}
-  while !thisBatch.isEmpty do
-    for n in thisBatch do
-      done := done.insert n
-      for m in (← getConstInfo n).getUsedConstantsAsSet do
-        if !done.contains m then
-          usedConstants := usedConstants.insert m
-          nextBatch := nextBatch.insert m
-    thisBatch := nextBatch
-    nextBatch := {}
+  let mut toProcess : NameSet := NameSet.empty.insert n
+  while !toProcess.isEmpty do
+    let current := toProcess.min.get!
+    toProcess := toProcess.erase current
+    usedConstants := usedConstants.insert current
+    for m in (← getConstInfo current).getUsedConstantsAsSet do
+      if !usedConstants.contains m then
+        toProcess := toProcess.insert m
   return usedConstants
 
 /--
@@ -62,15 +58,24 @@ in the specified declaration were defined.
 Note that this will *not* account for tactics and syntax used in the declaration,
 so the results may not suffice as imports.
 -/
-def Name.transitivelyRequiredModules (n : Name) : CoreM NameSet := do
-  let env ← getEnv
+def Name.transitivelyRequiredModules (n : Name) (env : Environment) : CoreM NameSet := do
   let mut requiredModules : NameSet := {}
   for m in (← n.transitivelyUsedConstants) do
-    match env.getModuleFor? m with
-    | some m =>
-      if ¬ (`Init).isPrefixOf m then
-        requiredModules := requiredModules.insert m
-    | none => pure ()
+    if let some module := env.getModuleFor? m then
+      requiredModules := requiredModules.insert module
+  return requiredModules
+
+/--
+Finds all constants defined in the specified module,
+and identifies all modules contains constants which are transitively required by those constants.
+-/
+def Environment.transitivelyRequiredModules (env : Environment) (module : Name) : CoreM NameSet := do
+  let mut requiredModules : NameSet := {}
+  for (_, ci) in env.constants.map₁.toList do
+    if let some m := env.getModuleFor? ci.name then
+      if m = module then
+        for r in (← ci.name.transitivelyRequiredModules env) do
+          requiredModules := requiredModules.insert r
   return requiredModules
 
 /--