[Object][DX] Parse DXContainer Parts

[beanz]

Summary

DXContainer files are structured as parts. This patch adds support for
parsing out the file part offsets and file part headers.

[beanz]

Adding @kuhar.

[beanz]

Fixing incorrect code coment.

[kuhar]

On high level, would it be possible to use an existing class like llvm::BinaryStreamReader instead that most of this data reading machinery?

And a general nit: I'm biased towards using ArrayRefs to represent data and StringRefs to represent text, but I don't think there's a consensus on this.

llvm/include/llvm/BinaryFormat/DXContainer.h

86

nit: swapBytes? LLVM prefers starting function names with verbs, especially the state-mutating ones: https://llvm.org/docs/CodingStandards.html#name-types-functions-variables-and-enumerators-properly

llvm/include/llvm/Object/DXContainer.h

39–42

nit: could we pull this class out of DXContainer? I think this would make it a bit more readable, especially since we have another level of nesting with the PartData struct.

58

Could you add a comment with a brief description of what this function does? It's impossible to tell based on the signature alone.

llvm/lib/Object/DXContainer.cpp

34

Can you give P a more descriptive name and/or add a function comment?
Also, why not return Expected<T>

39

What if P is not properly aligned for T*?
If it must be, can we add an assertion? If not, I think we would have to use a memcpy.

Also, shouldn't we check that T is trivially copyable in the first place?

76–77

Would it be possible for OffsetIt to equal the end iterator and end up with an uninitialized IteratorState?

llvm/unittests/Object/DXContainerTest.cpp

56

Could we have a test with an empty buffer, so that we know that this corner case was considered and whether this is supported or not?

[beanz]

libObject uses MemoryBufferRefs which play back and forth with StringRef, and while it has been pointed out over and over again that many of StringRef’s methods should probably be on MemoryBufferRef because they are more broadly useful, it is the way it is.

I think ArrayRef<char> is probably the least common pattern for representing arbitrary data buffers in LLVM, and given the needs of libObject to operate with MemoryBufferRef I think this code is best to operate on StringRefs.

llvm/lib/Object/DXContainer.cpp

34

Mostly just because the usage pattern for Error is cleaner in this case, but also it avoids a copy and is in-line able.

With Expected the code would be:

auto ExVal = readValue(…);
if (!ExVal)
  return ExVal.takeError();
Val = *ExVal;

Which is a bit less clean to me.

76–77

It shouldn’t be, but looking through I think there’s an edge case. I will update with a fix.

[kuhar]

In D124804#3526729, @beanz wrote:

libObject uses MemoryBufferRefs which play back and forth with StringRef, and while it has been pointed out over and over again that many of StringRef’s methods should probably be on MemoryBufferRef because they are more broadly useful, it is the way it is.
I think ArrayRef<char> is probably the least common pattern for representing arbitrary data buffers in LLVM, and given the needs of libObject to operate with MemoryBufferRef I think this code is best to operate on StringRefs.

I'm much more used to ArrayRef<uint8_t>. For example, you can see it used in BinaryStreamReader::readLongestContiguousChunk and there's also a helper conversion function ArrayRef<uint8_t> arrayRefFromStringRef(StringRef Input) in StringExtras.h. In total, I counted ~1500 uses of ArrayRef<uint8_t> in the monorepo, so I don't consider it an uncommon pattern.

That being said, I think sticking with StringRef is also a practical option and I don't oppose it.

[kuhar]

llvm/lib/Object/DXContainer.cpp

34

I'd disagree here. Let's look at a callsite and the code that follows:

uint32_t PartOffset;
if (Error Err = readValue(Data.getBuffer(), Current, PartOffset))
  return Err;
Current += sizeof(uint32_t);
if (PartOffset + sizeof(dxbc::PartHeader) > Data.getBufferSize())
  return parseFailed("Part offset points beyond boundary of the file");
PartOffsets.push_back(PartOffset);

In general, I would not say it's immediately that PartOffset is an output parameter. Without seeing the definition, we don't know if passing PartOffset to readValue will lead to uninitialized reads or not. I find it much more idiomatic to use return values to return values. With Expected<T>, the callsite and the surrounding code would look like this:

Expected<uint32_t> PartOffset = readValue(Data.getBuffer(), Current);
if (Error Err = PartOffset.takeError())
  return Err;
