From 20b1257411e93c1668cde40463afab1cf2255ba1 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Mon, 1 Nov 2021 13:50:42 +0000
Subject: [PATCH] Make sure that function constant members are fixup'd when
 patching

* We also add some extra debug asserts after the fixup path to ensure that all
  pointers are in the correct array (if we miss passing the old/new function for
  a value fixup like this, it could point to the old function and be silently
  ignored since that's conditional).
* A unit test harness is added too for convenience of iterating on a single
  shader's annotation.
---
 .../driver/d3d12/d3d12_shader_feedback.cpp    | 52 +++++++++++++++++++
 .../shaders/dxil/dxil_bytecode_editor.cpp     | 26 +++++++++-
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/renderdoc/driver/d3d12/d3d12_shader_feedback.cpp b/renderdoc/driver/d3d12/d3d12_shader_feedback.cpp
index 9946fa362..f6a27551d 100644
--- a/renderdoc/driver/d3d12/d3d12_shader_feedback.cpp
+++ b/renderdoc/driver/d3d12/d3d12_shader_feedback.cpp
@@ -1235,3 +1235,55 @@ void D3D12Replay::ClearFeedbackCache()
 {
   m_BindlessFeedback.Usage.clear();
 }
+
+#if ENABLED(ENABLE_UNIT_TESTS) && 0
+
+#include "catch/catch.hpp"
+
+TEST_CASE("DO NOT COMMIT - convenience test", "[dxbc]")
+{
+  // this test loads a file from disk and does a no-op edit pass on it then an annotation pass on
+  // it. Useful for when you are iterating on a shader and don't want to have to load a whole
+  // capture.
+  bytebuf buf;
+  FileIO::ReadAll("/path/to/container_file.dxbc", buf);
+  bytebuf editedBlob;
+
+  {
+    DXBC::DXBCContainer dxbc(buf, rdcstr(), GraphicsAPI::D3D12, ~0U, ~0U);
+
+    DXIL::ProgramEditor editor(&dxbc, 1234, editedBlob);
+  }
+
+  {
+    DXBC::DXBCContainer container(editedBlob, rdcstr(), GraphicsAPI::D3D12, ~0U, ~0U);
+
+    rdcstr disasm = container.GetDisassembly();
+
+    RDCLOG("no edits - %s", disasm.c_str());
+  }
+
+  {
+    WrappedID3D12Device device(NULL, D3D12InitParams(), false);
+
+    D3D12_SHADER_BYTECODE desc;
+    desc.BytecodeLength = buf.size();
+    desc.pShaderBytecode = buf.data();
+
+    WrappedID3D12PipelineState::ShaderEntry shad(desc, &device);
+
+    std::map<D3D12FeedbackKey, D3D12FeedbackSlot> slots;
+    uint32_t numSlots = 4;
+    AddArraySlots(&shad, 123456, 1000000, slots, numSlots, editedBlob, desc);
+  }
+
+  {
+    DXBC::DXBCContainer container(editedBlob, rdcstr(), GraphicsAPI::D3D12, ~0U, ~0U);
+
+    rdcstr disasm = container.GetDisassembly();
+
+    RDCLOG("annotated - %s", disasm.c_str());
+  }
+}
+
+#endif
diff --git a/renderdoc/driver/shaders/dxil/dxil_bytecode_editor.cpp b/renderdoc/driver/shaders/dxil/dxil_bytecode_editor.cpp
index b6b618cc3..04265644a 100644
--- a/renderdoc/driver/shaders/dxil/dxil_bytecode_editor.cpp
+++ b/renderdoc/driver/shaders/dxil/dxil_bytecode_editor.cpp
@@ -215,6 +215,10 @@ void ProgramEditor::Fixup(Function *&f)
       }
     }
   }
+
+#if ENABLED(RDOC_DEVEL)
+  RDCASSERT(IN_ARRAY(m_Functions, f));
+#endif
 }
 
 void ProgramEditor::Fixup(Type *&t)
@@ -229,6 +233,10 @@ void ProgramEditor::Fixup(Type *&t)
     if(result)
       t = result;
   }
+
+#if ENABLED(RDOC_DEVEL)
+  RDCASSERT(IN_ARRAY(m_Types, t));
+#endif
 }
 
 void ProgramEditor::Fixup(Block *&b, Function *oldf, Function *newf)
@@ -240,6 +248,10 @@ void ProgramEditor::Fixup(Block *&b, Function *oldf, Function *newf)
     if(result)
       b = result;
   }
+
+#if ENABLED(RDOC_DEVEL)
+  RDCASSERT(IN_ARRAY(newf->blocks, b));
+#endif
 }
 
 // this variant fixes up the pointer itself only, but doesn't recurse. We don't recurse because we
@@ -264,6 +276,10 @@ void ProgramEditor::Fixup(Instruction *&i, Function *oldf, Function *newf)
     if(result)
       i = result;
   }
+
+#if ENABLED(RDOC_DEVEL)
+  RDCASSERT(IN_ARRAY(newf->args, i) || IN_ARRAY(newf->instructions, i));
+#endif
 }
 
 void ProgramEditor::Fixup(Constant *&c, Function *oldf, Function *newf)
@@ -286,6 +302,10 @@ void ProgramEditor::Fixup(Constant *&c, Function *oldf, Function *newf)
     if(result)
       c = result;
   }
+
+#if ENABLED(RDOC_DEVEL)
+  RDCASSERT(IN_ARRAY(m_Constants, c) || (newf && IN_ARRAY(newf->constants, c)));
+#endif
 }
 
 void ProgramEditor::Fixup(Metadata *&m, Function *oldf, Function *newf)
@@ -308,6 +328,10 @@ void ProgramEditor::Fixup(Metadata *&m, Function *oldf, Function *newf)
     if(result)
       m = result;
   }
+
+#if ENABLED(RDOC_DEVEL)
+  RDCASSERT(IN_ARRAY(m_Metadata, m) || (newf && IN_ARRAY(newf->metadata, m)));
+#endif
 }
 
 ProgramEditor::~ProgramEditor()
@@ -407,7 +431,7 @@ ProgramEditor::~ProgramEditor()
       Fixup(c.type);
       Fixup(c.inner);
       for(size_t i = 0; i < c.members.size(); i++)
-        Fixup(c.members[i]);
+        Fixup(c.members[i], &oldf, &f);
     }
 
     for(Value &v : f.values)