chore: basic ExtraModUses

Kha · Kha · commit 641ea91cb7f7 · 2025-09-27T18:22:25.000+02:00
diff --git a/src/Lean.lean b/src/Lean.lean
@@ -47,3 +47,4 @@ public import Lean.ErrorExplanation
 public import Lean.ErrorExplanations
 public import Lean.DefEqAttrib
 public import Lean.Shell
+public import Lean.ExtraModUses
diff --git a/src/Lean/Elab/BuiltinTerm.lean b/src/Lean/Elab/BuiltinTerm.lean
@@ -12,6 +12,7 @@ public import Lean.Elab.Open
 public import Lean.Elab.SetOption
 public import Lean.Elab.Eval
 meta import Lean.Parser.Command
+import Lean.ExtraModUses
 
 public section
 
@@ -237,9 +238,11 @@ def elabScientificLit : TermElab := fun stx expectedType? => do
   | some val => pure $ toExpr val
   | none     => throwIllFormedSyntax
 
-@[builtin_term_elab doubleQuotedName] def elabDoubleQuotedName : TermElab := fun stx _ =>
+@[builtin_term_elab doubleQuotedName] def elabDoubleQuotedName : TermElab := fun stx _ => do
   -- Always allow quoting private names.
-  return toExpr (← withoutExporting <| realizeGlobalConstNoOverloadWithInfo stx[2])
+  let n ← withoutExporting <| realizeGlobalConstNoOverloadWithInfo stx[2]
+  recordExtraModUseFromDecl (isMeta := false) n
+  return toExpr n
 
 @[builtin_term_elab declName] def elabDeclName : TermElab := adaptExpander fun _ => do
   let some declName ← getDeclName?
diff --git a/src/Lean/Elab/Declaration.lean b/src/Lean/Elab/Declaration.lean
@@ -333,6 +333,11 @@ def elabMutual : CommandElab := fun stx => do
     Term.applyAttributes declName attrs
     for attrName in toErase do
       Attribute.erase declName attrName
+    if (← getEnv).isImportedConst declName && attrs.any (·.kind == .global) then
+      -- If an imported declaration is marked with a global attribute, there is no good way to track
+      -- its use generally and so Shake should conservatively preserve imports of the current
+      -- module.
+      recordExtraRevUseOfCurrentModule
 
 @[builtin_command_elab Lean.Parser.Command.«initialize»] def elabInitialize : CommandElab
   | stx@`($declModifiers:declModifiers $kw:initializeKeyword $[$id? : $type? ←]? $doSeq) => do
diff --git a/src/Lean/Elab/PreDefinition/WF/Main.lean b/src/Lean/Elab/PreDefinition/WF/Main.lean
@@ -16,6 +16,7 @@ public import Lean.Elab.PreDefinition.WF.Fix
 public import Lean.Elab.PreDefinition.WF.Unfold
 public import Lean.Elab.PreDefinition.WF.Preprocess
 public import Lean.Elab.PreDefinition.WF.GuessLex
+import Lean.ExtraModUses
 
 public section
 
@@ -72,6 +73,8 @@ def wfRecursion (docCtx : LocalContext × LocalInstances) (preDefs : Array PreDe
         eraseRecAppSyntaxExpr value
       /- `mkFix` invokes `decreasing_tactic` which may add auxiliary theorems to the environment. -/
       let value ← unfoldDeclsFrom envNew value
+      -- Make sure we remember invoked tactics
+      modifyEnv (copyExtraModUses envNew)
       return { unaryPreDefProcessed with value }
 
   trace[Elab.definition.wf] ">> {preDefNonRec.declName} :=\n{preDefNonRec.value}"
diff --git a/src/Lean/Elab/RecommendedSpelling.lean b/src/Lean/Elab/RecommendedSpelling.lean
@@ -9,6 +9,7 @@ prelude
 public import Lean.Parser.Term.Doc
 public import Lean.Parser.Command
 public import Lean.Elab.Command
+import Lean.ExtraModUses
 
 public section
 
@@ -20,6 +21,7 @@ open Lean.Parser.Command
 @[builtin_command_elab «recommended_spelling»] def elabRecommendedSpelling : CommandElab
   | `(«recommended_spelling»|$[$docs:docComment]? recommended_spelling $spelling:str for $«notation»:str in [$[$decls:ident],*]) => do
     let declNames ← liftTermElabM <| decls.mapM (fun decl => realizeGlobalConstNoOverloadWithInfo decl)
