From 0f541ea6ee276bbe050de5cbf3f5422d8b0411a1 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Mon, 8 Jan 2018 15:54:32 +0000
Subject: [PATCH] Add function and void type tracking, and helper for declaring
 structs

---
 .../driver/shaders/spirv/spirv_editor.cpp     | 68 +++++++++++++++++--
 renderdoc/driver/shaders/spirv/spirv_editor.h | 30 ++++++++
 2 files changed, 92 insertions(+), 6 deletions(-)

diff --git a/renderdoc/driver/shaders/spirv/spirv_editor.cpp b/renderdoc/driver/shaders/spirv/spirv_editor.cpp
index d28812bb8..d22088859 100644
--- a/renderdoc/driver/shaders/spirv/spirv_editor.cpp
+++ b/renderdoc/driver/shaders/spirv/spirv_editor.cpp
@@ -184,6 +184,25 @@ SPIRVEditor::SPIRVEditor(std::vector<uint32_t> &spirvWords) : spirv(spirvWords)
 
       pointerTypes[SPIRVPointer(it.word(3), (spv::StorageClass)it.word(2))] = id;
     }
+
+    if(opcode == spv::OpTypeFunction)
+    {
+      uint32_t id = it.word(1);
+      idOffsets[id] = it.offset;
+
+      std::vector<SPIRVId> args;
+
+      for(size_t i = 3; i < it.size(); i++)
+        args.push_back(it.word(i));
+
+      functionTypes[SPIRVFunction(it.word(2), args)] = id;
+    }
+
+    if(opcode == spv::OpTypeVoid)
+    {
+      voidType = it.word(1);
+      idOffsets[voidType] = it.offset;
+    }
   }
 }
 
@@ -294,6 +313,17 @@ SPIRVIterator SPIRVEditor::GetTypeInstructions()
   return SPIRVIterator(spirv, typeVarSection.startOffset);
 }
 
+SPIRVId SPIRVEditor::DeclareType(const SPIRVVoid &)
+{
+  if(voidType)
+    return voidType;
+
+  voidType = MakeId();
+  AddType(SPIRVOperation(spv::OpTypeVoid, {voidType}));
+
+  return voidType;
+}
+
 SPIRVId SPIRVEditor::DeclareType(const SPIRVScalar &scalar)
 {
   auto it = scalarTypes.lower_bound(scalar);
@@ -316,8 +346,7 @@ SPIRVId SPIRVEditor::DeclareType(const SPIRVVector &vector)
     return it->second;
 
   SPIRVId id = MakeId();
-  SPIRVOperation decl(spv::OpTypeVector, {id, DeclareType(vector.scalar), vector.count});
-  AddType(decl);
+  AddType(SPIRVOperation(spv::OpTypeVector, {id, DeclareType(vector.scalar), vector.count}));
 
   vectorTypes.insert(it, std::make_pair(vector, id));
 
@@ -331,8 +360,7 @@ SPIRVId SPIRVEditor::DeclareType(const SPIRVMatrix &matrix)
     return it->second;
 
   SPIRVId id = MakeId();
-  SPIRVOperation decl(spv::OpTypeVector, {id, DeclareType(matrix.vector), matrix.count});
-  AddType(decl);
+  AddType(SPIRVOperation(spv::OpTypeVector, {id, DeclareType(matrix.vector), matrix.count}));
 
   matrixTypes.insert(it, std::make_pair(matrix, id));
 
@@ -346,14 +374,42 @@ SPIRVId SPIRVEditor::DeclareType(const SPIRVPointer &pointer)
     return it->second;
 
   SPIRVId id = MakeId();
-  SPIRVOperation decl(spv::OpTypePointer, {id, (uint32_t)pointer.storage, pointer.baseId});
-  AddType(decl);
+  AddType(SPIRVOperation(spv::OpTypePointer, {id, (uint32_t)pointer.storage, pointer.baseId}));
 
   pointerTypes.insert(it, std::make_pair(pointer, id));
 
   return id;
 }
 
+SPIRVId SPIRVEditor::DeclareType(const SPIRVFunction &func)
+{
+  auto it = functionTypes.lower_bound(func);
+  if(it != functionTypes.end() && it->first == func)
+    return it->second;
+
+  SPIRVId id = MakeId();
+
+  std::vector<uint32_t> words;
+
+  words.push_back(id);
+  words.push_back(func.returnId);
+  for(SPIRVId id : func.argumentIds)
+    words.push_back(id);
+
+  AddType(SPIRVOperation(spv::OpTypeFunction, words));
+
+  functionTypes.insert(it, std::make_pair(func, id));
+
+  return id;
+}
+
+SPIRVId SPIRVEditor::DeclareStructType(std::vector<uint32_t> members)
+{
+  SPIRVId typeId = MakeId();
+  members.insert(members.begin(), typeId);
+  AddType(SPIRVOperation(spv::OpTypeStruct, members));
+  return typeId;
+}
 void SPIRVEditor::addWords(size_t offs, int32_t num)
 {
   // look through every section, any that are >= this point, adjust the offsets
diff --git a/renderdoc/driver/shaders/spirv/spirv_editor.h b/renderdoc/driver/shaders/spirv/spirv_editor.h
index eebe6873b..ce901fcbd 100644
--- a/renderdoc/driver/shaders/spirv/spirv_editor.h
+++ b/renderdoc/driver/shaders/spirv/spirv_editor.h
@@ -164,6 +164,10 @@ struct SPIRVEntry
   std::string name;
 };
 
+struct SPIRVVoid
+{
+};
+
 struct SPIRVScalar
 {
   constexpr SPIRVScalar(spv::Op t, uint32_t w, bool s) : type(t), width(w), signedness(s) {}
@@ -274,6 +278,26 @@ struct SPIRVPointer
   }
 };
 
+struct SPIRVFunction
+{
+  SPIRVFunction(SPIRVId ret, const std::vector<SPIRVId> &args) : returnId(ret), argumentIds(args) {}
+  SPIRVId returnId;
+  std::vector<SPIRVId> argumentIds;
+
+  bool operator<(const SPIRVFunction &o) const
+  {
+    if(returnId != o.returnId)
+      return returnId < o.returnId;
+    return argumentIds < o.argumentIds;
+  }
+
+  bool operator!=(const SPIRVFunction &o) const { return !operator==(o); }
+  bool operator==(const SPIRVFunction &o) const
+  {
+    return returnId == o.returnId && argumentIds == o.argumentIds;
+  }
+};
+
 class SPIRVEditor
 {
 public:
@@ -299,10 +323,14 @@ public:
 
   // fetches the id of this type. If it exists already the old ID will be returned, otherwise it
   // will be declared and the new ID returned
+  SPIRVId DeclareType(const SPIRVVoid &);
   SPIRVId DeclareType(const SPIRVScalar &scalar);
   SPIRVId DeclareType(const SPIRVVector &vector);
   SPIRVId DeclareType(const SPIRVMatrix &matrix);
   SPIRVId DeclareType(const SPIRVPointer &pointer);
+  SPIRVId DeclareType(const SPIRVFunction &func);
+
+  SPIRVId DeclareStructType(std::vector<uint32_t> members);
 
   // simple properties that are public.
   struct
@@ -341,6 +369,8 @@ private:
   std::map<SPIRVVector, SPIRVId> vectorTypes;
   std::map<SPIRVMatrix, SPIRVId> matrixTypes;
   std::map<SPIRVPointer, SPIRVId> pointerTypes;
+  std::map<SPIRVFunction, SPIRVId> functionTypes;
+  SPIRVId voidType;
 
   std::vector<uint32_t> &spirv;
 };