Current += sizeof(uint32_t);
if (*PartOffset + sizeof(dxbc::PartHeader) > Data.getBufferSize())
  return parseFailed("Part offset points beyond boundary of the file");
PartOffsets.push_back(*PartOffset);

Here it's very clear to me what is being returned and I would not worry about any uninitialized values, even without checking the implementation of readValue. And overall, the number of lines stays the same.

I also typically don't prefix/suffix expected values with anything special, I don't think it improves readability when you can see the variable type. This is also what the google style guide does: https://abseil.io/tips/181.

To be absolutely clear, this is not a big deal IMO either way and feel free to stick with the current API if you strongly prefer it.

[beanz]

In D124804#3526834, @kuhar wrote:

In total, I counted ~1500 uses of ArrayRef<uint8_t> in the monorepo, so I don't consider it an uncommon pattern.

Scanning the full monorepo is probably not the best measure for stylistic code patterns. The “golden-rule” of LLVM’s style guidelines is to match existing code patterns, and that is generally applied to the code in the area that you’re modifying (https://llvm.org/docs/CodingStandards.html#golden-rule).

[beanz]

Updates based on feedback from @kuhar. Thank you!

[beanz]

Updates uploaded.

[kuhar]

llvm/include/llvm/Object/DXContainer.h

39–42

Have you considered this suggestion?

llvm/lib/Object/DXContainer.cpp

37

nit: I think uintptr_t might be more appropriate

[beanz]

llvm/include/llvm/Object/DXContainer.h

39–42

Sorry, I had a comment on this that I failed to submit.

The downside to pulling the iterator out of the class is that trivial methods end up needing to be implemented in the implementation files. For example, if you put the iterator first in the header, the iterator's constructor and updateIterator method need to be defined in the implementation file because they depend on the DXContainer class. Conversely if you put the iterator after DXContainer, DXContainer begin and end need to be in the implementation file.

Moving trivial functions to the implementation file prevents them from being inlined.

I also experimented with replacing PartData with a std::pair, but found the first/second accessors to be less intuitive to read than the named members.

[beanz]

size_t -> uintptr_t

[kuhar]

LGTM

llvm/include/llvm/Object/DXContainer.h

39–42

You could define functions in the header if you are worried about performance in non-LTO builds. However, if you think that this change does not simplify take code or improve readability, I think that leaving it as-is is also fine.

llvm/include/llvm/BinaryFormat/DXContainer.h

struct Header {
   // Structure is followed by part offsets: uint32_t PartOffset[PartCount];
   // The offset is to a PartHeader, which is followed by the Part Data.
 };
 
 /// Use this type to describe the size and type of a DXIL container part.
 struct PartHeader {
   uint8_t Name[4];
   uint32_t Size;
+
+  void byteSwap() { sys::swapByteOrder(Size); }
   // Structure is followed directly by part data: uint8_t PartData[PartSize].
 };
 
 } // namespace dxbc
 } // namespace llvm
 
 #endif // LLVM_BINARYFORMAT_DXCONTAINER_H

