From ca9360a865fb4d28578b25da244b66d5d4483133 Mon Sep 17 00:00:00 2001
From: squidbus <175574877+squidbus@users.noreply.github.com>
Date: Tue, 27 Aug 2024 05:37:10 -0700
Subject: [PATCH] memory: Size direct memory based on requested flexible
 memory.

---
 src/core/libraries/kernel/memory_management.cpp |  9 +++++----
 src/core/libraries/kernel/memory_management.h   |  2 +-
 src/core/linker.cpp                             | 10 +++++-----
 src/core/memory.cpp                             | 16 ++++++++++++++--
 src/core/memory.h                               |  9 ++++++---
 5 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/src/core/libraries/kernel/memory_management.cpp b/src/core/libraries/kernel/memory_management.cpp
index ab6aa398..3fc837d9 100644
--- a/src/core/libraries/kernel/memory_management.cpp
+++ b/src/core/libraries/kernel/memory_management.cpp
@@ -15,7 +15,8 @@ namespace Libraries::Kernel {
 
 u64 PS4_SYSV_ABI sceKernelGetDirectMemorySize() {
     LOG_WARNING(Kernel_Vmm, "called");
-    return SCE_KERNEL_MAIN_DMEM_SIZE;
+    const auto* memory = Core::Memory::Instance();
+    return memory->GetTotalDirectSize();
 }
 
 int PS4_SYSV_ABI sceKernelAllocateDirectMemory(s64 searchStart, s64 searchEnd, u64 len,
@@ -52,8 +53,8 @@ int PS4_SYSV_ABI sceKernelAllocateDirectMemory(s64 searchStart, s64 searchEnd, u
 
 s32 PS4_SYSV_ABI sceKernelAllocateMainDirectMemory(size_t len, size_t alignment, int memoryType,
                                                    s64* physAddrOut) {
-    return sceKernelAllocateDirectMemory(0, SCE_KERNEL_MAIN_DMEM_SIZE, len, alignment, memoryType,
-                                         physAddrOut);
+    const auto searchEnd = static_cast<s64>(sceKernelGetDirectMemorySize());
+    return sceKernelAllocateDirectMemory(0, searchEnd, len, alignment, memoryType, physAddrOut);
 }
 
 s32 PS4_SYSV_ABI sceKernelCheckedReleaseDirectMemory(u64 start, size_t len) {
@@ -78,7 +79,7 @@ s32 PS4_SYSV_ABI sceKernelAvailableDirectMemorySize(u64 searchStart, u64 searchE
     if (physAddrOut == nullptr || sizeOut == nullptr) {
         return ORBIS_KERNEL_ERROR_EINVAL;
     }
-    if (searchEnd > SCE_KERNEL_MAIN_DMEM_SIZE) {
+    if (searchEnd > sceKernelGetDirectMemorySize()) {
         return ORBIS_KERNEL_ERROR_EINVAL;
     }
     if (searchEnd <= searchStart) {
diff --git a/src/core/libraries/kernel/memory_management.h b/src/core/libraries/kernel/memory_management.h
index c9dd86e1..9a447fe8 100644
--- a/src/core/libraries/kernel/memory_management.h
+++ b/src/core/libraries/kernel/memory_management.h
@@ -6,7 +6,7 @@
 #include "common/bit_field.h"
 #include "common/types.h"
 
-constexpr u64 SCE_KERNEL_MAIN_DMEM_SIZE = 4608_MB; // ~ 4.5GB
+constexpr u64 SCE_KERNEL_MAIN_DMEM_SIZE = 5056_MB; // ~ 5GB
 
 namespace Libraries::Kernel {
 
diff --git a/src/core/linker.cpp b/src/core/linker.cpp
index 6f46fdbe..37ac657d 100644
--- a/src/core/linker.cpp
+++ b/src/core/linker.cpp
@@ -68,11 +68,11 @@ void Linker::Execute() {
     }
 
     // Configure used flexible memory size.
-    // if (auto* mem_param = GetProcParam()->mem_param) {
-    //      if (u64* flexible_size = mem_param->flexible_memory_size) {
-    //          memory->SetTotalFlexibleSize(*flexible_size);
-    //      }
-    //  }
+    if (auto* mem_param = GetProcParam()->mem_param) {
+        if (u64* flexible_size = mem_param->flexible_memory_size) {
+            memory->SetupMemoryRegions(*flexible_size);
+        }
+    }
 
     // Init primary thread.
     Common::SetCurrentThreadName("GAME_MainThread");
diff --git a/src/core/memory.cpp b/src/core/memory.cpp
index 0feef260..420416bf 100644
--- a/src/core/memory.cpp
+++ b/src/core/memory.cpp
@@ -11,9 +11,11 @@
 
 namespace Core {
 
+constexpr u64 SCE_DEFAULT_FLEXIBLE_MEMORY_SIZE = 448_MB;
+
 MemoryManager::MemoryManager() {
-    // Insert an area that covers direct memory physical block.
-    dmem_map.emplace(0, DirectMemoryArea{0, SCE_KERNEL_MAIN_DMEM_SIZE});
+    // Set up the direct and flexible memory regions.
+    SetupMemoryRegions(SCE_DEFAULT_FLEXIBLE_MEMORY_SIZE);
 
     // Insert a virtual memory area that covers the entire area we manage.
     const VAddr system_managed_base = impl.SystemManagedVirtualBase();
@@ -35,6 +37,16 @@ MemoryManager::MemoryManager() {
 
 MemoryManager::~MemoryManager() = default;
 
+void MemoryManager::SetupMemoryRegions(u64 flexible_size) {
+    total_flexible_size = flexible_size;
+    total_direct_size = SCE_KERNEL_MAIN_DMEM_SIZE - flexible_size;
+
+    // Insert an area that covers direct memory physical block.
+    // Note that this should never be called after direct memory allocations have been made.
+    dmem_map.clear();
+    dmem_map.emplace(0, DirectMemoryArea{0, total_direct_size});
+}
+
 PAddr MemoryManager::Allocate(PAddr search_start, PAddr search_end, size_t size, u64 alignment,
                               int memory_type) {
     std::scoped_lock lk{mutex};
diff --git a/src/core/memory.h b/src/core/memory.h
index d5826967..c8638626 100644
--- a/src/core/memory.h
+++ b/src/core/memory.h
@@ -130,8 +130,8 @@ public:
         rasterizer = rasterizer_;
     }
 
-    void SetTotalFlexibleSize(u64 size) {
-        total_flexible_size = size;
+    u64 GetTotalDirectSize() const {
+        return total_direct_size;
     }
 
     u64 GetAvailableFlexibleSize() const {
@@ -142,6 +142,8 @@ public:
         return impl.SystemReservedVirtualBase();
     }
 
+    void SetupMemoryRegions(u64 flexible_size);
+
     PAddr Allocate(PAddr search_start, PAddr search_end, size_t size, u64 alignment,
                    int memory_type);
 
@@ -217,7 +219,8 @@ private:
     DMemMap dmem_map;
     VMAMap vma_map;
     std::recursive_mutex mutex;
-    size_t total_flexible_size = 448_MB;
+    size_t total_direct_size{};
+    size_t total_flexible_size{};
     size_t flexible_usage{};
     Vulkan::Rasterizer* rasterizer{};
 };