refactor(Tactic/Linter/Header): make the header linter stateful

Mathlib/Tactic/Linter/Header.lean

@@ -6,11 +6,9 @@ Authors: Michael Rothgang, Damiano Testa, Thomas R. Murrills
 module
 
 public meta import Lean.Elab.Command
-public meta import Std.Sync.Mutex
 public import Lean.Parser.Module
 public import Mathlib.Tactic.Linter.DirectoryDependency
 public meta import Lean.Linter.Basic
-public import Std.Sync.Mutex
 
 /-!
 # The "header" linter
@@ -43,6 +41,12 @@ The linter checks if it is linting the first command or not by (quickly) parsing
 checking whether the end position of the header is the same as the start position of the current
 command. If so, it re-parses the header more slowly and checks that the current (first) command is
 a module doc-string (unless the command is exempted for some reason).
+
+The linter is a stateful linter (`Lean.Elab.Command.registerStatefulLinter`). Its state records
+a verdict for the module: the header checks ran, or the module is exempt. Every later command
+takes a fast path that does no parsing at all. The state also stores the parsed module header.
+Other linters, such as `minImports` and `longLine`, read the parsed header from the state
+instead of a re-parse of the file.
 -/
 
 meta section
@@ -191,15 +195,16 @@ def isInLibraryRoot (modName : Name) : IO Bool := do
     return res.imports.any (·.module == modName)
   else return false
 
-/-- `inLibraryRootMutex` caches whether the current file is imported in the library root file
-(e.g. `Mathlib.lean`), as computed by `isInLibraryRoot`. It is
-* `none` at initialization time;
-* `some true` if the `header` linter has already discovered that the current file
-  is imported in the library root file;
-* `some false` if the `header` linter has already discovered that the current file
-  is *not* imported in the library root file.
--/
-initialize inLibraryRootMutex : Std.Mutex (Option Bool) ← Std.Mutex.new none
+/-- `HeaderState` is the persistent state of the `header` linter. -/
+public structure HeaderState where
+  /-- `true` after the linter reaches a verdict for the module: the header checks ran, or the
+  module is exempt from the checks. Every later command takes a fast path. -/
+  settled : Bool := false
+  /-- The syntax of the module header: the copyright comment and the `import` commands, with
+  source positions in the current file. The value is `.missing` when the header checks did not
+  run. Other linters read this field to avoid a re-parse of the header. -/
+  headerSyntax : Syntax := .missing
+  deriving Inhabited
 
 /--
 The "header" style linter checks that a file starts with
