kernel/tests: improve memory cleanup

1 files changed, 29 insertions, 6 deletions

diff --git a/kernel/src/test_support/kapi/memory.cpp b/kernel/src/test_support/kapi/memory.cpp
index 4f261ac..652a3d3 100644
--- a/kernel/src/test_support/kapi/memory.cpp
+++ b/kernel/src/test_support/kapi/memory.cpp
@@ -35,7 +35,9 @@ namespace kapi::memory
       auto release_many(std::pair<frame, std::size_t>) -> void override {}
 
       std::size_t next_free_frame{};
-    } bootstrap_allocator;
+    };
+
+    auto constinit bootstrap_allocator = std::optional<test_bootstrap_frame_allocator>{};
 
     struct test_page_mapper : page_mapper
     {
@@ -63,22 +65,31 @@ namespace kapi::memory
         kstd::println("trying to unmap page {}", page.number());
         return false;
       }
-    } test_mapper;
+    };
+
+    auto constinit test_mapper = std::optional<test_page_mapper>{};
 
     auto handoff_to_kernel_pmm(frame_allocator & new_allocator) -> void
     {
-      auto first_free_frame = bootstrap_allocator.next_free_frame;
+      auto first_free_frame = bootstrap_allocator->next_free_frame;
       auto number_of_free_frames = number_of_simulated_frames - first_free_frame;
       new_allocator.release_many({frame{first_free_frame}, number_of_free_frames});
     }
 
+    auto constinit previous_frame_allocator = std::optional<frame_allocator *>{};
+    auto constinit previous_page_mapper = std::optional<page_mapper *>{};
+
   }  // namespace
 
   auto init() -> void
   {
     kernel::tests::simulated_memory::init(simulate_memory_size);
-    set_frame_allocator(bootstrap_allocator);
-    set_page_mapper(test_mapper);
+
+    bootstrap_allocator.emplace();
+    test_mapper.emplace();
+
+    previous_frame_allocator = set_frame_allocator(*bootstrap_allocator);
+    previous_page_mapper = set_page_mapper(*test_mapper);
 
     init_pmm(simulate_memory_size / frame::size, handoff_to_kernel_pmm);
   }
@@ -86,7 +97,19 @@ namespace kapi::memory
   auto reset() -> void
   {
     kernel::tests::simulated_memory::reset();
-    bootstrap_allocator.next_free_frame = 0;
+
+    if (previous_frame_allocator && *previous_frame_allocator)
+    {
+      set_frame_allocator(**previous_frame_allocator);
+    }
+
+    if (previous_page_mapper && *previous_page_mapper)
+    {
+      set_page_mapper(**previous_page_mapper);
+    }
+
+    bootstrap_allocator.reset();
+    test_mapper.reset();
   }
 
 }  // namespace kapi::memory
\ No newline at end of file