diff --git a/ImportGraph/Cli.lean b/ImportGraph/Cli.lean index abc5963..b4ab70d 100644 --- a/ImportGraph/Cli.lean +++ b/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`, " ++ diff --git a/Main.lean b/Main.lean index 7e65269..2104d4f 100644 --- a/Main.lean +++ b/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. \