Add some helpers for manipulating BinaryStreamRefs.

llvm-svn: 303297

4 files changed, 62 insertions, 3 deletions

diff --git a/llvm/include/llvm/Support/BinaryStreamReader.h b/llvm/include/llvm/Support/BinaryStreamReader.h
index 0c1881d2e04..75e96a999a1 100644
--- a/llvm/include/llvm/Support/BinaryStreamReader.h
+++ b/llvm/include/llvm/Support/BinaryStreamReader.h
@@ -258,6 +258,8 @@ public:
   /// \returns the next byte in the stream.
   uint8_t peek() const;
 
+  Error padToAlignment(uint32_t Align);
+
   std::pair<BinaryStreamReader, BinaryStreamReader>
   split(uint32_t Offset) const;
 
diff --git a/llvm/include/llvm/Support/BinaryStreamRef.h b/llvm/include/llvm/Support/BinaryStreamRef.h
index 314668f331d..e3bd4bf0860 100644
--- a/llvm/include/llvm/Support/BinaryStreamRef.h
+++ b/llvm/include/llvm/Support/BinaryStreamRef.h
@@ -57,16 +57,35 @@ public:
     return Result;
   }
 
-  /// Return a new BinaryStreamRef with only the first \p N elements remaining.
-  RefType keep_front(uint32_t N) const {
+  /// Return a new BinaryStreamRef with the first \p N elements removed.
+  RefType drop_back(uint32_t N) const {
     if (!BorrowedImpl)
       return RefType();
+
     N = std::min(N, Length);
     RefType Result(static_cast<const RefType &>(*this));
-    Result.Length = N;
+    Result.Length -= N;
     return Result;
   }
 
+  /// Return a new BinaryStreamRef with only the first \p N elements remaining.
+  RefType keep_front(uint32_t N) const {
+    assert(N <= getLength());
+    return drop_back(getLength() - N);
+  }
+
+  /// Return a new BinaryStreamRef with only the last \p N elements remaining.
+  RefType keep_back(uint32_t N) const {
+    assert(N <= getLength());
+    return drop_front(getLength() - N);
+  }
+
+  /// Return a new BinaryStreamRef with the first and last \p N elements
+  /// removed.
+  RefType drop_symmetric(uint32_t N) const {
+    return drop_front(N).drop_back(N);
+  }
+
   /// Return a new BinaryStreamRef with the first \p Offset elements removed,
   /// and retaining exactly \p Len elements.
   RefType slice(uint32_t Offset, uint32_t Len) const {
diff --git a/llvm/lib/Support/BinaryStreamReader.cpp b/llvm/lib/Support/BinaryStreamReader.cpp
index 6b148903de9..5c277448a76 100644
--- a/llvm/lib/Support/BinaryStreamReader.cpp
+++ b/llvm/lib/Support/BinaryStreamReader.cpp
@@ -95,6 +95,11 @@ Error BinaryStreamReader::skip(uint32_t Amount) {
   return Error::success();
 }
 
+Error BinaryStreamReader::padToAlignment(uint32_t Align) {
+  uint32_t NewOffset = alignTo(Offset, Align);
+  return skip(NewOffset - Offset);
+}
+
 uint8_t BinaryStreamReader::peek() const {
   ArrayRef<uint8_t> Buffer;
   auto EC = Stream.readBytes(Offset, 1, Buffer);
diff --git a/llvm/unittests/Support/BinaryStreamTest.cpp b/llvm/unittests/Support/BinaryStreamTest.cpp
index ec3b0effc9e..1ce74cbb722 100644
--- a/llvm/unittests/Support/BinaryStreamTest.cpp
+++ b/llvm/unittests/Support/BinaryStreamTest.cpp
@@ -289,6 +289,39 @@ TEST_F(BinaryStreamTest, StreamRefBounds) {
   }
 }
 
+TEST_F(BinaryStreamTest, DropOperations) {
+  std::vector<uint8_t> InputData = {1, 2, 3, 4, 5, 4, 3, 2, 1};
+  auto RefData = makeArrayRef(InputData);
+  initializeInput(InputData, 1);
+
+  ArrayRef<uint8_t> Result;
+  BinaryStreamRef Original(InputData, support::little);
+  ASSERT_EQ(InputData.size(), Original.getLength());
+
+  EXPECT_NO_ERROR(Original.readBytes(0, InputData.size(), Result));
+  EXPECT_EQ(RefData, Result);
+
+  auto Dropped = Original.drop_front(2);
+  EXPECT_NO_ERROR(Dropped.readBytes(0, Dropped.getLength(), Result));
+  EXPECT_EQ(RefData.drop_front(2), Result);
+
+  Dropped = Original.drop_back(2);
+  EXPECT_NO_ERROR(Dropped.readBytes(0, Dropped.getLength(), Result));
+  EXPECT_EQ(RefData.drop_back(2), Result);
+
+  Dropped = Original.keep_front(2);
+  EXPECT_NO_ERROR(Dropped.readBytes(0, Dropped.getLength(), Result));
+  EXPECT_EQ(RefData.take_front(2), Result);
+
+  Dropped = Original.keep_back(2);
+  EXPECT_NO_ERROR(Dropped.readBytes(0, Dropped.getLength(), Result));
+  EXPECT_EQ(RefData.take_back(2), Result);
+
+  Dropped = Original.drop_symmetric(2);
+  EXPECT_NO_ERROR(Dropped.readBytes(0, Dropped.getLength(), Result));
+  EXPECT_EQ(RefData.drop_front(2).drop_back(2), Result);
+}
+
 // Test that we can write to a BinaryStream without a StreamWriter.
 TEST_F(BinaryStreamTest, MutableBinaryByteStreamBounds) {
   std::vector<uint8_t> InputData = {'T', 'e', 's', 't', '\0'};