From 0c4532f327edf0e49059c7a59ba49b3829ec8879 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
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         | 63 ++++++++++++++----------
 src/lake/Lake/Load/Lean/Elab.lean |  2 +-
 tests/lean/ctor_layout.lean       |  2 +-
 tests/lean/run/instances.lean     | 13 ++---
 tests/lean/run/instuniv.lean      |  2 +-
 tests/lean/run/meta1.lean         | 80 +++++++++++++++----------------
 tests/lean/run/meta3.lean         |  6 ---
 tests/pkg/user_attr_app/Main.lean |  2 +-
 9 files changed, 90 insertions(+), 83 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..0228b4d9d1c7 100644
--- a/src/Lean/Environment.lean
+++ b/src/Lean/Environment.lean
@@ -843,16 +843,13 @@ private def equivInfo (cinfo₁ cinfo₂ : ConstantInfo) : Bool := Id.run do
     && 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). -/
+Constructs environment from `importModulesCore` results.
+
+See also `importModules` for parameter documentation.
+-/
+
 def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0)
-    (leakEnv := false) : IO Environment := do
+    (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 +898,49 @@ 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
+may use the interpreter 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
+Creates environment object from imports and frees 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)
+    (act : Environment → IO α) (trustLevel : UInt32 := 0) : IO α := do
+  let env ← importModules (loadExts := false) 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
 
diff --git a/tests/lean/ctor_layout.lean b/tests/lean/ctor_layout.lean
index e9b7c43daa69..826a6e25c5d4 100644
--- a/tests/lean/ctor_layout.lean
+++ b/tests/lean/ctor_layout.lean
@@ -4,7 +4,7 @@ open Lean
 open Lean.IR
 
 unsafe def main : IO Unit :=
-withImportModules #[{module := `Lean.Compiler.IR.Basic}] {} 0 fun env => do
+withImportModules #[{module := `Lean.Compiler.IR.Basic}] {} fun env => do
    let ctorLayout ← IO.ofExcept $ getCtorLayout env `Lean.IR.Expr.reuse;
    ctorLayout.fieldInfo.forM $ fun finfo => IO.println (format finfo);
    IO.println "---";
diff --git a/tests/lean/run/instances.lean b/tests/lean/run/instances.lean
index b1ca83050a3b..b9b07e96cd1e 100644
--- a/tests/lean/run/instances.lean
+++ b/tests/lean/run/instances.lean
@@ -6,11 +6,12 @@ open Lean.Meta
 instance : ToFormat InstanceEntry where
   format e := format e.val
 
-unsafe def tst1 : IO Unit :=
-withImportModules #[{module := `Lean}] {} 0 fun env => do
-   let aux : MetaM Unit := do
-     let insts ← getGlobalInstancesIndex
-     IO.println (format insts)
-   discard <| aux.run |>.toIO { fileName := "", fileMap := default } { env := env }
+unsafe def tst1 : IO Unit := do
+  let env ← importModules (loadExts := true) #[{module := `Lean}] {}
+  let aux : MetaM Unit := do
+    let insts ← getGlobalInstancesIndex
+    assert! insts.size > 0
+    IO.println (format insts)
+  discard <| aux.run |>.toIO { fileName := "", fileMap := default } { env := env }
 
 #eval tst1
diff --git a/tests/lean/run/instuniv.lean b/tests/lean/run/instuniv.lean
index ce16db570562..8187d41b424b 100644
--- a/tests/lean/run/instuniv.lean
+++ b/tests/lean/run/instuniv.lean
@@ -3,7 +3,7 @@ import Lean
 open Lean
 
 unsafe def tst : IO Unit :=
