refactor(Tactic/Linter/MinImports): make the minImports linter stateful

Mathlib/Tactic/Linter/MinImports.lean

@@ -19,6 +19,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 +38,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 +104,85 @@ 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 (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
+    -- 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}"
+    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 _ _ _ => withSetOptionIn (minImportsPost · self) stx)
 
 end MinImports
 

MathlibTest/MinImports.lean

@@ -277,3 +277,61 @@ inductive ProposeToMoveThisInductive where
 theorem theorem_with_local_inductive : Nonempty ProposeToMoveThisInductive := ⟨.foo 0⟩
 
 end Linter.UpstreamableDecl
+
+section Linter.MinImports.Async
+
+/-!
+The `minImports` linter is a stateful linter: the elaborator threads the linter state through
+the commands of the file, so the accumulated imports are correct under parallel elaboration.
+The tests below run with `Elab.async` explicitly on, and check that `#import_bumps` keeps
+parallel elaboration enabled.
+-/
+
+set_option Elab.async true
+set_option linter.minImports.increases false
+
+-- Earlier sections of this file leave imports in the linter state: start from a clean state.
+set_option linter.minImports false in
+#reset_min_imports
+
+/-- info: Counting imports from here. -/
+#guard_msgs in
+#import_bumps
+
+-- `#import_bumps` must keep parallel elaboration on.
+set_option linter.minImports false in
+run_cmd do
+  unless Lean.Elab.async.get (← Lean.getOptions) do
+    throwError "`#import_bumps` must keep `Elab.async` on"
+
+/--
+warning: Imports increased to
+[Mathlib.Data.Int.Notation]
+
+New imports: [Mathlib.Data.Int.Notation]
+
+
+Note: This linter can be disabled with `set_option linter.minImports false`
+-/
+#guard_msgs in
+#guard (0 : ℤ) = 0
+
+#guard_msgs in
+-- no new imports needed here, so no message
+#guard (0 : ℤ) = 0
+
+/--
+warning: Imports increased to
+[Mathlib.Tactic.NormNum.Basic]
+
+New imports: [Mathlib.Tactic.NormNum.Basic]
+
+Now redundant: [Mathlib.Data.Int.Notation]
+
+
+Note: This linter can be disabled with `set_option linter.minImports false`
+-/
+#guard_msgs in
+lemma async_uses_norm_num : (0 + 1 : ℕ) = 1 := by norm_num
+
+end Linter.MinImports.Async