fix: freeing Environments in tests

src/Lean/Environment.lean

@@ -572,6 +572,13 @@ env ← finalizePersistentExtensions env;
 env ← mods.iterateM env $ fun _ mod env => performModifications env mod.serialized;
 pure env
 
+/--
+  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 : List Import) (trustLevel : UInt32 := 0) (x : Environment → IO α) : IO α := do
+env ← importModules imports trustLevel;
+finally (x env) env.freeRegions
+
 def regNamespacesExtension : IO (SimplePersistentEnvExtension Name NameSet) :=
 registerSimplePersistentEnvExtension {
   name            := `namespaces,

src/shell/lean.cpp

@@ -661,7 +661,9 @@ int main(int argc, char ** argv) {
         }
 
         if (run && ok) {
-            return ir::run_main(env, argc - optind, argv + optind);
+            uint32 ret = ir::run_main(env, argc - optind, argv + optind);
+            environment_free_regions(std::move(env));
+            return ret;
         }
         if (olean_fn && ok) {
             time_task t(".olean serialization",

tests/lean/ctor_layout.lean

@@ -2,8 +2,8 @@ import Lean.Compiler.IR
 open Lean
 open Lean.IR
 
-def tst : IO Unit :=
-do env ← importModules [{module := `Lean.Compiler.IR.Basic}];
+unsafe def main : IO Unit :=
+withImportModules [{module := `Lean.Compiler.IR.Basic}] 0 fun env => do
    ctorLayout ← IO.ofExcept $ getCtorLayout env `Lean.IR.Expr.reuse;
    ctorLayout.fieldInfo.forM $ fun finfo => IO.println (format finfo);
    IO.println "---";
@@ -14,4 +14,4 @@ do env ← importModules [{module := `Lean.Compiler.IR.Basic}];
    ctorLayout.fieldInfo.forM $ fun finfo => IO.println (format finfo);
    pure ()
 
-#eval tst
+#eval main

tests/lean/run/Reparen.lean

@@ -33,23 +33,23 @@ partial def unparen : Syntax → Syntax
   | `(level|($stx')) => unparenAux stx $ unparen stx'
   | _ => stx.modifyArgs $ Array.map unparen
 
-def main (args : List String) : IO Unit := do
+unsafe def main (args : List String) : IO Unit := do
 let (debug, f) : Bool × String := match args with
   | [f, "-d"] => (true, f)
   | [f]       => (false, f)
   | _         => panic! "usage: file [-d]";
 initSearchPath none;
-env ← importModules [{module := `Lean.Parser}];
-stx ← Lean.Parser.parseFile env args.head!;
-let header := stx.getArg 0;
-some s ← pure header.reprint | throw $ IO.userError "header reprint failed";
-IO.print s;
-let cmds := stx.getArgs.extract 1 stx.getArgs.size;
-cmds.forM $ fun cmd => do
-  let cmd := unparen cmd;
-  (cmd, _) ← IO.runMeta (PrettyPrinter.parenthesizeCommand cmd)
-    env { opts := KVMap.insert {} `trace.PrettyPrinter.parenthesize debug };
-  some s ← pure cmd.reprint | throw $ IO.userError "cmd reprint failed";
-  IO.print s
+withImportModules [{module := `Lean.Parser}] 0 fun env => do
+  stx ← Lean.Parser.parseFile env args.head!;
+  let header := stx.getArg 0;
+  some s ← pure header.reprint | throw $ IO.userError "header reprint failed";
+  IO.print s;
+  let cmds := stx.getArgs.extract 1 stx.getArgs.size;
+  cmds.forM $ fun cmd => do
+    let cmd := unparen cmd;
+    (cmd, _) ← IO.runMeta (PrettyPrinter.parenthesizeCommand cmd)
+      env { opts := KVMap.insert {} `trace.PrettyPrinter.parenthesize debug };
+    some s ← pure cmd.reprint | throw $ IO.userError "cmd reprint failed";
+    IO.print s
 
 #eval main ["../../../src/Init/Core.lean"]

tests/lean/run/instances.lean

@@ -2,11 +2,9 @@ import Lean
 open Lean
 open Lean.Meta
 
-def tst1 : IO Unit :=
-do let mods := [`Lean];
-   env ← importModules $ mods.map $ fun m => {module := m};
+unsafe def tst1 : IO Unit :=
+withImportModules [{module := `Lean}] 0 fun env => do
    let insts := env.getGlobalInstances;
-   IO.println (format insts);
-   pure ()
+   IO.println (format insts)
 
 #eval tst1

tests/lean/run/instuniv.lean

@@ -1,14 +1,13 @@
 import Lean
 open Lean
 
-def tst : IO Unit :=
-do env ← importModules [{module := `Init.Data.Array}];
+unsafe def tst : IO Unit :=
+withImportModules [{module := `Init.Data.Array}] 0 fun env => do
    match env.find? `Array.foldl with
    | some info => do
      IO.println (info.instantiateTypeLevelParams [levelZero, levelZero]);
      IO.println (info.instantiateValueLevelParams [levelZero, levelZero]);
      pure ()
-   | none      => IO.println ("Array.foldl not found");
-   pure ()
+   | none      => IO.println ("Array.foldl not found")
 
 #eval tst

tests/lean/run/meta1.lean

@@ -2,20 +2,20 @@ import Lean.Meta
 open Lean
 open Lean.Meta
 
-def tstInferType (mods : List Name) (e : Expr) : IO Unit :=
-do env ← importModules $ mods.map $ fun m => {module := m};
+unsafe def tstInferType (mods : List Name) (e : Expr) : IO Unit :=
+withImportModules (mods.map $ fun m => {module := m}) 0 fun env => do
    match inferType e { currRecDepth := 0, maxRecDepth := 100000 } { env := env } with
    | EStateM.Result.ok type s   => IO.println (toString e ++ " : " ++ toString type)
    | EStateM.Result.error err _ => throw (IO.userError (toString err))
 
-def tstWHNF (mods : List Name) (e : Expr) (t := TransparencyMode.default) : IO Unit :=
-do env ← importModules $ mods.map $ fun m => {module := m};
+unsafe def tstWHNF (mods : List Name) (e : Expr) (t := TransparencyMode.default) : IO Unit :=
+withImportModules (mods.map $ fun m => {module := m}) 0 fun env => do
    match whnf e { config := { transparency := t }, currRecDepth := 0, maxRecDepth := 100000  } { env := env } with
    | EStateM.Result.ok type s   => IO.println (toString e ++ " ==> " ++ toString type)
    | EStateM.Result.error err _ => throw (IO.userError (toString err))
 
-def tstIsProp (mods : List Name) (e : Expr) : IO Unit :=
-do env ← importModules $ mods.map $ fun m => {module := m};
+unsafe def tstIsProp (mods : List Name) (e : Expr) : IO Unit :=
+withImportModules (mods.map $ fun m => {module := m}) 0 fun env => do
    match isProp e { currRecDepth := 0, maxRecDepth := 100000 } { env := env } with
    | EStateM.Result.ok b s      => IO.println (toString e ++ ", isProp: " ++ toString b)
    | EStateM.Result.error err _ => throw (IO.userError (toString err))

tests/lean/run/meta3.lean

@@ -20,8 +20,8 @@ do v? ← getExprMVarAssignment? m.mvarId!;
    | some v => pure v
    | none   => throw $ Exception.other "metavariable is not assigned"
 
-def run (mods : List Name) (x : MetaM Unit) (opts : Options := dbgOpt) : IO Unit :=
-do env ← importModules $ mods.map $ fun m => {module := m};
+unsafe def run (mods : List Name) (x : MetaM Unit) (opts : Options := dbgOpt) : IO Unit :=
+withImportModules (mods.map $ fun m => {module := m}) 0 fun env => do
    match x { config := { opts := opts }, currRecDepth := 0, maxRecDepth := 100000 } { env := env } with
    | EStateM.Result.ok _ s    => do
      s.traceState.traces.forM $ fun m => IO.println $ format m;