diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean index cfcb8251d0..bbc90c2eb8 100644 --- a/src/Lean/Environment.lean +++ b/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 } }; diff --git a/src/library/module.cpp b/src/library/module.cpp index e1225de7bc..093c51d07f 100644 --- a/src/library/module.cpp +++ b/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(region))); + return set_io_result(mod_region); } catch (exception & ex) { return set_io_error((sstream() << "failed to read '" << olean_fn << "': " << ex.what()).str()); } diff --git a/src/runtime/compact.cpp b/src/runtime/compact.cpp index f5e2b9a059..1469a69bd4 100644 --- a/src/runtime/compact.cpp +++ b/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(region); + return lean_mk_io_result(lean_box(0)); +} } diff --git a/src/shell/lean.cpp b/src/shell/lean.cpp index b179a68144..a416515ad0 100644 --- a/src/shell/lean.cpp +++ b/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( diff --git a/tests/lean/ctor_layout.lean.expected.out b/tests/lean/ctor_layout.lean.expected.out index f3be927c15..763613e1b3 100644 --- a/tests/lean/ctor_layout.lean.expected.out +++ b/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 ◾