From 20915eb5b87ea7186fbcaa389100b99223b851ba Mon Sep 17 00:00:00 2001
From: squidbus <175574877+squidbus@users.noreply.github.com>
Date: Tue, 8 Oct 2024 00:29:05 -0700
Subject: [PATCH] core: Add support for Neo mode memory size. (#1196)

---
 src/core/address_space.cpp                    | 2 +-
 src/core/libraries/kernel/memory_management.h | 2 ++
 src/core/memory.cpp                           | 5 ++++-
 src/video_core/page_manager.cpp               | 3 +++
 4 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/core/address_space.cpp b/src/core/address_space.cpp
index 48748e4c..b2cb95a5 100644
--- a/src/core/address_space.cpp
+++ b/src/core/address_space.cpp
@@ -25,7 +25,7 @@ asm(".zerofill GUEST_SYSTEM,GUEST_SYSTEM,__guest_system,0xFBFC00000");
 
 namespace Core {
 
-static constexpr size_t BackingSize = SCE_KERNEL_MAIN_DMEM_SIZE;
+static constexpr size_t BackingSize = SCE_KERNEL_MAIN_DMEM_SIZE_PRO;
 
 #ifdef _WIN32
 
diff --git a/src/core/libraries/kernel/memory_management.h b/src/core/libraries/kernel/memory_management.h
index 38898aa5..6e90204c 100644
--- a/src/core/libraries/kernel/memory_management.h
+++ b/src/core/libraries/kernel/memory_management.h
@@ -7,6 +7,8 @@
 #include "common/types.h"
 
 constexpr u64 SCE_KERNEL_MAIN_DMEM_SIZE = 5056_MB; // ~ 5GB
+// TODO: Confirm this value on hardware.
+constexpr u64 SCE_KERNEL_MAIN_DMEM_SIZE_PRO = 5568_MB; // ~ 5.5GB
 
 namespace Libraries::Kernel {
 
diff --git a/src/core/memory.cpp b/src/core/memory.cpp
index d21ebae8..dd63522f 100644
--- a/src/core/memory.cpp
+++ b/src/core/memory.cpp
@@ -3,6 +3,7 @@
 
 #include "common/alignment.h"
 #include "common/assert.h"
+#include "common/config.h"
 #include "common/debug.h"
 #include "core/libraries/error_codes.h"
 #include "core/libraries/kernel/memory_management.h"
@@ -39,8 +40,10 @@ MemoryManager::MemoryManager() {
 MemoryManager::~MemoryManager() = default;
 
 void MemoryManager::SetupMemoryRegions(u64 flexible_size) {
+    const auto total_size =
+        Config::isNeoMode() ? SCE_KERNEL_MAIN_DMEM_SIZE_PRO : SCE_KERNEL_MAIN_DMEM_SIZE;
     total_flexible_size = flexible_size;
-    total_direct_size = SCE_KERNEL_MAIN_DMEM_SIZE - flexible_size;
+    total_direct_size = total_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.
diff --git a/src/video_core/page_manager.cpp b/src/video_core/page_manager.cpp
index fb09e70f..a49fff43 100644
--- a/src/video_core/page_manager.cpp
+++ b/src/video_core/page_manager.cpp
@@ -142,6 +142,9 @@ struct PageManager::Impl {
     }
 
     void Protect(VAddr address, size_t size, bool allow_write) {
+        ASSERT_MSG(owned_ranges.find(address) != owned_ranges.end(),
+                   "Attempted to track non-GPU memory at address {:#x}, size {:#x}.", address,
+                   size);
 #ifdef _WIN32
         DWORD prot = allow_write ? PAGE_READWRITE : PAGE_READONLY;
         DWORD old_prot{};