Merge pull request #29 from leanprover-community/eugster/direct-deps

feat: add --include-direct

ImportGraph/Cli.lean

@@ -17,18 +17,26 @@ 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.
-* Nodes which start with the `markedModule` and edges into them will be highlighted in green.
+* If `markedModule` is provided:
+  * Nodes which start with the `markedModule` will be highlighted in green and drawn closer together.
+  * Edges from `directDeps` into the module are highlighted in green
+  * Nodes in `directDeps` are marked with a green border and green text.
 -/
 def asDotGraph
     (graph : NameMap (Array Name))
     (unused : NameSet := {})
     (header := "import_graph")
-    (markedModule : Option Name := none) :
+    (markedModule : Option Name := none)
+    (directDeps : NameSet := {}) :
     String := Id.run do
 
   let mut lines := #[s!"digraph \"{header}\" " ++ "{"]
   for (n, is) in graph do
-    if unused.contains n then
+    if markedModule.isSome ∧ directDeps.contains n then
+      -- note: `fillcolor` defaults to `color` if not specified
+      let fill := if unused.contains n then "#e0e0e0" else "white"
+      lines := lines.push s!"  \"{n}\" [style=filled, fontcolor=\"#4b762d\", color=\"#71b144\", fillcolor=\"{fill}\", penwidth=2];"
+    else if unused.contains n then
       lines := lines.push s!"  \"{n}\" [style=filled, fillcolor=\"#e0e0e0\"];"
     else if isInModule markedModule n then
       -- mark node
@@ -77,13 +85,38 @@ 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" || args.hasFlag "mark-module" || includeStd
+    let includeDeps := args.hasFlag "include-deps" || includeStd
+    -- Note: `includeDirect` does not imply `includeDeps`!
+    -- e.g. if the module contains `import Lean`, the node `Lean` will be included with
+    -- `--include-direct`, but not included with `--include-deps`.
+    let includeDirect := args.hasFlag "include-direct"
+
+    -- `directDeps` contains files which are not in the module
+    -- but directly imported by a file in the module
+    let directDeps : NameSet := graph.fold (init := .empty) (fun acc n deps =>
+      if toModule.isPrefixOf n then
+        deps.filter (!toModule.isPrefixOf ·) |>.foldl (init := acc) NameSet.insert
+      else
+        acc)
+
     let filter (n : Name) : Bool :=
       toModule.isPrefixOf n ||
       bif isPrefixOf `Std n then includeStd else
       bif isPrefixOf `Lean n || isPrefixOf `Init n then includeLean else
       includeDeps
-    graph := graph.filterMap (fun n i => if filter n then (i.filter filter) else none)
+    let filterDirect (n : Name) : Bool :=
+      includeDirect ∧ directDeps.contains n
+
+    graph := graph.filterMap (fun n i =>
+      if filter n then
+        -- include node regularly
+        (i.filter (fun m => filterDirect m || filter m))
+      else if filterDirect n then
+        -- include node as direct dependency; drop any further deps.
+        some #[]
+      else
+        -- not included
+        none)
     if args.hasFlag "exclude-meta" then
       -- Mathlib-specific exclusion of tactics
       let filterMathlibMeta : Name → Bool := fun n => (
@@ -97,7 +130,7 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
 
     let markedModule : Option Name := if args.hasFlag "mark-module" then p else none
 
-    return asDotGraph graph (unused := unused) (markedModule := markedModule)
+    return asDotGraph graph (unused := unused) (markedModule := markedModule) (directDeps := directDeps)
   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`, " ++

Main.lean

@@ -10,7 +10,7 @@ open Cli
 
 /-- Setting up command line options and help text for `lake exe graph`. -/
 def graph : Cmd := `[Cli|
-  graph VIA importGraphCLI; ["0.0.2"]
+  graph VIA importGraphCLI; ["0.0.3"]
   "Generate representations of a Lean import graph. \
    By default generates the import graph up to `Mathlib`. \
    If you are working in a downstream project, use `lake exe graph --to MyProject`."
@@ -20,10 +20,11 @@ def graph : Cmd := `[Cli|
     to : ModuleName;      "Only show the upstream imports of the specified module."
     "from" : ModuleName;  "Only show the downstream dependencies of the specified module."
     "exclude-meta";       "Exclude any files starting with `Mathlib.[Tactic|Lean|Util|Mathport]`."
+    "include-direct";     "Include directly imported files from other libraries"
     "include-deps";       "Include used files from other libraries (not including Lean itself and `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`)"
+    "mark-module";        "Visually highlight the current module (used in combination with some `--include-XXX`)."
 
   ARGS:
     ...outputs : String;  "Filename(s) for the output. \