-  withImportModules #[{module := `Init.Data.Array}] {} 0 fun env =>
+  withImportModules #[{module := `Init.Data.Array}] {} fun env =>
     match env.find? `Array.foldl with
     | some info => do
       IO.println (info.instantiateTypeLevelParams [levelZero, levelZero])
diff --git a/tests/lean/run/meta1.lean b/tests/lean/run/meta1.lean
index dcba3d0f21cf..df4dde5d2263 100644
--- a/tests/lean/run/meta1.lean
+++ b/tests/lean/run/meta1.lean
@@ -3,18 +3,18 @@ import Lean.Meta
 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
-   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
-   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
+def tstInferType (e : Expr) : CoreM Unit := do
+  let env ← getEnv
+  let (type, _, _) ← (inferType e : MetaM _).toIO { fileName := "", fileMap := default } { env := env } {} {};
+  IO.println (toString e ++ " : " ++ toString type)
+
+def tstWHNF (e : Expr) : CoreM Unit := do
+  let env ← getEnv
+  let (s, _, _) ← (whnf e : MetaM _).toIO { fileName := "", fileMap := default } { env := env };
+  IO.println (toString e ++ " ==> " ++ toString s)
+
+unsafe def tstIsProp (e : Expr) : CoreM Unit := do
+  let env ← getEnv
   let (b, _, _) ← (isProp e : MetaM _).toIO { fileName := "", fileMap := default } { env := env };
   IO.println (toString e ++ ", isProp: " ++ toString b)
 
@@ -26,7 +26,7 @@ mkAppN map #[nat, bool]
 
 /-- info: List.map.{1, 1} Nat Bool : (Nat -> Bool) -> (List.{1} Nat) -> (List.{1} Bool) -/
 #guard_msgs in
-#eval tstInferType #[`Init.Data.List] t1
+#eval tstInferType t1
 
 def t2 : Expr :=
 let prop := mkSort levelZero;
@@ -34,7 +34,7 @@ mkForall `x BinderInfo.default prop prop
 
 /-- info: Prop -> Prop : Type -/
 #guard_msgs in
-#eval tstInferType #[`Init.Core] t2
+#eval tstInferType t2
 
 def t3 : Expr :=
 let nat   := mkConst `Nat [];
@@ -45,7 +45,7 @@ mkForall `x BinderInfo.default nat p
 
 /-- info: forall (x : Nat), Nat.le x 0 : Prop -/
 #guard_msgs in
-#eval tstInferType #[`Init.Data.Nat] t3
+#eval tstInferType t3
 
 def t4 : Expr :=
 let nat   := mkConst `Nat [];
@@ -54,7 +54,7 @@ mkLambda `x BinderInfo.default nat p
 
 /-- info: fun (x : Nat) => Nat.succ x : Nat -> Nat -/
 #guard_msgs in
-#eval tstInferType #[`Init.Core] t4
+#eval tstInferType t4
 
 def t5 : Expr :=
 let add   := mkConst `Nat.add [];
@@ -62,11 +62,7 @@ mkAppN add #[mkLit (Literal.natVal 3), mkLit (Literal.natVal 5)]
 
 /-- info: Nat.add 3 5 ==> 8 -/
 #guard_msgs in
-#eval tstWHNF #[`Init.Data.Nat] t5
-
-/-- info: Nat.add 3 5 ==> 8 -/
-#guard_msgs in
-#eval tstWHNF #[`Init.Data.Nat] t5 TransparencyMode.reducible
+#eval tstWHNF t5
 
 set_option pp.all true
 /-- info: @List.cons.{0} Nat : Nat → List.{0} Nat → List.{0} Nat -/
@@ -89,23 +85,23 @@ mkAppN map #[nat, nat, f, xs]
 info: List.map.{1, 1} Nat Nat (fun (x : Nat) => Nat.add x 1) (List.cons.{0} Nat 1 (List.cons.{0} Nat 4 (List.nil.{0} Nat))) : List.{1} Nat
 -/
 #guard_msgs in
-#eval tstInferType #[`Init.Data.List] t6
+#eval tstInferType t6
 
 /--
 info: List.map.{1, 1} Nat Nat (fun (x : Nat) => Nat.add x 1) (List.cons.{0} Nat 1 (List.cons.{0} Nat 4 (List.nil.{0} Nat))) ==> List.cons.{1} Nat ((fun (x : Nat) => Nat.add x 1) 1) (List.map.{1, 1} Nat Nat (fun (x : Nat) => Nat.add x 1) (List.cons.{0} Nat 4 (List.nil.{0} Nat)))
 -/
 #guard_msgs in
-#eval tstWHNF #[`Init.Data.List] t6
+#eval tstWHNF t6
 
 /-- info: Prop : Type -/
 #guard_msgs in
