optimize most important extensions

src/Lean/Class.lean

@@ -57,6 +57,7 @@ builtin_initialize classExtension : SimplePersistentEnvExtension ClassEntry Clas
   registerSimplePersistentEnvExtension {
     addEntryFn    := ClassState.addEntry
     addImportedFn := fun es => (mkStateFromImportedEntries ClassState.addEntry {} es).switch
+    asyncMode     := .mainOnly
   }
 
 /-- Return `true` if `n` is the name of type class in the given environment. -/

src/Lean/Environment.lean

@@ -759,6 +759,7 @@ end ConstantInfo
 inductive EnvExtension.AsyncMode where
   | sync
   | async
+  | local
   | mainOnly
   deriving Inhabited
 
@@ -827,17 +828,18 @@ def mkInitialExtStates : IO (Array EnvExtensionState) := do
 def modifyState {σ : Type} (ext : EnvExtension σ) (env : Environment) (f : σ → σ) : Environment :=
   -- safety: `ext`'s constructor is private, so we can assume the entry at `ext.idx` is of type `σ`
   match ext.asyncMode with
-  | .sync => env.setCheckedSync { env.checked.get with
-    extensions := unsafe ext.modifyStateImpl env.checked.get.extensions f }
-  | .async => env.modifyCheckedAsync fun env =>
-    { env with extensions := unsafe ext.modifyStateImpl env.extensions f }
   | .mainOnly =>
     if let some asyncCtx := env.asyncCtx? then
       let _ : Inhabited Environment := ⟨env⟩
       panic! s!"Environment.modifyState: environment extension is marked as `mainOnly` but used in \
         async context '{asyncCtx.declPrefix}'"
     else
       { env with checkedWithoutAsync.extensions := unsafe ext.modifyStateImpl env.checkedWithoutAsync.extensions f }
+  | .local =>
+    { env with checkedWithoutAsync.extensions := unsafe ext.modifyStateImpl env.checkedWithoutAsync.extensions f }
+  | _ =>
+    env.modifyCheckedAsync fun env =>
+      { env with extensions := unsafe ext.modifyStateImpl env.extensions f }
 
 def setState {σ : Type} (ext : EnvExtension σ) (env : Environment) (s : σ) : Environment :=
   inline <| modifyState ext env fun _ => s
@@ -848,9 +850,7 @@ private unsafe def getStateUnsafe {σ : Type} [Inhabited σ] (ext : EnvExtension
   -- safety: `ext`'s constructor is private, so we can assume the entry at `ext.idx` is of type `σ`
   match ext.asyncMode with
   | .sync     => ext.getStateImpl env.checked.get.extensions
-  | .async    =>
-    panic! "Environment.getState: environment extension is marked as `async`, use `findStateAsync`"
-  | .mainOnly => ext.getStateImpl env.checkedWithoutAsync.extensions
+  | _         => ext.getStateImpl env.checkedWithoutAsync.extensions
 
 @[implemented_by getStateUnsafe]
 opaque getState {σ : Type} [Inhabited σ] (ext : EnvExtension σ) (env : Environment) : σ

src/Lean/Meta/Match/MatcherInfo.lean

@@ -82,6 +82,7 @@ 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 :=

src/Lean/ScopedEnvExtension.lean

@@ -104,22 +104,21 @@ structure ScopedEnvExtension (α : Type) (β : Type) (σ : Type) where
 
 builtin_initialize scopedEnvExtensionsRef : IO.Ref (Array (ScopedEnvExtension EnvExtensionEntry EnvExtensionEntry EnvExtensionState)) ← IO.mkRef #[]
 
-unsafe def registerScopedEnvExtensionUnsafe (descr : Descr α β σ) : IO (ScopedEnvExtension α β σ) := do
+def registerScopedEnvExtension (descr : Descr α β σ)
+    (asyncMode : EnvExtension.AsyncMode := .local) : IO (ScopedEnvExtension α β σ) := do
   let ext ← registerPersistentEnvExtension {
     name            := descr.name
     mkInitial       := mkInitial descr
     addImportedFn   := addImportedFn descr
     addEntryFn      := addEntryFn descr
     exportEntriesFn := exportEntriesFn
     statsFn         := fun s => format "number of local entries: " ++ format s.newEntries.length
+    asyncMode
   }
   let ext := { descr := descr, ext := ext : ScopedEnvExtension α β σ }
-  scopedEnvExtensionsRef.modify fun exts => exts.push (unsafeCast ext)
+  unsafe scopedEnvExtensionsRef.modify fun exts => exts.push (unsafeCast ext)
   return ext
 
-@[implemented_by registerScopedEnvExtensionUnsafe]
-opaque registerScopedEnvExtension (descr : Descr α β σ) : IO (ScopedEnvExtension α β σ)
-
 def ScopedEnvExtension.pushScope (ext : ScopedEnvExtension α β σ) (env : Environment) : Environment :=
   let s := ext.ext.getState env
   match s.stateStack with