feat: add async support to more extensions and constructions (#7363)

src/Lean/Environment.lean

@@ -556,10 +556,18 @@ private def modifyCheckedAsync (env : Environment) (f : Kernel.Environment → K
 private def setCheckedSync (env : Environment) (newChecked : Kernel.Environment) : Environment :=
   { env with checked := .pure newChecked, base := newChecked }
 
+/-- The declaration prefix to which the environment is restricted to, if any. -/
+def asyncPrefix? (env : Environment) : Option Name :=
+  env.asyncCtx?.map (·.declPrefix)
+
 /-- True while inside `realizeConst`'s `realize`. -/
 def isRealizing (env : Environment) : Bool :=
   env.asyncCtx?.any (·.realizing)
 
+/-- Forgets about the asynchronous context restrictions. Used only for `withoutModifyingEnv`. -/
+def unlockAsync (env : Environment) : Environment :=
+  { env with asyncCtx? := none }
+
 /--
 Checks whether the given declaration name may potentially added, or have been added, to the current
 environment branch, which is the case either if this is the main branch or if the declaration name
@@ -1453,6 +1461,7 @@ def mkTagDeclarationExtension (name : Name := by exact decl_name%) : IO TagDecla
     addImportedFn := fun _ => {},
     addEntryFn    := fun s n => s.insert n,
     toArrayFn     := fun es => es.toArray.qsort Name.quickLt
+    asyncMode     := .async
   }
 
 namespace TagDeclarationExtension
@@ -1463,12 +1472,13 @@ instance : Inhabited TagDeclarationExtension :=
 def tag (ext : TagDeclarationExtension) (env : Environment) (declName : Name) : Environment :=
   have : Inhabited Environment := ⟨env⟩
   assert! env.getModuleIdxFor? declName |>.isNone -- See comment at `TagDeclarationExtension`
+  assert! env.asyncMayContain declName
   ext.addEntry env declName
 
 def isTagged (ext : TagDeclarationExtension) (env : Environment) (declName : Name) : Bool :=
   match env.getModuleIdxFor? declName with
   | some modIdx => (ext.getModuleEntries env modIdx).binSearchContains declName Name.quickLt
-  | none        => (ext.getState env).contains declName
+  | none        => (ext.findStateAsync env declName).contains declName
 
 end TagDeclarationExtension
 
@@ -1861,7 +1871,11 @@ def displayStats (env : Environment) : IO Unit := do
   IO.println ("number of extensions:                  " ++ toString env.base.extensions.size);
   pExtDescrs.forM fun extDescr => do
     IO.println ("extension '" ++ toString extDescr.name ++ "'")
-    let s := extDescr.toEnvExtension.getState env
+    -- get state from `checked` at the end if `async`; it would otherwise panic
+    let mut asyncMode := extDescr.toEnvExtension.asyncMode
+    if asyncMode matches .async then
+      asyncMode := .sync
+    let s := extDescr.toEnvExtension.getState (asyncMode := asyncMode) env
     let fmt := extDescr.statsFn s.state
     unless fmt.isNil do IO.println ("  " ++ toString (Format.nest 2 (extDescr.statsFn s.state)))
     IO.println ("  number of imported entries: " ++ toString (s.importedEntries.foldl (fun sum es => sum + es.size) 0))

src/Lean/Meta/Match/Match.lean

@@ -762,7 +762,8 @@ register_builtin_option bootstrap.genMatcherCode : Bool := {
   descr := "disable code generation for auxiliary matcher function"
 }
 
-builtin_initialize matcherExt : EnvExtension (PHashMap (Expr × Bool) Name) ← registerEnvExtension (pure {})
+builtin_initialize matcherExt : EnvExtension (PHashMap (Expr × Bool) Name) ←
+  registerEnvExtension (pure {}) (asyncMode := .local)  -- mere cache, keep it local
 
 /-- Similar to `mkAuxDefinition`, but uses the cache `matcherExt`.
    It also returns an Boolean that indicates whether a new matcher function was added to the environment or not. -/

src/Lean/Meta/Match/MatcherInfo.lean

@@ -82,13 +82,16 @@ builtin_initialize extension : SimplePersistentEnvExtension Entry State ←
   registerSimplePersistentEnvExtension {
     addEntryFn    := State.addEntry
     addImportedFn := fun es => (mkStateFromImportedEntries State.addEntry {} es).switch
+    asyncMode     := .async
   }
 
 def addMatcherInfo (env : Environment) (matcherName : Name) (info : MatcherInfo) : Environment :=
+  let _ : Inhabited Environment := ⟨env⟩
+  assert! env.asyncMayContain matcherName
   extension.addEntry env { name := matcherName, info := info }
 
 def getMatcherInfo? (env : Environment) (declName : Name) : Option MatcherInfo :=
-  (extension.getState env).map.find? declName
+  (extension.findStateAsync env declName).map.find? declName
 
 end Extension
 

src/Lean/Meta/Tactic/AuxLemma.lean

@@ -14,13 +14,16 @@ structure AuxLemmas where
   lemmas : PHashMap Expr (Name × List Name) := {}
   deriving Inhabited
 
-builtin_initialize auxLemmasExt : EnvExtension AuxLemmas ← registerEnvExtension (pure {})
+builtin_initialize auxLemmasExt : EnvExtension AuxLemmas ←
+  registerEnvExtension (pure {}) (asyncMode := .local)  -- a mere cache, keep local
 
 /--
   Helper method for creating auxiliary lemmas in the environment.
 
   It uses a cache that maps `type` to declaration name. The cache is not stored in `.olean` files.
-  It is useful to make sure the same auxiliary lemma is not created over and over again in the same file.
+  It is useful to make sure the same auxiliary lemma is not created over and over again in the same
+  environment branch. For expensive auxiliary lemmas that should be deduplicated even across
+  different environment branches, consider using `realizeConst` instead.
 
   This method is useful for tactics (e.g., `simp`) that may perform preprocessing steps to lemmas provided by
   users. For example, `simp` preprocessor may convert a lemma into multiple ones.
@@ -29,7 +32,7 @@ def mkAuxLemma (levelParams : List Name) (type : Expr) (value : Expr) : MetaM Na
   let env ← getEnv
   let s := auxLemmasExt.getState env
   let mkNewAuxLemma := do
-    let auxName := Name.mkNum (env.mainModule ++ `_auxLemma) s.idx
+    let auxName := Name.mkNum (env.asyncPrefix?.getD env.mainModule ++ `_auxLemma) s.idx
     addDecl <| Declaration.thmDecl {
       name := auxName
       levelParams, type, value

src/Lean/MonadEnv.lean

@@ -39,8 +39,9 @@ def isRec [Monad m] [MonadEnv m] (declName : Name) : m Bool :=
   return isRecCore (← getEnv) declName
 
 @[inline] def withoutModifyingEnv [Monad m] [MonadEnv m] [MonadFinally m] {α : Type} (x : m α) : m α := do
-  let env ← getEnv
-  try x finally setEnv env
+  -- Allow `x` to define new declarations even outside the asynchronous prefix (if any) as all
+  -- results will be discarded anway.
+  withEnv (← getEnv).unlockAsync x
 
 /-- Similar to `withoutModifyingEnv`, but also returns the updated environment -/
 @[inline] def withoutModifyingEnv' [Monad m] [MonadEnv m] [MonadFinally m] {α : Type} (x : m α) : m (α × Environment) := do

src/Lean/ReducibilityAttrs.lean

@@ -31,6 +31,7 @@ builtin_initialize reducibilityCoreExt : PersistentEnvExtension (Name × Reducib
       let r : Array (Name × ReducibilityStatus) := m.fold (fun a n p => a.push (n, p)) #[]
       r.qsort (fun a b => Name.quickLt a.1 b.1)
     statsFn         := fun s => "reducibility attribute core extension" ++ Format.line ++ "number of local entries: " ++ format s.size
+    asyncMode       := .async
   }
 
 builtin_initialize reducibilityExtraExt : SimpleScopedEnvExtension (Name × ReducibilityStatus) (SMap Name ReducibilityStatus) ←
@@ -51,7 +52,7 @@ def getReducibilityStatusCore (env : Environment) (declName : Name) : Reducibili
     match (reducibilityCoreExt.getModuleEntries env modIdx).binSearch (declName, .semireducible) (fun a b => Name.quickLt a.1 b.1) with
     | some (_, status) => status
     | none => .semireducible
-  | none => (reducibilityCoreExt.getState env).find? declName |>.getD .semireducible
+  | none => (reducibilityCoreExt.findStateAsync env declName).find? declName |>.getD .semireducible
 
 private def setReducibilityStatusCore (env : Environment) (declName : Name) (status : ReducibilityStatus) (attrKind : AttributeKind) (currNamespace : Name) : Environment :=
   if attrKind matches .global then
@@ -60,7 +61,8 @@ private def setReducibilityStatusCore (env : Environment) (declName : Name) (sta
       -- Trying to set the attribute of a declaration defined in an imported module.
       reducibilityExtraExt.addEntry env (declName, status)
     | none =>
-      --
+      let _ : Inhabited Environment := ⟨env⟩
+      assert! env.asyncMayContain declName
       reducibilityCoreExt.addEntry env (declName, status)
   else
     -- `scoped` and `local` must be handled by `reducibilityExtraExt`