fix

Author: kim-em

ImportGraph/RequiredModules.lean

@@ -123,31 +123,40 @@ partial def Environment.transitivelyRequiredModules' (env : Environment) (module
     CoreM (NameMap NameSet) := do
   let N := env.header.moduleNames.size
   let mut c2m : NameMap (BitVec N) := {}
+  let mut pushed : NameSet := {}
   let mut result : NameMap NameSet := {}
   for m in modules do
     if verbose then
       IO.println s!"Processing module {m}"
     let mut r : BitVec N := 0
     for n in env.header.moduleData[(env.header.moduleNames.getIdx? m).getD 0]!.constNames do
+      if ! n.isInternal then
       -- This is messy: Mathlib is big enough that writing a recursive function causes a stack overflow.
       -- So we use an explicit stack instead. We visit each constant twice:
       -- once to record the constants transitively used by it,
       -- and again to record the modules which defined those constants.
-      let mut stack : Array (ConstantInfo × Option NameSet) := #[⟨← getConstInfo n, none⟩]
+      let mut stack : List (Name × Option NameSet) := [⟨n, none⟩]
+      pushed := pushed.insert n
       while !stack.isEmpty do
-        let (ci, used?) := stack.back
-        stack := stack.pop
-        match used? with
-        | none =>
-          if !c2m.contains ci.name then
-            let used := ci.getUsedConstantsAsSet
-            stack := stack.push ⟨ci, some used⟩
-            for u in used do
-              if !c2m.contains u then
-                stack := stack.push ⟨← getConstInfo u, none⟩
-        | some used =>
-          let transitivelyUsed : BitVec N := used.fold (init := toBitVec used) (fun s u => s ||| ((c2m.find? u).getD 0))
-          c2m := c2m.insert ci.name transitivelyUsed
+        match stack with
+        | [] => panic! "Stack is empty"
+        | (c, used?) :: tail =>
+          stack := tail
+          match used? with
+          | none =>
+            if !c2m.contains c then
+              let used := (← getConstInfo c).getUsedConstantsAsSet
+              stack := ⟨c, some used⟩ :: stack
+              for u in used do
+                if !pushed.contains u then
+                  stack := ⟨u, none⟩ :: stack
+                  pushed := pushed.insert u
+          | some used =>
+            let usedModules : NameSet :=
+              used.fold (init := {}) (fun s u => if let some m := env.getModuleFor? u then s.insert m else s)
+            let transitivelyUsed : BitVec N :=
+              used.fold (init := toBitVec usedModules) (fun s u => s ||| ((c2m.find? u).getD 0))
+            c2m := c2m.insert c transitivelyUsed
       r := r ||| ((c2m.find? n).getD 0)
     result := result.insert m (toNameSet r)
   return result

ImportGraph/UnusedTransitiveImports.lean

@@ -0,0 +1,23 @@
+import ImportGraph.Imports
+
+open Lean
+
+def Core.withImportModules (modules : Array Name) {α} (f : CoreM α) : IO α := do
+  searchPathRef.set compile_time_search_path%
+  unsafe Lean.withImportModules (modules.map (fun m => {module := m})) {} (trustLevel := 1024)
+    fun env => Prod.fst <$> Core.CoreM.toIO
+        (ctx := { fileName := "<CoreM>", fileMap := default }) (s := { env := env }) do f
+
+/--
+`lake exe unused_transitive_imports m1 m2 ...`
+
+For each specified module `m`, prints those `n` from the argument list which are imported, but transitively unused by `m`.
+-/
+def main (args : List String) : IO UInt32 := do
+  let (flags, args) := args.partition (fun s => s.startsWith "-")
+  let mut modules := args.map (fun s => s.toName)
+  Core.withImportModules modules.toArray do
+    let r ← unusedTransitiveImports modules (verbose := flags.contains "-v" || flags.contains "--verbose")
+    for (n, u) in r do
+      IO.println s!"{n}: {u}"
+    return 0

ImportGraphTest/Imports.lean

@@ -1,5 +1,6 @@
 import ImportGraph.Imports
 import ImportGraph.RequiredModules
+import ImportGraphTest.Used
 
 open Lean
 
@@ -36,6 +37,7 @@ elab "#unused_transitive_imports" names:ident* : command => do
     logInfo <| s!"Transitively unused imports of {n}:\n{"\n".intercalate (u.map (fun i => s!"  {i}"))}"
 
 -- This test case can be removed after nightly-2024-10-24, because these imports have been cleaned up.
+-- It should be replaced with another test case!
 /--
 info: Transitively unused imports of Init.Control.StateRef:
   Init.System.IO
@@ -46,6 +48,13 @@ info: Transitively unused imports of Init.System.IO:
 #guard_msgs in
 #unused_transitive_imports Init.Control.StateRef Init.System.IO Init.Control.Reader Init.Control.Basic
 
+/--
+info: Transitively unused imports of ImportGraphTest.Used:
+  ImportGraphTest.Unused
+-/
+#guard_msgs in
+#unused_transitive_imports ImportGraphTest.Used ImportGraphTest.Unused Init.Control.Reader
+
 -- This is a spurious unused transitive import, because it relies on notation from `Init.Core`.
 /--
 info: Transitively unused imports of Init.Control.Basic:

lakefile.toml

@@ -28,5 +28,12 @@ root = "Main"
 # Remove this line if you do not need such functionality.
 supportInterpreter = true
 
+# `lake exe unused_transitive_imports` prints unused transitive imports from amongst a given list of modules.
+[[lean_exe]]
+name = "unused_transitive_imports"
+root = "ImportGraph.UnusedTransitiveImports"
+supportInterpreter = true
+
 [[lean_lib]]
 name = "ImportGraphTest"
+