feat: recordIndirectModUse (#11437)

This PR adds recording functionality such that `shake` can more
precisely track whether an import should be preserved solely for its
`attribute` commands.

Author: Kha

src/Lean/Elab/Declaration.lean

@@ -330,11 +330,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
+    if (← getEnv).isImportedConst declName then
+      if let some attr := attrs.find? (·.kind == .global) then
+        -- If an imported declaration is marked with a global attribute, Shake must be informed
+        -- about this indirect reference
+        recordIndirectModUse attr.name.toString declName
 
 @[builtin_command_elab Lean.Parser.Command.«initialize»] def elabInitialize : CommandElab
   | stx@`($declModifiers:declModifiers $kw:initializeKeyword $[$id? : $type? ←]? $doSeq) => do

src/Lean/Environment.lean

@@ -154,6 +154,7 @@ deriving Inhabited, BEq, Repr
 structure EffectiveImport extends Import where
   /-- Phases for which the import's IR is available. -/
   irPhases : IRPhases
+deriving Inhabited
 
 /-- Environment fields that are not used often. -/
 structure EnvironmentHeader where

src/Lean/ExtraModUses.lean

@@ -8,6 +8,7 @@ module
 prelude
 public import Lean.CoreM
 public import Lean.Compiler.MetaAttr  -- TODO: public because of specializations
+import Init.Data.Range.Polymorphic.Stream
 
 /-!
 Infrastructure for recording extra import dependencies not implied by the environment constants for
@@ -16,6 +17,41 @@ the benefit of `shake`.
 
 namespace Lean
 
+public structure IndirectModUse where
+  kind : String
+  declName : Name
+deriving BEq
+
+public builtin_initialize indirectModUseExt : SimplePersistentEnvExtension IndirectModUse (Std.HashMap Name (Array ModuleIdx)) ←
+  registerSimplePersistentEnvExtension {
+    addEntryFn s _ := s
+    addImportedFn es := Id.run do
+      let mut s := {}
+      for es in es, modIdx in 0...* do
+        for e in es do
+          s := s.alter e.declName (·.getD #[] |>.push modIdx)
+      return s
+    asyncMode := .sync
+  }
+
+public def getIndirectModUses (env : Environment) (modIdx : ModuleIdx) : Array IndirectModUse :=
+  indirectModUseExt.getModuleEntries env modIdx
+
+variable [Monad m] [MonadEnv m] [MonadTrace m] [MonadOptions m] [MonadRef m] [AddMessageContext m]
+
+/--
+Lets `shake` know that references to `declName` may also require importing the current module due to
+some additional metaprogramming dependency expressed by `kind`. Currently this is always the name of
+an attribute applied to `declName`, which is not from the current module, in the current module.
+`kind` is not currently used to further filter what references to `declName` require importing the
+current module but may in the future.
+-/
+public def recordIndirectModUse (kind : String) (declName : Name) : m Unit := do
+  -- We can assume this is called from the main thread only and that the list of entries is short
+  if !(indirectModUseExt.getEntries (asyncMode := .mainOnly) (← getEnv) |>.contains { kind, declName }) then
+    trace[extraModUses] "recording indirect mod use of `{declName}` ({kind})"
+    modifyEnv (indirectModUseExt.addEntry · { kind, declName })
+
 /-- Additional import dependency for elaboration. -/
 public structure ExtraModUse where
   /-- Dependency's module name. -/
@@ -49,8 +85,6 @@ public def copyExtraModUses (src dest : Environment) : Environment := Id.run do
       env := extraModUses.addEntry env entry
   env
 
-variable [Monad m] [MonadEnv m] [MonadTrace m] [MonadOptions m] [MonadRef m] [AddMessageContext m]
-
 def recordExtraModUseCore (mod : Name) (isMeta : Bool) (hint : Name := .anonymous) : m Unit := do
   let entry := { module := mod, isExported := (← getEnv).isExporting, isMeta }
   if !(extraModUses.getState (asyncMode := .local) (← getEnv)).contains entry then
@@ -62,6 +96,9 @@ def recordExtraModUseCore (mod : Name) (isMeta : Bool) (hint : Name := .anonymou
 /--
 Records an additional import dependency for the current module, using `Environment.isExporting` as
 the visibility level.
+
+NOTE: Directly recording a module name does not trigger including indirect dependencies recorded via
+`recordIndirectModUse`, prefer `recordExtraModUseFromDecl` instead.
 -/
 public def recordExtraModUse (modName : Name) (isMeta : Bool) : m Unit := do
   if modName != (← getEnv).mainModule then
@@ -78,6 +115,8 @@ public def recordExtraModUseFromDecl (declName : Name) (isMeta : Bool) : m Unit
     -- If the declaration itself is already `meta`, no need to mark the import.
     let isMeta := isMeta && !isMarkedMeta (← getEnv) declName
     recordExtraModUseCore mod.module isMeta (hint := declName)
+    for modIdx in (indirectModUseExt.getState (asyncMode := .local) env)[declName]?.getD #[] do
+      recordExtraModUseCore env.header.modules[modIdx]!.module (isMeta := false) (hint := declName)
 
 builtin_initialize isExtraRevModUseExt : SimplePersistentEnvExtension Unit Unit ←
   registerSimplePersistentEnvExtension {

stage0/src/stdlib_flags.h

@@ -1,5 +1,7 @@
 #include "util/options.h"
 
+// Dear CI, please update stage 0
+
 namespace lean {
 options get_default_options() {
     options opts;

tests/lean/run/extraModUses.lean

@@ -69,7 +69,7 @@ Is rev mod use: false
 recommended_spelling "id" for "id" in [id]
 
 /--
-info: Entries: [import Init.Prelude]
+info: Entries: [import Init.Core, import Init.Prelude]
 Is rev mod use: false
 -/
 #guard_msgs in #eval showExtraModUses
@@ -174,7 +174,7 @@ attribute [grind =] List.append
 
 /--
 info: Entries: [import Init.Grind.Attr, public import Init.Prelude]
-Is rev mod use: true
+Is rev mod use: false
 -/
 #guard_msgs in #eval showExtraModUses
 
@@ -201,7 +201,7 @@ Simp theorems (especially defeq ones) are tracked (here `Nat.pow_succ` from Init
 def test7 : 2 ^ 8 = 256 := by simp [Nat.pow_succ]
 
 /--
-info: Entries: [import Init.Tactics, import Init.Data.Nat.Basic, import Init.SimpLemmas, import Init.Notation]
+info: Entries: [import Init.Tactics, import Init.Data.Nat.Lemmas, import Init.Data.Nat.Basic, import Init.SimpLemmas, import Init.Notation]
 Is rev mod use: false
 -/
 #guard_msgs in #eval showExtraModUses