feat: unsafe functions for freeing compacted regions

src/Lean/Environment.lean

@@ -29,13 +29,24 @@ structure Import :=
 instance : HasToString Import :=
 ⟨fun imp => toString imp.module ++ if imp.runtimeOnly then " (runtime)" else ""⟩
 
+/--
+  A compacted region holds multiple Lean objects in a contiguous memory region, which can be read/written to/from disk.
+  Objects inside the region do not have reference counters and cannot be freed individually. The contents of .olean
+  files are compacted regions. -/
+def CompactedRegion := USize
+
+/-- Free a compacted region and its contents. No live references to the contents may exist at the time of invocation. -/
+@[extern 2 "lean_compacted_region_free"]
+unsafe constant CompactedRegion.free : CompactedRegion → IO Unit := arbitrary _
+
 /- Environment fields that are not used often. -/
 structure EnvironmentHeader :=
 (trustLevel   : UInt32       := 0)
 (quotInit     : Bool         := false)
 (mainModule   : Name         := arbitrary _)
 (imports      : Array Import := #[]) -- direct imports
-(moduleNames  : NameSet      := {})  -- all imported .lean modules
+(regions      : Array CompactedRegion := #[]) -- compacted regions of all imported modules
+(moduleNames  : NameSet      := {})  -- names of all imported modules
 
 open Std (HashMap)
 
@@ -447,7 +458,29 @@ instance ModuleData.inhabited : Inhabited ModuleData :=
 @[extern 3 "lean_save_module_data"]
 constant saveModuleData (fname : @& String) (m : ModuleData) : IO Unit := arbitrary _
 @[extern 2 "lean_read_module_data"]
-constant readModuleData (fname : @& String) : IO ModuleData := arbitrary _
+constant readModuleData (fname : @& String) : IO (ModuleData × CompactedRegion) := arbitrary _
+
+/--
+  Free compacted regions of imports. No live references to imported objects may exist at the time of invocation; in
+  particular, `env` should be the last reference to any `Environment` derived from these imports. -/
+@[noinline, export lean_environment_free_regions]
+unsafe def Environment.freeRegions (env : Environment) : IO Unit :=
+/-
+  NOTE: This assumes `env` is not inferred as a borrowed parameter, and is freed after extracting the `header` field.
+  Otherwise, we would encounter undefined behavior when the constant map in `env`, which may reference objects in
+  compacted regions, is freed after the regions.
+
+  In the currently produced IR, we indeed see:
+  ```
+    def Lean.Environment.freeRegions (x_1 : obj) (x_2 : obj) : obj :=
+      let x_3 : obj := proj[3] x_1;
+      inc x_3;
+      dec x_1;
+      ...
+  ```
+
+  TODO: statically check for this. -/
+env.header.regions.forM CompactedRegion.free
 
 def mkModuleData (env : Environment) : IO ModuleData := do
 pExts ← persistentEnvExtensionsRef.get;
@@ -470,20 +503,21 @@ pure {
 def writeModule (env : Environment) (fname : String) : IO Unit := do
 modData ← mkModuleData env; saveModuleData fname modData
 
-partial def importModulesAux : List Import → (NameSet × Array ModuleData) → IO (NameSet × Array ModuleData)
+partial def importModulesAux : List Import → (NameSet × Array ModuleData × Array CompactedRegion) → IO (NameSet × Array ModuleData × Array CompactedRegion)
 | [],    r         => pure r
-| i::is, (s, mods) =>
+| i::is, (s, mods, regions) =>
   if i.runtimeOnly || s.contains i.module then
-    importModulesAux is (s, mods)
+    importModulesAux is (s, mods, regions)
   else do
     let s := s.insert i.module;
     mFile ← findOLean i.module;
     unlessM (IO.fileExists mFile) $
       throw $ IO.userError $ "object file '" ++ mFile ++ "' of module " ++ toString i.module ++ " does not exist";
-    mod ← readModuleData mFile;
-    (s, mods) ← importModulesAux mod.imports.toList (s, mods);
+    (mod, region) ← readModuleData mFile;
+    (s, mods, regions) ← importModulesAux mod.imports.toList (s, mods, regions);
     let mods := mods.push mod;
-    importModulesAux is (s, mods)
+    let regions := regions.push region;
+    importModulesAux is (s, mods, regions)
 
 private partial def getEntriesFor (mod : ModuleData) (extId : Name) : Nat → Array EnvExtensionState
 | i =>
@@ -510,7 +544,7 @@ pExtDescrs.iterateM env $ fun _ extDescr env => do
 
 @[export lean_import_modules]
 def importModules (imports : List Import) (trustLevel : UInt32 := 0) : IO Environment := do
-(moduleNames, mods) ← importModulesAux imports ({}, #[]);
+(moduleNames, mods, regions) ← importModulesAux imports ({}, #[], #[]);
 let const2ModIdx := mods.iterate {} $ fun (modIdx) (mod : ModuleData) (m : HashMap Name ModuleIdx) =>
   mod.constants.iterate m $ fun _ cinfo m =>
     m.insert cinfo.name modIdx.val;
@@ -529,6 +563,7 @@ let env : Environment := {
     quotInit     := !imports.isEmpty, -- We assume `core.lean` initializes quotient module
     trustLevel   := trustLevel,
     imports      := imports.toArray,
+    regions      := regions,
     moduleNames  := moduleNames
   }
 };

src/library/module.cpp

@@ -86,16 +86,12 @@ extern "C" object * lean_read_module_data(object * fname, object *) {
             return set_io_error((sstream() << "failed to read file '" << olean_fn << "'").str());
         }
         in.close();
-        /* We don't free compacted_region objects */
         compacted_region * region = new compacted_region(size - header_size, buffer);
-#if defined(__has_feature)
-#if __has_feature(address_sanitizer)
-        // do not report as leak
-        __lsan_ignore_object(region);
-#endif
-#endif
         object * mod = region->read();
-        return set_io_result(mod);
+        object * mod_region = alloc_cnstr(0, 2, 0);
+        cnstr_set(mod_region, 0, mod);
+        cnstr_set(mod_region, 1, box_size_t(reinterpret_cast<size_t>(region)));
+        return set_io_result(mod_region);
     } catch (exception & ex) {
         return set_io_error((sstream() << "failed to read '" << olean_fn << "': " << ex.what()).str());
     }

src/runtime/compact.cpp

@@ -482,4 +482,9 @@ object * compacted_region::read() {
         }
     }
 }
+
+extern "C" obj_res lean_compacted_region_free(usize region, object *) {
+    delete reinterpret_cast<compacted_region *>(region);
+    return lean_mk_io_result(lean_box(0));
+}
 }

src/shell/lean.cpp

@@ -373,6 +373,11 @@ extern "C" object* lean_print_deps(object* deps, object* w);
 void print_deps(object_ref const & deps) {
     consume_io_result(lean_print_deps(deps.to_obj_arg(), io_mk_world()));
 }
+
+extern "C" object* lean_environment_free_regions(object * env, object * w);
+void environment_free_regions(environment && env) {
+    consume_io_result(lean_environment_free_regions(env.steal(), io_mk_world()));
+}
 }
 
 void check_optarg(char const * option_name) {
@@ -681,6 +686,7 @@ int main(int argc, char ** argv) {
         if (!json_output)
             display_cumulative_profiling_times(std::cerr);
 
+        environment_free_regions(std::move(env));
         return ok ? 0 : 1;
     } catch (lean::throwable & ex) {
         std::cerr << lean::message_builder(env, ios, mod_fn, lean::pos_info(1, 1), lean::ERROR).set_exception(

tests/lean/ctor_layout.lean.expected.out

@@ -8,6 +8,7 @@ scalar#1@4:u8
 obj@0
 obj@1
 obj@2
+obj@3
 ---
 obj@0
 ◾