chore: performance of lake builtin-lint (experiment)

src/Lean/Linter/EnvLinter/Frontend.lean

@@ -195,4 +195,20 @@ def getDeclsInPackage (pkg : Name) : CoreM (Array Name) := do
       decls.push declName
     else decls
 
+/--
+Get the list of declarations whose defining module is exactly `modName`.
+
+Unlike `getDeclsInPackage`, this does *not* match modules by prefix. Used by the
+per-module linting flow, where each module is imported as the root so its
+private decls are loaded (`Root ≥ all`), and we restrict linting to the decls
+that module itself defined.
+-/
+def getDeclsInModule (modName : Name) : CoreM (Array Name) := do
+  let env ← getEnv
+  let some idx := env.getModuleIdx? modName | return #[]
+  return env.constants.map₁.fold (init := #[]) fun decls declName _ =>
+    if env.getModuleIdxFor? declName == some idx then
+      decls.push declName
+    else decls
+
 end Lean.Linter.EnvLinter

src/lake/Lake/CLI/BuiltinLint.lean

@@ -23,6 +23,14 @@ public structure Args where
   only : Array Name := #[]
   /-- The list of root modules to lint. -/
   mods : Array Name := #[]
+  /--
+  Parallel to `mods`: for each top-level target, the buildable modules in its
+  library. Used only by the experimental per-module flow, which imports every
+  package module as a root in a single `importModules` call so each module's
+  private decls are visible (`Root ≥ all`). Populated by Lake from
+  `LeanLib.getModuleArray`; the barrel flow ignores this field.
+  -/
+  pkgMods : Array (Array Name) := #[]
 
 public def leanOptOverrides (args : Args) : LeanOptions :=
   let enableAll : Array LeanOption :=
@@ -52,22 +60,27 @@ private def collectTextLints
     else
       acc
 
-public def run (args : Args) : IO UInt32 := do
-  let mods := args.mods
-  if mods.isEmpty then
-    IO.eprintln "lake lint: no modules specified for builtin linting"
-    return 1
+/--
+For benchmarking, set `LEAN_LINT_SKIP_TEXT=1` to skip text-linter collection
+in both the barrel and per-module flows. Text-linter results come from
+`lintLogExt` populated at build time, so they're essentially free, but
+omitting them removes one cross-flow variable from timing comparisons.
+-/
+private def skipTextLints : IO Bool :=
+  return (← IO.getEnv "LEAN_LINT_SKIP_TEXT").any (· == "1")
 
+private def runBarrel (args : Args) : IO UInt32 := do
   let runOnly := if args.only.isEmpty then none else some args.only.toList
   let scope := args.scope
   let envLinterModule : Import := { module := `Lean.Linter.EnvLinter }
+  let skipText ← skipTextLints
 
   let mut anyFailed := false
-  for mod in mods do
+  for mod in args.mods do
     unsafe Lean.enableInitializersExecution
     let env ← importModules #[{ module := mod }, envLinterModule] {} (trustLevel := 1024) (loadExts := true)
 
-    let textEntries := collectTextLints env args mod.getRoot
+    let textEntries := if skipText then #[] else collectTextLints env args mod.getRoot
     let textFailed := !textEntries.isEmpty
     if textFailed then
       IO.println s!"-- Text linter diagnostics in {mod}:"
@@ -97,4 +110,117 @@ public def run (args : Args) : IO UInt32 := do
 
   return if anyFailed then 1 else 0
 
+/--
+Parse `LEAN_LINT_LEVEL` for the experimental per-module flow.
+
+Accepts `exported`, `server`, `private`. Defaults to `.exported` (the
+public-only setting that motivates the experiment); `private` lets the loop
+run on non-module-system roots that `.exported` would reject.
+-/
+private def parseLintLevel : IO OLeanLevel := do
+  match (← IO.getEnv "LEAN_LINT_LEVEL") with
+  | none | some "exported" => return .exported
+  | some "server"          => return .server
+  | some "private"         => return .private
+  | some other =>
+    IO.eprintln s!"lake lint: unknown LEAN_LINT_LEVEL={other} (expected exported|server|private)"
+    return .exported
+
+/--
+Experimental per-module flow gated on `LEAN_LINT_PER_MODULE=1`.
+
+For each top-level target in `args.mods`, we import every module in the
+target's library (provided by Lake as `args.pkgMods`) in a single
+`importModules` call, listing each module as a root. `importModulesCore`
+deduplicates underlying olean reads via the DAG walk, and every module
+listed as a root gets `importAll := true` — i.e. `Root ≥ all`, with private
+decls visible. Transitive non-package deps (Init, Lean libs, …) load at
+`LEAN_LINT_LEVEL` (default `.exported`).
+
+Then a single linter pass walks each module's decls (filtered by
+`getModuleIdxFor?`) and runs the configured env-linter set against them.
+
+Compared to the original per-module-loop, this design:
+  * runs *one* `importModules` call instead of N, avoiding the dynlib
+    re-init segfaults observed on Mathlib;
+  * keeps the same coverage (every module's privates visible);
+  * still benefits from `.exported` for transitive non-package deps.
+
+`LEAN_LINT_LEVEL` selects the transitive level (default `.exported`). Use
+`private` for non-module-system roots that `.exported`/`.server` reject.
+-/
+private def runPerModule (args : Args) : IO UInt32 := do
+  let runOnly := if args.only.isEmpty then none else some args.only.toList
+  let scope := args.scope
+  let envLinterModule : Import := { module := `Lean.Linter.EnvLinter }
+  let level ← parseLintLevel
+  let skipText ← skipTextLints
+  let trace := (← IO.getEnv "LEAN_LINT_TRACE").any (· == "1")
+
+  if args.pkgMods.size != args.mods.size then
+    IO.eprintln s!"lake lint: per-module flow requires `pkgMods` (size {args.pkgMods.size}) to match `mods` (size {args.mods.size}). \
+      This field must be populated by Lake's CLI; see `runBuiltinLint` in `Main.lean`."
+    return 1
+
+  let mut anyFailed := false
+  for h : i in [:args.mods.size] do
+    let topMod := args.mods[i]
+    let pkgModules := args.pkgMods[i]!
+
+    if trace then
+      let t ← IO.monoMsNow
+      IO.eprintln s!"-- [+{t}ms] importing {pkgModules.size} modules as roots (target {topMod})"
+
+    -- Single `importModules` call with every package module as a root.
+    -- `importAll := true` for each → `Root ≥ all` → privates loaded.
+    unsafe Lean.enableInitializersExecution
+    let roots : Array Import :=
+      pkgModules.map (fun m => ({ module := m } : Import)) |>.push envLinterModule
+    let env ← importModules roots {} (trustLevel := 1024) (loadExts := true) (level := level)
+
+    if trace then
+      let t ← IO.monoMsNow
+      IO.eprintln s!"-- [+{t}ms] import done; running linters"
+
+    let textEntries :=
+      if skipText then #[] else collectTextLints env args topMod.getRoot
+    let textFailed := !textEntries.isEmpty
+    if textFailed then
+      IO.println s!"-- Text linter diagnostics in {topMod}:"
+      for e in textEntries do
+        IO.print e.message.toString
+
+    -- One CoreM pass: enumerate each module's decls and lint them.
+    let (perModFailed, _) ← CoreM.toIO (ctx := { fileName := "", fileMap := default }) (s := { env }) do
+      let linters ← Linter.EnvLinter.getChecks (scope := scope) (runOnly := runOnly)
+      if linters.isEmpty then return false
+      let mut anyMod := false
+      for mod in pkgModules do
+        let decls ← Linter.EnvLinter.getDeclsInModule mod
+        if decls.isEmpty then continue
+        let results ← Linter.EnvLinter.lintCore decls linters
+        let failed := results.any (!·.2.isEmpty)
+        if failed then
+          anyMod := true
+          let fmtResults ←
+            Linter.EnvLinter.formatLinterResults results decls
+              (groupByFilename := true) (useErrorFormat := true)
+              s!"in {mod}" (scope := if args.only.isEmpty then scope else .all) .medium linters.size
+          IO.print (← fmtResults.toString)
+      return anyMod
+
+    if !textFailed && !perModFailed then
+      IO.println s!"-- Linting passed for {topMod}."
+    if textFailed || perModFailed then
+      anyFailed := true
+
+  return if anyFailed then 1 else 0
+
+public def run (args : Args) : IO UInt32 := do
+  if args.mods.isEmpty then
+    IO.eprintln "lake lint: no modules specified for builtin linting"
+    return 1
+  let perModule := (← IO.getEnv "LEAN_LINT_PER_MODULE").any (· == "1")
+  if perModule then runPerModule args else runBarrel args
+
 end Lake.BuiltinLint

src/lake/Lake/CLI/Main.lean

@@ -977,8 +977,17 @@ private def runBuiltinLint
   let mods := if specifiedMods.isEmpty then ws.defaultTargetRoots else specifiedMods
   if mods.isEmpty then
     error "no modules specified and there are no applicable default targets"
+  -- Resolve each top-level target to the full list of buildable modules in its
+  -- library. The per-module flow in `BuiltinLint.run` consumes this to do a
+  -- single `importModules` call per target with every module as a root, which
+  -- gives each module `Root ≥ all` (private decls visible) and avoids the
+  -- segfault/freeze of doing one `importModules` per module.
+  let pkgMods : Array (Array Lean.Name) ← mods.mapM fun modName => do
+    match ws.findTargetModule? modName with
+    | some mod => return (← mod.lib.getModuleArray).map (·.name)
+    | none     => return #[modName]
   let args := opts.builtinLint
-  let args := {args with mods}
+  let args := {args with mods, pkgMods}
   let specs ← parseTargetSpecs ws (mods.map (s!"+{·}") |>.toList)
   let lintOpts := BuiltinLint.leanOptOverrides args
   let overrides : Lean.NameMap Lean.LeanOptions :=

tests/lake/tests/builtin-lint/Clean.lean

@@ -1,3 +1,5 @@
+module
+
 -- No linter violations here.
 theorem correctTheorem : 1 = 1 := rfl
 

tests/lake/tests/builtin-lint/ExtraViolations.lean

@@ -1,3 +1,5 @@
+module
+
 import Linters
 
 -- This name ends with 'Extra' — the dummyExtra linter should flag it.

tests/lake/tests/builtin-lint/Linters.lean

@@ -1,4 +1,7 @@
-import Lean.Linter.EnvLinter
+module
+
+public import Lean.Linter.EnvLinter
+public import Lean.Elab.Command
 
 open Lean Meta Lean.Linter Lean.Elab.Command
 

tests/lake/tests/builtin-lint/Main.lean

@@ -1,3 +1,5 @@
+module
+
 import Main.Sub
 
 -- This uses `def` for a Prop — the `defLemma` linter should flag this.

tests/lake/tests/builtin-lint/Main/Sub.lean

@@ -1,3 +1,5 @@
+module
+
 -- Env-linter violation: `def` on a Prop — should be caught by `defLemma`
 -- regardless of build-time options, since env linters run post-build.
 def shouldBeTheoremInSub : 1 = 1 := rfl

tests/lake/tests/builtin-lint/TextLints.lean

@@ -1,3 +1,5 @@
+module
+
 -- `linter.unusedVariables` has `defValue := true`, so it fires during the normal
 -- build that backs `lake lint --builtin-lint`. The warning is captured in the
 -- olean via `lintLogExt` and re-emitted by `lake lint`.

tests/lake/tests/builtin-lint/dep/Dep.lean

@@ -1,3 +1,5 @@
+module
+
 -- A public def without a docstring. `linter.missingDocs` (defValue=false) only
 -- fires when `linter.all=true` is injected during the build. This file lives in
 -- a non-root package (`dep`), so the option override must be keyed by the