From 807a398126f98030c6270bb6caaafe21af943327 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Wed, 15 Feb 2023 17:29:22 +0000
Subject: [PATCH] Handle COMDAT entries in DXIL

---
 .../driver/shaders/dxil/dxil_bytecode.cpp     | 15 +++++++++++
 renderdoc/driver/shaders/dxil/dxil_bytecode.h |  2 ++
 .../shaders/dxil/dxil_bytecode_editor.cpp     | 11 ++++++--
 .../driver/shaders/dxil/dxil_disassemble.cpp  | 26 +++++++++++++++++++
 renderdoc/driver/shaders/dxil/llvm_common.h   |  1 +
 5 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/renderdoc/driver/shaders/dxil/dxil_bytecode.cpp b/renderdoc/driver/shaders/dxil/dxil_bytecode.cpp
index 46574c672..c09c4fb96 100644
--- a/renderdoc/driver/shaders/dxil/dxil_bytecode.cpp
+++ b/renderdoc/driver/shaders/dxil/dxil_bytecode.cpp
@@ -490,7 +490,17 @@ Program::Program(const byte *bytes, size_t length) : alloc(32 * 1024)
 
         // ignore rest of properties, assert that if present they are 0
         for(size_t p = 6; p < rootchild.ops.size(); p++)
+        {
+          // 12, if present, is the comdat index
+          if(p == 12 && rootchild.ops[p] > 0)
+          {
+            RDCASSERT(rootchild.ops[p] - 1 < m_Comdats.size(), rootchild.ops[p], m_Comdats.size());
+            f->comdatIdx = uint32_t(rootchild.ops[p] - 1);
+            continue;
+          }
+
           RDCASSERT(rootchild.ops[p] == 0, p, rootchild.ops[p]);
+        }
 
         if(!f->external)
           functionDecls.push_back(m_Functions.size());
@@ -517,6 +527,11 @@ Program::Program(const byte *bytes, size_t length) : alloc(32 * 1024)
       {
         m_Sections.push_back(rootchild.getString(0));
       }
+      else if(IS_KNOWN(rootchild.id, ModuleRecord::COMDAT))
+      {
+        // can ignore the length for now, it's implicit anyway as there's nothing after the string
+        m_Comdats.push_back({rootchild.ops[0], rootchild.getString(2)});
+      }
       else
       {
         RDCERR("Unknown record ID %u encountered at module scope", rootchild.id);
diff --git a/renderdoc/driver/shaders/dxil/dxil_bytecode.h b/renderdoc/driver/shaders/dxil/dxil_bytecode.h
index 39e855e1a..1767b3c75 100644
--- a/renderdoc/driver/shaders/dxil/dxil_bytecode.h
+++ b/renderdoc/driver/shaders/dxil/dxil_bytecode.h
@@ -998,6 +998,7 @@ struct Function : public Value
   bool external = false;
   bool internalLinkage = false;
   bool sortedSymtab = true;
+  uint32_t comdatIdx = ~0U;
   const AttributeSet *attrs = NULL;
 
   uint64_t align = 0;
@@ -1093,6 +1094,7 @@ protected:
   rdcarray<Function *> m_Functions;
   rdcarray<Alias *> m_Aliases;
   rdcarray<rdcstr> m_Sections;
+  rdcarray<rdcpair<uint64_t, rdcstr>> m_Comdats;
   uint32_t m_directHeapAccessCount = 0;
 
   rdcarray<rdcstr> m_Kinds;
diff --git a/renderdoc/driver/shaders/dxil/dxil_bytecode_editor.cpp b/renderdoc/driver/shaders/dxil/dxil_bytecode_editor.cpp
index 2c1b3acdc..6ab9a0324 100644
--- a/renderdoc/driver/shaders/dxil/dxil_bytecode_editor.cpp
+++ b/renderdoc/driver/shaders/dxil/dxil_bytecode_editor.cpp
@@ -658,7 +658,14 @@ bytebuf ProgramEditor::EncodeProgram()
     writer.EndBlock();
   }
 
-  // COMDAT would be next, but we don't read these (DXIL seems not to use them...?)
+  for(size_t i = 0; i < m_Comdats.size(); i++)
+  {
+    rdcarray<uint64_t> vals;
+    vals.push_back(m_Comdats[i].first);
+    for(char c : m_Comdats[i].second)
+      vals.push_back(c);
+    writer.Record(LLVMBC::ModuleRecord::COMDAT, vals);
+  }
 
   if(!m_Triple.empty())
     writer.Record(LLVMBC::ModuleRecord::TRIPLE, m_Triple);
@@ -761,7 +768,7 @@ bytebuf ProgramEditor::EncodeProgram()
                       // dllstorageclass
                       0U,
                       // comdat
-                      0U,
+                      uint64_t(f.comdatIdx != ~0U ? 1U + f.comdatIdx : 0U),
                       // prefixdata
                       0U,
                       // personality
diff --git a/renderdoc/driver/shaders/dxil/dxil_disassemble.cpp b/renderdoc/driver/shaders/dxil/dxil_disassemble.cpp
index b3f7de9c8..39c4e860f 100644
--- a/renderdoc/driver/shaders/dxil/dxil_disassemble.cpp
+++ b/renderdoc/driver/shaders/dxil/dxil_disassemble.cpp
@@ -313,6 +313,28 @@ void Program::MakeDisassemblyString()
 
   int instructionLine = 6;
 
+  for(const rdcpair<uint64_t, rdcstr> &comdat : m_Comdats)
+  {
+    rdcstr type = "unknown";
+    switch(comdat.first)
+    {
+      case 1: type = "any"; break;
+      case 2: type = "exactmatch"; break;
+      case 3: type = "largest"; break;
+      case 4: type = "noduplicates"; break;
+      case 5: type = "samesize"; break;
+    }
+    m_Disassembly += StringFormat::Fmt("$%s = comdat %s\n",
+                                       escapeStringIfNeeded(comdat.second).c_str(), type.c_str());
+    instructionLine++;
+  }
+
+  if(!m_Comdats.empty())
+  {
+    m_Disassembly += "\n";
+    instructionLine++;
+  }
+
   LLVMOrderAccumulator accum;
   accum.processGlobals(this);
 
@@ -551,6 +573,10 @@ void Program::MakeDisassemblyString()
     m_Disassembly +=
         func.type->declFunction("@" + escapeStringIfNeeded(func.name), func.args, func.attrs);
 
+    if(func.comdatIdx < m_Comdats.size())
+      m_Disassembly += StringFormat::Fmt(
+          " comdat($%s)", escapeStringIfNeeded(m_Comdats[func.comdatIdx].second).c_str());
+
     if(func.attrs && func.attrs->functionSlot)
       m_Disassembly += StringFormat::Fmt(" #%u", funcAttrGroups.indexOf(func.attrs->functionSlot));
 
diff --git a/renderdoc/driver/shaders/dxil/llvm_common.h b/renderdoc/driver/shaders/dxil/llvm_common.h
index 5ec8be68b..1880b0250 100644
--- a/renderdoc/driver/shaders/dxil/llvm_common.h
+++ b/renderdoc/driver/shaders/dxil/llvm_common.h
@@ -96,6 +96,7 @@ enum class ModuleRecord : uint32_t
   GLOBALVAR = 7,
   FUNCTION = 8,
   ALIAS = 9,
+  COMDAT = 12,
 };
 
 enum class ConstantsRecord : uint32_t