@@ -284,6 +289,7 @@ def headerTestFiles : NameSet := .ofList [
   `MathlibTest.Linter.Header.Basic,
   `MathlibTest.Linter.Header.Fail,
   `MathlibTest.Linter.Header.Verso,
+  `MathlibTest.Linter.Header.MinImportsPayload,
   `MathlibTest.DirectoryDependencyLinter.Test]
 
 /-- Check the `Syntax` `imports` for broad imports:
@@ -332,47 +338,40 @@ def duplicateImportsCheck (imports : Array ImportRef) : CommandElabM Unit := do
       importsSoFar := importsSoFar.insert imp.toImport
 
 @[inherit_doc Mathlib.Linter.linter.style.header]
-def headerLinter : Linter where run := withSetOptionIn fun stx ↦ do
+def headerPost (stx : Syntax) (self : HeaderState) : CommandElabM HeaderState := do
   unless getLinterValue linter.style.header (← getLinterOptions) do
-    return
+    return self
   if (← get).messages.hasErrors then
-    return
+    return self
   let map ← getFileMap
   -- This is essentially the same as calling `parseImports'`, but we need access to `s.pos`.
   let s := ParseImports.main map.source (ParseImports.whitespace map.source {})
   unless (← read).cmdPos == s.pos do
-    return
+    -- Only a command that starts at the end of the header can be the first command of the
+    -- module. Later commands start further down, so the linter has no work left in this module.
+    return { self with settled := true }
   let mainModule ← getMainModule
   if Parser.isTerminalCommand stx ||
     -- Deprecated module files are exempt from all header style checks (copyright, doc-string,
     -- directory dependency, etc.) since they are just import-redirect stubs.
     stx.isOfKind ``Parser.Command.deprecated_module ||
     -- Skip linting the library root file itself.
-    -- In practice, the `inLibraryRoot?` check above already covers this (a well-formed
+    -- In practice, the `isInLibraryRoot` check below already covers this (a well-formed
     -- `<root>.lean` does not import itself), but a root module could appear in `headerTestFiles`.
     mainModule matches .str .anonymous _libroot
-  then return
+  then return { self with settled := true }
   unless stx.isOfKind ``Parser.Command.moduleDoc do
     Linter.logLint linter.style.header stx
       m!"The module doc-string for a file should be the first command after the imports.\n\
         Please, add a module doc-string (`/-! ... -/`) before `{stx}`.\
         {.hint' m!"Type `m(odule docstring) + [tab]` to insert a template via snippet."}"
-  let inLibraryRoot? ← inLibraryRootMutex.atomically do
-    match ← get with
-    | some d => return d
-    | none =>
-      let val ← isInLibraryRoot mainModule
-      -- We cache the answer to avoid recomputing it on every command. The fill runs under the mutex
-      -- so that concurrent (async) linter runs don't all miss the cache and each redundantly parse
-      -- the library root file; `mainModule` is fixed for the duration of the elaboration.
-      set (some val)
-      return val
   -- The linter skips files not imported in their library root (e.g. `Mathlib.lean`), to avoid
   -- linting "scratch files". It is however active in the test files for the linter itself.
-  unless inLibraryRoot? || headerTestFiles.contains mainModule do return
+  unless (← isInLibraryRoot mainModule) || headerTestFiles.contains mainModule do
+    return { self with settled := true }
   -- Re-parse the header slowly, to get the leading whitespace and nice source locations for imports
-  let (headerStx, s, log) ← Parser.parseHeader (Parser.mkInputContext map.source (← getFileName))
-  if log.hasErrors then return
+  let (headerStx, _, log) ← Parser.parseHeader (Parser.mkInputContext map.source (← getFileName))
+  if log.hasErrors then return { self with settled := true }
   let importRefs := headerToImportRefs headerStx
   -- Report on broad or duplicate imports.
   broadImportsCheck importRefs mainModule
@@ -389,8 +388,17 @@ def headerLinter : Linter where run := withSetOptionIn fun stx ↦ do
       | _ => ""
     for (stx, m) in copyrightHeaderChecks copyright expectedLicense do
       Linter.logLint linter.style.header stx m!"* `{stx.getAtomVal}`:\n{m}\n"
+  return { self with settled := true, headerSyntax := headerStx.raw }
 
-initialize addLinter headerLinter
+/--
+The typed handle of the `header` linter. Other stateful linters can read the previous
+`HeaderState` of the linter through this handle.
+-/
+public initialize headerLinter : StatefulLinter HeaderState Unit ←
+  registerStatefulLinter {}
+    (post := fun stx self _ _ _ => do
+      if self.settled then return self
+      withSetOptionIn (headerPost · self) stx)
 
 end Style.Header
 

Mathlib/Tactic/Linter/MinImports.lean

@@ -7,6 +7,9 @@ module
 
 public meta import ImportGraph.Imports.ImportGraph
 public meta import ImportGraph.Graph.TransitiveClosure
+-- The `headerLinter` handle comes from `Mathlib.Tactic.Linter.Header`, which `Mathlib.Init`
+-- imports: `linter.checkInitImports` requires the import of `Mathlib.Init` here.
+public meta import Mathlib.Init
 public import Mathlib.Tactic.MinImports
 
 /-! # The `minImports` linter
@@ -19,6 +22,10 @@ Unlike the related `#min_imports` command, the linter takes into account notatio
 information.
 It also works incrementally, accumulating increasing import information.
 This is better suited, for instance, to split files.
+
+The linter is a stateful linter (`Lean.Elab.Command.registerStatefulLinter`). It stores the
+cumulative import information in its linter state, and the elaborator threads the state
+through the commands of the file.
 -/
 
 meta section
@@ -34,28 +41,28 @@ The "minImports" linter tracks information about minimal imports over several co
 namespace Mathlib.Linter
 
 /--
-`ImportState` is the structure keeping track of the data that the `minImports` linter uses.
+`ImportState` is the persistent state of the `minImports` linter.
 * `transClosure` is the import graph of the current file.
 * `minImports` is the `NameSet` of minimal imports to build the file up to the current command.
 * `importSize` is the number of transitive imports to build the file up to the current command.
 -/
-structure ImportState where
-  /-- The transitive closure of the import graph of the current file.  The value is `none` only at
-  initialization time, as the linter immediately sets it to its value for the current file. -/
+public structure ImportState where
+  /-- The transitive closure of the import graph of the current file. The value is `none` only
+  before the linter processes the first command. The linter then sets it to the value for the
+  current file. -/
   transClosure : Option (NameMap NameSet) := none
   /-- The minimal imports needed to build the file up to the current command. -/
   minImports   : NameSet := {}
   /-- The number of transitive imports needed to build the file up to the current command. -/
   importSize   : Nat := 0
   deriving Inhabited
 
-/--
-`minImportsRef` keeps track of cumulative imports across multiple commands, using `ImportState`.
--/
-initialize minImportsRef : IO.Ref ImportState ← IO.mkRef {}
-
-/-- `#reset_min_imports` sets to empty the current list of cumulative imports. -/
-elab "#reset_min_imports" : command => minImportsRef.set {}
+-- TODO: we could give `#import_bumps` reset semantics and deprecate `#reset_min_imports`.
+/-- `#reset_min_imports` makes the `minImports` linter start again from empty cumulative
+imports. -/
+-- Only the linter can write the linter state. The elaborator does nothing: the `minImports`
+-- linter detects the `resetMinImports` syntax kind and clears its state.
+elab (name := resetMinImports) "#reset_min_imports" : command => pure ()
 
 /--
 The `minImports` linter incrementally computes the minimal imports needed for each file to build.
@@ -100,76 +107,92 @@ def importsBelow (tc : NameMap NameSet) (ms : NameSet) : NameSet :=
 macro "#import_bumps" : command => `(
   -- We emit a message to prevent the `#`-command linter from flagging `#import_bumps`.
   run_cmd logInfo "Counting imports from here."
-  set_option Elab.async false
   set_option linter.minImports true)
 
 @[inherit_doc Mathlib.Linter.linter.minImports]
-def minImportsLinter : Linter where run := withSetOptionIn fun stx ↦ do
-    unless getLinterValue linter.minImports (← getLinterOptions) do
-      return
-    if (← get).messages.hasErrors then
-      return
-    if stx == (← `(command| #import_bumps)) then return
-    if stx == (← `(command| set_option $(mkIdent `linter.minImports) true)) then
-      logInfo "Try using '#import_bumps', instead of manually setting the linter option: \
-              the linter works best with linear parsing of the file and '#import_bumps' \
-              also sets the `Elab.async` option to `false`."
-      return
-    let env ← getEnv
-    -- the first time `minImportsRef` is read, it has `transClosure = none`;
-    -- in this case, we set it to be the `transClosure` for the file.
-    if (← minImportsRef.get).transClosure.isNone then
-      minImportsRef.modify ({· with transClosure := env.importGraph.transitiveClosure})
-    let impState ← minImportsRef.get
-    let (importsSoFar, oldCumulImps) := (impState.minImports, impState.importSize)
-    -- when the linter reaches the end of the file or `#exit`, it gives a report
-    if #[``Parser.Command.eoi, ``Lean.Parser.Command.exit].contains stx.getKind then
-      let explicitImportsInFile : NameSet :=
-        .ofArray ((env.imports.map (·.module)).filter (!isInitImport ·))
-      let newImps := importsSoFar \ explicitImportsInFile
-      let currentlyUnneededImports := explicitImportsInFile \ importsSoFar
-      -- we read the current file, to do a custom parsing of the imports:
-      -- this is a hack to obtain some `Syntax` information for the `import X` commands
-      let fname ← getFileName
-      let contents ← IO.FS.readFile fname
-      -- `impMods` is the syntax for the modules imported in the current file
-      let (impMods, _) ← Parser.parseHeader (Parser.mkInputContext contents fname)
-      for i in currentlyUnneededImports do
-        match impMods.raw.find? (·.getId == i) with
-          | some impPos => logWarningAt impPos m!"unneeded import '{i}'"
-          | _ => dbg_trace f!"'{i}' not found"  -- this should be unreachable
-      -- if the linter found new imports that should be added (likely to *reduce* the dependencies)
-      if !newImps.isEmpty then
-        -- format the imports prepending `import ` to each module name
-        let withImport := (newImps.toArray.qsort Name.lt).map (s!"import {·}")
-        -- log a warning at the first `import`, if there is one.
-        logWarningAt ((impMods.raw.find? (·.isOfKind `import)).getD default)
-          m!"-- missing imports\n{"\n".intercalate withImport.toList}"
-    let id ← getId stx
-    let newImports := (getIrredundantImports env (← getAllImports stx id)).filter (!isInitImport ·)
-    let tot := (newImports.append importsSoFar)
-    let redundant := env.findRedundantImports tot.toArray
-    let currImports := tot \ redundant
-    let currImpArray := currImports.toArray.qsort Name.lt
-    if currImpArray != #[] &&
-       currImpArray ≠ importsSoFar.toArray.qsort Name.lt then
-      let newCumulImps := -- We should always be in the situation where `getD` finds something
-        (importsBelow (impState.transClosure.getD env.importGraph.transitiveClosure) tot).size
-      minImportsRef.modify ({· with minImports := currImports, importSize := newCumulImps})
-      let new := currImpArray.filter (!importsSoFar.contains ·)
-      let redundant := importsSoFar.toArray.filter (!currImports.contains ·)
-      -- to make `test` files more stable, we suppress the exact count of import changes if
-      -- the `linter.minImports.increases` option is `false`
-      let byCount :=  if getLinterValue linter.minImports.increases (← getLinterOptions) then
-                      m!"by {newCumulImps - oldCumulImps} "
-                    else
-                      m!""
-      Linter.logLint linter.minImports stx <|
-        m!"Imports increased {byCount}to\n{currImpArray}\n\n\
-          New imports: {new}\n" ++
-            if redundant.isEmpty then m!"" else m!"\nNow redundant: {redundant}\n"
-
-initialize addLinter minImportsLinter
+def minImportsPost (readPrev : PrevStateFn) (stx : Syntax) (self : ImportState) :
+    CommandElabM ImportState := do
+  -- The reset applies also when the linter option is off. The linter then lints the
+  -- `#reset_min_imports` command itself, against the empty state.
+  let self := if stx.isOfKind ``resetMinImports then {} else self
+  unless getLinterValue linter.minImports (← getLinterOptions) do
+    return self
+  if (← get).messages.hasErrors then
+    return self
+  if stx == (← `(command| #import_bumps)) then return self
+  if stx == (← `(command| set_option $(mkIdent `linter.minImports) true)) then
+    logInfo "Consider using '#import_bumps' instead of setting the linter option: \
+            it also logs the position where the import count starts."
+    return self
+  let env ← getEnv
+  -- On the first command that the linter processes, it computes the transitive closure of the
+  -- imports of the file.
+  let self := if self.transClosure.isNone then
+      { self with transClosure := env.importGraph.transitiveClosure }
+    else self
+  let (importsSoFar, oldCumulImps) := (self.minImports, self.importSize)
+  -- when the linter reaches the end of the file or `#exit`, it gives a report
+  if #[``Parser.Command.eoi, ``Lean.Parser.Command.exit].contains stx.getKind then
+    let explicitImportsInFile : NameSet :=
+      .ofArray ((env.imports.map (·.module)).filter (!isInitImport ·))
+    let newImps := importsSoFar \ explicitImportsInFile
+    let currentlyUnneededImports := explicitImportsInFile \ importsSoFar
+    -- `impMods` is the syntax for the modules imported in the current file. The state of the
+    -- `header` linter provides it when the header checks ran. Otherwise, we read the current
+    -- file and do a custom parsing of the imports: this is a hack to obtain some `Syntax`
+    -- information for the `import X` commands.
+    let headerStx := (readPrev Style.Header.headerLinter).headerSyntax
+    let impMods ← if headerStx.isMissing then do
+        let fname ← getFileName
+        let contents ← IO.FS.readFile fname
+        let (impMods, _) ← Parser.parseHeader (Parser.mkInputContext contents fname)
+        pure impMods.raw
+      else
+        pure headerStx
+    for i in currentlyUnneededImports do
+      match impMods.find? (·.getId == i) with
+        | some impPos => logWarningAt impPos m!"unneeded import '{i}'"
+        | _ => dbg_trace f!"'{i}' not found"  -- this should be unreachable
+    -- if the linter found new imports that should be added (likely to *reduce* the dependencies)
+    if !newImps.isEmpty then
+      -- format the imports prepending `import ` to each module name
+      let withImport := (newImps.toArray.qsort Name.lt).map (s!"import {·}")
+      -- log a warning at the first `import`, if there is one.
+      logWarningAt ((impMods.find? (·.isOfKind `import)).getD default)
+        m!"-- missing imports\n{"\n".intercalate withImport.toList}"
+    return self
+  let id ← getId stx
+  let newImports := (getIrredundantImports env (← getAllImports stx id)).filter (!isInitImport ·)
+  let tot := (newImports.append importsSoFar)
+  let redundant := env.findRedundantImports tot.toArray
+  let currImports := tot \ redundant
+  let currImpArray := currImports.toArray.qsort Name.lt
+  if currImpArray != #[] &&
+     currImpArray ≠ importsSoFar.toArray.qsort Name.lt then
+    let newCumulImps := -- We should always be in the situation where `getD` finds something
+      (importsBelow (self.transClosure.getD env.importGraph.transitiveClosure) tot).size
+    let new := currImpArray.filter (!importsSoFar.contains ·)
+    let redundant := importsSoFar.toArray.filter (!currImports.contains ·)
+    -- to make `test` files more stable, we suppress the exact count of import changes if
+    -- the `linter.minImports.increases` option is `false`
+    let byCount :=  if getLinterValue linter.minImports.increases (← getLinterOptions) then
+                    m!"by {newCumulImps - oldCumulImps} "
+                  else
+                    m!""
+    Linter.logLint linter.minImports stx <|
+      m!"Imports increased {byCount}to\n{currImpArray}\n\n\
+        New imports: {new}\n" ++
+          if redundant.isEmpty then m!"" else m!"\nNow redundant: {redundant}\n"
+    return { self with minImports := currImports, importSize := newCumulImps }
+  return self
+
+/--
+The typed handle of the `minImports` linter. Other stateful linters can read the previous
+`ImportState` of the linter through this handle.
+-/
+public initialize minImportsLinter : StatefulLinter ImportState Unit ←
+  registerStatefulLinter {}
+    (post := fun stx self _ readPrev _ => withSetOptionIn (minImportsPost readPrev · self) stx)
 
 end MinImports