+    declNames.forM (recordExtraModUseFromDecl (isMeta := false))
     let recommendedSpelling : RecommendedSpelling := {
       «notation» := «notation».getString
       recommendedSpelling := spelling.getString
diff --git a/src/Lean/Elab/Syntax.lean b/src/Lean/Elab/Syntax.lean
@@ -10,6 +10,7 @@ public import Lean.Elab.Command
 public import Lean.Parser.Syntax
 public import Lean.Elab.Util
 public meta import Lean.Parser.Syntax
+import Lean.ExtraModUses
 
 public section
 
@@ -394,7 +395,10 @@ def elabSyntax (stx : Syntax) : CommandElabM Name := do
       syntax%$tk $[: $prec? ]? $[(name := $name?)]? $[(priority := $prio?)]? $[$ps:stx]* : $catStx) := stx
     | throwUnsupportedSyntax
   let cat := catStx.getId.eraseMacroScopes
-  unless (Parser.isParserCategory (← getEnv) cat) do
+  if let some cat := Parser.getParserCategory? (← getEnv) cat then
+    -- The category must be imported but is not directly referenced afterwards.
+    recordExtraModUseFromDecl (isMeta := true) cat.declName
+  else
     throwErrorAt catStx "unknown category `{cat}`"
   liftTermElabM <| Term.addCategoryInfo catStx cat
   let syntaxParser := mkNullNode ps
diff --git a/src/Lean/Elab/Tactic/Basic.lean b/src/Lean/Elab/Tactic/Basic.lean
@@ -8,6 +8,7 @@ module
 prelude
 public import Lean.Meta.Tactic.Util
 public import Lean.Elab.Term
+import Lean.ExtraModUses
 
 public section
 
@@ -235,6 +236,7 @@ where
           withReader ({ · with elaborator := m.declName }) do
             withTacticInfoContext stx do
               let stx' ← adaptMacro m.value stx
+              recordExtraModUseFromDecl (isMeta := true) m.declName
               -- Support incrementality; see also Note [Incremental Macros]
               if evalFns.isEmpty && ms.isEmpty then  -- Only try incrementality in one branch
                 if let some snap := (← readThe Term.Context).tacSnap? then
diff --git a/src/Lean/Elab/Tactic/Config.lean b/src/Lean/Elab/Tactic/Config.lean
@@ -9,6 +9,7 @@ prelude
 public import Lean.Meta.Eval
 public import Lean.Elab.Tactic.Basic
 public import Lean.Elab.SyntheticMVars
+public import Lean.ExtraModUses
 import Lean.Linter.MissingDocs
 meta import Lean.Parser.Tactic
 
@@ -142,6 +143,7 @@ private meta def mkConfigElaborator
           return $empty
         unless (← getEnv).contains ``$type do
           throwError m!"Error evaluating configuration: Environment does not yet contain type {``$type}"
+        recordExtraModUseFromDecl (isMeta := true) ``$type
         let c ← elabConfig recover ``$type items
         if c.hasSyntheticSorry then
           -- An error is already logged, return the default.
diff --git a/src/Lean/Elab/Tactic/Omega/Frontend.lean b/src/Lean/Elab/Tactic/Omega/Frontend.lean
@@ -675,6 +675,8 @@ open Lean Elab Tactic Parser.Tactic
 
 /-- The `omega` tactic, for resolving integer and natural linear arithmetic problems. -/
 def omegaTactic (cfg : OmegaConfig) : TacticM Unit := do
