From 570f66908b84b0a66c6209c225ce1e1b8ac3c5b1 Mon Sep 17 00:00:00 2001 From: Sebastian Ullrich Date: Fri, 6 Dec 2024 10:20:59 +0100 Subject: [PATCH] feat: `importModules` without loading environment extensions --- src/Lean/Elab/Import.lean | 3 +- src/Lean/Environment.lean | 56 ++++++++++++++++++------------- src/lake/Lake/Load/Lean/Elab.lean | 2 +- 3 files changed, 35 insertions(+), 26 deletions(-) diff --git a/src/Lean/Elab/Import.lean b/src/Lean/Elab/Import.lean index 88a013c49040..7b78108d8ec0 100644 --- a/src/Lean/Elab/Import.lean +++ b/src/Lean/Elab/Import.lean @@ -21,7 +21,8 @@ def processHeader (header : Syntax) (opts : Options) (messages : MessageLog) (inputCtx : Parser.InputContext) (trustLevel : UInt32 := 0) (leakEnv := false) : IO (Environment × MessageLog) := do try - let env ← importModules (leakEnv := leakEnv) (headerToImports header) opts trustLevel + let env ← + importModules (leakEnv := leakEnv) (loadExts := true) (headerToImports header) opts trustLevel pure (env, messages) catch e => let env ← mkEmptyEnvironment diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean index 7c28cedbf7c1..fd35a4e444ed 100644 --- a/src/Lean/Environment.lean +++ b/src/Lean/Environment.lean @@ -842,17 +842,8 @@ private def equivInfo (cinfo₁ cinfo₂ : ConstantInfo) : Bool := Id.run do && tval₁.levelParams == tval₂.levelParams && tval₁.all == tval₂.all -/-- - Construct environment from `importModulesCore` results. - - If `leakEnv` is true, we mark the environment as persistent, which means it - will not be freed. We set this when the object would survive until the end of - the process anyway. In exchange, RC updates are avoided, which is especially - important when they would be atomic because the environment is shared across - threads (potentially, storing it in an `IO.Ref` is sufficient for marking it - as such). -/ -def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0) - (leakEnv := false) : IO Environment := do +private def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0) + (leakEnv loadExts : Bool): IO Environment := do let numConsts := s.moduleData.foldl (init := 0) fun numConsts mod => numConsts + mod.constants.size + mod.extraConstNames.size let mut const2ModIdx : Std.HashMap Name ModuleIdx := Std.HashMap.empty (capacity := numConsts) @@ -901,29 +892,46 @@ 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 - if leakEnv then - /- Ensure the final environment including environment extension states is - marked persistent as documented. - - 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 + if loadExts then + env ← finalizePersistentExtensions env s.moduleData opts + if leakEnv then + /- Ensure the final environment including environment extension states is + marked persistent as documented. + + 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 -@[export lean_import_modules] +/-- +Creates environment object from given imports. + +If `leakEnv` is true, we mark the environment as persistent, which means it will not be freed. We +set this when the object would survive until the end of the process anyway. In exchange, RC updates +are avoided, which is especially important when they would be atomic because the environment is +shared across threads (potentially, storing it in an `IO.Ref` is sufficient for marking it as such). + +If `loadExts` is true, we initialize the environment extensions using the imported data. Doing so +can run arbitrary code and thus is only safe to do after calling `enableInitializersExecution`; see +also caveats there. If not set, every extension will have its initial value as its state. While the +environment's constant map can be accessed without `loadExts`, many functions that take +`Environment` or are in a monad carrying it such as `CoreM` may not function properly without it. +-/ def importModules (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0) - (leakEnv := false) : IO Environment := profileitIO "import" opts do + (leakEnv := false) (loadExts := false) : IO Environment := profileitIO "import" opts do for imp in imports do if imp.module matches .anonymous then throw <| IO.userError "import failed, trying to import module with anonymous name" withImporting do let (_, s) ← importModulesCore imports |>.run - finalizeImport (leakEnv := leakEnv) s imports opts trustLevel + finalizeImport (leakEnv := leakEnv) (loadExts := loadExts) s imports opts trustLevel /-- - Create environment object from imports and free compacted regions after calling `act`. No live references to the - environment object or imported objects may exist after `act` finishes. -/ +Creates environment object from imports and free compacted regions after calling `act`. No live +references to the environment object or imported objects may exist after `act` finishes. As this +cannot be ruled out after loading environment extensions, `importModules`'s `loadExts` is always +unset using this function. +-/ unsafe def withImportModules {α : Type} (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0) (act : Environment → IO α) : IO α := do let env ← importModules imports opts trustLevel try act env finally env.freeRegions diff --git a/src/lake/Lake/Load/Lean/Elab.lean b/src/lake/Lake/Load/Lean/Elab.lean index c83fe093df16..21d2161fbac8 100644 --- a/src/lake/Lake/Load/Lean/Elab.lean +++ b/src/lake/Lake/Load/Lean/Elab.lean @@ -30,7 +30,7 @@ initialize importEnvCache : IO.Ref (Std.HashMap (Array Import) Environment) ← def importModulesUsingCache (imports : Array Import) (opts : Options) (trustLevel : UInt32) : IO Environment := do if let some env := (← importEnvCache.get)[imports]? then return env - let env ← importModules imports opts trustLevel + let env ← importModules (loadExts := true) imports opts trustLevel importEnvCache.modify (·.insert imports env) return env