fix: codegen was allowed improper env ext accesses

src/Lean/Compiler/ClosedTermCache.lean

@@ -13,7 +13,8 @@ structure ClosedTermCache where
   constNames : NameSet := {}
   deriving Inhabited
 
-builtin_initialize closedTermCacheExt : EnvExtension ClosedTermCache ← registerEnvExtension (pure {})
+builtin_initialize closedTermCacheExt : EnvExtension ClosedTermCache ←
+  registerEnvExtension (pure {}) (asyncMode := .sync)  -- compilation is non-parallel anyway
 
 @[export lean_cache_closed_term_name]
 def cacheClosedTermName (env : Environment) (e : Expr) (n : Name) : Environment :=

src/Lean/Compiler/IR/ElimDeadBranches.lean

@@ -142,6 +142,7 @@ builtin_initialize functionSummariesExt : SimplePersistentEnvExtension (FunId ×
     addImportedFn := fun _ => {}
     addEntryFn := fun s ⟨e, n⟩ => s.insert e n
     toArrayFn := fun s => sortEntries s.toArray
+    asyncMode := .sync  -- compilation is non-parallel anyway
   }
 
 def addFunctionSummary (env : Environment) (fid : FunId) (v : Value) : Environment :=

src/Lean/CoreM.lean

@@ -527,16 +527,17 @@ partial def compileDecls (decls : List Name) (ref? : Option Declaration := none)
     doCompile
     return
   let env ← getEnv
-  let (postEnv, prom) ← env.promiseChecked
+  let res ← env.promiseChecked
+  setEnv res.mainEnv
   let checkAct ← Core.wrapAsyncAsSnapshot fun _ => do
+    setEnv res.asyncEnv
     try
       doCompile
     finally
-      prom.resolve (← getEnv)
+      res.commitChecked (← getEnv)
   let t ← BaseIO.mapTask (fun _ => checkAct) env.checked
   let endRange? := (← getRef).getTailPos?.map fun pos => ⟨pos, pos⟩
   Core.logSnapshotTask { range? := endRange?, task := t }
-  setEnv postEnv
 where doCompile := do
   -- don't compile if kernel errored; should be converted into a task dependency when compilation
   -- is made async as well

src/Lean/Environment.lean

@@ -462,10 +462,6 @@ 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 }
 
-def promiseChecked (env : Environment) : BaseIO (Environment × IO.Promise Environment) := do
-  let prom ← IO.Promise.new
-  return ({ env with checked := prom.result?.bind (sync := true) (·.getD env |>.checked) }, prom)
-
 /--
 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
@@ -595,6 +591,47 @@ def dbgFormatAsyncState (env : Environment) : BaseIO String :=
 def dbgFormatCheckedSyncState (env : Environment) : BaseIO String :=
   return s!"checked.get.constants.map₂: {repr <| env.checked.get.constants.map₂.toList.map (·.1)}"
 
+/-- Result of `Lean.Environment.promiseChecked`. -/
+structure PromiseCheckedResult where
+  /--
+  Resulting "main branch" environment. Accessing the kernel environment will block until
+  `PromiseCheckedResult.commitChecked` has been called.
+  -/
+  mainEnv : Environment
+  /--
+  Resulting "async branch" environment which should be used in a new task and then to call
+  `PromiseCheckedResult.commitChecked` to commit results back to the main environment. If it is not
+  called and the `PromiseCheckedResult` object is dropped, the kernel environment will be left
+  unchanged.
+  -/
+  asyncEnv : Environment
+  private checkedEnvPromise : IO.Promise Kernel.Environment
+
+/--
+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
+`PromiseCheckedResult.commitChecked` has been called on the "async" environment branch.
+-/
+def promiseChecked (env : Environment) : BaseIO PromiseCheckedResult := do
+  let checkedEnvPromise ← IO.Promise.new
+  return {
+    mainEnv := { env with
+      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 }
+    }
+    checkedEnvPromise
+  }
+
+/-- Commits the kernel environment of the given environment back to the main branch. -/
+def PromiseCheckedResult.commitChecked (res : PromiseCheckedResult) (env : Environment) :
+    BaseIO Unit :=
+  assert! env.asyncCtx?.isSome
+  res.checkedEnvPromise.resolve env.toKernelEnv
+
 /--
 Result of `Lean.Environment.addConstAsync` which is necessary to complete the asynchronous addition.
 -/