1 files changed, 170 insertions, 33 deletions

diff --git a/src/core/hle/kernel/k_memory_manager.h b/src/core/hle/kernel/k_memory_manager.h
index abd6c8ace..18775b262 100644
--- a/src/core/hle/kernel/k_memory_manager.h
+++ b/src/core/hle/kernel/k_memory_manager.h
@@ -5,10 +5,12 @@
 #pragma once
 
 #include <array>
-#include <mutex>
 #include <tuple>
 
+#include "common/common_funcs.h"
 #include "common/common_types.h"
+#include "core/hle/kernel/k_light_lock.h"
+#include "core/hle/kernel/k_memory_layout.h"
 #include "core/hle/kernel/k_page_heap.h"
 #include "core/hle/result.h"
 
@@ -20,8 +22,11 @@ namespace Kernel {
 
 class KPageLinkedList;
 
-class KMemoryManager final : NonCopyable {
+class KMemoryManager final {
 public:
+    YUZU_NON_COPYABLE(KMemoryManager);
+    YUZU_NON_MOVEABLE(KMemoryManager);
+
     enum class Pool : u32 {
         Application = 0,
         Applet = 1,
@@ -48,22 +53,33 @@ public:
 
     explicit KMemoryManager(Core::System& system_);
 
-    constexpr std::size_t GetSize(Pool pool) const {
-        return managers[static_cast<std::size_t>(pool)].GetSize();
+    void Initialize(VAddr management_region, size_t management_region_size);
+
+    constexpr size_t GetSize(Pool pool) const {
+        constexpr Direction GetSizeDirection = Direction::FromFront;
+        size_t total = 0;
+        for (auto* manager = this->GetFirstManager(pool, GetSizeDirection); manager != nullptr;
+             manager = this->GetNextManager(manager, GetSizeDirection)) {
+            total += manager->GetSize();
+        }
+        return total;
     }
 
-    void InitializeManager(Pool pool, u64 start_address, u64 end_address);
+    PAddr AllocateAndOpenContinuous(size_t num_pages, size_t align_pages, u32 option);
+    ResultCode AllocateAndOpen(KPageLinkedList* out, size_t num_pages, u32 option);
+    ResultCode AllocateAndOpenForProcess(KPageLinkedList* out, size_t num_pages, u32 option,
+                                         u64 process_id, u8 fill_pattern);
+
+    static constexpr size_t MaxManagerCount = 10;
 
-    VAddr AllocateAndOpenContinuous(size_t num_pages, size_t align_pages, u32 option);
-    ResultCode Allocate(KPageLinkedList& page_list, std::size_t num_pages, Pool pool, Direction dir,
-                        u32 heap_fill_value = 0);
-    ResultCode Free(KPageLinkedList& page_list, std::size_t num_pages, Pool pool, Direction dir,
-                    u32 heap_fill_value = 0);
+    void Close(PAddr address, size_t num_pages);
+    void Close(const KPageLinkedList& pg);
 
-    static constexpr std::size_t MaxManagerCount = 10;
+    void Open(PAddr address, size_t num_pages);
+    void Open(const KPageLinkedList& pg);
 
 public:
-    static std::size_t CalculateManagementOverheadSize(std::size_t region_size) {
+    static size_t CalculateManagementOverheadSize(size_t region_size) {
         return Impl::CalculateManagementOverheadSize(region_size);
     }
 
@@ -88,38 +104,34 @@ public:
     }
 
 private:
-    class Impl final : NonCopyable {
-    private:
-        using RefCount = u16;
-
-    private:
-        KPageHeap heap;
-        Pool pool{};
-
+    class Impl final {
     public:
-        static std::size_t CalculateManagementOverheadSize(std::size_t region_size);
-
-        static constexpr std::size_t CalculateOptimizedProcessOverheadSize(
-            std::size_t region_size) {
-            return (Common::AlignUp((region_size / PageSize), Common::BitSize<u64>()) /
-                    Common::BitSize<u64>()) *
-                   sizeof(u64);
-        }
+        YUZU_NON_COPYABLE(Impl);
+        YUZU_NON_MOVEABLE(Impl);
 
-    public:
         Impl() = default;
+        ~Impl() = default;
 
-        std::size_t Initialize(Pool new_pool, u64 start_address, u64 end_address);
+        size_t Initialize(PAddr address, size_t size, VAddr management, VAddr management_end,
+                          Pool p);
 
         VAddr AllocateBlock(s32 index, bool random) {
             return heap.AllocateBlock(index, random);
         }
 
-        void Free(VAddr addr, std::size_t num_pages) {
+        void Free(VAddr addr, size_t num_pages) {
             heap.Free(addr, num_pages);
         }
 
-        constexpr std::size_t GetSize() const {
+        void SetInitialUsedHeapSize(size_t reserved_size) {
+            heap.SetInitialUsedSize(reserved_size);
+        }
+
+        constexpr Pool GetPool() const {
+            return pool;
+        }
+
+        constexpr size_t GetSize() const {
             return heap.GetSize();
         }
 
@@ -130,12 +142,137 @@ private:
         constexpr VAddr GetEndAddress() const {
             return heap.GetEndAddress();
         }
+
+        constexpr size_t GetPageOffset(PAddr address) const {
+            return heap.GetPageOffset(address);
+        }
+
+        constexpr size_t GetPageOffsetToEnd(PAddr address) const {
+            return heap.GetPageOffsetToEnd(address);
+        }
+
+        constexpr void SetNext(Impl* n) {
+            next = n;
+        }
+
+        constexpr void SetPrev(Impl* n) {
+            prev = n;
+        }
+
+        constexpr Impl* GetNext() const {
+            return next;
+        }
+
+        constexpr Impl* GetPrev() const {
+            return prev;
+        }
+
+        void OpenFirst(PAddr address, size_t num_pages) {
+            size_t index = this->GetPageOffset(address);
+            const size_t end = index + num_pages;
+            while (index < end) {
+                const RefCount ref_count = (++page_reference_counts[index]);
+                ASSERT(ref_count == 1);
+
+                index++;
+            }
+        }
+
+        void Open(PAddr address, size_t num_pages) {
+            size_t index = this->GetPageOffset(address);
+            const size_t end = index + num_pages;
+            while (index < end) {
+                const RefCount ref_count = (++page_reference_counts[index]);
+                ASSERT(ref_count > 1);
+
+                index++;
+            }
+        }
+
+        void Close(PAddr address, size_t num_pages) {
+            size_t index = this->GetPageOffset(address);
+            const size_t end = index + num_pages;
+
+            size_t free_start = 0;
+            size_t free_count = 0;
+            while (index < end) {
+                ASSERT(page_reference_counts[index] > 0);
+                const RefCount ref_count = (--page_reference_counts[index]);
+
+                // Keep track of how many zero refcounts we see in a row, to minimize calls to free.
+                if (ref_count == 0) {
+                    if (free_count > 0) {
+                        free_count++;
+                    } else {
+                        free_start = index;
+                        free_count = 1;
+                    }
+                } else {
+                    if (free_count > 0) {
+                        this->Free(heap.GetAddress() + free_start * PageSize, free_count);
+                        free_count = 0;
+                    }
+                }
+
+                index++;
+            }
+
+            if (free_count > 0) {
+                this->Free(heap.GetAddress() + free_start * PageSize, free_count);
+            }
+        }
+
+        static size_t CalculateManagementOverheadSize(size_t region_size);
+
+        static constexpr size_t CalculateOptimizedProcessOverheadSize(size_t region_size) {
+            return (Common::AlignUp((region_size / PageSize), Common::BitSize<u64>()) /
+                    Common::BitSize<u64>()) *
+                   sizeof(u64);
+        }
+
+    private:
+        using RefCount = u16;
+
+        KPageHeap heap;
+        std::vector<RefCount> page_reference_counts;
+        VAddr management_region{};
+        Pool pool{};
+        Impl* next{};
+        Impl* prev{};
     };
 
 private:
+    Impl& GetManager(const KMemoryLayout& memory_layout, PAddr address) {
+        return managers[memory_layout.GetPhysicalLinearRegion(address).GetAttributes()];
+    }
+
+    const Impl& GetManager(const KMemoryLayout& memory_layout, PAddr address) const {
+        return managers[memory_layout.GetPhysicalLinearRegion(address).GetAttributes()];
+    }
+
+    constexpr Impl* GetFirstManager(Pool pool, Direction dir) const {
+        return dir == Direction::FromBack ? pool_managers_tail[static_cast<size_t>(pool)]
+                                          : pool_managers_head[static_cast<size_t>(pool)];
+    }
+
+    constexpr Impl* GetNextManager(Impl* cur, Direction dir) const {
+        if (dir == Direction::FromBack) {
+            return cur->GetPrev();
+        } else {
+            return cur->GetNext();
+        }
+    }
+
+    ResultCode AllocatePageGroupImpl(KPageLinkedList* out, size_t num_pages, Pool pool,
+                                     Direction dir, bool random);
+
+private:
     Core::System& system;
-    std::array<std::mutex, static_cast<std::size_t>(Pool::Count)> pool_locks;
+    std::array<KLightLock, static_cast<size_t>(Pool::Count)> pool_locks;
+    std::array<Impl*, MaxManagerCount> pool_managers_head{};
+    std::array<Impl*, MaxManagerCount> pool_managers_tail{};
     std::array<Impl, MaxManagerCount> managers;
+    size_t num_managers{};
 };
 
 } // namespace Kernel