adding root parameter view

Author: joaosaffran

llvm/include/llvm/Object/DXContainer.h

@@ -18,6 +18,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/Object/Error.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/TargetParser/Triple.h"
@@ -122,7 +123,107 @@ struct RootParameter {
   union {
     dxbc::RootConstants Constants;
   };
+
+  RootParameter() = default;
+};
+
+struct ViewRootParameter {
+  ViewArray<dxbc::RootParameterHeader> HeaderView;
+  StringRef Data;
+
+  ViewRootParameter() = default;
+  ViewRootParameter(ViewArray<dxbc::RootParameterHeader> HV, StringRef D)
+      : HeaderView(HV), Data(D) {}
+
+  using value_type = DirectX::RootParameter;
+  struct iterator {
+    const ViewArray<dxbc::RootParameterHeader> HeaderView;
+    ViewArray<dxbc::RootParameterHeader>::iterator Current;
+    StringRef Data;
+
+    iterator(const ViewRootParameter &VRP,
+             ViewArray<dxbc::RootParameterHeader>::iterator C)
+        : HeaderView(VRP.HeaderView), Current(C), Data(VRP.Data) {}
+    iterator(const iterator &) = default;
+
+    template <typename T>
+    static Error readParameter(StringRef Buffer, const char *Src, T &Struct) {
+      // Don't read before the beginning or past the end of the file
+      if (Src < Buffer.begin() || Src + sizeof(T) > Buffer.end())
+        return make_error<GenericBinaryError>(
+            "Reading structure out of file bounds", object_error::parse_failed);
+
+      memcpy(&Struct, Src, sizeof(T));
+      // DXContainer is always little endian
+      if (sys::IsBigEndianHost)
+        Struct.swapBytes();
+      return Error::success();
+    }
+
+    static Error validationFailed(const Twine &Msg) {
+      return make_error<StringError>(Msg.str(), inconvertibleErrorCode());
+    }
+
+    llvm::Expected<value_type> operator*() {
+      value_type Parameter;
+      std::memset(&Parameter.Header, 0, sizeof(dxbc::RootParameterHeader));
+      memcpy(&Parameter.Header, Current.Current,
+             sizeof(dxbc::RootParameterHeader));
+
+      switch (Parameter.Header.ParameterType) {
+      case dxbc::RootParameterType::Constants32Bit:
+        std::memset(&Parameter.Constants, 0, sizeof(dxbc::RootConstants));
+        if (Error Err = readParameter(
+                Data, Data.begin() + Parameter.Header.ParameterOffset,
+                Parameter.Constants))
+          return Err;
+        break;
+      case dxbc::RootParameterType::Empty:
+        break;
+      }
+
+      return Parameter;
+    }
+
+    iterator operator++() {
+      if (Current != HeaderView.end())
+        Current++;
+      return *this;
+    }
+
+    iterator operator++(int I) {
+      iterator Tmp = *this;
+      for (; I > 0; I--)
+        ++Tmp;
+      return Tmp;
+    }
+
+    iterator operator--() {
+      if (Current != HeaderView.begin())
+        Current--;
+      return *this;
+    }
+
+    iterator operator--(int I) {
+      iterator Tmp = *this;
+      for (; I > 0; I--)
+        --Tmp;
+      return Tmp;
+    }
+
+    bool operator==(iterator I) { return I.Current == Current; }
+    bool operator!=(const iterator I) { return !(*this == I); }
+  };
+
+  iterator begin() const { return iterator(*this, this->HeaderView.begin()); }
+
+  iterator end() const { return iterator(*this, this->HeaderView.end()); }
+
+  size_t size() const { return HeaderView.size(); }
+
+  bool isEmpty() const { return HeaderView.isEmpty(); }
 };