-#eval tstInferType #[] $ mkSort levelZero
+#eval tstInferType $ mkSort levelZero
 
 /--
 info: fun {a : Type} (x : a) (xs : List.{0} a) => xs : forall {a : Type}, a -> (List.{0} a) -> (List.{0} a)
 -/
 #guard_msgs in
-#eval tstInferType #[`Init.Data.List] $ mkLambda `a BinderInfo.implicit (mkSort levelOne) (mkLambda `x BinderInfo.default (mkBVar 0) (mkLambda `xs BinderInfo.default (mkApp (mkConst `List [levelZero]) (mkBVar 1)) (mkBVar 0)))
+#eval tstInferType $ mkLambda `a BinderInfo.implicit (mkSort levelOne) (mkLambda `x BinderInfo.default (mkBVar 0) (mkLambda `xs BinderInfo.default (mkApp (mkConst `List [levelZero]) (mkBVar 1)) (mkBVar 0)))
 
 def t7 : Expr :=
 let nat  := mkConst `Nat [];
@@ -114,11 +110,11 @@ mkLet `x nat one one
 
 /-- info: let x : Nat := 1; 1 : Nat -/
 #guard_msgs in
-#eval tstInferType #[`Init.Core] $ t7
+#eval tstInferType $ t7
 
 /-- info: let x : Nat := 1; 1 ==> 1 -/
 #guard_msgs in
-#eval tstWHNF #[`Init.Core] $ t7
+#eval tstWHNF $ t7
 
 def t8 : Expr :=
 let nat  := mkConst `Nat [];
@@ -128,11 +124,11 @@ mkLet `x nat one (mkAppN add #[one, mkBVar 0])
 
 /-- info: let x : Nat := 1; Nat.add 1 x : Nat -/
 #guard_msgs in
-#eval tstInferType #[`Init.Core] $ t8
+#eval tstInferType $ t8
 
 /-- info: let x : Nat := 1; Nat.add 1 x ==> 2 -/
 #guard_msgs in
-#eval tstWHNF #[`Init.Core] $ t8
+#eval tstWHNF $ t8
 
 def t9 : Expr :=
 let nat  := mkConst `Nat [];
@@ -140,23 +136,23 @@ mkLet `a (mkSort levelOne) nat (mkForall `x BinderInfo.default (mkBVar 0) (mkBVa
 
 /-- info: let a : Type := Nat; a -> a : Type -/
 #guard_msgs in
-#eval tstInferType #[`Init.Core] $ t9
+#eval tstInferType $ t9
 
 /-- info: let a : Type := Nat; a -> a ==> Nat -> Nat -/
 #guard_msgs in
-#eval tstWHNF #[`Init.Core] $ t9
+#eval tstWHNF $ t9
 
 /-- info: 10 : Nat -/
 #guard_msgs in
-#eval tstInferType #[`Init.Core] $ mkLit (Literal.natVal 10)
+#eval tstInferType $ mkLit (Literal.natVal 10)
 
 /-- info: "hello" : String -/
 #guard_msgs in
-#eval tstInferType #[`Init.Core] $ mkLit (Literal.strVal "hello")
+#eval tstInferType $ mkLit (Literal.strVal "hello")
 
 /-- info: [mdata 10] : Nat -/
 #guard_msgs in
-#eval tstInferType #[`Init.Core] $ mkMData {} $ mkLit (Literal.natVal 10)
+#eval tstInferType $ mkMData {} $ mkLit (Literal.natVal 10)
 
 def t10 : Expr :=
 let nat  := mkConst `Nat [];
@@ -165,39 +161,39 @@ mkLambda `a BinderInfo.default nat (mkApp refl (mkBVar 0))
 
 /-- info: fun (a : Nat) => Eq.refl.{1} Nat a : forall (a : Nat), Eq.{1} Nat a a -/
 #guard_msgs in
-#eval tstInferType #[`Init.Core] t10
+#eval tstInferType t10
 
 /-- info: fun (a : Nat) => Eq.refl.{1} Nat a, isProp: false -/
 #guard_msgs in
-#eval tstIsProp #[`Init.Core] t10
+#eval tstIsProp t10
 
 /-- info: And True True, isProp: true -/
 #guard_msgs in
-#eval tstIsProp #[`Init.Core] (mkAppN (mkConst `And []) #[mkConst `True [], mkConst `True []])
+#eval tstIsProp (mkAppN (mkConst `And []) #[mkConst `True [], mkConst `True []])
 
 /-- info: And, isProp: false -/
 #guard_msgs in
-#eval tstIsProp #[`Init.Core] (mkConst `And [])
+#eval tstIsProp (mkConst `And [])
 
 -- Example where isPropQuick fails
 /-- info: id.{0} Prop (And True True), isProp: true -/
 #guard_msgs in
-#eval tstIsProp #[`Init.Core] (mkAppN (mkConst `id [levelZero]) #[mkSort levelZero, mkAppN (mkConst `And []) #[mkConst `True [], mkConst
+#eval tstIsProp (mkAppN (mkConst `id [levelZero]) #[mkSort levelZero, mkAppN (mkConst `And []) #[mkConst `True [], mkConst
  `True []]])
 
 /-- info: Eq.{1} Nat 0 1, isProp: true -/
 #guard_msgs in
-#eval tstIsProp #[`Init.Core] (mkAppN (mkConst `Eq [levelOne]) #[mkConst `Nat [], mkLit (Literal.natVal 0), mkLit (Literal.natVal 1)])
+#eval tstIsProp (mkAppN (mkConst `Eq [levelOne]) #[mkConst `Nat [], mkLit (Literal.natVal 0), mkLit (Literal.natVal 1)])
 
 /-- info: forall (x : Nat), Eq.{1} Nat x 1, isProp: true -/
 #guard_msgs in
