feat: Environment.realizeConst (#7076)

This PR introduces the central parallelism API for ensuring that helper
declarations can be generated lazily without duplicating work or
creating conflicts across threads.

src/Init/System/IO.lean

@@ -57,6 +57,11 @@ def EIO.catchExceptions (act : EIO ε α) (h : ε → BaseIO α) : BaseIO α :=
   | EStateM.Result.ok a s     => EStateM.Result.ok a s
   | EStateM.Result.error ex s => h ex s
 
+def EIO.ofExcept (e : Except ε α) : EIO ε α :=
+  match e with
+  | Except.ok a    => pure a
+  | Except.error e => throw e
+
 open IO (Error) in
 abbrev IO : Type → Type := EIO Error
 

src/Init/System/IOError.lean

@@ -48,7 +48,9 @@ inductive IO.Error where
 
   | unexpectedEof
   | userError (msg : String)
-  deriving Inhabited
+
+instance : Inhabited IO.Error where
+  default := .userError "(`Inhabited.default` for `IO.Error`)"
 
 @[export lean_mk_io_user_error]
 def IO.userError (s : String) : IO.Error :=

src/Init/System/Promise.lean

@@ -73,5 +73,5 @@ def Promise.result := @Promise.result!
 /--
 Like `Promise.result`, but resolves to `dflt` if the promise is dropped without ever being resolved.
 -/
-def Promise.resultD (promise : Promise α) (dflt : α): Task α :=
+@[macro_inline] def Promise.resultD (promise : Promise α) (dflt : α) : Task α :=
   promise.result?.map (sync := true) (·.getD dflt)

src/Lean/AddDecl.lean

@@ -8,12 +8,6 @@ import Lean.CoreM
 
 namespace Lean
 
-register_builtin_option debug.skipKernelTC : Bool := {
-  defValue := false
-  group    := "debug"
-  descr    := "skip kernel type checker. WARNING: setting this option to true may compromise soundness because your proofs will not be checked by the Lean kernel"
-}
-
 /-- Adds given declaration to the environment, respecting `debug.skipKernelTC`. -/
 def Kernel.Environment.addDecl (env : Environment) (opts : Options) (decl : Declaration)
     (cancelTk? : Option IO.CancelToken := none) : Except Exception Environment :=

src/Lean/CoreM.lean

@@ -194,7 +194,7 @@ protected def withFreshMacroScope (x : CoreM α) : CoreM α := do
 
 instance : MonadQuotation CoreM where
   getCurrMacroScope   := return (← read).currMacroScope
-  getMainModule       := return (← get).env.mainModule
+  getMainModule       := return (← getEnv).mainModule
   withFreshMacroScope := Core.withFreshMacroScope
 
 instance : Elab.MonadInfoTree CoreM where
@@ -413,6 +413,26 @@ register_builtin_option stderrAsMessages : Bool := {
   descr    := "(server) capture output to the Lean stderr channel (such as from `dbg_trace`) during elaboration of a command as a diagnostic message"
 }
 
+/--
+Creates snapshot reporting given `withIsolatedStreams` output and diagnostics and traces from the
+given state.
+-/
+def mkSnapshot (output : String) (ctx : Context) (st : State)
+    (desc : String := by exact decl_name%.toString) : BaseIO Language.SnapshotTree := do
+  let mut msgs := st.messages
+  if !output.isEmpty then
+    msgs := msgs.add {
+      fileName := ctx.fileName
+      severity := MessageSeverity.information
+      pos      := ctx.fileMap.toPosition <| ctx.ref.getPos?.getD 0
+      data     := output
+    }
+  return .mk {
+    desc
+    diagnostics := (← Language.Snapshot.Diagnostics.ofMessageLog msgs)
+    traces := st.traceState
+  } st.snapshotTasks
+
 open Language in
 /--
 Wraps the given action for use in `BaseIO.asTask` etc., discarding its final state except for
@@ -443,20 +463,7 @@ def wrapAsyncAsSnapshot (act : Unit → CoreM Unit) (cancelTk? : Option IO.Cance
   let ctx ← readThe Core.Context
   return do
     match (← t.toBaseIO) with
-    | .ok (output, st) =>
-      let mut msgs := st.messages
-      if !output.isEmpty then
-        msgs := msgs.add {
-          fileName := ctx.fileName
-          severity := MessageSeverity.information
-          pos      := ctx.fileMap.toPosition <| ctx.ref.getPos?.getD 0
-          data     := output
-        }
-      return .mk {
-        desc
-        diagnostics := (← Language.Snapshot.Diagnostics.ofMessageLog msgs)
-        traces := st.traceState
-      } st.snapshotTasks
+    | .ok (output, st) => mkSnapshot output ctx st desc
     -- interrupt or abort exception as `try catch` above should have caught any others
     | .error _ => default
 
@@ -646,6 +653,11 @@ def logMessageKind (kind : Name) : CoreM Bool := do
     modify fun s => { s with messages.loggedKinds := s.messages.loggedKinds.insert kind }
     return true
 
+@[inherit_doc Environment.enableRealizationsForConst]
+def enableRealizationsForConst (n : Name) : CoreM Unit := do
+  let env ← (← getEnv).enableRealizationsForConst (← getOptions) n
+  setEnv env
+
 builtin_initialize
   registerTraceClass `Elab.async
   registerTraceClass `Elab.block

src/Lean/Elab/MutualInductive.lean

@@ -931,6 +931,7 @@ private def mkInductiveDecl (vars : Array Expr) (elabs : Array InductiveElabStep
         for ctor in view.ctors do
           if (ctor.declId.getPos? (canonicalOnly := true)).isSome then
             Term.addTermInfo' ctor.declId (← mkConstWithLevelParams ctor.declName) (isBinder := true)
+            enableRealizationsForConst ctor.declName
     return res
 
 private def mkAuxConstructions (declNames : Array Name) : TermElabM Unit := do

src/Lean/Elab/PreDefinition/Basic.lean

@@ -161,6 +161,7 @@ private def addNonRecAux (preDef : PreDefinition) (compile : Bool) (all : List N
     if compile && shouldGenCodeFor preDef then
       compileDecl decl
     if applyAttrAfterCompilation then
+      enableRealizationsForConst preDef.declName
       generateEagerEqns preDef.declName
       applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation
 

src/Lean/Elab/PreDefinition/Mutual.lean

@@ -82,6 +82,7 @@ Assign final attributes to the definitions. Assumes the EqnInfos to be already p
 def addPreDefAttributes (preDefs : Array PreDefinition) : TermElabM Unit := do
   for preDef in preDefs do
     markAsRecursive preDef.declName
+    enableRealizationsForConst preDef.declName
     generateEagerEqns preDef.declName
     applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation
     -- Unless the user asks for something else, mark the definition as irreducible

src/Lean/Elab/PreDefinition/Nonrec/Eqns.lean

@@ -20,18 +20,23 @@ Simple, coarse-grained equation theorem for nonrecursive definitions.
 -/
 private def mkSimpleEqThm (declName : Name) (suffix := Name.mkSimple unfoldThmSuffix) : MetaM (Option Name) := do
   if let some (.defnInfo info) := (← getEnv).find? declName then
+    let name := declName ++ suffix
+    -- determinism: `name` and `info` are dependent only on `declName`, not any later env
+    -- modifications
+    realizeConst declName name (doRealize name info)
+    return some name
+  else
+    return none
+where
+  doRealize name info :=
     lambdaTelescope (cleanupAnnotations := true) info.value fun xs body => do
       let lhs := mkAppN (mkConst info.name <| info.levelParams.map mkLevelParam) xs
       let type  ← mkForallFVars xs (← mkEq lhs body)
       let value ← mkLambdaFVars xs (← mkEqRefl lhs)
-      let name := declName ++ suffix
-      addDecl <| Declaration.thmDecl {
+      addDecl <| .thmDecl {
         name, type, value
         levelParams := info.levelParams
       }
-      return some name
-  else
-    return none
 
 def getEqnsFor? (declName : Name) : MetaM (Option (Array Name)) := do
   if (← isRecursiveDefinition declName) then

src/Lean/Elab/PreDefinition/Structural/Main.lean

@@ -193,6 +193,10 @@ def structuralRecursion (preDefs : Array PreDefinition) (termMeasure?s : Array (
         registerEqnsInfo preDef (preDefs.map (·.declName)) recArgPos numFixed
     addSmartUnfoldingDef preDef recArgPos
     markAsRecursive preDef.declName
+  for preDef in preDefs do
+    -- must happen in separate loop so realizations can see eqnInfos of all other preDefs
+    enableRealizationsForConst preDef.declName
+    -- must happen after `enableRealizationsForConst`
     generateEagerEqns preDef.declName
   applyAttributesOf preDefsNonRec AttributeApplicationTime.afterCompilation
 

src/Lean/Elab/PreDefinition/WF/Main.lean

@@ -68,6 +68,7 @@ def wfRecursion (preDefs : Array PreDefinition) (termMeasure?s : Array (Option T
       unless (← isProp preDef.type) do
         WF.mkUnfoldEq preDef preDefNonRec.declName wfPreprocessProof
   Mutual.addPreDefAttributes preDefs
+  enableRealizationsForConst preDefNonRec.declName
 
 builtin_initialize registerTraceClass `Elab.definition.wf
 

src/Lean/Elab/PreDefinition/WF/Unfold.lean

@@ -100,4 +100,7 @@ def mkUnfoldEq (preDef : PreDefinition) (unaryPreDefName : Name) (wfPreprocessPr
       }
       trace[Elab.definition.wf] "mkUnfoldEq defined {.ofConstName name}"
 
+builtin_initialize
+  registerTraceClass `Elab.definition.wf.eqns
+
 end Lean.Elab.WF

src/Lean/Elab/Structure.lean

@@ -839,6 +839,9 @@ private def addProjections (r : ElabHeaderResult) (fieldInfos : Array StructFiel
   for fieldInfo in fieldInfos do
     if fieldInfo.isSubobject then
       addDeclarationRangesFromSyntax fieldInfo.declName r.view.ref fieldInfo.ref
+  for n in projNames do
+    -- projections may generate equation theorems
+    enableRealizationsForConst n
 
 private def registerStructure (structName : Name) (infos : Array StructFieldInfo) : TermElabM Unit := do
   let fields ← infos.filterMapM fun info => do

src/Lean/Environment.lean

@@ -18,8 +18,10 @@ import Lean.Util.Path
 import Lean.Util.FindExpr
 import Lean.Util.Profile
 import Lean.Util.InstantiateLevelParams
+import Lean.Util.FoldConsts
 import Lean.PrivateName
 import Lean.LoadDynlib
+import Init.Dynamic
 
 /-!
 # Note [Environment Branches]
@@ -65,6 +67,12 @@ paths back together.
 -/
 
 namespace Lean
+register_builtin_option debug.skipKernelTC : Bool := {
+  defValue := false
+  group    := "debug"
+  descr    := "skip kernel type checker. WARNING: setting this option to true may compromise soundness because your proofs will not be checked by the Lean kernel"
+}
+
 /-- Opaque environment extension state. -/
 opaque EnvExtensionStateSpec : (α : Type) × Inhabited α := ⟨Unit, ⟨()⟩⟩
 def EnvExtensionState : Type := EnvExtensionStateSpec.fst
@@ -252,6 +260,28 @@ inductive Exception where
   | excessiveMemory
   | deepRecursion
   | interrupted
+deriving Nonempty
+
+/-- Basic `Exception` formatting without `MessageData` dependency. -/
+private def Exception.toRawString : Kernel.Exception → String
+  | unknownConstant _ constName       => s!"(kernel) unknown constant '{constName}'"
+  | alreadyDeclared _ constName       => s!"(kernel) constant has already been declared '{constName}'"
+  | declTypeMismatch _ _ _ => s!"(kernel) declaration type mismatch"
+  | declHasMVars _ constName _        => s!"(kernel) declaration has metavariables '{constName}'"
+  | declHasFVars _ constName _        => s!"(kernel) declaration has free variables '{constName}'"
+  | funExpected _ _ e                 => s!"(kernel) function expected: {e}"
+  | typeExpected _ _ e                => s!"(kernel) type expected: {e}"
+  | letTypeMismatch  _ _ n _ _        => s!"(kernel) let-declaration type mismatch '{n}'"
+  | exprTypeMismatch _ _ e _          => s!"(kernel) type mismatch at {e}"
+  | appTypeMismatch  _ _ e fnType argType =>
+    s!"application type mismatch: {e}\nargument has type {argType}\nbut function has type {fnType}"
+  | invalidProj _ _ e                 => s!"(kernel) invalid projection {e}"
+  | thmTypeIsNotProp _ constName type => s!"(kernel) type of theorem '{constName}' is not a proposition: {type}"
+  | other msg                         => s!"(kernel) {msg}"
+  | deterministicTimeout              => "(kernel) deterministic timeout"
+  | excessiveMemory                   => "(kernel) excessive memory consumption detected"
+  | deepRecursion                     => "(kernel) deep recursion detected"
+  | interrupted                       => "(kernel) interrupted"
 
 namespace Environment
 
@@ -346,6 +376,7 @@ structure AsyncConstantInfo where
   sig       : Task ConstantVal
   /-- The final, complete constant info, potentially filled asynchronously. -/
   constInfo : Task ConstantInfo
+deriving Inhabited
 
 namespace AsyncConstantInfo
 
@@ -365,21 +396,25 @@ end AsyncConstantInfo
 
 /--
 Information about the current branch of the environment representing asynchronous elaboration.
+
+Use `Environment.enterAsync` instead of `mkRaw`.
 -/
-structure AsyncContext where
+private structure AsyncContext where mkRaw ::
   /--
   Name of the declaration asynchronous elaboration was started for. All constants added to this
   environment branch must have the name as a prefix, after erasing macro scopes and private name
   prefixes.
   -/
   declPrefix : Name
+  /-- Whether we are in `realizeConst`, used to restrict env ext modifications. -/
+  realizing  : Bool
 deriving Nonempty
 
 /--
 Checks whether a declaration named `n` may be added to the environment in the given context. See
 also `AsyncContext.declPrefix`.
 -/
-def AsyncContext.mayContain (ctx : AsyncContext) (n : Name) : Bool :=
+private def AsyncContext.mayContain (ctx : AsyncContext) (n : Name) : Bool :=
   ctx.declPrefix.isPrefixOf <| privateToUserName n.eraseMacroScopes
 
 /--
@@ -394,29 +429,47 @@ structure AsyncConst where
   exts?     : Option (Task (Array EnvExtensionState))
 
 /-- Data structure holding a sequence of `AsyncConst`s optimized for efficient access. -/
-structure AsyncConsts where
-  toArray : Array AsyncConst := #[]
+private structure AsyncConsts where
+  size : Nat
+  revList : List AsyncConst
   /-- Map from declaration name to const for fast direct access. -/
-  private map : NameMap AsyncConst := {}
+  map : NameMap AsyncConst
   /-- Trie of declaration names without private name prefixes for fast longest-prefix access. -/
-  private normalizedTrie : NameTrie AsyncConst := {}
+  normalizedTrie : NameTrie AsyncConst
 deriving Inhabited
 
-def AsyncConsts.add (aconsts : AsyncConsts) (aconst : AsyncConst) : AsyncConsts :=
+private def AsyncConsts.add (aconsts : AsyncConsts) (aconst : AsyncConst) : AsyncConsts :=
   { aconsts with
-    toArray := aconsts.toArray.push aconst
+    size := aconsts.size + 1
+    revList := aconst :: aconsts.revList
     map := aconsts.map.insert aconst.constInfo.name aconst
     normalizedTrie := aconsts.normalizedTrie.insert (privateToUserName aconst.constInfo.name) aconst
   }
 
-def AsyncConsts.find? (aconsts : AsyncConsts) (declName : Name) : Option AsyncConst :=
+private def AsyncConsts.find? (aconsts : AsyncConsts) (declName : Name) : Option AsyncConst :=
   aconsts.map.find? declName
 
 /-- Finds the constant in the collection that is a prefix of `declName`, if any. -/
-def AsyncConsts.findPrefix? (aconsts : AsyncConsts) (declName : Name) : Option AsyncConst :=
+private def AsyncConsts.findPrefix? (aconsts : AsyncConsts) (declName : Name) : Option AsyncConst :=
   -- as macro scopes are a strict suffix,
   aconsts.normalizedTrie.findLongestPrefix? (privateToUserName declName.eraseMacroScopes)
 
+/-- Context for `realizeConst` established by `enableRealizationsForConst`. -/
+private structure RealizationContext where
+  /--
+  Saved `Environment`, untyped to avoid cyclic reference. Import environment for imported constants.
+  -/
+  env       : NonScalar
+  /-- Saved options. Empty for imported constants. -/
+  opts      : Options
+  /--
+  `realizeConst _ c ..` adds a mapping from `c` to a task of the realization results: the newly
+  added constants (incl. extension data in `AsyncConst.exts?`),  a function for replaying the
+  changes onto a derived kernel environment, and auxiliary data (always `SnapshotTree` in builtin
+  uses, but untyped to avoid cyclic module references).
+  -/
+  constsRef : IO.Ref (NameMap (Task (List AsyncConst × (Kernel.Environment → Kernel.Environment) × Dynamic)))
+
 /--
 Elaboration-specific extension of `Kernel.Environment` that adds tracking of asynchronously
 elaborated declarations.
@@ -443,19 +496,32 @@ structure Environment where
   -/
   checked             : Task Kernel.Environment := .pure checkedWithoutAsync
   /--
-  Container of asynchronously elaborated declarations, i.e.
-  `checked = checkedWithoutAsync ⨃ asyncConsts`.
+  Container of asynchronously elaborated declarations. For consistency, `updateBaseAfterKernelAdd`
+  makes sure this contains constants added even synchronously, i.e. this is a superset of
+  `checkedWithoutAsync` except for imported constants.
   -/
-  private asyncConsts : AsyncConsts := {}
+  private asyncConsts : AsyncConsts := default
   /-- Information about this asynchronous branch of the environment, if any. -/
   private asyncCtx?   : Option AsyncContext := none
+  /--
+  Realized constants belonging to imported declarations. `none` only from `Environment.ofKernelEnv`,
+  which should never leak into general elaboration.
+  -/
+  private realizedImportedConsts? : Option RealizationContext
+  /--
+  Realized constants belonging to local declarations. This is a map from local declarations, which
+  need to be registered synchronously using `enableRealizationsForConst`, to their realization
+  context incl. a ref of realized constants.
+  -/
+  private realizedLocalConsts  : NameMap RealizationContext := {}
 deriving Nonempty
 
 namespace Environment
 
+-- used only when the kernel calls into the interpreter, and in `Lean.Kernel.Exception.mkCtx`
 @[export lean_elab_environment_of_kernel_env]
 def ofKernelEnv (env : Kernel.Environment) : Environment :=
-  { checkedWithoutAsync := env }
+  { checkedWithoutAsync := env, realizedImportedConsts? := none }
 
 @[export lean_elab_environment_to_kernel_env]
 def toKernelEnv (env : Environment) : Kernel.Environment :=
@@ -469,6 +535,10 @@ private def modifyCheckedAsync (env : Environment) (f : Kernel.Environment → K
 private def setCheckedSync (env : Environment) (newChecked : Kernel.Environment) : Environment :=
   { env with checked := .pure newChecked, checkedWithoutAsync := newChecked }
 
+/-- True while inside `realizeConst`'s `realize`. -/
+def isRealizing (env : Environment) : Bool :=
+  env.asyncCtx?.any (·.realizing)
+
 /--
 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
@@ -574,6 +644,30 @@ def findConstVal? (env : Environment) (n : Name) : Option ConstantVal := do
     return asyncConst.constInfo.toConstantVal
   else env.findNoAsync n |>.map (·.toConstantVal)
 
+/--
+Allows `realizeConst` calls for the given declaration in all derived environment branches.
+Realizations will run using the given environment and options to ensure deterministic results. Note
+that while we check that the function isn't called too *early*, i.e. before the declaration is
+actually added to the environment, we cannot automatically check that it isn't called too *late*,
+i.e. before all environment extensions that may be relevant to realizations have been set. We do
+check that we are not calling it from a different branch than `c` was added on, which would be
+definitely too late.
+-/
+def enableRealizationsForConst (env : Environment) (opts : Options) (c : Name) :
+    BaseIO Environment := do
+  if env.findAsync? c |>.isNone then
+    panic! s!"Environment.enableRealizationsForConst: declaration {c} not found in environment"
+  if let some asyncCtx := env.asyncCtx? then
+    if !asyncCtx.mayContain c then
+      panic! s!"Environment.enableRealizationsForConst: {c} is outside current context {asyncCtx.declPrefix}"
+  if env.realizedLocalConsts.contains c then
+    return env
+  return { env with realizedLocalConsts := env.realizedLocalConsts.insert c {
+    -- safety: `RealizationContext` is private
+    env := unsafe unsafeCast env
+    opts
+    constsRef := (← IO.mkRef {}) } }
+
 /--
 Looks up the given declaration name in the environment, blocking on the corresponding elaboration
 task if not yet complete.
@@ -590,9 +684,14 @@ def find? (env : Environment) (n : Name) : Option ConstantInfo :=
 def dbgFormatAsyncState (env : Environment) : BaseIO String :=
   return s!"\
     asyncCtx.declPrefix: {repr <| env.asyncCtx?.map (·.declPrefix)}\
-  \nasyncConsts: {repr <| env.asyncConsts.toArray.map (·.constInfo.name)}\
-  \ncheckedWithoutAsync.constants.map₂: {repr <|
-      env.checkedWithoutAsync.constants.map₂.toList.map (·.1)}"
+  \nasyncConsts: {repr <| env.asyncConsts.revList.reverse.map (·.constInfo.name)}\
+  \nrealizedLocalConsts: {repr (← env.realizedLocalConsts.toList.mapM fun (n, ctx) => do
+    let consts := (← ctx.constsRef.get).toList
+    return (n, consts.map (·.1)))}
+  \nrealizedImportedConsts?: {repr <| (← env.realizedImportedConsts?.mapM fun ctx => do
+    return (← ctx.constsRef.get).toList.map fun (n, m?) =>
+      (n, m?.get.1.map (fun c : AsyncConst => c.constInfo.name.toString) |> toString))}
+  \ncheckedWithoutAsync.constants.map₂: {repr <| env.checkedWithoutAsync.constants.map₂.toList.map (·.1)}"
 
 /-- Returns debug output about the synchronous state of the environment. -/
 def dbgFormatCheckedSyncState (env : Environment) : BaseIO String :=
@@ -614,6 +713,13 @@ structure PromiseCheckedResult where
   asyncEnv : Environment
   private checkedEnvPromise : IO.Promise Kernel.Environment
 
+/-- Creates an async context for the given declaration name, normalizing it for use as a prefix. -/
+private def enterAsync (declName : Name) (realizing := false) (env : Environment) : Environment :=
+  { env with asyncCtx? := some {
+    declPrefix := privateToUserName declName.eraseMacroScopes
+    -- `realizing` is sticky
+    realizing := realizing || env.asyncCtx?.any (·.realizing) } }
+
 /--
 Starts an asynchronous modification of the kernel environment. The environment is split into a
 "main" branch that will block on access to the kernel environment until
@@ -626,10 +732,8 @@ def promiseChecked (env : Environment) : BaseIO PromiseCheckedResult := do
       checked := checkedEnvPromise.result?.bind (sync := true) fun
         | some kenv => .pure kenv
         | none      => env.checked }
-    asyncEnv := { env with
-      -- Do not allow adding new constants
-      asyncCtx? := some { declPrefix := `__reserved__Environment_promiseChecked }
-    }
+    -- Do not allow adding new constants
+    asyncEnv := env.enterAsync `__reserved__Environment_promiseChecked
     checkedEnvPromise
   }
 
@@ -664,28 +768,14 @@ structure AddConstAsyncResult where
   private extensionsPromise : IO.Promise (Array EnvExtensionState)
   private checkedEnvPromise : IO.Promise Kernel.Environment
 
-/--
-Starts the asynchronous addition of a constant to the environment. The environment is split into a
-"main" branch that holds a reference to the constant to be added but will block on access until the
-corresponding information has been added on the "async" environment branch and committed there; see
-the respective fields of `AddConstAsyncResult` as well as the [Environment Branches] note for more
-information.
--/
-def addConstAsync (env : Environment) (constName : Name) (kind : ConstantKind) (reportExts := true) :
-    IO AddConstAsyncResult := do
-  assert! env.asyncMayContain constName
-  let sigPromise ← IO.Promise.new
-  let infoPromise ← IO.Promise.new
-  let extensionsPromise ← IO.Promise.new
-  let checkedEnvPromise ← IO.Promise.new
-
-  -- fallback info in case promises are dropped unfulfilled
-  let fallbackVal := {
+/-- Creates fallback info to be used in case promises are dropped unfulfilled. -/
+private def mkFallbackConstInfo (constName : Name) (kind : ConstantKind) : ConstantInfo :=
+  let fallbackVal : ConstantVal := {
     name := constName
     levelParams := []
-    type := mkApp2 (mkConst ``sorryAx [0]) (mkSort 0) (mkConst ``true)
+    type := mkApp2 (mkConst ``sorryAx [1]) (mkSort 0) (mkConst ``true)
   }
-  let fallbackInfo := match kind with
+  match kind with
     | .defn => .defnInfo { fallbackVal with
       value := mkApp2 (mkConst ``sorryAx [0]) fallbackVal.type (mkConst ``true)
       hints := .abbrev
@@ -697,16 +787,38 @@ def addConstAsync (env : Environment) (constName : Name) (kind : ConstantKind) (
     | .axiom  => .axiomInfo { fallbackVal with
       isUnsafe := false
     }
-    | k => panic! s!"AddConstAsyncResult.addConstAsync: unsupported constant kind {repr k}"
+    | k => panic! s!"Environment.mkFallbackConstInfo: unsupported constant kind {repr k}"
+
+/--
+Starts the asynchronous addition of a constant to the environment. The environment is split into a
+"main" branch that holds a reference to the constant to be added but will block on access until the
+corresponding information has been added on the "async" environment branch and committed there; see
+the respective fields of `AddConstAsyncResult` as well as the [Environment Branches] note for more
+information.
+-/
+def addConstAsync (env : Environment) (constName : Name) (kind : ConstantKind) (reportExts := true)
+    (checkMayContain := true) :
+    IO AddConstAsyncResult := do
+  if checkMayContain then
+    if let some ctx := env.asyncCtx? then
+      if !ctx.mayContain constName then
+        throw <| .userError s!"cannot add declaration {constName} to environment as it is \
+          restricted to the prefix {ctx.declPrefix}"
+  let sigPromise ← IO.Promise.new
+  let infoPromise ← IO.Promise.new
+  let extensionsPromise ← IO.Promise.new
+  let checkedEnvPromise ← IO.Promise.new
+
+  let fallbackConstInfo := mkFallbackConstInfo constName kind
 
   let asyncConst := {
     constInfo := {
       name := constName
       kind
-      sig := sigPromise.resultD fallbackVal
-      constInfo := infoPromise.resultD fallbackInfo
+      sig := sigPromise.resultD fallbackConstInfo.toConstantVal
+      constInfo := infoPromise.resultD fallbackConstInfo
     }
-    exts? := guard reportExts *> some (extensionsPromise.resultD #[])
+    exts? := guard reportExts *> some (extensionsPromise.resultD env.toKernelEnv.extensions)
   }
   return {
     constName, kind
@@ -715,9 +827,7 @@ def addConstAsync (env : Environment) (constName : Name) (kind : ConstantKind) (
       checked := checkedEnvPromise.result?.bind (sync := true) fun
         | some kenv => .pure kenv
         | none      => env.checked }
-    asyncEnv := { env with
-      asyncCtx? := some { declPrefix := privateToUserName constName.eraseMacroScopes }
-    }
+    asyncEnv := env.enterAsync constName
     sigPromise, infoPromise, extensionsPromise, checkedEnvPromise
   }
 
@@ -880,6 +990,9 @@ inductive EnvExtension.AsyncMode where
   | async
   deriving Inhabited
 
+abbrev ReplayFn (σ : Type) :=
+  (oldState : σ) → (newState : σ) → (newConsts : List Name) → σ → σ
+
 /--
 Environment extension, can only be generated by `registerEnvExtension` that allocates a unique index
 for this extension into each environment's extension state's array.
@@ -888,6 +1001,13 @@ structure EnvExtension (σ : Type) where private mk ::
   idx       : Nat
   mkInitial : IO σ
   asyncMode : EnvExtension.AsyncMode
+  /--
+  Optional function that, given state before and after realization and newly added constants,
+  replays this change onto a state from another (derived) environment. This function is used only
+  when making changes to an extension inside a `realizeConst` call, in which case it must be
+  present.
+  -/
+  replay?   : Option (ReplayFn σ)
   deriving Inhabited
 
 namespace EnvExtension
@@ -949,19 +1069,24 @@ from different environment branches are reconciled.
 Note that in modes `sync` and `async`, `f` will be called twice, on the local and on the `checked`
 state.
 -/
-def modifyState {σ : Type} (ext : EnvExtension σ) (env : Environment) (f : σ → σ) : Environment :=
+def modifyState {σ : Type} (ext : EnvExtension σ) (env : Environment) (f : σ → σ) : Environment := Id.run do
   -- safety: `ext`'s constructor is private, so we can assume the entry at `ext.idx` is of type `σ`
   match ext.asyncMode with
   | .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 }
+        {if asyncCtx.realizing then "realization" else "async"} context '{asyncCtx.declPrefix}'"
+    return { env with checkedWithoutAsync.extensions := unsafe ext.modifyStateImpl env.checkedWithoutAsync.extensions f }
   | .local =>
-    { env with checkedWithoutAsync.extensions := unsafe ext.modifyStateImpl env.checkedWithoutAsync.extensions f }
+    if let some asyncCtx := env.asyncCtx?.filter (·.realizing) then
+      panic! s!"Environment.modifyState: environment extension is marked as `local` but used in \
+        realization context '{asyncCtx.declPrefix}'"
+    return { env with checkedWithoutAsync.extensions := unsafe ext.modifyStateImpl env.checkedWithoutAsync.extensions f }
   | _ =>
+    if ext.replay?.isNone then
+      if let some asyncCtx := env.asyncCtx?.filter (·.realizing) then
+        panic! s!"Environment.modifyState: environment extension must set `replay?` field to be \
+          used in realization context '{asyncCtx.declPrefix}'"
     env.modifyCheckedAsync fun env =>
       { env with extensions := unsafe ext.modifyStateImpl env.extensions f }
 
@@ -992,6 +1117,24 @@ recommended and should be considered only for important optimizations.
 opaque getState {σ : Type} [Inhabited σ] (ext : EnvExtension σ) (env : Environment)
   (asyncMode := ext.asyncMode) : σ
 
+-- `unsafe` fails to infer `Nonempty` here
+private unsafe def findStateAsyncUnsafe {σ : Type} [Inhabited σ]
+    (ext : EnvExtension σ) (env : Environment) (declPrefix : Name) : σ :=
+  -- safety: `ext`'s constructor is private, so we can assume the entry at `ext.idx` is of type `σ`
+  if let some { exts? := some exts, .. } := env.asyncConsts.findPrefix? declPrefix then
+    ext.getStateImpl exts.get
+  else
+    ext.getStateImpl env.checkedWithoutAsync.extensions
+
+/--
+Returns the final extension state on the environment branch corresponding to the passed declaration
+name, if any, or otherwise the state on the current branch. In other words, at most one environment
+branch will be blocked on.
+-/
+@[implemented_by findStateAsyncUnsafe]
+opaque findStateAsync {σ : Type} [Inhabited σ] (ext : EnvExtension σ)
+  (env : Environment) (declPrefix : Name) : σ
+
 end EnvExtension
 
 /-- Environment extensions can only be registered during initialization.
@@ -1002,12 +1145,13 @@ end EnvExtension
    Note that by default, extension state is *not* stored in .olean files and will not propagate across `import`s.
    For that, you need to register a persistent environment extension. -/
 def registerEnvExtension {σ : Type} (mkInitial : IO σ)
+    (replay? : Option (ReplayFn σ) := none)
     (asyncMode : EnvExtension.AsyncMode := .mainOnly) : IO (EnvExtension σ) := do
   unless (← initializing) do
     throw (IO.userError "failed to register environment, extensions can only be registered during initialization")
   let exts ← EnvExtension.envExtensionsRef.get
   let idx := exts.size
-  let ext : EnvExtension σ := { idx, mkInitial, asyncMode }
+  let ext : EnvExtension σ := { idx, mkInitial, asyncMode, replay? }
   -- safety: `EnvExtensionState` is opaque, so we can upcast to it
   EnvExtension.envExtensionsRef.modify fun exts => exts.push (unsafe unsafeCast ext)
   pure ext
@@ -1019,7 +1163,7 @@ def mkEmptyEnvironment (trustLevel : UInt32 := 0) : IO Environment := do
   let initializing ← IO.initializing
   if initializing then throw (IO.userError "environment objects cannot be created during initialization")
   let exts ← mkInitialExtensionStates
-  pure {
+  return {
     checkedWithoutAsync := {
       const2ModIdx    := {}
       constants       := {}
@@ -1027,6 +1171,7 @@ def mkEmptyEnvironment (trustLevel : UInt32 := 0) : IO Environment := do
       extraConstNames := {}
       extensions      := exts
     }
+    realizedImportedConsts? := none
   }
 
 structure PersistentEnvExtensionState (α : Type) (σ : Type) where
@@ -1117,8 +1262,9 @@ def addEntry {α β σ : Type} (ext : PersistentEnvExtension α β σ) (env : En
     { s with state := state }
 
 /-- Get the current state of the given extension in the given environment. -/
-def getState {α β σ : Type} [Inhabited σ] (ext : PersistentEnvExtension α β σ) (env : Environment) : σ :=
-  (ext.toEnvExtension.getState env).state
+def getState {α β σ : Type} [Inhabited σ] (ext : PersistentEnvExtension α β σ) (env : Environment)
+    (asyncMode := ext.toEnvExtension.asyncMode) : σ :=
+  (ext.toEnvExtension.getState (asyncMode := asyncMode) env).state
 
 /-- Set the current state of the given extension in the given environment. -/
 def setState {α β σ : Type} (ext : PersistentEnvExtension α β σ) (env : Environment) (s : σ) : Environment :=
@@ -1128,23 +1274,11 @@ def setState {α β σ : Type} (ext : PersistentEnvExtension α β σ) (env : En
 def modifyState {α β σ : Type} (ext : PersistentEnvExtension α β σ) (env : Environment) (f : σ → σ) : Environment :=
   ext.toEnvExtension.modifyState env fun ps => { ps with state := f (ps.state) }
 
--- `unsafe` fails to infer `Nonempty` here
-private unsafe def findStateAsyncUnsafe {α β σ : Type} [Inhabited σ]
-    (ext : PersistentEnvExtension α β σ) (env : Environment) (declPrefix : Name) : σ :=
-  -- safety: `ext`'s constructor is private, so we can assume the entry at `ext.idx` is of type `σ`
-  if let some { exts? := some exts, .. } := env.asyncConsts.findPrefix? declPrefix then
-    ext.toEnvExtension.getStateImpl exts.get |>.state
-  else
-    ext.toEnvExtension.getStateImpl env.checkedWithoutAsync.extensions |>.state
+@[inherit_doc EnvExtension.findStateAsync]
+def findStateAsync {α β σ : Type} [Inhabited σ] (ext : PersistentEnvExtension α β σ)
+    (env : Environment) (declPrefix : Name) : σ :=
+  ext.toEnvExtension.findStateAsync env declPrefix |>.state
 
-/--
-Returns the final extension state on the environment branch corresponding to the passed declaration
-name, if any, or otherwise the state on the current branch. In other words, at most one environment
-branch will be blocked on.
--/
-@[implemented_by findStateAsyncUnsafe]
-opaque findStateAsync {α β σ : Type} [Inhabited σ] (ext : PersistentEnvExtension α β σ)
-  (env : Environment) (declPrefix : Name) : σ
 
 end PersistentEnvExtension
 
@@ -1158,11 +1292,14 @@ structure PersistentEnvExtensionDescr (α β σ : Type) where
   exportEntriesFn : σ → Array α
   statsFn         : σ → Format := fun _ => Format.nil
   asyncMode       : EnvExtension.AsyncMode := .mainOnly
+  replay?         : Option (ReplayFn σ) := none
 
 unsafe def registerPersistentEnvExtensionUnsafe {α β σ : Type} [Inhabited σ] (descr : PersistentEnvExtensionDescr α β σ) : IO (PersistentEnvExtension α β σ) := do
   let pExts ← persistentEnvExtensionsRef.get
   if pExts.any (fun ext => ext.name == descr.name) then throw (IO.userError s!"invalid environment extension, '{descr.name}' has already been used")
-  let ext ← registerEnvExtension (asyncMode := descr.asyncMode) do
+  let replay? := descr.replay?.map fun replay =>
+    fun oldState newState newConsts s => { s with state := replay oldState.state newState.state newConsts s.state }
+  let ext ← registerEnvExtension (asyncMode := descr.asyncMode) (replay? := replay?) do
     let initial ← descr.mkInitial
     let s : PersistentEnvExtensionState α σ := {
       importedEntries := #[],
@@ -1206,6 +1343,9 @@ def registerSimplePersistentEnvExtension {α σ : Type} [Inhabited σ] (descr :
     exportEntriesFn := fun s => descr.toArrayFn s.1.reverse,
     statsFn := fun s => format "number of local entries: " ++ format s.1.length
     asyncMode := descr.asyncMode
+    replay? := some fun oldState newState _ (entries, s) =>
+      let newEntries := newState.1.drop oldState.1.length
+      (newEntries ++ entries, newEntries.foldl descr.addEntryFn s)
   }
 
 namespace SimplePersistentEnvExtension
@@ -1219,8 +1359,9 @@ def getEntries {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension
   (PersistentEnvExtension.getState ext env).1
 
 /-- Get the current state of the given `SimplePersistentEnvExtension`. -/
-def getState {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ) (env : Environment) : σ :=
-  (PersistentEnvExtension.getState ext env).2
+def getState {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ) (env : Environment)
+    (asyncMode := ext.toEnvExtension.asyncMode) : σ :=
+  (PersistentEnvExtension.getState (asyncMode := asyncMode) ext env).2
 
 /-- Set the current state of the given `SimplePersistentEnvExtension`. This change is *not* persisted across files. -/
 def setState {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env : Environment) (s : σ) : Environment :=
@@ -1230,6 +1371,11 @@ def setState {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env : En
 def modifyState {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env : Environment) (f : σ → σ) : Environment :=
   PersistentEnvExtension.modifyState ext env (fun ⟨entries, s⟩ => (entries, f s))
 
+@[inherit_doc PersistentEnvExtension.findStateAsync]
+def findStateAsync {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ)
+    (env : Environment) (declPrefix : Name) : σ :=
+  PersistentEnvExtension.findStateAsync ext env declPrefix |>.2
+
 end SimplePersistentEnvExtension
 
 /-- Environment extension for tagging declarations.
@@ -1329,8 +1475,12 @@ unsafe def Environment.freeRegions (env : Environment) : IO Unit :=
 
 def mkModuleData (env : Environment) : IO ModuleData := do
   let pExts ← persistentEnvExtensionsRef.get
-  let entries := pExts.map fun pExt =>
-    let state := pExt.getState env
+  let entries := pExts.map fun pExt => Id.run do
+    -- get state from `checked` at the end if `async`; it would otherwise panic
+    let mut asyncMode := pExt.toEnvExtension.asyncMode
+    if asyncMode matches .async then
+      asyncMode := .sync
+    let state := pExt.getState (asyncMode := asyncMode) env
     (pExt.name, pExt.exportEntriesFn state)
   let kenv := env.toKernelEnv
   let constNames := kenv.constants.foldStage2 (fun names name _ => names.push name) #[]
@@ -1403,7 +1553,9 @@ where
     let pExtDescrs ← persistentEnvExtensionsRef.get
     if h : i < pExtDescrs.size then
       let extDescr := pExtDescrs[i]
-      let s := extDescr.toEnvExtension.getState env
+      -- `local` as `async` does not allow for `getState` but it's all safe here as there is only
+      -- one branch so far.
+      let s := extDescr.toEnvExtension.getState (asyncMode := .local) env
       let prevSize := (← persistentEnvExtensionsRef.get).size
       let prevAttrSize ← getNumBuiltinAttributes
       let newState ← extDescr.addImportedFn s.importedEntries { env := env, opts := opts }
@@ -1522,6 +1674,7 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
         moduleData   := s.moduleData
       }
     }
+    realizedImportedConsts? := none
   }
   env ← setImportedEntries env s.moduleData
   if leakEnv then
@@ -1539,6 +1692,14 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
        Safety: There are no concurrent accesses to `env` at this point. -/
     env ← unsafe Runtime.markPersistent env
   env ← finalizePersistentExtensions env s.moduleData opts
+  env := { env with
+    realizedImportedConsts? := some {
+      -- safety: `RealizationContext` is private
+      env := unsafe unsafeCast env
+      opts
+      constsRef := (← IO.mkRef {})
+    }
+  }
   if leakEnv then
     /- Ensure the final environment including environment extension states is
        marked persistent as documented.
@@ -1583,6 +1744,9 @@ builtin_initialize namespacesExt : SimplePersistentEnvExtension Name NameSSet
       let map := mkStateFromImportedEntries (fun map name => map.insert name ()) map as
       SMap.fromHashMap map |>.switch
     addEntryFn      := fun s n => s.insert n
+    -- Namespaces from local helper constants can be disregarded in other environment branches. We
+    -- do *not* want `getNamespaceSet` to have to wait on all prior branches.
+    asyncMode       := .local
   }
 
 @[inherit_doc Kernel.Environment.enableDiag]
@@ -1616,8 +1780,18 @@ def getNamespaceSet (env : Environment) : NameSSet :=
   namespacesExt.getState env
 
 @[export lean_elab_environment_update_base_after_kernel_add]
-private def updateBaseAfterKernelAdd (env : Environment) (kernel : Kernel.Environment) : Environment :=
-  { env with checked := .pure kernel, checkedWithoutAsync := { kernel with extensions := env.checkedWithoutAsync.extensions } }
+private def updateBaseAfterKernelAdd (env : Environment) (kernel : Kernel.Environment) (decl : Declaration) : Environment :=
+  { env with
+    checked := .pure kernel
+    checkedWithoutAsync := { kernel with extensions := env.checkedWithoutAsync.extensions }
+    -- make constants available in `asyncConsts` as well; see its docstring
+    asyncConsts := decl.getNames.foldl (init := env.asyncConsts) fun asyncConsts n =>
+      if asyncConsts.find? n |>.isNone then
+        asyncConsts.add {
+          constInfo := .ofConstantInfo (kernel.find? n |>.get!)
+          exts? := none
+        }
+      else asyncConsts }
 
 @[export lean_display_stats]
 def displayStats (env : Environment) : IO Unit := do
@@ -1666,6 +1840,101 @@ def hasUnsafe (env : Environment) (e : Expr) : Bool :=
     | _ => false;
   c?.isSome
 
+/-- Plumbing function for `Lean.Meta.realizeConst`; see documentation there. -/
+def realizeConst (env : Environment) (forConst : Name) (constName : Name)
+    (realize : Environment → Options → BaseIO (Environment × Dynamic)) :
+    IO (Environment × Dynamic) := do
+  let mut env := env
+  -- find `RealizationContext` for `forConst` in `realizedImportedConsts?` or `realizedLocalConsts`
+  let ctx ← if env.checkedWithoutAsync.const2ModIdx.contains forConst then
+    env.realizedImportedConsts?.getDM <|
+      throw <| .userError s!"Environment.realizeConst: `realizedImportedConsts` is empty"
+  else
+    match env.realizedLocalConsts.find? forConst with
+    | some ctx => pure ctx
+    | none =>
+      throw <| .userError s!"trying to realize {constName} but `enableRealizationsForConst` must be called for '{forConst}' first"
+  let prom ← IO.Promise.new
+  -- ensure `prom` is not left unresolved from stray exceptions
+  BaseIO.toIO do
+    -- atomically check whether we are the first branch to realize `constName`
+    let existingConsts? ← ctx.constsRef.modifyGet fun m => match m.find? constName with
+      | some prom' => (some prom', m)
+      | none       => (none, m.insert constName prom.result!)
+    let (consts, replay, dyn) ← if let some existingConsts := existingConsts? then
+      pure existingConsts.get
+    else
+      -- safety: `RealizationContext` is private
+      let realizeEnv : Environment := unsafe unsafeCast ctx.env
+      let realizeEnv := { realizeEnv with
+        -- allow realizations to recursively realize other constants for `forConst`. Do note that
+        -- this allows for recursive realization of `constName` itself, which will deadlock.
+        realizedLocalConsts := realizeEnv.realizedLocalConsts.insert forConst ctx
+      }
+      -- ensure realized constants are nested below `forConst` and that environment extension
+      -- modifications know they are in an async context
+      let realizeEnv := realizeEnv.enterAsync (realizing := true) forConst
+      -- skip kernel in `realize`, we'll re-typecheck anyway
+      let realizeOpts := debug.skipKernelTC.set ctx.opts true
+      let (realizeEnv', dyn) ← realize realizeEnv realizeOpts
+      -- We could check that `c` was indeed added here but in practice `realize` has already
+      -- reported an error so we don't.
+
+      -- find new constants incl. nested realizations, add current extension state, and compute
+      -- closure
+      let consts := realizeEnv'.asyncConsts.revList.take (realizeEnv'.asyncConsts.size - realizeEnv.asyncConsts.size)
+      let consts := consts.map fun c =>
+        if c.exts?.isNone then
+          { c with exts? := some <| .pure realizeEnv'.checkedWithoutAsync.extensions }
+        else c
+      let exts ← EnvExtension.envExtensionsRef.get
+      let replay := (maybeAddToKernelEnv realizeEnv realizeEnv' consts · exts)
+      prom.resolve (consts, replay, dyn)
+      pure (consts, replay, dyn)
+    return ({ env with
+      asyncConsts := consts.foldl (·.add) env.asyncConsts
+      checked := env.checked.map replay
+    }, dyn)
+where
+  -- Adds `consts` if they haven't already been added by a previous branch. Note that this
+  -- conditional is deterministic because of the linearizing effect of `env.checked`.
+  maybeAddToKernelEnv (oldEnv newEnv : Environment) (consts : List AsyncConst)
+      (kenv : Kernel.Environment)
+      (exts : Array (EnvExtension EnvExtensionState)) : Kernel.Environment := Id.run do
+    let mut kenv := kenv
+    for c in consts do
+      if kenv.find? c.constInfo.name |>.isSome then
+        continue
+      let info := c.constInfo.toConstantInfo
+      if info.isUnsafe then
+        -- Checking unsafe declarations is not necessary for consistency, and it is necessary to
+        -- avoid checking them in the case of the old code generator, which adds ill-typed constants
+        -- to the kernel environment. We can delete this branch after removing the old code
+        -- generator.
+        kenv := kenv.add info
+        continue
+      let decl := match info with
+        | .thmInfo thm   => .thmDecl thm
+        | .defnInfo defn => .defnDecl defn
+        | _              => panic! s!"Environment.realizeConst: {c.constInfo.name} must be definition/theorem"
+      -- realized kernel additions cannot be interrupted - which would be bad anyway as they can be
+      -- reused between snapshots
+      match kenv.addDeclCore 0 decl none with
+      | .ok kenv' => kenv := kenv'
+      | .error e =>
+        let _ : Inhabited Kernel.Environment := ⟨kenv⟩
+        panic! s!"Environment.realizeConst: failed to add {c.constInfo.name} to environment\n{e.toRawString}"
+    for ext in exts do
+      if let some replay := ext.replay? then
+        kenv := { kenv with
+          -- safety: like in `modifyState`, but that one takes an elab env instead of a kernel env
+          extensions := unsafe (ext.modifyStateImpl kenv.extensions <|
+            replay
+              (ext.getStateImpl oldEnv.toKernelEnv.extensions)
+              (ext.getStateImpl newEnv.toKernelEnv.extensions)
+              (consts.map (·.constInfo.name))) }
+    return kenv
+
 end Environment
 
 namespace Kernel
@@ -1721,4 +1990,13 @@ def mkDefinitionValInferrringUnsafe [Monad m] [MonadEnv m] (name : Name) (levelP
   let safety := if env.hasUnsafe type || env.hasUnsafe value then DefinitionSafety.unsafe else DefinitionSafety.safe
   return { name, levelParams, type, value, hints, safety }
 
+def getMaxHeight (env : Environment) (e : Expr) : UInt32 :=
+  e.foldConsts 0 fun constName max =>
+    match env.find? constName with
+    | ConstantInfo.defnInfo val =>
+      match val.hints with
+      | ReducibilityHints.regular h => if h > max then h else max
+      | _                           => max
+    | _ => max
+
 end Lean

src/Lean/Language/Basic.lean

@@ -184,7 +184,7 @@ structure SnapshotTree where
   element : Snapshot
   /-- The asynchronously available children of the snapshot tree node. -/
   children : Array (SnapshotTask SnapshotTree)
-deriving Inhabited
+deriving Inhabited, TypeName
 
 /--
   Helper class for projecting a heterogeneous hierarchy of snapshot classes to a homogeneous

src/Lean/Meta/Basic.lean

@@ -2203,10 +2203,84 @@ def instantiateMVarsIfMVarApp (e : Expr) : MetaM Expr := do
   else
     return e
 
+private partial def setAllDiagRanges (snap : Language.SnapshotTree) (pos endPos : Position) :
+    BaseIO Language.SnapshotTree := do
+  let msgLog := snap.element.diagnostics.msgLog
+  let msgLog := { msgLog with unreported := msgLog.unreported.map fun diag =>
+      { diag with pos, endPos } }
+  return {
+    element.diagnostics := (← Language.Snapshot.Diagnostics.ofMessageLog msgLog)
+    children := (← snap.children.mapM fun task => return { task with
+      stx? := none
+      task := (← BaseIO.mapTask (t := task.task) (setAllDiagRanges · pos endPos)) })
+  }
+
+/--
+Makes the helper constant `constName` that is derived from `forConst` available in the environment.
+`enableRealizationsForConst forConst` must have been called first on this environment branch. If
+this is the first environment branch requesting `constName` to be realized (atomically), `realize`
+is called with the environment and options at the time of calling `enableRealizationsForConst` if
+`forConst` is from the current module and the state just after importing otherwise, thus helping
+achieve deterministic results despite the non-deterministic choice of which thread is tasked with
+realization. In other words, the state after calling `realizeConst` is *as if* `realize` had been
+called immediately after `enableRealizationsForConst forConst`, though the effects of this call are
+visible only after calling `realizeConst`. See below for more details on the replayed effects.
+
+`realizeConst` cannot check what other data is captured in the `realize` closure,
+so it is best practice to extract it into a separate function and pay close attention to the passed
+arguments, if any. `realize` must return with `constName` added to the environment,
+at which point all callers of `realizeConst` with this `constName` will be unblocked
+and have access to an updated version of their own environment containing any new constants
+`realize` added, including recursively realized constants. Traces, diagnostics, and raw std stream
+output are reported at all callers via `Core.logSnapshotTask` (so that the location of generated
+diagnostics is deterministic). Note that, as `realize` is run using the options at declaration time
+of `forConst`, trace options must be set prior to that (or, for imported constants, on the cmdline)
+in order to be active. The environment extension state at the end of `realize` is available to each
+caller via `EnvExtension.findStateAsync` for `constName`. If `realize` throws an exception or fails
+to add `constName` to the environment, an appropriate diagnostic is reported to all callers but no
+constants are added to the environment.
+-/
+def realizeConst (forConst : Name) (constName : Name) (realize : MetaM Unit) :
+    MetaM Unit := withTraceNode `Meta.realizeConst (fun _ => return constName) do
+  let env ← getEnv
+  let coreCtx ← readThe Core.Context
+  -- these fields should be invariant throughout the file
+  let coreCtx := { fileName := coreCtx.fileName, fileMap := coreCtx.fileMap }
+  let (env, dyn) ← env.realizeConst forConst constName (realizeAndReport coreCtx)
+  if let some snap := dyn.get? Language.SnapshotTree then
+    let mut snap := snap
+    -- localize diagnostics
+    if let some range := (← getRef).getRange? then
+      let fileMap ← getFileMap
+      snap ← setAllDiagRanges snap (fileMap.toPosition range.start) (fileMap.toPosition range.stop)
+    Core.logSnapshotTask <| .finished (stx? := none) snap
+  setEnv env
+where
+  -- similar to `wrapAsyncAsSnapshot` but not sufficiently so to share code
+  realizeAndReport (coreCtx : Core.Context) env opts := do
+    let coreCtx := { coreCtx with options := opts }
+    let act :=
+      IO.FS.withIsolatedStreams (isolateStderr := Core.stderrAsMessages.get opts) do
+        -- catch all exceptions
+        let _ : MonadExceptOf _ MetaM := MonadAlwaysExcept.except
+        try
+          realize
+          if !(← getEnv).contains constName then
+            throwError "Lean.Meta.realizeConst: {constName} was not added to the environment"
+        catch e : Exception =>
+          logError e.toMessageData
+        finally
+          addTraceAsMessages
+    let res? ← act |>.run' |>.run coreCtx { env } |>.toBaseIO
+    match res? with
+    | .ok ((output, ()), st) => pure (st.env, .mk (← Core.mkSnapshot output coreCtx st))
+    | .error _e => unreachable!; pure (env, .mk ({ diagnostics := .empty : Language.SnapshotLeaf}))
+
 end Meta
 
 builtin_initialize
   registerTraceClass `Meta.isLevelDefEq.postponed
+  registerTraceClass `Meta.realizeConst
 
 export Meta (MetaM)
 

src/Lean/Meta/Match/Match.lean

@@ -784,6 +784,7 @@ def mkMatcherAuxDefinition (name : Name) (type : Expr) (value : Expr) : MetaM (E
       modifyEnv fun env => matcherExt.modifyState env fun s => s.insert (result.value, compile) name
       addMatcherInfo name mi
       setInlineAttribute name
+      enableRealizationsForConst name
       if compile then
         compileDecl decl
     return (mkMatcherConst name, some addMatcher)

src/Lean/Util/FoldConsts.lean

@@ -5,7 +5,8 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Lean.Expr
-import Lean.Environment
+import Lean.Util.PtrSet
+import Lean.Declaration
 
 namespace Lean
 namespace Expr
@@ -70,14 +71,4 @@ def getUsedConstantsAsSet (c : ConstantInfo) : NameSet :=
     | _ => {}
 
 end ConstantInfo
-
-def getMaxHeight (env : Environment) (e : Expr) : UInt32 :=
-  e.foldConsts 0 fun constName max =>
-    match env.find? constName with
-    | ConstantInfo.defnInfo val =>
-      match val.hints with
-      | ReducibilityHints.regular h => if h > max then h else max
-      | _                           => max
-    | _ => max
-
 end Lean

src/library/elab_environment.cpp

@@ -11,17 +11,18 @@ Authors: Leonardo de Moura, Sebastian Ullrich
 #include "library/compiler/ir_interpreter.h"
 
 namespace lean {
-/* updateBaseAfterKernelAdd (env : Environment) (base : Kernel.Environment) : Environment
+/* updateBaseAfterKernelAdd (env : Environment) (base : Kernel.Environment) (decl : Declaration) : Environment
 
    Updates an elab environment with a given kernel environment.
 
    NOTE: Ideally this language switching would not be necessary and we could do all this in Lean
-   only but the old code generator and `mk_projections` still need a C++ `elab_environment::add`. */
-extern "C" obj_res lean_elab_environment_update_base_after_kernel_add(obj_arg env, obj_arg kenv);
+   only but the old code generator and `mk_projections` still need a C++ `elab_environment::add`
+   that throws C++ exceptions. */
+extern "C" obj_res lean_elab_environment_update_base_after_kernel_add(obj_arg env, obj_arg kenv, obj_arg decl);
 
 elab_environment elab_environment::add(declaration const & d, bool check) const {
     environment kenv = to_kernel_env().add(d, check);
-    return elab_environment(lean_elab_environment_update_base_after_kernel_add(this->to_obj_arg(), kenv.to_obj_arg()));
+    return elab_environment(lean_elab_environment_update_base_after_kernel_add(this->to_obj_arg(), kenv.to_obj_arg(), d.to_obj_arg()));
 }
 
 extern "C" LEAN_EXPORT object * lean_elab_add_decl(object * env, size_t max_heartbeat, object * decl,

src/runtime/object.cpp

@@ -102,8 +102,8 @@ extern "C" LEAN_EXPORT void lean_set_panic_messages(bool flag) {
     g_panic_messages = flag;
 }
 
-static void panic_eprintln(char const * line) {
-    if (g_exit_on_panic || should_abort_on_panic()) {
+static void panic_eprintln(char const * line, bool force_stderr) {
+    if (force_stderr || g_exit_on_panic || should_abort_on_panic()) {
         // If we are about to kill the process, we should skip the Lean stderr buffer
         std::cerr << line << "\n";
     } else {
@@ -111,37 +111,33 @@ static void panic_eprintln(char const * line) {
     }
 }
 
-static void print_backtrace() {
+static void print_backtrace(bool force_stderr) {
 #if LEAN_SUPPORTS_BACKTRACE
     void * bt_buf[100];
     int nptrs = backtrace(bt_buf, sizeof(bt_buf) / sizeof(void *));
     if (char ** symbols = backtrace_symbols(bt_buf, nptrs)) {
         for (int i = 0; i < nptrs; i++) {
-            panic_eprintln(symbols[i]);
+            panic_eprintln(symbols[i], force_stderr);
         }
         // According to `man backtrace`, each `symbols[i]` should NOT be freed
         free(symbols);
         if (nptrs == sizeof(bt_buf)) {
-            panic_eprintln("...");
+            panic_eprintln("...", force_stderr);
         }
     }
 #else
-    panic_eprintln("(stack trace unavailable)");
+    panic_eprintln("(stack trace unavailable)", force_stderr);
 #endif
 }
 
 extern "C" LEAN_EXPORT void lean_panic(char const * msg, bool force_stderr = false) {
     if (g_panic_messages) {
-        if (force_stderr) {
-            std::cerr << msg << "\n";
-        } else {
-            panic_eprintln(msg);
-        }
+        panic_eprintln(msg, force_stderr);
 #if LEAN_SUPPORTS_BACKTRACE
         char * bt_env = getenv("LEAN_BACKTRACE");
         if (!bt_env || strcmp(bt_env, "0") != 0) {
-            panic_eprintln("backtrace:");
-            print_backtrace();
+            panic_eprintln("backtrace:", force_stderr);
+            print_backtrace(force_stderr);
         }
 #endif
     }
@@ -2558,7 +2554,7 @@ extern "C" LEAN_EXPORT object * lean_dbg_trace_if_shared(obj_arg s, obj_arg a) {
 }
 
 extern "C" LEAN_EXPORT object * lean_dbg_stack_trace(obj_arg fn) {
-    print_backtrace();
+    print_backtrace(/* force_stderr */ false);
     return lean_apply_1(fn, lean_box(0));
 }