From bb5b36145abd30a1579239037b365e86cacee8a3 Mon Sep 17 00:00:00 2001
From: WinterSnowfall <WinterSnowfall@users.noreply.github.com>
Date: Wed, 27 Nov 2024 13:14:50 +0200
Subject: [PATCH] [d3d9] Fixes for state block specific behavior

---
 src/d3d9/d3d9_device.cpp     | 21 ++++++++++++---------
 src/d3d9/d3d9_device.h       | 12 +++++++-----
 src/d3d9/d3d9_stateblock.cpp |  9 +++++++--
 src/d3d9/d3d9_stateblock.h   |  8 +-------
 4 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/src/d3d9/d3d9_device.cpp b/src/d3d9/d3d9_device.cpp
index 548e0d80288..0870676680e 100644
--- a/src/d3d9/d3d9_device.cpp
+++ b/src/d3d9/d3d9_device.cpp
@@ -2599,6 +2599,10 @@ namespace dxvk {
           IDirect3DStateBlock9** ppSB) {
     D3D9DeviceLock lock = LockDevice();
 
+    // A state block can not be created while another is being recorded.
+    if (unlikely(ShouldRecord()))
+      return D3DERR_INVALIDCALL;
+
     InitReturnPtr(ppSB);
 
     if (unlikely(ppSB == nullptr))
@@ -2622,7 +2626,8 @@ namespace dxvk {
   HRESULT STDMETHODCALLTYPE D3D9DeviceEx::BeginStateBlock() {
     D3D9DeviceLock lock = LockDevice();
 
-    if (unlikely(m_recorder != nullptr))
+    // Only one state block can be recorded at a given time.
+    if (unlikely(ShouldRecord()))
       return D3DERR_INVALIDCALL;
 
     m_recorder = new D3D9StateBlock(this, D3D9StateBlockType::None);
@@ -2634,11 +2639,12 @@ namespace dxvk {
   HRESULT STDMETHODCALLTYPE D3D9DeviceEx::EndStateBlock(IDirect3DStateBlock9** ppSB) {
     D3D9DeviceLock lock = LockDevice();
 
-    InitReturnPtr(ppSB);
-
-    if (unlikely(ppSB == nullptr || m_recorder == nullptr))
+    // Recording a state block can't end if recording hasn't been started.
+    if (unlikely(ppSB == nullptr || !ShouldRecord()))
       return D3DERR_INVALIDCALL;
 
+    InitReturnPtr(ppSB);
+
     *ppSB = m_recorder.ref();
     if (!m_isD3D8Compatible)
       m_losableResourceCounter++;
@@ -4639,21 +4645,18 @@ namespace dxvk {
   }
 
 
-  inline bool D3D9DeviceEx::ShouldRecord() {
-    return m_recorder != nullptr && !m_recorder->IsApplying();
-  }
-
-
   D3D9_VK_FORMAT_MAPPING D3D9DeviceEx::LookupFormat(
     D3D9Format            Format) const {
     return m_adapter->GetFormatMapping(Format);
   }
 
+
   const DxvkFormatInfo* D3D9DeviceEx::UnsupportedFormatInfo(
     D3D9Format            Format) const {
     return m_adapter->GetUnsupportedFormatInfo(Format);
   }
 
+
   bool D3D9DeviceEx::WaitForResource(
     const DxvkPagedResource&                Resource,
           uint64_t                          SequenceNumber,
diff --git a/src/d3d9/d3d9_device.h b/src/d3d9/d3d9_device.h
index c9c75a8b516..ca7362173df 100644
--- a/src/d3d9/d3d9_device.h
+++ b/src/d3d9/d3d9_device.h
@@ -1020,6 +1020,13 @@ namespace dxvk {
      */
     void RemoveMappedTexture(D3D9CommonTexture* pTexture);
 
+    /**
+     * \brief Returns whether the device is currently recording a StateBlock
+     */
+    bool ShouldRecord() const {
+      return m_recorder != nullptr;
+    }
+
     bool IsD3D8Compatible() const {
       return m_isD3D8Compatible;
     }
@@ -1160,11 +1167,6 @@ namespace dxvk {
      */
     void WaitStagingBuffer();
 
-    /**
-     * \brief Returns whether the device is currently recording a StateBlock
-     */
-    inline bool ShouldRecord();
-
     HRESULT               CreateShaderModule(
             D3D9CommonShader*     pShaderModule,
             uint32_t*             pLength,
diff --git a/src/d3d9/d3d9_stateblock.cpp b/src/d3d9/d3d9_stateblock.cpp
index 85154e209c9..9998fa94d22 100644
--- a/src/d3d9/d3d9_stateblock.cpp
+++ b/src/d3d9/d3d9_stateblock.cpp
@@ -46,6 +46,10 @@ namespace dxvk {
 
 
   HRESULT STDMETHODCALLTYPE D3D9StateBlock::Capture() {
+    // A state block can't capture state while another is being recorded.
+    if (unlikely(m_parent->ShouldRecord()))
+      return D3DERR_INVALIDCALL;
+
     if (m_captures.flags.test(D3D9CapturedStateFlag::VertexDecl))
       SetVertexDeclaration(m_deviceState->vertexDecl.ptr());
 
@@ -56,13 +60,14 @@ namespace dxvk {
 
 
   HRESULT STDMETHODCALLTYPE D3D9StateBlock::Apply() {
-    m_applying = true;
+    // A state block can't be applied while another is being recorded.
+    if (unlikely(m_parent->ShouldRecord()))
+      return D3DERR_INVALIDCALL;
 
     if (m_captures.flags.test(D3D9CapturedStateFlag::VertexDecl) && m_state.vertexDecl != nullptr)
       m_parent->SetVertexDeclaration(m_state.vertexDecl.ptr());
 
     ApplyOrCapture<D3D9StateFunction::Apply, false>();
-    m_applying = false;
 
     return D3D_OK;
   }
diff --git a/src/d3d9/d3d9_stateblock.h b/src/d3d9/d3d9_stateblock.h
index e21cc26df87..74284d4f7ad 100644
--- a/src/d3d9/d3d9_stateblock.h
+++ b/src/d3d9/d3d9_stateblock.h
@@ -388,10 +388,6 @@ namespace dxvk {
     HRESULT SetVertexBoolBitfield(uint32_t idx, uint32_t mask, uint32_t bits);
     HRESULT SetPixelBoolBitfield (uint32_t idx, uint32_t mask, uint32_t bits);
 
-    inline bool IsApplying() {
-      return m_applying;
-    }
-
   private:
 
     void CapturePixelRenderStates();
@@ -407,9 +403,7 @@ namespace dxvk {
     D3D9CapturableState  m_state;
     D3D9StateCaptures    m_captures;
 
-    D3D9DeviceState* m_deviceState;
-
-    bool                 m_applying = false;
+    D3D9DeviceState*     m_deviceState;
 
   };