+  -- Conservatively expect all of `Init.Omega` to be imported.
+  recordExtraModUse (isMeta := false) `Init.Omega
   liftMetaFinishingTactic fun g => do
     if debug.terminalTacticsAsSorry.get (← getOptions) then
       g.admit
diff --git a/src/Lean/EnvExtension.lean b/src/Lean/EnvExtension.lean
@@ -66,8 +66,9 @@ instance {α σ : Type} [Inhabited σ] : Inhabited (SimplePersistentEnvExtension
 
 /-- Get the list of values used to update the state of the given
 `SimplePersistentEnvExtension` in the current file. -/
-def getEntries {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ) (env : Environment) : List α :=
-  (PersistentEnvExtension.getState ext env).1
+def getEntries {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ)
+    (env : Environment) (asyncMode := ext.toEnvExtension.asyncMode) : List α :=
+  (PersistentEnvExtension.getState (asyncMode := asyncMode) ext env).1
 
 /-- Get the current state of the given `SimplePersistentEnvExtension`. -/
 def getState {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ) (env : Environment)
diff --git a/src/Lean/ExtraModUses.lean b/src/Lean/ExtraModUses.lean
@@ -0,0 +1,99 @@
+/-
+Copyright (c) 2025 Lean FRO. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sebastian Ullrich
+-/
+module
+
+prelude
+public import Lean.CoreM
+public import Lean.Compiler.MetaAttr  -- TODO: public because of specializations
+
+/-!
+Infrastructure for recording extra import dependencies not implied by the environment constants for
+the benefit of `shake`.
+-/
+
+namespace Lean
+
+/-- Additional import dependency for elaboration. -/
+public structure ExtraModUse where
+  /-- Dependency's module name. -/
+  module : Name
+  /-- Whether dependency must be imported as `public`. -/
+  isExported : Bool
+  /-- Whether dependency must be imported as `meta`. -/
+  isMeta : Bool
+deriving BEq, Hashable, Repr
+
+builtin_initialize extraModUses : SimplePersistentEnvExtension ExtraModUse (PHashSet ExtraModUse) ←
+  registerSimplePersistentEnvExtension {
+    addEntryFn m k := m.insert k
+    addImportedFn _ := {}
+    exportEntriesFnEx? := some fun _ _ entries => fun
+      | .private => entries.toArray
+      | _        => #[]
+    asyncMode := .sync
+    replay? := some <| SimplePersistentEnvExtension.replayOfFilter (·.contains ·) (·.insert ·)
+  }
+
+/-- Returns additional recorded import dependencies of the given module. -/
+public def getExtraModUses (env : Environment) (modIdx : ModuleIdx) : Array ExtraModUse :=
+  extraModUses.getModuleEntries env modIdx
+
+/-- Copies additional recorded import dependencies from one environment to another. -/
+public def copyExtraModUses (src dest : Environment) : Environment := Id.run do
+  let mut env := dest
+  for entry in extraModUses.getEntries src do
+    if !(extraModUses.getState (asyncMode := .local) env).contains entry then
+      env := extraModUses.addEntry env entry
+  env
+
+variable [Monad m] [MonadEnv m] [MonadTrace m] [MonadOptions m] [MonadRef m] [AddMessageContext m]
+
+def recordExtraModUseCore (modIdx : ModuleIdx) (isMeta : Bool) (hint : Name := .anonymous) : m Unit := do
+  let name := (← getEnv).header.moduleNames[modIdx]!
+  let entry := { module := name, isExported := (← getEnv).isExporting, isMeta }
+  if !(extraModUses.getState (asyncMode := .local) (← getEnv)).contains entry then
+    trace[extraModUses] "recording {if entry.isExported then "public" else "private"} \
+      {if isMeta then "meta" else "regular"} extra mod use {name}\
+      {if hint.isAnonymous then m!"" else m!" of {hint}"}"
+    modifyEnv (extraModUses.addEntry · entry)
+
+/--
+Records an additional import dependency for the current module, using `Environment.isExporting` as
+the visibility level.
+-/
+public def recordExtraModUse (modName : Name) (isMeta : Bool) : m Unit := do
+  if let some modIdx := (← getEnv).getModuleIdx? modName then
+    recordExtraModUseCore modIdx isMeta
+
+/--
+Records the module of the given declaration as an additional import dependency for the current
+module, using `Environment.isExporting` as the visibility level. If the declaration itself is
+already `meta`, the module dependency is recorded as a non-`meta` dependency.
+-/
+public def recordExtraModUseFromDecl (declName : Name) (isMeta : Bool) : m Unit := do
+  if let some modIdx := (← getEnv).getModuleIdxFor? declName then
+    -- If the declaration itself is already `meta`, no need to mark the import.
+    let isMeta := isMeta && !Lean.isMeta (← getEnv) declName
+    recordExtraModUseCore modIdx isMeta (hint := declName)
+
+builtin_initialize isExtraRevModUseExt : SimplePersistentEnvExtension Unit Unit ←
+  registerSimplePersistentEnvExtension {
+    addEntryFn s e := ()
+    addImportedFn _ := ()
+    asyncMode := .sync
+  }
+
+/-- Checks whether this module should be preserved as an import by `shake`. -/
+public def isExtraRevModUse (env : Environment) (modIdx : ModuleIdx) : Bool :=
+  !(isExtraRevModUseExt.getModuleEntries env modIdx |>.isEmpty)
+
+/-- Records this module to be preserved as an import by `shake`. -/
+public def recordExtraRevUseOfCurrentModule : m Unit := do
+  if isExtraRevModUseExt.getEntries (← getEnv) |>.isEmpty then
+    modifyEnv fun env => isExtraRevModUseExt.addEntry env ()
+
+builtin_initialize
+  registerTraceClass `extraModUses
diff --git a/src/Lean/KeyedDeclsAttribute.lean b/src/Lean/KeyedDeclsAttribute.lean
@@ -56,6 +56,7 @@ structure OLeanEntry where
   deriving Inhabited
 
 structure AttributeEntry (γ : Type) extends OLeanEntry where
