From b33e5e018a11df8c7c22d8a96fb3a15311bf000a Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Thu, 20 Dec 2018 13:32:08 +0000
Subject: [PATCH] Add tracking of which extensions are in the module

---
 .../driver/shaders/spirv/spirv_editor.cpp     | 34 +++++++++++++++++++
 renderdoc/driver/shaders/spirv/spirv_editor.h |  2 ++
 2 files changed, 36 insertions(+)

diff --git a/renderdoc/driver/shaders/spirv/spirv_editor.cpp b/renderdoc/driver/shaders/spirv/spirv_editor.cpp
index 57385ebf1..8351d582b 100644
--- a/renderdoc/driver/shaders/spirv/spirv_editor.cpp
+++ b/renderdoc/driver/shaders/spirv/spirv_editor.cpp
@@ -373,6 +373,30 @@ void SPIRVEditor::AddCapability(spv::Capability cap)
   addWords(FirstRealWord, op.size());
 }
 
+void SPIRVEditor::AddExtension(const std::string &extension)
+{
+  // don't add duplicate extensions
+  if(extensions.find(extension) != extensions.end())
+    return;
+
+  // start at the beginning
+  SPIRVIterator it(spirv, FirstRealWord);
+
+  // skip past any capabilities
+  while(it.opcode() == spv::OpCapability)
+    it++;
+
+  // insert the extension instruction
+  size_t sz = extension.size();
+  std::vector<uint32_t> uintName((sz / 4) + 1);
+  memcpy(&uintName[0], extension.c_str(), sz);
+
+  SPIRVOperation op(spv::OpExtension, uintName);
+  spirv.insert(spirv.begin() + it.offset, op.begin(), op.end());
+  RegisterOp(it);
+  addWords(it.offset, op.size());
+}
+
 SPIRVId SPIRVEditor::ImportExtInst(const char *setname)
 {
   SPIRVId ret = extSets[setname];
@@ -578,6 +602,11 @@ void SPIRVEditor::RegisterOp(SPIRVIterator it)
   {
     capabilities.insert((spv::Capability)it.word(1));
   }
+  else if(opcode == spv::OpExtension)
+  {
+    const char *name = (const char *)&it.word(1);
+    extensions.insert(name);
+  }
   else if(opcode == spv::OpExtInstImport)
   {
     SPIRVId id = it.word(1);
@@ -716,6 +745,11 @@ void SPIRVEditor::UnregisterOp(SPIRVIterator it)
   {
     capabilities.erase((spv::Capability)it.word(1));
   }
+  else if(opcode == spv::OpExtension)
+  {
+    const char *name = (const char *)&it.word(1);
+    extensions.erase(name);
+  }
   else if(opcode == spv::OpExtInstImport)
   {
     const char *name = (const char *)&it.word(2);
diff --git a/renderdoc/driver/shaders/spirv/spirv_editor.h b/renderdoc/driver/shaders/spirv/spirv_editor.h
index a41bfe9db..2fc337d4a 100644
--- a/renderdoc/driver/shaders/spirv/spirv_editor.h
+++ b/renderdoc/driver/shaders/spirv/spirv_editor.h
@@ -384,6 +384,7 @@ public:
   void SetName(uint32_t id, const char *name);
   void AddDecoration(const SPIRVOperation &op);
   void AddCapability(spv::Capability cap);
+  void AddExtension(const std::string &extension);
   SPIRVId ImportExtInst(const char *setname);
   SPIRVId AddType(const SPIRVOperation &op);
   SPIRVId AddVariable(const SPIRVOperation &op);
@@ -487,6 +488,7 @@ private:
 
   std::vector<SPIRVEntry> entries;
   std::vector<SPIRVId> functions;
+  std::set<std::string> extensions;
   std::set<spv::Capability> capabilities;
 
   std::map<std::string, SPIRVId> extSets;