llvm/include/llvm/Object/DXContainer.h

 // interface for DXContainer files.
 //
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_OBJECT_DXCONTAINER_H
 #define LLVM_OBJECT_DXCONTAINER_H
 
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBufferRef.h"
 
 namespace llvm {
 namespace object {
 class DXContainer {
 private:
   DXContainer(MemoryBufferRef O);
 
   MemoryBufferRef Data;
   dxbc::Header Header;
+  SmallVector<uint32_t, 4> PartOffsets;
 
   Error parseHeader();
+  Error parsePartOffsets();
+  friend class PartIterator;
 
 public:
+  // The PartIterator is a wrapper around the iterator for the PartOffsets
+  // member of the DXContainer. It contains a refernce to the container, and the
+  // current iterator value, as well as storage for a parsed part header.
+  class PartIterator {
+    const DXContainer &Container;
+    SmallVectorImpl<uint32_t>::const_iterator OffsetIt;
+    struct PartData {
+      dxbc::PartHeader Part;
+      StringRef Data;
+    } IteratorState;
+
+    friend class DXContainer;
+
+    PartIterator(const DXContainer &C,
+                 SmallVectorImpl<uint32_t>::const_iterator It)
+        : Container(C), OffsetIt(It) {
+      updateIterator();
+    }
+
+    void updateIterator();
+
+  public:
+    PartIterator &operator++() {
+      if (OffsetIt == Container.PartOffsets.end())
+        return *this;
+      ++OffsetIt;
+      updateIterator();
+      return *this;
+    }
+
+    PartIterator operator++(int) {
+      PartIterator Tmp = *this;
+      ++(*this);
+      return Tmp;
+    }
+
+    PartIterator &operator--() {
+      if (OffsetIt == Container.PartOffsets.begin())
+        return *this;
+      --OffsetIt;
+      updateIterator();
+      return *this;
+    }
+
+    PartIterator operator--(int) {
+      PartIterator Tmp = *this;
+      --(*this);
+      return Tmp;
+    }
+
+    bool operator==(const PartIterator &RHS) const {
+      return OffsetIt == RHS.OffsetIt;
+    }
+
+    bool operator!=(const PartIterator &RHS) const {
+      return OffsetIt != RHS.OffsetIt;
+    }
+
+    const PartData &operator*() { return IteratorState; }
+    const PartData *operator->() { return &IteratorState; }
+  };
+
+  PartIterator begin() const {
+    return PartIterator(*this, PartOffsets.begin());
+  }
+
+  PartIterator end() const { return PartIterator(*this, PartOffsets.end()); }
+
   StringRef getData() const { return Data.getBuffer(); }
   static Expected<DXContainer> create(MemoryBufferRef Object);
 
   const dxbc::Header &getHeader() const { return Header; }
 };
 
 } // namespace object
 } // namespace llvm
 
 #endif // LLVM_OBJECT_DXCONTAINERFILE_H

llvm/lib/Object/DXContainer.cpp

static Error readStruct(StringRef Buffer, const char *P, T &Struct) {
 
   memcpy(&Struct, P, sizeof(T));
   // DXContainer is always little endian
   if (sys::IsBigEndianHost)
     Struct.byteSwap();
   return Error::success();
 }
 
+template <typename T>
+static Error readValue(StringRef Buffer, const char *P, T &Val) {
+  // Don't read before the beginning or past the end of the file
+  if (P < Buffer.begin() || P + sizeof(T) > Buffer.end())
+    return parseFailed("Reading structure out of file bounds");
+
+  Val = *reinterpret_cast<const T *>(P);
+  // DXContainer is always little endian
+  if (sys::IsBigEndianHost)
+    sys::swapByteOrder(Val);
+  return Error::success();
+}
+
 DXContainer::DXContainer(MemoryBufferRef O) : Data(O) {}
 
 Error DXContainer::parseHeader() {
   return readStruct(Data.getBuffer(), Data.getBuffer().data(), Header);
 }
 
+Error DXContainer::parsePartOffsets() {
+  const char *Current = Data.getBuffer().data() + sizeof(dxbc::Header);
+  for (uint32_t Part = 0; Part < Header.PartCount; ++Part) {
+    uint32_t PartOffset;
+    if (Error Err = readValue(Data.getBuffer(), Current, PartOffset))
+      return Err;
+    Current += sizeof(uint32_t);
+    if (PartOffset + sizeof(dxbc::PartHeader) > Data.getBufferSize())
+      return parseFailed("Part offset points beyond boundary of the file");
+    PartOffsets.push_back(PartOffset);
+  }
+  return Error::success();
+}
+
 Expected<DXContainer> DXContainer::create(MemoryBufferRef Object) {
   DXContainer Container(Object);
   if (Error Err = Container.parseHeader())
     return std::move(Err);
+  if (Error Err = Container.parsePartOffsets())
+    return std::move(Err);
   return Container;
 }
+
+void DXContainer::PartIterator::updateIterator() {
+  if (OffsetIt == Container.PartOffsets.end())
+    return;
+  StringRef Buffer = Container.Data.getBuffer();
+  const char *Current = Buffer.data() + *OffsetIt;
+  // Offsets are validated during parsing, so all offsets in the container are
+  // valid and contain enough readable data to read a header.
+  cantFail(readStruct(Buffer, Current, IteratorState.Part));
+  IteratorState.Data =
+      StringRef(Current + sizeof(dxbc::PartHeader), IteratorState.Part.Size);
+}

llvm/unittests/Object/DXContainerTest.cpp

