From 895cdce9bca3b20dfd5f8e9535e996d3342a5480 Mon Sep 17 00:00:00 2001 From: Sebastian Ullrich Date: Mon, 10 Feb 2025 16:08:02 +0100 Subject: [PATCH] fix: codegen was allowed improper env ext accesses (#7023) --- src/Lean/Compiler/ClosedTermCache.lean | 3 +- src/Lean/Compiler/IR/ElimDeadBranches.lean | 1 + src/Lean/Compiler/LCNF/BaseTypes.lean | 2 +- src/Lean/Compiler/LCNF/ElimDeadBranches.lean | 1 + src/Lean/Compiler/LCNF/MonoTypes.lean | 2 +- src/Lean/CoreM.lean | 7 +-- src/Lean/Environment.lean | 45 ++++++++++++++++++-- 7 files changed, 51 insertions(+), 10 deletions(-) diff --git a/src/Lean/Compiler/ClosedTermCache.lean b/src/Lean/Compiler/ClosedTermCache.lean index 3456bf46d1..a620e09663 100644 --- a/src/Lean/Compiler/ClosedTermCache.lean +++ b/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 := diff --git a/src/Lean/Compiler/IR/ElimDeadBranches.lean b/src/Lean/Compiler/IR/ElimDeadBranches.lean index ce7c9726ac..bc77895958 100644 --- a/src/Lean/Compiler/IR/ElimDeadBranches.lean +++ b/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 := diff --git a/src/Lean/Compiler/LCNF/BaseTypes.lean b/src/Lean/Compiler/LCNF/BaseTypes.lean index afe90bf39c..70a6a1266a 100644 --- a/src/Lean/Compiler/LCNF/BaseTypes.lean +++ b/src/Lean/Compiler/LCNF/BaseTypes.lean @@ -20,7 +20,7 @@ structure BaseTypeExtState where deriving Inhabited builtin_initialize baseTypeExt : EnvExtension BaseTypeExtState ← - registerEnvExtension (pure {}) + registerEnvExtension (pure {}) (asyncMode := .sync) -- compilation is non-parallel anyway def getOtherDeclBaseType (declName : Name) (us : List Level) : CoreM Expr := do let info ← getConstInfo declName diff --git a/src/Lean/Compiler/LCNF/ElimDeadBranches.lean b/src/Lean/Compiler/LCNF/ElimDeadBranches.lean index d06ec02ba1..3147620d59 100644 --- a/src/Lean/Compiler/LCNF/ElimDeadBranches.lean +++ b/src/Lean/Compiler/LCNF/ElimDeadBranches.lean @@ -248,6 +248,7 @@ builtin_initialize functionSummariesExt : SimplePersistentEnvExtension (Name × addImportedFn := fun _ => {} addEntryFn := fun s ⟨e, n⟩ => s.insert e n toArrayFn := fun s => s.toArray.qsort decLt + asyncMode := .sync -- compilation is non-parallel anyway } /-- diff --git a/src/Lean/Compiler/LCNF/MonoTypes.lean b/src/Lean/Compiler/LCNF/MonoTypes.lean index 88b65701d8..cd38edcf88 100644 --- a/src/Lean/Compiler/LCNF/MonoTypes.lean +++ b/src/Lean/Compiler/LCNF/MonoTypes.lean @@ -117,7 +117,7 @@ structure MonoTypeExtState where deriving Inhabited builtin_initialize monoTypeExt : EnvExtension MonoTypeExtState ← - registerEnvExtension (pure {}) + registerEnvExtension (pure {}) (asyncMode := .sync) -- compilation is non-parallel anyway def getOtherDeclMonoType (declName : Name) : CoreM Expr := do match monoTypeExt.getState (← getEnv) |>.mono.find? declName with diff --git a/src/Lean/CoreM.lean b/src/Lean/CoreM.lean index cd4817682d..9d4e524b5c 100644 --- a/src/Lean/CoreM.lean +++ b/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 diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean index bbe640f45f..01afdc2ef4 100644 --- a/src/Lean/Environment.lean +++ b/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. -/