+
 class RootSignature {
 private:
   uint32_t Version = 2;
@@ -132,9 +233,9 @@ class RootSignature {
   uint32_t StaticSamplersOffset = 0;
   uint32_t Flags = 0;
 
-  SmallVector<DirectX::RootParameter> Parameters;
+  ViewRootParameter Parameters;
 
-  using param_iterator = SmallVector<DirectX::RootParameter>::iterator;
+  using param_iterator = ViewRootParameter::iterator;
 
 public:
   RootSignature() {}
@@ -145,7 +246,7 @@ class RootSignature {
   uint32_t getRootParametersOffset() const { return RootParametersOffset; }
   uint32_t getNumStaticSamplers() const { return NumStaticSamplers; }
   uint32_t getStaticSamplersOffset() const { return StaticSamplersOffset; }
-  llvm::iterator_range<const RootParameter *> params() const {
+  llvm::iterator_range<param_iterator> params() const {
     return llvm::make_range(Parameters.begin(), Parameters.end());
   }
   param_iterator param_begin() { return Parameters.begin(); }

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -87,24 +87,7 @@ 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) {
-
-    case dxbc::RootParameterType::Constants32Bit:
-      Constants.NumOfConstants = Parameter.Constants.NumOfConstants;
-      Constants.Space = Parameter.Constants.Space;
-      Constants.Register = Parameter.Constants.Register;
-      break;
-    case dxbc::RootParameterType::Empty:
-      llvm_unreachable("Invalid Root Parameter Type. It should be verified "
-                       "before reaching here.");
-      break;
-    }
-  }
+
   union {
     RootConstantsYaml Constants;
   };

llvm/lib/Object/DXContainer.cpp

@@ -305,58 +305,26 @@ Error DirectX::RootSignature::parse(StringRef Data) {
   Flags = FValue;
 
   assert(Current == Begin + RootParametersOffset);
-  for (uint32_t I = 0; I < NumParameters; I++) {
-    DirectX::RootParameter NewParam;
 
-    uint32_t MaybeParamType =
-        support::endian::read<uint32_t, llvm::endianness::little>(Current);
+  Parameters.HeaderView.Data = Data.substr(
+      RootParametersOffset, NumParameters * sizeof(dxbc::RootParameterHeader));
+  Parameters.Data = Data;
 
-    if (Error Err =
-            safeConvertEnum(MaybeParamType, NewParam.Header.ParameterType))
-      return Err;
+  for (auto P : params()) {
+    if (!P)
+      return P.takeError();
 
     if (!dxbc::RootSignatureValidations::isValidParameterType(
-            NewParam.Header.ParameterType))
-      return validationFailed(
-          "unsupported parameter type value read: " +
-          llvm::Twine((uint32_t)NewParam.Header.ParameterType));
-
-    Current += sizeof(dxbc::RootParameterType);
-
-    uint32_t MaybeVisibility =
-        support::endian::read<uint32_t, llvm::endianness::little>(Current);
-
-    if (Error Err =
-            safeConvertEnum(MaybeVisibility, NewParam.Header.ShaderVisibility))
-      return Err;
+            P->Header.ParameterType))
+      return validationFailed("unsupported parameter type value read: " +
+                              llvm::Twine((uint32_t)P->Header.ParameterType));
 
     if (!dxbc::RootSignatureValidations::isValidShaderVisibility(
-            NewParam.Header.ShaderVisibility))
+            P->Header.ShaderVisibility))
       return validationFailed(
           "unsupported shader visility flag value read: " +
-          llvm::Twine((uint32_t)NewParam.Header.ShaderVisibility));
-
-    Current += sizeof(dxbc::ShaderVisibility);
-
-    NewParam.Header.ParameterOffset =
-        support::endian::read<uint32_t, llvm::endianness::little>(Current);
-    Current += sizeof(uint32_t);
-
-    switch (NewParam.Header.ParameterType) {
-
-    case dxbc::RootParameterType::Constants32Bit:
-      if (Error Err = readStruct(Data, Begin + NewParam.Header.ParameterOffset,
-                                 NewParam.Constants))
-        return Err;
-      break;
-    case dxbc::RootParameterType::Empty:
-      // unreachable  because it was validated and assigned before this point.
-      llvm_unreachable("Invalid value for RootParameterType");
-    }
-
-    Parameters.push_back(NewParam);
+          llvm::Twine((uint32_t)P->Header.ShaderVisibility));
   }
