chore: unconditionally re-enable realizeConst (#7334)

To be merged when Mathlib adaption passes

src/Lean/Environment.lean

@@ -692,21 +692,6 @@ def find? (env : Environment) (n : Name) (skipRealize := false) : Option Constan
     return c
   env.findAsyncCore? n (skipRealize := skipRealize) |>.map (·.toConstantInfo)
 
-/--
-Allows `realizeConst` calls for imported declarations in all derived environment branches.
-Realizations will run using the given environment and options to ensure deterministic results.
-This function should be called directly after `setMainModule` to ensure that all realized constants
-use consistent private prefixes.
--/
-def enableRealizationsForImports (env : Environment) (opts : Options) : BaseIO Environment :=
-  return { env with realizedImportedConsts? := some {
-      -- safety: `RealizationContext` is private
-      env := unsafe unsafeCast env
-      opts
-      constsRef := (← IO.mkRef {})
-    }
-  }
-
 /--
 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
@@ -977,10 +962,13 @@ def allImportedModuleNames (env : Environment) : Array Name :=
   env.header.moduleNames
 
 def setMainModule (env : Environment) (m : Name) : Environment := Id.run do
-  if env.realizedImportedConsts?.isSome then
-    let _ : Inhabited Environment := ⟨env⟩
-    return panic! "cannot set after `enableRealizationsForImports`"
-  env.modifyCheckedAsync ({ · with header.mainModule := m })
+  let env := env.modifyCheckedAsync ({ · with
+    header.mainModule := m
+  })
+  { env with realizedImportedConsts? := env.realizedImportedConsts?.map ({ · with
+      -- safety: `RealizationContext` is private
+      env := unsafe unsafeCast env
+    }) }
 
 def mainModule (env : Environment) : Name :=
   env.header.mainModule
@@ -1473,7 +1461,7 @@ def registerSimplePersistentEnvExtension {α σ : Type} [Inhabited σ] (descr :
     replay? := descr.replay?.map fun replay oldState newState _ (entries, s) =>
       let newEntries := newState.1.take (newState.1.length - oldState.1.length)
       let (newEntries, s) := replay newEntries newState.2 s
-      (entries ++ newEntries, s)
+      (newEntries ++ entries, s)
   }
 
 namespace SimplePersistentEnvExtension
@@ -1834,7 +1822,12 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
        Safety: There are no concurrent accesses to `env` at this point, assuming
        extensions' `addImportFn`s did not spawn any unbound tasks. -/
     env ← unsafe Runtime.markPersistent env
-  pure env
+  return { env with realizedImportedConsts? := some {
+    -- safety: `RealizationContext` is private
+    env := unsafe unsafeCast env
+    opts
+    constsRef := (← IO.mkRef {})
+  } }
 
 @[export lean_import_modules]
 def importModules (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0)
@@ -1917,6 +1910,64 @@ unsafe opaque evalConst (α) (env : @& Environment) (opts : @& Options) (constNa
 private def throwUnexpectedType {α} (typeName : Name) (constName : Name) : ExceptT String Id α :=
   throw ("unexpected type at '" ++ toString constName ++ "', `" ++ toString typeName ++ "` expected")
 
+/--
+Replays the difference between `newEnv` and `oldEnv` onto `dest`: the set of constants in `newEnv`
+but not `oldEnv` and the environment extension state for extensions defining `replay?`. If
+`skipExisting` is true, constants that are already in `dest` are not added. If `newEnv` and `dest`
+are not derived from `oldEnv`, the result is undefined.
+-/
+def replayConsts (oldEnv newEnv : Environment) (dest : Environment) (skipExisting := false) :
+    BaseIO Environment := do
+  let numNewConsts := newEnv.asyncConsts.size - oldEnv.asyncConsts.size
+  let consts := newEnv.asyncConsts.revList.take numNewConsts |>.reverse
+  let exts ← EnvExtension.envExtensionsRef.get
+  return { dest with
+    asyncConsts := consts.foldl (init := dest.asyncConsts) fun consts c =>
+      if skipExisting && (consts.find? c.constInfo.name).isSome then
+        consts
+      else
+        consts.add c
+    checked := dest.checked.map (replayKernel exts consts)
+  }
+where
+  replayKernel (exts : Array (EnvExtension EnvExtensionState)) (consts : List AsyncConst)
+      (kenv : Kernel.Environment) : Kernel.Environment := Id.run do
+    let mut kenv := kenv
+    for c in consts do
+      if skipExisting && (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
+      -- for panics
+      let _ : Inhabited Kernel.Environment := ⟨kenv⟩
+      let decl ← match info with
+        | .thmInfo thm   => .thmDecl thm
+        | .defnInfo defn => .defnDecl defn
+        | _              =>
+          return panic! s!"{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 =>
+        return panic! s!"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
+
 /-- Like `evalConst`, but first check that `constName` indeed is a declaration of type `typeName`.
     Note that this function cannot guarantee that `typeName` is in fact the name of the type `α`. -/
 unsafe def evalConstCheck (α) (env : Environment) (opts : Options) (typeName : Name) (constName : Name) : ExceptT String Id α :=
@@ -1988,7 +2039,7 @@ def realizeConst (env : Environment) (forConst : Name) (constName : Name)
           { c with exts? := some <| .pure realizeEnv'.base.extensions }
         else c
       let exts ← EnvExtension.envExtensionsRef.get
-      let replay := (maybeAddToKernelEnv realizeEnv realizeEnv' consts · exts)
+      let replay := replayConsts.replayKernel (skipExisting := true) realizeEnv realizeEnv' exts consts
       prom.resolve (consts, replay, dyn)
       pure (consts, replay, dyn)
     return ({ env with
@@ -2002,47 +2053,6 @@ def realizeConst (env : Environment) (forConst : Name) (constName : Name)
         consts.foldl (init := allRealizations) fun allRealizations c =>
           allRealizations.insert c.constInfo.name c
     }, 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
-      -- for panics
-      let _ : Inhabited Kernel.Environment := ⟨kenv⟩
-      let decl ← match info with
-        | .thmInfo thm   => .thmDecl thm
-        | .defnInfo defn => .defnDecl defn
-        | _              =>
-          return panic! s!"{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 =>
-        return panic! s!"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
 

src/Lean/Language/Lean.lean

@@ -467,7 +467,6 @@ where
         return { diagnostics, result? := none }
 
       let headerEnv := headerEnv.setMainModule setup.mainModuleName
-      let headerEnv ← headerEnv.enableRealizationsForImports setup.opts
       let mut traceState := default
       if trace.profiler.output.get? setup.opts |>.isSome then
         traceState := {

src/Lean/Meta/Basic.lean

@@ -2242,12 +2242,11 @@ Makes the helper constant `constName` that is derived from `forConst` available
 `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 (when
-`enableRealizationsForImports` should be called) 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.
+`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
@@ -2271,10 +2270,6 @@ def realizeConst (forConst : Name) (constName : Name) (realize : MetaM Unit) :
   -- the relevant local environment extension state when accessed on this branch.
   if env.containsOnBranch constName then
     return
-  -- TODO: remove when Mathlib passes without it
-  if !Elab.async.get (← getOptions) then
-    realize
-    return
   withTraceNode `Meta.realizeConst (fun _ => return constName) do
     let coreCtx ← readThe Core.Context
     let coreCtx := {

src/lake/Lake/Load/Lean/Elab.lean

@@ -58,7 +58,6 @@ def elabConfigFile (pkgDir : FilePath) (lakeOpts : NameMap String)
   let (header, parserState, messages) ← Parser.parseHeader inputCtx
   let (env, messages) ← processHeader header leanOpts inputCtx messages
   let env := env.setMainModule configModuleName
-  let env ← env.enableRealizationsForImports leanOpts
 
   -- Configure extensions
   let env := dirExt.setState env pkgDir

tests/lean/run/issue7322.lean

@@ -51,12 +51,10 @@ def bar (distance : Nat) (idx : Nat) (a : Fin distance) (fuel : Nat) :
 termination_by fuel
 decreasing_by sorry
 
-set_option Elab.async false -- for stable message ordering in guard_msgs
-
 /--
-warning: declaration uses 'sorry'
----
 info: bar.induct (motive : Nat → Prop) (case1 : ∀ (x : Nat), (∀ (n : Nat), motive n) → motive x) (fuel : Nat) : motive fuel
+---
+warning: declaration uses 'sorry'
 -/
 #guard_msgs in
 #check bar.induct