From 5cf9a17c36d4cf9b6dbc49b0afa09994c2d82550 Mon Sep 17 00:00:00 2001
From: praydog <praydog@praydog.com>
Date: Sat, 1 Jun 2024 21:12:26 -0700
Subject: [PATCH] D3D12Hook: More frame generation fixes

---
 src/D3D12Hook.cpp   | 20 ++++++++++++++------
 src/D3D12Hook.hpp   |  4 +++-
 src/REFramework.hpp |  2 ++
 3 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/D3D12Hook.cpp b/src/D3D12Hook.cpp
index 0d5cacab9..26369b083 100644
--- a/src/D3D12Hook.cpp
+++ b/src/D3D12Hook.cpp
@@ -29,16 +29,24 @@ void* D3D12Hook::Streamline::link_swapchain_to_cmd_queue(void* rcx, void* rdx, v
     spdlog::info("[Streamline] linkSwapchainToCmdQueue: {:x}", (uintptr_t)_ReturnAddress());
 
     g_framework->on_reset(); // Needed to prevent a crash due to resources hanging around
-    g_d3d12_hook->unhook(); // Removes all vtable hooks
 
-    auto& hook = g_d3d12_hook->m_streamline.link_swapchain_to_cmd_queue_hook;
+    if (g_d3d12_hook != nullptr) {
+        g_d3d12_hook->unhook(); // Removes all vtable hooks
+    }
+
+    auto& hook = D3D12Hook::s_streamline.link_swapchain_to_cmd_queue_hook;
     const auto result = hook->get_original<decltype(link_swapchain_to_cmd_queue)>()(rcx, rdx, r8, r9);
 
+    // Re-hooks present after the above function creates the swapchain
+    // This allows the hook to immediately still function
+    // rather than waiting on the hook monitor to notice the hook isn't working
+    g_framework->hook_d3d12();
+
     return result;
 }
 
 void D3D12Hook::hook_streamline() {
-    if (m_streamline.setup) {
+    if (D3D12Hook::s_streamline.setup) {
         return;
     }
 
@@ -72,15 +80,15 @@ void D3D12Hook::hook_streamline() {
         return;
     }
 
-    m_streamline.link_swapchain_to_cmd_queue_hook = std::make_unique<FunctionHook>(*fn, (uintptr_t)&Streamline::link_swapchain_to_cmd_queue);
+    D3D12Hook::s_streamline.link_swapchain_to_cmd_queue_hook = std::make_unique<FunctionHook>(*fn, (uintptr_t)&Streamline::link_swapchain_to_cmd_queue);
 
-    if (m_streamline.link_swapchain_to_cmd_queue_hook->create()) {
+    if (D3D12Hook::s_streamline.link_swapchain_to_cmd_queue_hook->create()) {
         spdlog::info("[Streamline] Hooked linkSwapchainToCmdQueue");
     } else {
         spdlog::error("[Streamline] Failed to hook linkSwapchainToCmdQueue");
     }
 
-    m_streamline.setup = true;
+    D3D12Hook::s_streamline.setup = true;
 }
 
 bool D3D12Hook::hook() {
diff --git a/src/D3D12Hook.hpp b/src/D3D12Hook.hpp
index cc46005b1..d65a85552 100644
--- a/src/D3D12Hook.hpp
+++ b/src/D3D12Hook.hpp
@@ -131,7 +131,9 @@ class D3D12Hook
 
         std::unique_ptr<FunctionHook> link_swapchain_to_cmd_queue_hook{};
         bool setup{ false };
-    } m_streamline{};
+    };
+
+    static inline Streamline s_streamline{};
     
     OnPresentFn m_on_present{ nullptr };
     OnPresentFn m_on_post_present{ nullptr };
diff --git a/src/REFramework.hpp b/src/REFramework.hpp
index c51dc8131..aa77c8a0c 100644
--- a/src/REFramework.hpp
+++ b/src/REFramework.hpp
@@ -144,9 +144,11 @@ class REFramework {
     void draw_ui();
     void draw_about();
 
+public:
     bool hook_d3d11();
     bool hook_d3d12();
 
+private:
     bool initialize();
     bool initialize_game_data();
     bool initialize_windows_message_hook();