From 9eccd341f0dd07265437ad2f4a656159370c5197 Mon Sep 17 00:00:00 2001
From: Jon Eugster <eugster.jon@gmail.com>
Date: Mon, 2 Sep 2024 11:13:26 +0200
Subject: [PATCH] split into --include-direct-deps and --mark-module

---
 ImportGraph/Cli.lean | 16 ++++++++++------
 Main.lean            | 19 ++++++++++---------
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/ImportGraph/Cli.lean b/ImportGraph/Cli.lean
index c1cac5b..ed051f5 100644
--- a/ImportGraph/Cli.lean
+++ b/ImportGraph/Cli.lean
@@ -30,7 +30,7 @@ def asDotGraph
 
   let mut lines := #[s!"digraph \"{header}\" " ++ "{"]
   for (n, is) in graph do
-    if directDeps.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];"
@@ -84,21 +84,25 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
     let includeLean := args.hasFlag "include-lean"
     let includeStd := args.hasFlag "include-std" || includeLean
     let includeDeps := args.hasFlag "include-deps" || includeStd
+    let includeDirectDeps := args.hasFlag "include-direct-deps" || args.hasFlag "mark-module"
 
-    -- `directDeps` are exempt from being filtered out.
-    let directDeps := if args.hasFlag "mark-module" then
+    -- If the flag is set, `directDeps` contains files which are not in the module
+    -- but directly imported by a file in the module
+    let directDeps := if includeDirectDeps then
         graph.fold (fun acc n deps => if toModule.isPrefixOf n then
-          -- append all dependencies of `n` which are not in the module
           deps.filter (!toModule.isPrefixOf ·) |>.foldl NameSet.insert acc
         else acc) NameSet.empty
       else NameSet.empty
 
     let filter (n : Name) : Bool :=
-      toModule.isPrefixOf n || directDeps.contains n ||
+      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)
+    graph := graph.filterMap (fun n i => if filter n then
+      (i.filter (fun m => directDeps.contains m || filter m)) else
+        -- include direct dep but without any further deps
+        if directDeps.contains n then some #[] else none)
     if args.hasFlag "exclude-meta" then
       -- Mathlib-specific exclusion of tactics
       let filterMathlibMeta : Name → Bool := fun n => (
diff --git a/Main.lean b/Main.lean
index c84f5d6..b314193 100644
--- a/Main.lean
+++ b/Main.lean
@@ -10,20 +10,21 @@ 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`."
 
   FLAGS:
-    "show-transitive";    "Show transitively redundant edges."
-    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-deps";       "Include used files from other libraries (not including Lean itself and `std`)"
-    "mark-module";        "Include directly imported files from other libraries and visually highlight the current module"
-    "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`)"
+    "show-transitive";     "Show transitively redundant edges."
+    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-deps"; "Include directly imported files from other libraries"
+    "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-direct-deps`)"
 
   ARGS:
     ...outputs : String;  "Filename(s) for the output. \