feat: add --mark-module to hightlight current module when using --include-deps

ImportGraph/Cli.lean

@@ -14,20 +14,38 @@ open Cli
 open Lean
 open ImportGraph
 
-/-- 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. -/
+/--
+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.
+* Nodes which start with the `markedModule` and edges into them will be highlighted in green.
+-/
 def asDotGraph
-    (graph : NameMap (Array Name)) (unused : NameSet := {}) (header := "import_graph") :
+    (graph : NameMap (Array Name))
+    (unused : NameSet := {})
+    (header := "import_graph")
+    (markedModule : Option Name := none) :
     String := Id.run do
+
   let mut lines := #[s!"digraph \"{header}\" " ++ "{"]
   for (n, is) in graph do
     if unused.contains n then
       lines := lines.push s!"  \"{n}\" [style=filled, fillcolor=\"#e0e0e0\"];"
+    else if isInModule markedModule n then
+      -- mark node
+      lines := lines.push s!"  \"{n}\" [style=filled, fillcolor=\"#96ec5b\"];"
     else
       lines := lines.push s!"  \"{n}\";"
     -- Then add edges
     for i in is do
-      lines := lines.push s!"  \"{i}\" -> \"{n}\";"
+      if isInModule markedModule n then
+        if isInModule markedModule i then
+          -- draw the main project close together
+          lines := lines.push s!"  \"{i}\" -> \"{n}\" [weight=100];"
+        else
+          -- mark edges into the main project
+          lines := lines.push s!"  \"{i}\" -> \"{n}\" [penwidth=2, color=\"#71b144\"];"
+      else
+        lines := lines.push s!"  \"{i}\" -> \"{n}\";"
   lines := lines.push "}"
   return "\n".intercalate lines.toList
 
@@ -44,6 +62,7 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
   | none => none
   searchPathRef.set compile_time_search_path%
   let dotFile ← try unsafe withImportModules #[{module := to}] {} (trustLevel := 1024) fun env => do
+    let p := ImportGraph.getModule to
     let mut graph := env.importGraph
     let unused ←
       match args.flag? "to"  with
@@ -58,7 +77,7 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
     let toModule := ImportGraph.getModule to
     let includeLean := args.hasFlag "include-lean"
     let includeStd := args.hasFlag "include-std" || includeLean
-    let includeDeps := args.hasFlag "include-deps" || includeStd
+    let includeDeps := args.hasFlag "include-deps" || args.hasFlag "mark-module" || includeStd
     let filter (n : Name) : Bool :=
       toModule.isPrefixOf n ||
       bif isPrefixOf `Std n then includeStd else
@@ -75,12 +94,16 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
       graph := graph.filterGraph filterMathlibMeta (replacement := `«Mathlib.Tactics»)
     if !args.hasFlag "show-transitive" then
       graph := graph.transitiveReduction
-    return asDotGraph graph (unused := unused)
+
+    let markedModule : Option Name := if args.hasFlag "mark-module" then p else none
+
+    return asDotGraph graph (unused := unused) (markedModule := markedModule)
   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`, " ++
       s!"try if `lake build {to}` fixes the issue"
     throw err
+
   match args.variableArgsAs! String with
   | #[] => writeFile "import_graph.dot" dotFile
   | outputs => for o in outputs do

ImportGraph/CurrentModule.lean

@@ -22,3 +22,11 @@ def getCurrentModule : IO Name := do
     -- Would be better to read the `.defaultTargets` from the
     -- `← getRootPackage` from `Lake`, but I can't make that work with the monads involved.
     return manifest.name.capitalize
+
+/--
+Helper which only returns `true` if the `module` is provided and the name `n` lies
+inside it.
+ -/
+def isInModule (module : Option Name) (n : Name) := match module with
+  | some m => m.isPrefixOf n
+  | none => false

Main.lean

@@ -21,8 +21,9 @@ def graph : Cmd := `[Cli|
     "from" : ModuleName;  "Only show the downstream dependencies of the specified module."
     "exclude-meta";       "Exclude any files starting with `Mathlib.[Tactic|Lean|Util|Mathport]`."
     "include-deps";       "Include used files from other libraries (not including Lean itself and `std`)"
-    "include-std";        "Include used files from the Lean standard library (implies 'include-deps')"
-    "include-lean";       "Include used files from Lean itself (implies 'include-deps' and 'include-std')"
+    "mark-module";        "Visually highlight the current module. (implies `--include-deps`)"
+    "include-std";        "Include used files from the Lean standard library (implies `--include-deps`)"
+    "include-lean";       "Include used files from Lean itself (implies `--include-deps` and `--include-std`)"
 
   ARGS:
     ...outputs : String;  "Filename(s) for the output. \