From 5d3f966cb89791d8a7a9851c51a1daa7ca87d9d3 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Thu, 14 Nov 2024 15:56:35 +0000
Subject: [PATCH] Check BLASs comprising a TLAS are actually BLASs and are
 older (by id)

* If the same AS is re-used for a new BLAS, consider the TLAS invalid/stale and
  0 out its address. Similarly if a TLAS somehow ends up there (though it should
  always fail the age check first)
---
 renderdoc/driver/d3d12/d3d12_initstate.cpp | 26 ++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/renderdoc/driver/d3d12/d3d12_initstate.cpp b/renderdoc/driver/d3d12/d3d12_initstate.cpp
index 889e2e395..a9b2f32f4 100644
--- a/renderdoc/driver/d3d12/d3d12_initstate.cpp
+++ b/renderdoc/driver/d3d12/d3d12_initstate.cpp
@@ -1447,10 +1447,11 @@ bool D3D12ResourceManager::Serialise_InitialState(SerialiserType &ser, ResourceI
               UINT64 blasOffs;
               m_Device->GetResIDFromOrigAddr(instances[i].AccelerationStructure, blasId, blasOffs);
 
-              WrappedID3D12Resource *blas = GetLiveAs<WrappedID3D12Resource>(blasId);
+              WrappedID3D12Resource *blasASB = GetLiveAs<WrappedID3D12Resource>(blasId);
 
-              D3D12AccelerationStructure *as = NULL;
-              if(blasId == ResourceId() || blas == NULL || !blas->GetAccStructIfExist(blasOffs, &as))
+              D3D12AccelerationStructure *blasCheck = NULL;
+              if(blasId == ResourceId() || blasASB == NULL ||
+                 !blasASB->GetAccStructIfExist(blasOffs, &blasCheck))
               {
                 RDCWARN(
                     "  %u: BLAS referenced by TLAS is not available on replay - possibly stale "
@@ -1460,7 +1461,24 @@ bool D3D12ResourceManager::Serialise_InitialState(SerialiserType &ser, ResourceI
                 continue;
               }
 
-              instances[i].AccelerationStructure = blas->GetGPUVirtualAddress() + blasOffs;
+              if(id < GetOriginalID(blasCheck->GetResourceID()))
+              {
+                RDCWARN("  %u: BLAS referenced by TLAS is newer than TLAS - possibly stale TLAS", i);
+                instances[i].AccelerationStructure = 0;
+                continue;
+              }
+
+              if(blasCheck->Type() != D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL)
+              {
+                RDCWARN("  %u: BLAS is not of correct type - possibly stale TLAS", i);
+                instances[i].AccelerationStructure = 0;
+                continue;
+              }
+
+              RDCASSERTEQUAL(blasCheck->GetVirtualAddress(),
+                             blasASB->GetGPUVirtualAddress() + blasOffs);
+
+              instances[i].AccelerationStructure = blasASB->GetGPUVirtualAddress() + blasOffs;
             }
 
             void *upload = mappedBuffer->Map();