+  isBuiltin : Bool
   /-- Recall that we cannot store `γ` into .olean files because it is a closure.
      Given `OLeanEntry.declName`, we convert it into a `γ` by using the unsafe function `evalConstCheck`. -/
   value : γ
@@ -96,7 +97,7 @@ def ExtensionState.insert (s : ExtensionState γ) (v : AttributeEntry γ) :  Ext
 }
 
 def addBuiltin (attr : KeyedDeclsAttribute γ) (key : Key) (declName : Name) (value : γ) : IO Unit :=
-  attr.tableRef.modify fun m => m.insert { key, declName, value }
+  attr.tableRef.modify fun m => m.insert { key, declName, value, isBuiltin := true }
 
 def mkStateOfTable (table : Table γ) : ExtensionState γ := {
   table
@@ -116,7 +117,7 @@ protected unsafe def init {γ} (df : Def γ) (attrDeclName : Name := by exact de
     ofOLeanEntry := fun _ entry => do
       let ctx ← read
       match ctx.env.evalConstCheck γ ctx.opts df.valueTypeName entry.declName with
-      | Except.ok f     => return { toOLeanEntry := entry, value := f }
+      | Except.ok f     => return { toOLeanEntry := entry, value := f, isBuiltin := false }
       | Except.error ex => throw (IO.userError ex)
     addEntry     := fun s e => s.insert e
     toOLeanEntry := (·.toOLeanEntry)
@@ -157,7 +158,7 @@ protected unsafe def init {γ} (df : Def γ) (attrDeclName : Name := by exact de
       match IR.getSorryDep (← getEnv) declName with
       | none =>
         let val ← evalConstCheck γ df.valueTypeName declName
-        ext.add { key := key, declName := declName, value := val } attrKind
+        ext.add { key := key, declName := declName, value := val, isBuiltin := true } attrKind
         df.onAdded false declName key
       | _ =>
         -- If the declaration contains `sorry`, we skip `evalConstCheck` to avoid unnecessary bizarre error message
diff --git a/src/Lean/Linter/Deprecated.lean b/src/Lean/Linter/Deprecated.lean
@@ -9,6 +9,7 @@ prelude
 public import Lean.Meta.Basic
 import Lean.Linter.Basic
 import Lean.Elab.InfoTree.Main
+import Lean.ExtraModUses
 
 public section
 
@@ -33,6 +34,8 @@ builtin_initialize deprecatedAttr : ParametricAttribute DeprecationEntry ←
       let `(attr| deprecated $[$id?]? $[$text?]? $[(since := $since?)]?) := stx
         | throwError "Invalid `[deprecated]` attribute syntax"
       let newName? ← id?.mapM Elab.realizeGlobalConstNoOverloadWithInfo
+      if let some newName := newName? then
+        recordExtraModUseFromDecl (isMeta := false) newName
       let text? := text?.map TSyntax.getString
       let since? := since?.map TSyntax.getString
       if id?.isNone && text?.isNone then
diff --git a/src/Lean/Meta/Tactic/Grind/Attr.lean b/src/Lean/Meta/Tactic/Grind/Attr.lean
@@ -9,6 +9,8 @@ public import Lean.Meta.Tactic.Grind.EMatchTheorem
 public import Lean.Meta.Tactic.Grind.Injective
 public import Lean.Meta.Tactic.Grind.Cases
 public import Lean.Meta.Tactic.Grind.ExtAttr
+import Lean.ExtraModUses
+
 public section
 namespace Lean.Meta.Grind
 
@@ -97,6 +99,7 @@ private def registerGrindAttr (minIndexable : Bool) (showInfo : Bool) : IO Unit
       `grind` will add an instance of this theorem to the local context whenever it encounters the pattern `foo (foo x)`."
     applicationTime := .afterCompilation
     add := fun declName stx attrKind => MetaM.run' do
+      recordExtraModUseFromDecl (isMeta := false) declName
       match (← getAttrKindFromOpt stx) with
       | .ematch .user => throwInvalidUsrModifier
       | .ematch k => addEMatchAttr declName attrKind k (← getGlobalSymbolPriorities) (minIndexable := minIndexable) (showInfo := showInfo)
diff --git a/src/Lean/Meta/Tactic/Grind/EMatchTheorem.lean b/src/Lean/Meta/Tactic/Grind/EMatchTheorem.lean
@@ -19,6 +19,7 @@ import Lean.Message
 import Lean.Meta.Tactic.FVarSubst
 import Lean.Meta.Match.Basic
 import Lean.Meta.Tactic.TryThis
+import Lean.ExtraModUses
 public section
 namespace Lean.Meta.Grind
 /--
@@ -235,6 +236,9 @@ pbind (Grind.genPattern h x (some a)) f = f a h
 ```
 -/
 private def detectGeneralizedPatterns? (type : Expr) : MetaM Expr := do
+  -- `genPattern` is not necessarily referenced in the result but always importing it for `grind`
+  -- uses is a solid approximation.
+  recordExtraModUseFromDecl (isMeta := false) ``Grind.genPattern
   forallTelescopeReducing type fun hs resultType => do
     let isTarget? (lhs : Expr) (rhs : Expr) (s : FVarSubst) : Option (FVarId × Expr) := Id.run do
       let .fvar fvarId := lhs | return none
diff --git a/src/Lean/Parser/Extension.lean b/src/Lean/Parser/Extension.lean
@@ -351,8 +351,11 @@ builtin_initialize parserExtension : ParserExtension ←
     ofOLeanEntry    := ParserExtension.OLeanEntry.toEntry
   }
 
+def getParserCategory? (env : Environment) (catName : Name) : Option ParserCategory :=
+  (parserExtension.getState env).categories.find? catName
+
 def isParserCategory (env : Environment) (catName : Name) : Bool :=
-  (parserExtension.getState env).categories.contains catName
+  getParserCategory? env catName |>.isSome
 
 def addParserCategory (env : Environment) (catName declName : Name) (behavior : LeadingIdentBehavior) : Except String Environment := do
   if isParserCategory env catName then
