From d2b418bfe2f568065d98b428e54d94cadb4d2437 Mon Sep 17 00:00:00 2001
From: Zachary Turner <zturner@google.com>
Date: Wed, 17 May 2017 20:42:52 +0000
Subject: [PATCH] Add some helpers for manipulating BinaryStreamRefs.

llvm-svn: 303297
---
 llvm/include/llvm/Support/BinaryStreamReader.h |  2 ++
 llvm/include/llvm/Support/BinaryStreamRef.h    | 25 ++++++++++++++++---
 llvm/lib/Support/BinaryStreamReader.cpp        |  5 ++++
 llvm/unittests/Support/BinaryStreamTest.cpp    | 33 ++++++++++++++++++++++++++
 4 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/Support/BinaryStreamReader.h b/llvm/include/llvm/Support/BinaryStreamReader.h
index 0c1881d..75e96a9 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 314668f..e3bd4bf 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 6b14890..5c27744 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 ec3b0ef..1ce74cb 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'};
-- 
2.7.4