-#eval tstIsProp #[`Init.Core] $
+#eval tstIsProp $
   mkForall `x BinderInfo.default (mkConst `Nat [])
     (mkAppN (mkConst `Eq [levelOne]) #[mkConst `Nat [], mkBVar 0, mkLit (Literal.natVal 1)])
 
 /-- info: (fun (x : Nat) => Eq.{1} Nat x 1) 0, isProp: true -/
 #guard_msgs in
-#eval tstIsProp #[`Init.Core] $
+#eval tstIsProp $
   mkApp
     (mkLambda `x BinderInfo.default (mkConst `Nat [])
       (mkAppN (mkConst `Eq [levelOne]) #[mkConst `Nat [], mkBVar 0, mkLit (Literal.natVal 1)]))
diff --git a/tests/lean/run/meta3.lean b/tests/lean/run/meta3.lean
index 9a7097bdb3a2..79d3b5b03180 100644
--- a/tests/lean/run/meta3.lean
+++ b/tests/lean/run/meta3.lean
@@ -21,12 +21,6 @@ do let v? ← getExprMVarAssignment? m.mvarId!;
    | some v => pure v
    | none   => throwError "metavariable is not assigned")
 
-unsafe def run (mods : Array Name) (x : MetaM Unit) (opts : Options := dbgOpt) : IO Unit :=
- withImportModules (mods.map $ fun m => {module := m}) {} 0 fun env => do
-   let x : MetaM Unit := do { x; printTraces };
-   discard $ x.toIO { options := opts, fileName := "", fileMap := default } { env := env };
-   pure ()
-
 def nat  := mkConst `Nat
 def succ := mkConst `Nat.succ
 def add  := mkAppN (mkConst `Add.add [levelZero]) #[nat, mkConst `Nat.add]
diff --git a/tests/pkg/user_attr_app/Main.lean b/tests/pkg/user_attr_app/Main.lean
index 2f373ecdc64d..d53f1f300f45 100644
--- a/tests/pkg/user_attr_app/Main.lean
+++ b/tests/pkg/user_attr_app/Main.lean
@@ -13,4 +13,4 @@ def tst : MetaM Unit := do
 
 unsafe def main : IO Unit := do
   initSearchPath (← Lean.findSysroot) [".lake/build/lib"]
-  withImportModules #[{ module := `UserAttr.Tst : Import }] {} 0 fun env => pure ()
+  withImportModules #[{ module := `UserAttr.Tst : Import }] {} fun env => pure ()