From 0b9ed7bd8e13b322b789d7bb7edf12b24fd61a34 Mon Sep 17 00:00:00 2001
From: Jake Turner <jake@evansturner.co.uk>
Date: Tue, 2 Jul 2024 14:16:34 +0100
Subject: [PATCH] DXIL disassembly parsing track labels in functions and blocks

Per Function compute:
labelToBlockIndex container : to go from a label to its block index
blockIndexToLabel container : to go from block index to its label

For each Block store the starting instruction index
---
 renderdoc/driver/shaders/dxil/dxil_bytecode.h |  3 +++
 .../driver/shaders/dxil/dxil_disassemble.cpp  | 24 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/renderdoc/driver/shaders/dxil/dxil_bytecode.h b/renderdoc/driver/shaders/dxil/dxil_bytecode.h
index 5d409ec7c..5035bfead 100644
--- a/renderdoc/driver/shaders/dxil/dxil_bytecode.h
+++ b/renderdoc/driver/shaders/dxil/dxil_bytecode.h
@@ -1327,6 +1327,7 @@ struct Block : public ForwardReferencableValue<Block>
   rdcinflexiblestr name;
   rdcarray<const Block *> preds;
   uint32_t slot = ~0U;
+  uint32_t startInstructionIdx = ~0U;
 };
 
 struct UselistEntry
@@ -1360,6 +1361,8 @@ struct Function : public Value
   rdcarray<Block *> blocks;
 
   rdcarray<UselistEntry> uselist;
+  std::map<rdcstr, uint32_t> labelToBlockIndex;
+  rdcarray<rdcstr> blockIndexToLabel;
 
   AttachedMetadata attachedMeta;
 };
diff --git a/renderdoc/driver/shaders/dxil/dxil_disassemble.cpp b/renderdoc/driver/shaders/dxil/dxil_disassemble.cpp
index 5e4fda341..11fa458e5 100644
--- a/renderdoc/driver/shaders/dxil/dxil_disassemble.cpp
+++ b/renderdoc/driver/shaders/dxil/dxil_disassemble.cpp
@@ -1065,7 +1065,7 @@ void Program::SettleIDs()
   {
     m_Accum.processFunction(m_Functions[i]);
 
-    const Function &func = *m_Functions[i];
+    Function &func = *m_Functions[i];
 
     auto argMetaSlot = [this, &metaSlots, &nextMetaSlot](const Value *v) {
       if(const Metadata *meta = cast<Metadata>(v))
@@ -1094,6 +1094,15 @@ void Program::SettleIDs()
 
     if(!func.external)
     {
+      func.labelToBlockIndex.clear();
+      func.blockIndexToLabel.clear();
+
+      size_t curBlock = 0;
+      RDCASSERT(!func.blocks.empty());
+      rdcstr labelName = StringFormat::Fmt("_label%u", func.blocks[curBlock]->slot);
+      func.labelToBlockIndex[labelName] = (uint32_t)curBlock;
+      func.blocks[curBlock]->startInstructionIdx = 0;
+      func.blockIndexToLabel.push_back(labelName);
       for(size_t funcIdx = 0; funcIdx < func.instructions.size(); funcIdx++)
       {
         Instruction &inst = *func.instructions[funcIdx];
@@ -1234,6 +1243,19 @@ void Program::SettleIDs()
           for(size_t m = 0; m < attachedMeta.size(); m++)
             AssignMetaSlot(metaSlots, nextMetaSlot, attachedMeta[m].second);
         }
+        // add labels before the final instruction in the function
+        if(funcIdx < func.instructions.size() - 1)
+        {
+          if(inst.op == Operation::Branch || inst.op == Operation::Unreachable ||
+             inst.op == Operation::Switch || inst.op == Operation::Ret)
+          {
+            curBlock++;
+            labelName = StringFormat::Fmt("_label%u", func.blocks[curBlock]->slot);
+            func.labelToBlockIndex[labelName] = (uint32_t)curBlock;
+            func.blocks[curBlock]->startInstructionIdx = (uint32_t)(funcIdx + 1);
+            func.blockIndexToLabel.push_back(labelName);
+          }
+        }
       }
     }
     m_Accum.exitFunction();