2 files changed, 92 insertions, 1 deletions

diff --git a/llvm/include/llvm/Bitcode/BitstreamReader.h b/llvm/include/llvm/Bitcode/BitstreamReader.h
index 801df459a99..a7b1392483f 100644
--- a/llvm/include/llvm/Bitcode/BitstreamReader.h
+++ b/llvm/include/llvm/Bitcode/BitstreamReader.h
@@ -171,7 +171,7 @@ public:
     if (BitsInCurWord != 0)
       return false;
     if (Size != 0)
-      return Size == NextChar;
+      return Size <= NextChar;
     fillCurWord();
     return BitsInCurWord == 0;
   }
@@ -351,6 +351,28 @@ public:
 
   /// Skip to the end of the file.
   void skipToEnd() { NextChar = R->getBitcodeBytes().getExtent(); }
+
+  /// Prevent the cursor from reading past a byte boundary.
+  ///
+  /// Prevent the cursor from requesting byte reads past \c Limit.  This is
+  /// useful when working with a cursor on a StreamingMemoryObject, when it's
+  /// desirable to avoid invalidating the result of getPointerToByte().
+  ///
+  /// If \c Limit is on a word boundary, AtEndOfStream() will return true if
+  /// the cursor position reaches or exceeds \c Limit, regardless of the true
+  /// number of available bytes.  Otherwise, AtEndOfStream() returns true when
+  /// it reaches or exceeds the next word boundary.
+  void setArtificialByteLimit(uint64_t Limit) {
+    assert(getCurrentByteNo() < Limit && "Move cursor before lowering limit");
+
+    // Round to word boundary.
+    if (Limit & (sizeof(word_t) - 1))
+      Limit += sizeof(word_t) - Limit & (sizeof(word_t) - 1);
+
+    // Only change size if the new one is lower.
+    if (!Size || Size > Limit)
+      Size = Limit;
+  }
 };
 
 /// When advancing through a bitstream cursor, each advance can discover a few
diff --git a/llvm/unittests/Bitcode/BitstreamReaderTest.cpp b/llvm/unittests/Bitcode/BitstreamReaderTest.cpp
index 80285b84b32..9d3bf7bf14d 100644
--- a/llvm/unittests/Bitcode/BitstreamReaderTest.cpp
+++ b/llvm/unittests/Bitcode/BitstreamReaderTest.cpp
@@ -96,4 +96,73 @@ TEST(BitstreamReaderTest, jumpToPointer) {
   }
 }
 
+TEST(BitstreamReaderTest, setArtificialByteLimit) {
+  uint8_t Bytes[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+                     0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f};
+  BitstreamReader Reader(std::begin(Bytes), std::end(Bytes));
+  SimpleBitstreamCursor Cursor(Reader);
+
+  Cursor.setArtificialByteLimit(8);
+  while (!Cursor.AtEndOfStream())
+    (void)Cursor.Read(1);
+
+  EXPECT_EQ(8u, Cursor.getCurrentByteNo());
+}
+
+TEST(BitstreamReaderTest, setArtificialByteLimitNotWordBoundary) {
+  uint8_t Bytes[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+                     0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f};
+  BitstreamReader Reader(std::begin(Bytes), std::end(Bytes));
+  SimpleBitstreamCursor Cursor(Reader);
+
+  Cursor.setArtificialByteLimit(5);
+  while (!Cursor.AtEndOfStream())
+    (void)Cursor.Read(1);
+
+  EXPECT_EQ(8u, Cursor.getCurrentByteNo());
+}
+
+TEST(BitstreamReaderTest, setArtificialByteLimitNot4ByteBoundary) {
+  uint8_t Bytes[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+                     0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f};
+  BitstreamReader Reader(std::begin(Bytes), std::end(Bytes));
+  SimpleBitstreamCursor Cursor(Reader);
+
+  Cursor.setArtificialByteLimit(5);
+  while (!Cursor.AtEndOfStream())
+    (void)Cursor.Read(1);
+
+  EXPECT_EQ(8u, Cursor.getCurrentByteNo());
+}
+
+TEST(BitstreamReaderTest, setArtificialByteLimitPastTheEnd) {
+  uint8_t Bytes[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+                     0x08, 0x09, 0x0a, 0x0b};
+  BitstreamReader Reader(std::begin(Bytes), std::end(Bytes));
+  SimpleBitstreamCursor Cursor(Reader);
+
+  // The size of the memory object isn't known yet.  Set it too high and
+  // confirm that we don't read too far.
+  Cursor.setArtificialByteLimit(20);
+  while (!Cursor.AtEndOfStream())
+    (void)Cursor.Read(1);
+
+  EXPECT_EQ(12u, Cursor.getCurrentByteNo());
+}
+
+TEST(BitstreamReaderTest, setArtificialByteLimitPastTheEndKnown) {
+  uint8_t Bytes[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+                     0x08, 0x09, 0x0a, 0x0b};
+  BitstreamReader Reader(std::begin(Bytes), std::end(Bytes));
+  SimpleBitstreamCursor Cursor(Reader);
+
+  // Save the size of the memory object in the cursor.
+  while (!Cursor.AtEndOfStream())
+    (void)Cursor.Read(1);
+  EXPECT_EQ(12u, Cursor.getCurrentByteNo());
+
+  Cursor.setArtificialByteLimit(20);
+  EXPECT_TRUE(Cursor.AtEndOfStream());
+}
+
 } // end anonymous namespace