feat(Tactic/Linter/DirectoryDependency): add an allow-list for directories with few imports (leanprover-community#24109)

Add an allow-list mechanism to the `directoryDependency` linter: modules covered by that list may only have imports listed in that list 

Co-authored-by: Michael Rothgang <rothgang@math.uni-bonn.de>
Co-authored-by: grunweg <rothgami@math.hu-berlin.de>

Author: Vierkantor

Mathlib/Tactic/Linter/DirectoryDependency.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Anne Baanen
 -/
 import Lean.Elab.Command
+import Lean.Elab.ParseImportsFast
 
 /-! # The `directoryDependency` linter
 
@@ -12,6 +13,14 @@ independent. By specifying that one directory does not import from another, we c
 modularity of Mathlib.
 -/
 
+-- XXX: is there a better long-time place for this
+/-- Parse all imports in a text file at `path` and return just their names:
+this is just a thin wrapper around `Lean.parseImports'`.
+Omit `Init` (which is part of the prelude). -/
+def findImports (path : System.FilePath) : IO (Array Lean.Name) := do
+  return (← Lean.parseImports' (← IO.FS.readFile path) path.toString).imports
+    |>.map (fun imp ↦ imp.module) |>.erase `Init
+
 /-- Find the longest prefix of `n` such that `f` returns `some` (or return `none` otherwise). -/
 def Lean.Name.findPrefix {α} (f : Name → Option α) (n : Name) : Option α := do
   f n <|> match n with
@@ -26,6 +35,13 @@ def Lean.Name.prefixes (n : Name) : NameSet :=
     | str n' _ => n'.prefixes
     | num n' _ => n'.prefixes
 
+/-- Return the immediate prefix of `n` (if any). -/
+def Lean.Name.prefix? (n : Name) : Option Name :=
+  match n with
+    | anonymous => none
+    | str n' _ => some n'
+    | num n' _ => some n'
+
 /-- Collect all prefixes of names in `ns` into a single `NameSet`. -/
 def Lean.Name.collectPrefixes (ns : Array Name) : NameSet :=
   ns.foldl (fun ns n => ns.union n.prefixes) ∅
@@ -117,11 +133,95 @@ def findAny (r : NamePrefixRel) (n₁ : Name) (ns : Array Name) : Option (Name 
         pure ()
     none
 
+/-- Does `r` contain any entries with key `n`? -/
+def containsKey (r : NamePrefixRel) (n : Name) : Bool := NameMap.contains r n
+
 /-- Is a prefix of `n₁` related to a prefix of `n₂`? -/
 def contains (r : NamePrefixRel) (n₁ n₂ : Name) : Bool := (r.find n₁ n₂).isSome
 
+/-- Look up all names `m` which are values of some prefix of `n` under this relation. -/
+def getAllLeft (r : NamePrefixRel) (n : Name) : NameSet := Id.run do
+  let matchingPrefixes := n.prefixes.filter (fun prf ↦ r.containsKey prf)
+  let mut allRules := NameSet.empty
+  for prfix in matchingPrefixes do
+    let some rules := RBMap.find? r prfix | unreachable!
+    allRules := allRules.append rules
+  allRules
+
 end NamePrefixRel
 
+-- TODO: add/extend tests for this linter, to ensure the allow-list works
+-- TODO: move the following three lists to a JSON file, for easier evolution over time!
+-- Future: enforce that allowed and forbidden keys are disjoint
+-- Future: move further directories to use this allow-list instead of the blocklist
+
+/-- `allowedImportDirs` relates module prefixes, specifying that modules with the first prefix
+are only allowed to import modules in the second directory.
+For directories which are low in the import hierarchy, this opt-out approach is both more ergonomic
+(fewer updates needed) and needs less configuration.
+
+We always allow imports of `Init`, `Lean`, `Std`, `Qq` and
+`Mathlib.Init` (as well as their transitive dependencies.)
+-/
+def allowedImportDirs : NamePrefixRel := .ofArray #[
+  -- This is used to test the linter.
+  (`MathlibTest.DirectoryDependencyLinter, `Mathlib.Lean),
+  -- Mathlib.Tactic has large transitive imports: just allow all of mathlib,
+  -- as we don't care about the details a lot.
+  (`MathlibTest.Header, `Mathlib),
+  (`MathlibTest.Header, `Aesop),
+  (`MathlibTest.Header, `ImportGraph),
+  (`MathlibTest.Header, `LeanSearchClient),
+  (`MathlibTest.Header, `Plausible),
+  (`MathlibTest.Header, `ProofWidgets),
+  (`MathlibTest.Header, `Qq),
+  -- (`MathlibTest.Header, `Mathlib.Tactic),
+  -- (`MathlibTest.Header, `Mathlib.Deprecated),
+  (`MathlibTest.Header, `Batteries),
+  (`MathlibTest.Header, `Lake),
+
+  (`Mathlib.Util, `Batteries),
+  (`Mathlib.Util, `Mathlib.Lean),
+  (`Mathlib.Util, `Mathlib.Tactic),
+  -- TODO: reduce this dependency by upstreaming `Data.String.Defs to batteries
+  (`Mathlib.Util.FormatTable, `Mathlib.Data.String.Defs),
+
+  (`Mathlib.Lean, `Batteries.CodeAction),
+  (`Mathlib.Lean, `Batteries.Tactic.Lint),
+  -- TODO: decide if this is acceptable or should be split in a more fine-grained way
+  (`Mathlib.Lean, `Batteries),
+  (`Mathlib.Lean.Expr, `Mathlib.Util),
+  (`Mathlib.Lean.Meta.RefinedDiscrTree, `Mathlib.Util),
+  -- Fine-grained exceptions: TODO decide if these are fine, or should be scoped more broadly.
+  (`Mathlib.Lean.CoreM, `Mathlib.Tactic.ToExpr),
+  (`Mathlib.Lean.CoreM, `Mathlib.Util.WhatsNew),
+  (`Mathlib.Lean.Meta.RefinedDiscrTree, `Mathlib.Tactic.Lemma),
+  (`Mathlib.Lean.Meta.RefinedDiscrTree, `Mathlib.Tactic.TypeStar),
+  (`Mathlib.Lean.Meta.RefinedDiscrTree, `Mathlib.Tactic.ToAdditive),
+  (`Mathlib.Lean.Meta.RefinedDiscrTree, `Mathlib.Tactic), -- split this up further?
+  (`Mathlib.Lean.Meta.RefinedDiscrTree, `Mathlib.Data), -- split this up further?
+  (`Mathlib.Lean.Meta.RefinedDiscrTree, `Mathlib.Algebra.Notation),
+  (`Mathlib.Lean.Meta.RefinedDiscrTree, `Mathlib.Data.Notation),
+  (`Mathlib.Lean.Meta.RefinedDiscrTree, `Mathlib.Data.Array),
+
+  (`Mathlib.Lean.Meta.CongrTheorems, `Mathlib.Data),
+  (`Mathlib.Lean.Meta.CongrTheorems, `Mathlib.Logic),
+  (`Mathlib.Lean.Meta.CongrTheorems, `Mathlib.Order.Defs),
+  (`Mathlib.Lean.Meta.CongrTheorems, `Mathlib.Tactic),
+
+  (`Mathlib.Lean.Expr.ExtraRecognizers, `Mathlib.Data),
+  (`Mathlib.Lean.Expr.ExtraRecognizers, `Mathlib.Order),
+  (`Mathlib.Lean.Expr.ExtraRecognizers, `Mathlib.Logic),
+  (`Mathlib.Lean.Expr.ExtraRecognizers, `Mathlib.Tactic),
+
+  (`Mathlib.Tactic.Linter, `Batteries),
+  -- The Mathlib.Tactic.Linter *module* imports all linters, hence requires all the imports.
+  -- For more fine-grained exceptions of the next two imports, one needs to rename that file.
+  (`Mathlib.Tactic.Linter, `ImportGraph),
+  (`Mathlib.Tactic.Linter, `Mathlib.Tactic.MinImports),
+  (`Mathlib.Tactic.Linter.TextBased, `Mathlib.Data.Nat.Notation),
+]
+
 /-- `forbiddenImportDirs` relates module prefixes, specifying that modules with the first prefix
 should not import modules with the second prefix (except if specifically allowed in
 `overrideAllowedImportDirs`).
@@ -496,25 +596,81 @@ end DirectoryDependency
 
 open DirectoryDependency
 
-@[inherit_doc Mathlib.Linter.linter.directoryDependency]
-def directoryDependencyCheck (mainModule : Name) : CommandElabM (Option MessageData) := do
-  unless Linter.getLinterValue linter.directoryDependency (← getOptions) do
-    return none
-  let env ← getEnv
-  let imports := env.allImportedModuleNames
+/-- Check if one of the imports `imports` to `mainModule` is forbidden by `forbiddenImportDirs`;
+if so, return an error describing how the import transitively arises. -/
+private def checkBlocklist (env : Environment) (mainModule : Name) (imports : Array Name) : Option MessageData := Id.run do
   match forbiddenImportDirs.findAny mainModule imports with
   | some (n₁, n₂) => do
     if let some imported := n₂.prefixToName imports then
       if !overrideAllowedImportDirs.contains mainModule imported then
-        let mut msg := m!"Modules starting with {n₁} are not allowed to import modules starting with {n₂}.
-This module depends on {imported}\n"
+        let mut msg := m!"Modules starting with {n₁} are not allowed to import modules starting with {n₂}. \
+        This module depends on {imported}\n"
         for dep in env.importPath imported do
           msg := msg ++ m!"which is imported by {dep},\n"
-        return some <| msg ++ m!"which is imported by this module.
-(Exceptions can be added to `overrideAllowedImportDirs`.)"
+        return some (msg ++ m!"which is imported by this module. \
+          (Exceptions can be added to `overrideAllowedImportDirs`.)")
+      else none
     else
-      return some m!"Internal error in `directoryDependency` linter: this module claims to depend on a module starting with {n₂} but a module with that prefix was not found in the import graph."
-  | none => pure ()
-  return none
+      return some m!"Internal error in `directoryDependency` linter: this module claims to depend \
+      on a module starting with {n₂} but a module with that prefix was not found in the import graph."
+  | none => none
+
+@[inherit_doc Mathlib.Linter.linter.directoryDependency]
+def directoryDependencyCheck (mainModule : Name) : CommandElabM (Array MessageData) := do
+  unless Linter.getLinterValue linter.directoryDependency (← getOptions) do
+    return #[]
+  let env ← getEnv
+  let imports := env.allImportedModuleNames
+
+  -- If this module is in the allow-list, we only allow imports from directories specified there.
+  -- Collect all prefixes which have a matching entry.
+  let matchingPrefixes := mainModule.prefixes.filter (fun prf ↦ allowedImportDirs.containsKey prf)
+  if matchingPrefixes.isEmpty then
+    -- Otherwise, we fall back to the blocklist `forbiddenImportDirs`.
+    if let some msg := checkBlocklist env mainModule imports then return #[msg] else return #[]
+  else
+    -- We always allow imports in the same directory (for each matching prefix),
+    -- and from `Init`, `Lean`, `Std` and `Qq`.
+    -- We also allow transitive imports of Mathlib.Init, as well as Mathlib.Init itself.
+    let initImports := (← findImports ("Mathlib" / "Init.lean")).append
+      #[`Mathlib.Init, `Mathlib.Tactic.DeclarationNames]
+    let exclude := [
+      `Init, `Std, `Lean, `Qq,
+    ]
+    let importsToCheck := imports.filter (fun imp ↦ !exclude.any (·.isPrefixOf imp))
+      |>.filter (fun imp ↦ !matchingPrefixes.any (·.isPrefixOf imp))
+      |>.filter (!initImports.contains ·)
+
+    -- Find all prefixes which are allowed for one of these directories.
+    let allRules := allowedImportDirs.getAllLeft mainModule
+    -- Error about those imports which are not covered by allowedImportDirs.
+    let mut messages := #[]
+    for imported in importsToCheck do
+      if !allowedImportDirs.contains mainModule imported then
+        let importPath := env.importPath imported
+        let mut msg := m!"Module {mainModule} depends on {imported},\n\
+        but is only allowed to import modules starting with one of {allRules.toArray.qsort (·.toString < ·.toString)}.\n\
+        Note: module {imported}"
+        let mut superseded := false
+        match importPath.toList with
+        | [] => msg := msg ++ " is directly imported by this module"
+        | a :: rest =>
+          -- Only add messages about imports that aren't themselves transitive imports of
+          -- forbidden imports.
+          -- This should prevent redundant messages.
+          if !allowedImportDirs.contains mainModule a then
+            superseded := true
+          else
+            msg := msg ++ s!" is imported by {a},\n"
+            for dep in rest do
+              if !allowedImportDirs.contains mainModule dep then
+                superseded := true
+                break
+              msg := msg ++ m!"which is imported by {dep},\n"
+            msg := msg ++ m!"which is imported by this module."
+            msg := msg ++ "(Exceptions can be added to `allowedImportDirs`.)"
+        if !superseded then
+          messages := messages.push msg
+    return messages
 
 end Mathlib.Linter

Mathlib/Tactic/Linter/Header.lean

@@ -304,7 +304,8 @@ def headerLinter : Linter where run := withSetOptionIn fun stx ↦ do
       return val
   -- The linter skips files not imported in `Mathlib.lean`, to avoid linting "scratch files".
   -- It is however active in the test file `MathlibTest.Header` for the linter itself.
-  unless inMathlib? || mainModule == `MathlibTest.Header do return
+  unless inMathlib? ||
+    mainModule == `MathlibTest.Header || mainModule == `MathlibTest.DirectoryDependencyLinter.Test do return
   unless Linter.getLinterValue linter.style.header (← getOptions) do
     return
   if (← get).messages.hasErrors then
@@ -332,9 +333,12 @@ def headerLinter : Linter where run := withSetOptionIn fun stx ↦ do
   -- Report on broad or duplicate imports.
   broadImportsCheck importIds mainModule
   duplicateImportsCheck importIds
-  if let some msg ← directoryDependencyCheck mainModule then
-    Linter.logLint linter.directoryDependency stx msg
-
+  let errors ← directoryDependencyCheck mainModule
+  if errors.size > 0 then
+    let mut msgs := ""
+    for msg in errors do
+      msgs := msgs ++ "\n\n" ++ (← msg.toString)
+    Linter.logLint linter.directoryDependency stx msgs.trimLeft
   let afterImports := firstNonImport? upToStx
   if afterImports.isNone then return
   let copyright := match upToStx.getHeadInfo with

MathlibTest/DirectoryDependencyLinter/Test.lean

@@ -0,0 +1,34 @@
+/-
+Copyright (c) 2025 Michael Rothgang. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Michael Rothgang
+-/
+import Mathlib.Init
+import Qq
+import Mathlib.Util.AssertExists
+
+/--
+warning: Module MathlibTest.DirectoryDependencyLinter.Test depends on Mathlib.Util.AssertExists,
+but is only allowed to import modules starting with one of [Mathlib.Lean].
+Note: module Mathlib.Util.AssertExists is directly imported by this module
+note: this linter can be disabled with `set_option linter.directoryDependency false`
+---
+warning: The module doc-string for a file should be the first command after the imports.
+Please, add a module doc-string before `/-!# Tests for the `directoryDependency` linter
+-/
+`.
+note: this linter can be disabled with `set_option linter.style.header false`
+-/
+#guard_msgs in
+set_option linter.style.header true in
+
+/-!
+# Tests for the `directoryDependency` linter
+-/
+
+-- Some unit-tests for internal functions.
+#guard Lean.Name.isPrefixOf `Mathlib.Util `Mathlib.Util.Basic == true
+#guard Lean.Name.isPrefixOf `Mathlib.Util `Mathlib.Util.Nested.Basic == true
+#guard Lean.Name.isPrefixOf `Mathlib.Util `Mathlib.Utils.Basic == false
+#guard Lean.Name.isPrefixOf `Mathlib.Foo `Mathlib.Util.Foo == false
+#guard Lean.Name.isPrefixOf `Mathlib.Util `Mathlib.Utils == false

MathlibTest/Header.lean

@@ -24,12 +24,6 @@ note: this linter can be disabled with `set_option linter.style.header false`
 warning: Duplicate imports: 'Mathlib.Tactic.Linter.Header' already imported
 note: this linter can be disabled with `set_option linter.style.header false`
 ---
-warning: Modules starting with MathlibTest.Header are not allowed to import modules starting with Mathlib.Deprecated.
-This module depends on Mathlib.Deprecated.Aliases
-which is imported by this module.
-(Exceptions can be added to `overrideAllowedImportDirs`.)
-note: this linter can be disabled with `set_option linter.directoryDependency false`
----
 warning: The module doc-string for a file should be the first command after the imports.
 Please, add a module doc-string before `/-!# Tests for the `docModule` linter
 -/

scripts/lint-style.lean

@@ -23,13 +23,6 @@ In addition, this checks that
 
 open Cli Lean.Linter Mathlib.Linter.TextBased System.FilePath
 
-/-- Parse all imports in a text file at `path` and return just their names:
-this is just a thin wrapper around `Lean.parseImports'`.
-Omit `Init (which is part of the prelude). -/
-def findImports (path : System.FilePath) : IO (Array Lean.Name) := do
-  return (← Lean.parseImports' (← IO.FS.readFile path) path.toString).imports
-    |>.map (fun imp ↦ imp.module) |>.erase `Init
-
 /-- Additional imports generated by `mk_all`. -/
 def explicitImports : Array Lean.Name := #[`Batteries, `Std]