adding samplers_iterator

Author: joaosaffran

llvm/include/llvm/Object/DXContainer.h

@@ -223,6 +223,28 @@ static Error parseFailed(const Twine &Msg) {
 
 class RootSignature {
 private:
+  struct samplers_iterator {
+    const RootSignature *Parent = nullptr;
+    uint32_t Index = 0;
+
+    llvm::Expected<dxbc::RTS0::v3::StaticSampler> operator*() const {
+      return Parent->getSampler(Index);
+    }
+
+    samplers_iterator &operator++() {
+      ++Index;
+      return *this;
+    }
+
+    bool operator==(const samplers_iterator &Other) const {
+      return Parent == Other.Parent && Index == Other.Index;
+    }
+
+    bool operator!=(const samplers_iterator &Other) const {
+      return !(*this == Other);
+    }
+  };
+
   uint32_t Version;
   uint32_t NumParameters;
   uint32_t RootParametersOffset;
@@ -233,23 +255,6 @@ class RootSignature {
   StringRef PartData;
   ViewArray<dxbc::RTS0::v1::StaticSampler> StaticSamplers;
 
-  using param_header_iterator =
-      ViewArray<dxbc::RTS0::v1::RootParameterHeader>::iterator;
-  using samplers_iterator = ViewArray<dxbc::RTS0::v1::StaticSampler>::iterator;
-
-public:
-  RootSignature(StringRef PD) : PartData(PD) {}
-
-  LLVM_ABI Error parse();
-  uint32_t getVersion() const { return Version; }
-  uint32_t getNumParameters() const { return NumParameters; }
-  uint32_t getRootParametersOffset() const { return RootParametersOffset; }
-  uint32_t getNumStaticSamplers() const { return NumStaticSamplers; }
-  uint32_t getStaticSamplersOffset() const { return StaticSamplersOffset; }
-  uint32_t getNumRootParameters() const { return ParametersHeaders.size(); }
-  llvm::iterator_range<param_header_iterator> param_headers() const {
-    return llvm::make_range(ParametersHeaders.begin(), ParametersHeaders.end());
-  }
   llvm::Expected<dxbc::RTS0::v3::StaticSampler> getSampler(uint32_t Loc) const {
     if (Loc >= getNumStaticSamplers())
       return parseFailed("Static sampler index out of range");
@@ -271,6 +276,37 @@ class RootSignature {
                                             object_error::parse_failed);
     return readParameter<dxbc::RTS0::v3::StaticSampler>(Buff);
   }
+
+  using param_header_iterator =
+      ViewArray<dxbc::RTS0::v1::RootParameterHeader>::iterator;
+
+public:
+  RootSignature(StringRef PD) : PartData(PD) {}
+
+  LLVM_ABI Error parse();
+  uint32_t getVersion() const { return Version; }
+  uint32_t getNumParameters() const { return NumParameters; }
+  uint32_t getRootParametersOffset() const { return RootParametersOffset; }
+  uint32_t getNumStaticSamplers() const { return NumStaticSamplers; }
+  uint32_t getStaticSamplersOffset() const { return StaticSamplersOffset; }
+  uint32_t getNumRootParameters() const { return ParametersHeaders.size(); }
+
+  samplers_iterator samplers_begin() const {
+    return samplers_iterator{this, 0};
+  }
+
+  samplers_iterator samplers_end() const {
+    return samplers_iterator{this, getNumStaticSamplers()};
+  }
+
+  llvm::iterator_range<samplers_iterator> samplers() const {
+    return llvm::make_range(samplers_begin(), samplers_end());
+  }
+
+  llvm::iterator_range<param_header_iterator> param_headers() const {
+    return llvm::make_range(ParametersHeaders.begin(), ParametersHeaders.end());
+  }
+
   uint32_t getFlags() const { return Flags; }
 
   llvm::Expected<RootParameterView>

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -163,9 +163,8 @@ DXContainerYAML::RootSignatureYamlDesc::create(
     }
   }
 
-  for (uint32_t Loc = 0; Loc < Data.getNumStaticSamplers(); ++Loc) {
-    llvm::Expected<dxbc::RTS0::v3::StaticSampler> MaybeSampler =
-        Data.getSampler(Loc);
+  for (llvm::Expected<dxbc::RTS0::v3::StaticSampler> MaybeSampler :
+       Data.samplers()) {
     if (Error E = MaybeSampler.takeError())
       return std::move(E);
     const llvm::dxbc::RTS0::v3::StaticSampler &S = *MaybeSampler;

llvm/unittests/Object/DXContainerTest.cpp

@@ -1184,7 +1184,7 @@ TEST(RootSignature, ParseStaticSamplers) {
     ASSERT_EQ(RS.getStaticSamplersOffset(), 24u);
     ASSERT_EQ(RS.getFlags(), 17u);
 
-    auto MaybeSamplerView = RS.getSampler(0);
+    auto MaybeSamplerView = *RS.samplers().begin();
     ASSERT_THAT_ERROR(MaybeSamplerView.takeError(), Succeeded());
     const auto &Sampler = *MaybeSamplerView;