From a72e907f98934524fbaea744fcf59db581a81068 Mon Sep 17 00:00:00 2001
From: yamashi <yamashi@live.com>
Date: Mon, 14 Dec 2020 11:12:09 +0100
Subject: [PATCH] Add memory pool patch

---
 cyberpunk_amd_patch/src/Options.cpp    | 18 +++++-----
 cyberpunk_amd_patch/src/Options.h      |  1 +
 cyberpunk_amd_patch/src/dllmain.cpp    |  4 +++
 cyberpunk_amd_patch/src/pool_patch.cpp | 47 ++++++++++++++++++++++++++
 cyberpunk_amd_patch/xmake.lua          |  4 +--
 5 files changed, 63 insertions(+), 11 deletions(-)
 create mode 100644 cyberpunk_amd_patch/src/pool_patch.cpp

diff --git a/cyberpunk_amd_patch/src/Options.cpp b/cyberpunk_amd_patch/src/Options.cpp
index 58041560..5d6c8a11 100644
--- a/cyberpunk_amd_patch/src/Options.cpp
+++ b/cyberpunk_amd_patch/src/Options.cpp
@@ -30,15 +30,15 @@ Options::Options(HMODULE aModule)
         this->PatchAVX = config.value("avx", this->PatchAVX);
         this->PatchSMT = config.value("smt", this->PatchSMT);
         this->PatchSpectre = config.value("spectre", this->PatchSpectre);
+        this->PatchMemoryPool = config.value("memory_pool", true);
     }
-    else
-    {
-        nlohmann::json config;
-        config["avx"] = this->PatchAVX;
-        config["smt"] = this->PatchSMT;
-        config["spectre"] = this->PatchSpectre;
 
-        std::ofstream o(configPath);
-        o << config.dump(4) << std::endl;
-    }
+    nlohmann::json config;
+    config["avx"] = this->PatchAVX;
+    config["smt"] = this->PatchSMT;
+    config["spectre"] = this->PatchSpectre;
+    config["memory_pool"] = this->PatchMemoryPool;
+
+    std::ofstream o(configPath);
+    o << config.dump(4) << std::endl;
 }
diff --git a/cyberpunk_amd_patch/src/Options.h b/cyberpunk_amd_patch/src/Options.h
index ec4ba138..4ccc2d4a 100644
--- a/cyberpunk_amd_patch/src/Options.h
+++ b/cyberpunk_amd_patch/src/Options.h
@@ -11,5 +11,6 @@ struct Options
 	bool PatchSpectre { true };
 	bool PatchSMT{ true };
 	bool PatchAVX{ false };
+	bool PatchMemoryPool{ true };
 	std::filesystem::path Path;
 };
\ No newline at end of file
diff --git a/cyberpunk_amd_patch/src/dllmain.cpp b/cyberpunk_amd_patch/src/dllmain.cpp
index ce4808cb..85f60924 100644
--- a/cyberpunk_amd_patch/src/dllmain.cpp
+++ b/cyberpunk_amd_patch/src/dllmain.cpp
@@ -13,6 +13,7 @@
 #pragma comment( lib, "dbghelp.lib" )
 #pragma comment(linker, "/DLL")
 
+void PoolPatch(Image* apImage);
 void PatchAmd(Image* apImage);
 void PatchAvx(Image* apImage);
 void HotPatchFix(Image* apImage);
@@ -34,6 +35,9 @@ void Initialize(HMODULE mod)
     if (options.PatchAVX)
         PatchAvx(&image);
 
+    if(options.PatchMemoryPool)
+        PoolPatch(&image);
+
     spdlog::default_logger()->flush();
 }
 
diff --git a/cyberpunk_amd_patch/src/pool_patch.cpp b/cyberpunk_amd_patch/src/pool_patch.cpp
new file mode 100644
index 00000000..87afe4bf
--- /dev/null
+++ b/cyberpunk_amd_patch/src/pool_patch.cpp
@@ -0,0 +1,47 @@
+#include "Image.h"
+#include <spdlog/spdlog.h>
+#include <mhook-lib/mhook.h>
+
+using TRegisterPoolOptions = void(void*, const char*, uint64_t);
+TRegisterPoolOptions* RealRegisterPoolOptions = nullptr;
+
+void RegisterPoolOptions(void* apThis, const char* acpName, uint64_t aSize)
+{
+    if (strcmp(acpName, "PoolCPU") == 0)
+    {
+        MEMORYSTATUSEX statex;
+        statex.dwLength = sizeof(statex);
+        GlobalMemoryStatusEx(&statex);
+
+        if (statex.ullTotalPhys)
+        {
+            const auto gigsInstalled = statex.ullTotalPhys >> 30;
+            aSize = (gigsInstalled - 4) << 30;
+
+            spdlog::info("\t\tDetected RAM: {}GB, using {}GB", gigsInstalled, aSize >> 30);
+        }
+    }
+    else if (strcmp(acpName, "PoolGPU") == 0)
+    {
+        const auto fourGigs = 4ull << 30; // Assume at least 4 gigs of vram is available
+        aSize = aSize > fourGigs ? aSize : fourGigs;
+
+        spdlog::info("\t\tUsing {}GB of VRAM", aSize >> 30);
+    }
+
+    RealRegisterPoolOptions(apThis, acpName, aSize);
+}
+
+void PoolPatch(Image* apImage)
+{
+    if (apImage->version == Image::MakeVersion(1, 4))
+        RealRegisterPoolOptions = reinterpret_cast<TRegisterPoolOptions*>(0x1AD0F0 + apImage->base_address);
+
+    if (RealRegisterPoolOptions)
+    {
+        auto result = Mhook_SetHook(reinterpret_cast<PVOID*>(&RealRegisterPoolOptions), &RegisterPoolOptions);
+        spdlog::info("\tPool patch: {}", result ? "success":"error");
+    }
+    else
+        spdlog::info("\tPool patch: failed");
+}
diff --git a/cyberpunk_amd_patch/xmake.lua b/cyberpunk_amd_patch/xmake.lua
index a4b04e1e..560140ad 100644
--- a/cyberpunk_amd_patch/xmake.lua
+++ b/cyberpunk_amd_patch/xmake.lua
@@ -1,6 +1,6 @@
 set_languages("cxx17")
 
-add_requires("zlib", "spdlog", "nlohmann_json")
+add_requires("zlib", "spdlog", "nlohmann_json", "mhook")
 
 add_rules("mode.debug", "mode.release")
 add_cxflags("-flto")
@@ -14,4 +14,4 @@ target("performance_overhaul")
     add_files("src/**.cpp")
     add_includedirs("src/")
     add_syslinks("User32")
-    add_packages("zlib", "spdlog", "nlohmann_json")
+    add_packages("zlib", "spdlog", "nlohmann_json", "mhook")