EXPECT_THAT_EXPECTED(
       DXContainer::create(getMemoryBuffer<4>(Buffer)),
       FailedWithMessage("Reading structure out of file bounds"));
 }
 
 TEST(DXCFile, ParseHeader) {
   uint8_t Buffer[] = {0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00,
                       0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                       0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
-                      0x70, 0x0D, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00};
+                      0x70, 0x0D, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
   DXContainer C =
       llvm::cantFail(DXContainer::create(getMemoryBuffer<32>(Buffer)));
   EXPECT_TRUE(memcmp(C.getHeader().Magic, "DXBC", 4) == 0);
   EXPECT_TRUE(memcmp(C.getHeader().FileHash.Digest,
                      "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0", 16) == 0);
   EXPECT_EQ(C.getHeader().Version.Major, 1u);
   EXPECT_EQ(C.getHeader().Version.Minor, 0u);
 }
+
+TEST(DXCFile, ParsePartMissingOffsets) {
+  uint8_t Buffer[] = {
+      0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00,
+      0x00, 0x00, 0x70, 0x0D, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+  };
+  EXPECT_THAT_EXPECTED(
+      DXContainer::create(getMemoryBuffer<32>(Buffer)),
+      FailedWithMessage("Reading structure out of file bounds"));
+}
+
+TEST(DXCFile, ParsePartInvalidOffsets) {
+  uint8_t Buffer[] = {
+      0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+      0x70, 0x0D, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF,
+  };
+  EXPECT_THAT_EXPECTED(
+      DXContainer::create(getMemoryBuffer<36>(Buffer)),
+      FailedWithMessage("Part offset points beyond boundary of the file"));
+}
+
+TEST(DXCFile, ParseEmptyParts) {
+  uint8_t Buffer[] = {
+      0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+      0x70, 0x0D, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00, 0x3C, 0x00, 0x00, 0x00,
+      0x44, 0x00, 0x00, 0x00, 0x4C, 0x00, 0x00, 0x00, 0x54, 0x00, 0x00, 0x00,
+      0x5C, 0x00, 0x00, 0x00, 0x64, 0x00, 0x00, 0x00, 0x6C, 0x00, 0x00, 0x00,
+      0x53, 0x46, 0x49, 0x30, 0x00, 0x00, 0x00, 0x00, 0x49, 0x53, 0x47, 0x31,
+      0x00, 0x00, 0x00, 0x00, 0x4F, 0x53, 0x47, 0x31, 0x00, 0x00, 0x00, 0x00,
+      0x50, 0x53, 0x56, 0x30, 0x00, 0x00, 0x00, 0x00, 0x53, 0x54, 0x41, 0x54,
+      0x00, 0x00, 0x00, 0x00, 0x44, 0x58, 0x49, 0x4C, 0x00, 0x00, 0x00, 0x00,
+      0x44, 0x45, 0x41, 0x44, 0x00, 0x00, 0x00, 0x00,
+  };
+  DXContainer C =
+      llvm::cantFail(DXContainer::create(getMemoryBuffer<116>(Buffer)));
+  EXPECT_EQ(C.getHeader().PartCount, 7u);
+
+  // All the part sizes are 0, which makes a nice test of the range based for
+  int ElementsVisited = 0;
+  for (auto Part : C) {
+    EXPECT_EQ(Part.Part.Size, 0u);
+    EXPECT_EQ(Part.Data.size(), 0u);
+    ++ElementsVisited;
+  }
+  EXPECT_EQ(ElementsVisited, 7);
+
+  auto It = C.begin();
+  EXPECT_TRUE(memcmp(It->Part.Name, "SFI0", 4) == 0);
+  --It; // Don't decrement past begin.
+  EXPECT_TRUE(memcmp(It->Part.Name, "SFI0", 4) == 0);
+  ++It;
+  EXPECT_TRUE(memcmp(It->Part.Name, "ISG1", 4) == 0);
+  ++It;
+  EXPECT_TRUE(memcmp(It->Part.Name, "OSG1", 4) == 0);
+  ++It;
+  EXPECT_TRUE(memcmp(It->Part.Name, "PSV0", 4) == 0);
+  ++It;
+  EXPECT_TRUE(memcmp(It->Part.Name, "STAT", 4) == 0);
+  ++It;
+  EXPECT_TRUE(memcmp(It->Part.Name, "DXIL", 4) == 0);
+  ++It;
+  EXPECT_TRUE(memcmp(It->Part.Name, "DEAD", 4) == 0);
+  ++It; // Don't increment past the end
+  EXPECT_TRUE(memcmp(It->Part.Name, "DEAD", 4) == 0);
+}