feat: importModules without loading environment extensions

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

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,31 +892,47 @@ 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. -/
-unsafe def withImportModules {α : Type} (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0) (act : Environment → IO α) : IO α := do
-  let env ← importModules imports opts trustLevel
+Like `importModules` but frees compacted regions after calling `act`. No live references to the
+environment object or imported objects may exist after `act` finishes.
+-/
+unsafe def withImportModules {α : Type} (imports : Array Import) (opts : Options)
+    (act : Environment → IO α) (trustLevel : UInt32 := 0) (loadExts := false) : IO α := do
+  let env ← importModules (loadExts := loadExts) imports opts trustLevel
   try act env finally env.freeRegions
 
 /--

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
 

tests/lean/run/meta1.lean

@@ -4,17 +4,17 @@ open Lean
 open Lean.Meta
 
 unsafe def tstInferType (mods : Array Name) (e : Expr) : IO Unit :=
-withImportModules (mods.map $ fun m => {module := m}) {} 0 fun env => do
+withImportModules (loadExts := true) (mods.map $ fun m => {module := m}) {} fun env => do
    let (type, _, _) ← (inferType e : MetaM _).toIO { fileName := "", fileMap := default } { env := env } {} {};
    IO.println (toString e ++ " : " ++ toString type)
 
 unsafe def tstWHNF (mods : Array Name) (e : Expr) (t := TransparencyMode.default) : IO Unit :=
-withImportModules (mods.map $ fun m => {module := m}) {} 0 fun env => do
+withImportModules (loadExts := true) (mods.map $ fun m => {module := m}) {} fun env => do
    let (s, _, _) ← (whnf e : MetaM _).toIO { fileName := "", fileMap := default } { env := env };
    IO.println (toString e ++ " ==> " ++ toString s)
 
 unsafe def tstIsProp (mods : Array Name) (e : Expr) : IO Unit :=
-withImportModules (mods.map $ fun m => {module := m}) {} 0 fun env => do
+withImportModules (loadExts := true) (mods.map $ fun m => {module := m}) {} fun env => do
   let (b, _, _) ← (isProp e : MetaM _).toIO { fileName := "", fileMap := default } { env := env };
   IO.println (toString e ++ ", isProp: " ++ toString b)