From 2b779038afb08519282ca0d939c4bf6801350ad9 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Tue, 10 Oct 2023 11:51:43 +0100
Subject: [PATCH] Add a SPIR-V editor helper to add or re-use a builtin input

---
 .../driver/shaders/spirv/spirv_editor.cpp     | 84 +++++++++++++++++++
 renderdoc/driver/shaders/spirv/spirv_editor.h | 16 ++++
 2 files changed, 100 insertions(+)

diff --git a/renderdoc/driver/shaders/spirv/spirv_editor.cpp b/renderdoc/driver/shaders/spirv/spirv_editor.cpp
index d9ac46bef..19b5367c6 100644
--- a/renderdoc/driver/shaders/spirv/spirv_editor.cpp
+++ b/renderdoc/driver/shaders/spirv/spirv_editor.cpp
@@ -403,6 +403,54 @@ Iter Editor::GetEntry(Id id)
   return Iter();
 }
 
+rdcpair<Id, Id> Editor::AddBuiltinInputLoad(OperationList &ops, ShaderStage stage, BuiltIn builtin,
+                                            Id type)
+{
+  BuiltinInputData &data = builtinInputs[builtin];
+
+  Id ptrType = DeclareType(Pointer(type, StorageClass::Input));
+
+  if(data.variable == Id())
+  {
+    const DataType &dataType = dataTypes[type];
+
+    Id var = AddVariable(OpVariable(ptrType, MakeId(), StorageClass::Input));
+    AddDecoration(OpDecorate(var, DecorationParam<Decoration::BuiltIn>(builtin)));
+
+    // Fragment shader inputs that are signed or unsigned integers, integer vectors, or any
+    // double-precision floating-point type must be decorated with Flat.
+    if(dataType.scalar().type != Op::TypeFloat && stage == ShaderStage::Pixel)
+      AddDecoration(OpDecorate(var, Decoration::Flat));
+
+    builtinInputs[builtin] = {var, ptrType, {}};
+
+    return {ops.add(OpLoad(type, MakeId(), var)), var};
+  }
+
+  Id varType = dataTypes[data.type].InnerType();
+
+  Id ret;
+
+  if(data.chain.empty())
+  {
+    ret = ops.add(OpLoad(varType, MakeId(), data.variable));
+  }
+  else
+  {
+    rdcarray<rdcspv::Id> chain;
+    for(uint32_t accessIdx : data.chain)
+      chain.push_back(AddConstantImmediate<uint32_t>(accessIdx));
+    Id subElement = ops.add(OpAccessChain(ptrType, MakeId(), data.variable, chain));
+
+    ret = ops.add(rdcspv::OpLoad(varType, MakeId(), subElement));
+  }
+
+  if(varType != type)
+    ret = ops.add(rdcspv::OpBitcast(type, MakeId(), ret));
+
+  return {ret, Id()};
+}
+
 Id Editor::DeclareStructType(const rdcarray<Id> &members)
 {
   Id typeId = MakeId();
@@ -544,6 +592,42 @@ void Editor::UnregisterOp(Iter it)
   }
 }
 
+void Editor::RegisterBuiltinMembers(rdcspv::Id baseId, const rdcarray<uint32_t> chainSoFar,
+                                    const DataType *type)
+{
+  uint32_t i = 0;
+  for(const DataType::Child &member : type->children)
+  {
+    if(member.decorations.flags & Decorations::HasBuiltIn)
+    {
+      rdcarray<uint32_t> chain = chainSoFar;
+      chain.push_back(i);
+      builtinInputs[member.decorations.builtIn] = {baseId, member.type, chain};
+    }
+    i++;
+  }
+}
+
+void Editor::PostParse()
+{
+  Processor::PostParse();
+
+  for(const Variable &v : globals)
+  {
+    if(v.storage == StorageClass::Input)
+    {
+      if(decorations[v.id].flags & Decorations::HasBuiltIn)
+      {
+        builtinInputs[decorations[v.id].builtIn] = {v.id, v.type, {}};
+      }
+      else
+      {
+        RegisterBuiltinMembers(v.id, {}, &dataTypes[v.type]);
+      }
+    }
+  }
+}
+
 void Editor::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 70037d928..4246fc599 100644
--- a/renderdoc/driver/shaders/spirv/spirv_editor.h
+++ b/renderdoc/driver/shaders/spirv/spirv_editor.h
@@ -174,6 +174,9 @@ public:
     return it->second;
   }
 
+  rdcpair<Id, Id> AddBuiltinInputLoad(OperationList &ops, ShaderStage stage, BuiltIn builtin,
+                                      Id type);
+
   Id DeclareStructType(const rdcarray<Id> &members);
 
   // helper for AddConstant
@@ -252,6 +255,19 @@ private:
 
   virtual void RegisterOp(Iter iter);
   virtual void UnregisterOp(Iter iter);
+  virtual void PostParse();
+
+  void RegisterBuiltinMembers(rdcspv::Id baseId, const rdcarray<uint32_t> chainSoFar,
+                              const DataType *type);
+
+  struct BuiltinInputData
+  {
+    Id variable;
+    Id type;
+    rdcarray<uint32_t> chain;
+  };
+
+  std::map<BuiltIn, BuiltinInputData> builtinInputs;
 
   std::map<Id, Binding> bindings;