addressing pr comments

Author: joaosaffran

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -7,6 +7,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/BinaryFormat/DXContainer.h"
+#include <cstddef>
+#include <cstdint>
 
 namespace llvm {
 
@@ -28,6 +30,10 @@ struct RootSignatureDesc {
             2, 0, sizeof(dxbc::RootSignatureHeader), 0, 0, 0}) {}
 
   void write(raw_ostream &OS) const;
+
+  size_t getSize() const {
+    return sizeof(dxbc::RootSignatureHeader) + Parameters.size_in_bytes();
+  }
 };
 } // namespace mcdxbc
 } // namespace llvm

llvm/include/llvm/Object/DXContainer.h

@@ -133,6 +133,8 @@ class RootSignature {
 
   SmallVector<DirectX::RootParameter> Parameters;
 
+  using ParamsIter = SmallVector<DirectX::RootParameter>::iterator;
+
 public:
   RootSignature() {}
 
@@ -142,9 +144,9 @@ class RootSignature {
   uint32_t getRootParametersOffset() const { return RootParametersOffset; }
   uint32_t getNumStaticSamplers() const { return NumStaticSamplers; }
   uint32_t getStaticSamplersOffset() const { return StaticSamplersOffset; }
-  const SmallVector<DirectX::RootParameter> &getParameters() const {
-    return Parameters;
-  }
+  ParamsIter beginParams() { return Parameters.begin(); }
+
+  ParamsIter endParams() { return Parameters.end(); }
   uint32_t getFlags() const { return Flags; }
 };
 

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -88,16 +88,16 @@ struct RootParameterYamlDesc {
   uint32_t Offset;
 
   RootParameterYamlDesc() = default;
-  RootParameterYamlDesc(const object::DirectX::RootParameter &Parameter) {
-    Type = Parameter.Header.ParameterType;
-    Visibility = Parameter.Header.ShaderVisibility;
-    Offset = Parameter.Header.ParameterOffset;
-    switch (Parameter.Header.ParameterType) {
+  RootParameterYamlDesc(object::DirectX::RootParameter *Parameter) {
+    Type = Parameter->Header.ParameterType;
+    Visibility = Parameter->Header.ShaderVisibility;
+    Offset = Parameter->Header.ParameterOffset;
+    switch (Parameter->Header.ParameterType) {
 
     case dxbc::RootParameterType::Constants32Bit: {
-      Constants.Num32BitValues = Parameter.Constants.Num32BitValues;
-      Constants.RegisterSpace = Parameter.Constants.RegisterSpace;
-      Constants.ShaderRegister = Parameter.Constants.ShaderRegister;
+      Constants.Num32BitValues = Parameter->Constants.Num32BitValues;
+      Constants.RegisterSpace = Parameter->Constants.RegisterSpace;
+      Constants.ShaderRegister = Parameter->Constants.ShaderRegister;
     } break;
     case dxbc::RootParameterType::Empty:
       llvm_unreachable("Invalid Root Parameter Type. It should be verified "
@@ -112,7 +112,7 @@ struct RootParameterYamlDesc {
 
 struct RootSignatureYamlDesc {
   RootSignatureYamlDesc() = default;
-  RootSignatureYamlDesc(const object::DirectX::RootSignature &Data);
+  RootSignatureYamlDesc(object::DirectX::RootSignature Data);
 
   uint32_t Version;
   uint32_t NumStaticSamplers;

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -7,41 +7,47 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/MC/DXContainerRootSignature.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Support/EndianStream.h"
+#include <cstdint>
 
 using namespace llvm;
 using namespace llvm::mcdxbc;
 
-static uint32_t writePlaceholder(raw_ostream &Stream) {
+static uint32_t writePlaceholder(raw_svector_ostream &Stream) {
   const uint32_t DummyValue = std::numeric_limits<uint32_t>::max();
   uint32_t Offset = Stream.tell();
   support::endian::write(Stream, DummyValue, llvm::endianness::little);
   return Offset;
 }
 
-static void rewriteOffset(buffer_ostream &Stream, uint32_t Offset) {
+static void rewriteOffset(raw_svector_ostream &Stream, uint32_t Offset) {
   uint32_t Value =
       support::endian::byte_swap<uint32_t, llvm::endianness::little>(
           Stream.tell());
   Stream.pwrite(reinterpret_cast<const char *>(&Value), sizeof(Value), Offset);
 }
 
 void RootSignatureDesc::write(raw_ostream &OS) const {
-  buffer_ostream BOS(OS);
+  SmallString<256> Storage;
+  raw_svector_ostream BOS(Storage);
+  BOS.reserveExtraSpace(getSize());
   const uint32_t NumParameters = Parameters.size();
   const uint32_t Zero = 0;
 
   support::endian::write(BOS, Header.Version, llvm::endianness::little);
   support::endian::write(BOS, NumParameters, llvm::endianness::little);
 
-  uint32_t HeaderPoint = writePlaceholder(BOS);
+  // Root Parameters offset should always start after the
+  support::endian::write(BOS, (uint32_t)sizeof(dxbc::RootSignatureHeader),
+                         llvm::endianness::little);
 
   support::endian::write(BOS, Zero, llvm::endianness::little);
   support::endian::write(BOS, Zero, llvm::endianness::little);
   support::endian::write(BOS, Header.Flags, llvm::endianness::little);
 
-  rewriteOffset(BOS, HeaderPoint);
-
   SmallVector<uint32_t> ParamsOffsets;
   for (const auto &P : Parameters) {
     support::endian::write(BOS, P.Header.ParameterType,
@@ -70,4 +76,6 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
       llvm_unreachable("Invalid RootParameterType");
     }
   }
+
+  OS.write(Storage.data(), Storage.size());
 }

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -14,6 +14,7 @@
 #include "llvm/ObjectYAML/DXContainerYAML.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/Object/DXContainer.h"
 #include "llvm/Support/ScopedPrinter.h"
 
 namespace llvm {
@@ -30,12 +31,13 @@ DXContainerYAML::ShaderFeatureFlags::ShaderFeatureFlags(uint64_t FlagData) {
 }
 
 DXContainerYAML::RootSignatureYamlDesc::RootSignatureYamlDesc(
-    const object::DirectX::RootSignature &Data)
+    object::DirectX::RootSignature Data)
     : Version(Data.getVersion()),
       NumStaticSamplers(Data.getNumStaticSamplers()),
       StaticSamplersOffset(Data.getStaticSamplersOffset()) {
   uint32_t Flags = Data.getFlags();
-  for (const auto &P : Data.getParameters()) {
+  for (auto *P = Data.beginParams(); P != Data.endParams(); ++P) {
+
     Parameters.push_back(RootParameterYamlDesc(P));
   }
 #define ROOT_ELEMENT_FLAG(Num, Val)                                            \

llvm/unittests/Object/DXContainerTest.cpp

@@ -909,7 +909,7 @@ TEST(RootSignature, ParseRootConstant) {
     DXContainer C =
         llvm::cantFail(DXContainer::create(getMemoryBuffer<133>(Buffer)));
 
-    const auto &RS = C.getRootSignature();
+    auto RS = C.getRootSignature();
     ASSERT_TRUE(RS.has_value());
     ASSERT_EQ(RS->getVersion(), 2u);
     ASSERT_EQ(RS->getNumParameters(), 1);
@@ -918,12 +918,12 @@ TEST(RootSignature, ParseRootConstant) {
     ASSERT_EQ(RS->getStaticSamplersOffset(), 44u);
     ASSERT_EQ(RS->getFlags(), 17u);
 
-    const auto RootParam = RS->getParameters()[0];
-    ASSERT_EQ((uint32_t)RootParam.Header.ParameterType, 1u);
-    ASSERT_EQ((uint32_t)RootParam.Header.ShaderVisibility, 2u);
-    ASSERT_EQ(RootParam.Constants.ShaderRegister, 15u);
-    ASSERT_EQ(RootParam.Constants.RegisterSpace, 14u);
-    ASSERT_EQ(RootParam.Constants.Num32BitValues, 16u);
+    object::DirectX::RootParameter *RootParam = RS->beginParams();
+    ASSERT_EQ((uint32_t)RootParam->Header.ParameterType, 1u);
+    ASSERT_EQ((uint32_t)RootParam->Header.ShaderVisibility, 2u);
+    ASSERT_EQ(RootParam->Constants.ShaderRegister, 15u);
+    ASSERT_EQ(RootParam->Constants.RegisterSpace, 14u);
+    ASSERT_EQ(RootParam->Constants.Num32BitValues, 16u);
   }
   {
     // ParameterType has been set to an invalid value