-
   return Error::success();
 }
 

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/Support/ErrorHandling.h"
 #include "llvm/Support/ScopedPrinter.h"
 
 namespace llvm {
@@ -35,9 +36,27 @@ DXContainerYAML::RootSignatureYamlDesc::RootSignatureYamlDesc(
       NumStaticSamplers(Data.getNumStaticSamplers()),
       StaticSamplersOffset(Data.getStaticSamplersOffset()) {
   uint32_t Flags = Data.getFlags();
-  for (auto const &P : Data.params()) {
-
-    Parameters.push_back(RootParameterYamlDesc(P));
+  for (auto P : Data.params()) {
+    if (!P)
+      llvm_unreachable("Parameters are validated on read.");
+
+    auto NewP = RootParameterYamlDesc();
+    NewP.Offset = P->Header.ParameterOffset;
+    NewP.Type = P->Header.ParameterType;
+    NewP.Visibility = P->Header.ShaderVisibility;
+
+    switch (NewP.Type) {
+
+    case dxbc::RootParameterType::Constants32Bit:
+      NewP.Constants.NumOfConstants = P->Constants.NumOfConstants;
+      NewP.Constants.Register = P->Constants.Register;
+      NewP.Constants.Space = P->Constants.Space;
+      break;
+    case dxbc::RootParameterType::Empty:
+      llvm_unreachable("Invalid RootParameterType");
+      break;
+    }
+    Parameters.push_back(NewP);
   }
 #define ROOT_ELEMENT_FLAG(Num, Val)                                            \
   Val = (Flags & (uint32_t)dxbc::RootElementFlag::Val) > 0;

llvm/unittests/Object/DXContainerTest.cpp

@@ -909,22 +909,26 @@ TEST(RootSignature, ParseRootConstant) {
     DXContainer C =
         llvm::cantFail(DXContainer::create(getMemoryBuffer<133>(Buffer)));
 
-    auto RS = C.getRootSignature();
-    ASSERT_TRUE(RS.has_value());
-    ASSERT_EQ(RS->getVersion(), 2u);
-    ASSERT_EQ(RS->getNumParameters(), 1u);
-    ASSERT_EQ(RS->getRootParametersOffset(), 24u);
-    ASSERT_EQ(RS->getNumStaticSamplers(), 0u);
-    ASSERT_EQ(RS->getStaticSamplersOffset(), 44u);
-    ASSERT_EQ(RS->getFlags(), 17u);
-
-    for (auto const &RootParam : RS->params()) {
-      ASSERT_EQ((uint32_t)RootParam.Header.ParameterType, 1u);
-      ASSERT_EQ((uint32_t)RootParam.Header.ShaderVisibility, 2u);
-      ASSERT_EQ(RootParam.Constants.Register, 15u);
-      ASSERT_EQ(RootParam.Constants.Space, 14u);
-      ASSERT_EQ(RootParam.Constants.NumOfConstants, 16u);
-    }
+    auto MaybeRS = C.getRootSignature();
+    ASSERT_TRUE(MaybeRS.has_value());
+    const auto &RS = MaybeRS.value();
+    ASSERT_EQ(RS.getVersion(), 2u);
+    ASSERT_EQ(RS.getNumParameters(), 1u);
+    ASSERT_EQ(RS.getRootParametersOffset(), 24u);
+    ASSERT_EQ(RS.getNumStaticSamplers(), 0u);
+    ASSERT_EQ(RS.getStaticSamplersOffset(), 44u);
+    ASSERT_EQ(RS.getFlags(), 17u);
+
+    auto RootParam = *RS.params().begin();
+    ASSERT_FALSE(!RootParam);
+    auto X = (unsigned)RootParam->Header.ParameterType;
+    auto Y = (unsigned)RootParam->Header.ShaderVisibility;
+
+    ASSERT_EQ(X, 1u);
+    ASSERT_EQ(Y, 2u);
+    ASSERT_EQ(RootParam->Constants.Register, 15u);
+    ASSERT_EQ(RootParam->Constants.Space, 14u);
+    ASSERT_EQ(RootParam->Constants.NumOfConstants, 16u);
   }
   {
     // ParameterType has been set to an invalid value
@@ -943,7 +947,7 @@ TEST(RootSignature, ParseRootConstant) {
         0x00};
     EXPECT_THAT_EXPECTED(
         DXContainer::create(getMemoryBuffer<133>(Buffer)),
-        FailedWithMessage("Value is not within range for enum"));
+        FailedWithMessage("unsupported parameter type value read: 255"));
   }
   {
     // ShaderVisibility has been set to an invalid value
@@ -962,7 +966,7 @@ TEST(RootSignature, ParseRootConstant) {
         0x00};
     EXPECT_THAT_EXPECTED(
         DXContainer::create(getMemoryBuffer<133>(Buffer)),
-        FailedWithMessage("Value is not within range for enum"));
+        FailedWithMessage("unsupported shader visility flag value read: 255"));
   }
   {
     // Offset has been set to an invalid value