Support: Implement StreamingMemoryObject::getPointer

The implementation is fairly obvious.  This is preparation for using
some blobs in bitcode.

For clarity (and perhaps future-proofing?), I moved the call to
JumpToBit in BitstreamCursor::readRecord ahead of calling
MemoryObject::getPointer, since JumpToBit can theoretically (a) read
bytes, which (b) invalidates the blob pointer.

This isn't strictly necessary the two memory objects we have:

  - The return of RawMemoryObject::getPointer is valid until the memory
    object is destroyed.

  - StreamingMemoryObject::getPointer is valid until the next chunk is
    read from the stream.  Since the JumpToBit call is only going ahead
    to a word boundary, we'll never load another chunk.

However, reordering makes it clear by inspection that the blob returned
by BitstreamCursor::readRecord will be valid.

I added some tests for StreamingMemoryObject::getPointer and
BitstreamCursor::readRecord.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@264549 91177308-0d34-0410-b5e6-96231b3b80d8

7 files changed, 116 insertions, 17 deletions

diff --git a/include/llvm/Support/StreamingMemoryObject.h b/include/llvm/Support/StreamingMemoryObject.h
index a5980c23594..1ab85372cd2 100644
--- a/include/llvm/Support/StreamingMemoryObject.h
+++ b/include/llvm/Support/StreamingMemoryObject.h
@@ -28,15 +28,7 @@ public:
   uint64_t getExtent() const override;
   uint64_t readBytes(uint8_t *Buf, uint64_t Size,
                      uint64_t Address) const override;
-  const uint8_t *getPointer(uint64_t address, uint64_t size) const override {
-    // FIXME: This could be fixed by ensuring the bytes are fetched and
-    // making a copy, requiring that the bitcode size be known, or
-    // otherwise ensuring that the memory doesn't go away/get reallocated,
-    // but it's not currently necessary. Users that need the pointer (any
-    // that need Blobs) don't stream.
-    report_fatal_error("getPointer in streaming memory objects not allowed");
-    return nullptr;
-  }
+  const uint8_t *getPointer(uint64_t Address, uint64_t Size) const override;
   bool isValidAddress(uint64_t address) const override;
 
   /// Drop s bytes from the front of the stream, pushing the positions of the
diff --git a/lib/Bitcode/Reader/BitstreamReader.cpp b/lib/Bitcode/Reader/BitstreamReader.cpp
index fe3f6e8b59a..60360d2ef78 100644
--- a/lib/Bitcode/Reader/BitstreamReader.cpp
+++ b/lib/Bitcode/Reader/BitstreamReader.cpp
@@ -260,7 +260,10 @@ unsigned BitstreamCursor::readRecord(unsigned AbbrevID,
       break;
     }
 
-    // Otherwise, inform the streamer that we need these bytes in memory.
+    // Otherwise, inform the streamer that we need these bytes in memory.  Skip
+    // over tail padding first, in case jumping to NewEnd invalidates the Blob
+    // pointer.
+    JumpToBit(NewEnd);
     const char *Ptr = (const char *)getPointerToBit(CurBitPos, NumElts);
 
     // If we can return a reference to the data, do so to avoid copying it.
@@ -271,8 +274,6 @@ unsigned BitstreamCursor::readRecord(unsigned AbbrevID,
       for (; NumElts; --NumElts)
         Vals.push_back((unsigned char)*Ptr++);
     }
-    // Skip over tail padding.
-    JumpToBit(NewEnd);
   }
 
   return Code;
diff --git a/lib/Support/StreamingMemoryObject.cpp b/lib/Support/StreamingMemoryObject.cpp
index 5a44e624eb8..fb566179486 100644
--- a/lib/Support/StreamingMemoryObject.cpp
+++ b/lib/Support/StreamingMemoryObject.cpp
@@ -104,6 +104,12 @@ uint64_t StreamingMemoryObject::readBytes(uint8_t *Buf, uint64_t Size,
   return Size;
 }
 
+const uint8_t *StreamingMemoryObject::getPointer(uint64_t Address,
+                                                 uint64_t Size) const {
+  fetchToPos(Address + Size - 1);
+  return &Bytes[Address + BytesSkipped];
+}
+
 bool StreamingMemoryObject::dropLeadingBytes(size_t s) {
   if (BytesRead < s) return true;
   BytesSkipped = s;
diff --git a/test/Bitcode/Inputs/invalid-fixme-streaming-blob.bc b/test/Bitcode/Inputs/invalid-fixme-streaming-blob.bc
deleted file mode 100644
index 7e32f8b0774..00000000000
--- a/test/Bitcode/Inputs/invalid-fixme-streaming-blob.bc
+++ /dev/null
diff --git a/test/Bitcode/invalid.test b/test/Bitcode/invalid.test
index 3425adc8410..499320140d7 100644
--- a/test/Bitcode/invalid.test
+++ b/test/Bitcode/invalid.test
@@ -168,11 +168,6 @@ RUN:   FileCheck --check-prefix=INVALID-ARGUMENT-TYPE %s
 
 INVALID-ARGUMENT-TYPE: Invalid function argument type
 
-RUN: not llvm-dis -disable-output %p/Inputs/invalid-fixme-streaming-blob.bc 2>&1 | \
-RUN:   FileCheck --check-prefix=STREAMING-BLOB %s
-
-STREAMING-BLOB: getPointer in streaming memory objects not allowed
-
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-function-comdat-id.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=INVALID-FCOMDAT-ID %s
 
diff --git a/unittests/Bitcode/BitstreamReaderTest.cpp b/unittests/Bitcode/BitstreamReaderTest.cpp
index 9d3bf7bf14d..935980fc4df 100644
--- a/unittests/Bitcode/BitstreamReaderTest.cpp
+++ b/unittests/Bitcode/BitstreamReaderTest.cpp
@@ -7,13 +7,31 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Bitcode/BitstreamReader.h"
+#include "llvm/Bitcode/BitstreamWriter.h"
+#include "llvm/Support/StreamingMemoryObject.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 namespace {
 
+class BufferStreamer : public DataStreamer {
+  StringRef Buffer;
+
+public:
+  BufferStreamer(StringRef Buffer) : Buffer(Buffer) {}
+  size_t GetBytes(unsigned char *OutBuffer, size_t Length) override {
+    if (Length >= Buffer.size())
+      Length = Buffer.size();
+
+    std::copy(Buffer.begin(), Buffer.begin() + Length, OutBuffer);
+    Buffer = Buffer.drop_front(Length);
+    return Length;
+  }
+};
+
 TEST(BitstreamReaderTest, AtEndOfStream) {
   uint8_t Bytes[4] = {
     0x00, 0x01, 0x02, 0x03
@@ -165,4 +183,65 @@ TEST(BitstreamReaderTest, setArtificialByteLimitPastTheEndKnown) {
   EXPECT_TRUE(Cursor.AtEndOfStream());
 }
 
+TEST(BitstreamReaderTest, readRecordWithBlobWhileStreaming) {
+  SmallVector<uint8_t, 1> BlobData;
+  for (unsigned I = 0, E = 1024; I != E; ++I)
+    BlobData.push_back(I);
+
+  // Try a bunch of different sizes.
+  const unsigned Magic = 0x12345678;
+  const unsigned BlockID = bitc::FIRST_APPLICATION_BLOCKID;
+  const unsigned RecordID = 1;
+  for (unsigned I = 0, BlobSize = 0, E = BlobData.size(); BlobSize < E;
+       BlobSize += ++I) {
+    StringRef BlobIn((const char *)BlobData.begin(), BlobSize);
+
+    // Write the bitcode.
+    SmallVector<char, 1> Buffer;
+    unsigned AbbrevID;
+    {
+      BitstreamWriter Stream(Buffer);
+      Stream.Emit(Magic, 32);
+      Stream.EnterSubblock(BlockID, 3);
+
+      BitCodeAbbrev *Abbrev = new BitCodeAbbrev();
+      Abbrev->Add(BitCodeAbbrevOp(RecordID));
+      Abbrev->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Blob));
+      AbbrevID = Stream.EmitAbbrev(Abbrev);
+      unsigned Record[] = {RecordID};
+      Stream.EmitRecordWithBlob(AbbrevID, makeArrayRef(Record), BlobIn);
+
+      Stream.ExitBlock();
+    }
+
+    // Stream the buffer into the reader.
+    BitstreamReader R(make_unique<StreamingMemoryObject>(
+        make_unique<BufferStreamer>(StringRef(Buffer.begin(), Buffer.size()))));
+    BitstreamCursor Stream(R);
+
+    // Header.  Included in test so that we can run llvm-bcanalyzer to debug
+    // when there are problems.
+    ASSERT_EQ(Magic, Stream.Read(32));
+
+    // Block.
+    BitstreamEntry Entry =
+        Stream.advance(BitstreamCursor::AF_DontAutoprocessAbbrevs);
+    ASSERT_EQ(BitstreamEntry::SubBlock, Entry.Kind);
+    ASSERT_EQ(BlockID, Entry.ID);
+    ASSERT_FALSE(Stream.EnterSubBlock(BlockID));
+
+    // Abbreviation.
+    Entry = Stream.advance();
+    ASSERT_EQ(BitstreamEntry::Record, Entry.Kind);
+    ASSERT_EQ(AbbrevID, Entry.ID);
+
+    // Record.
+    StringRef BlobOut;
+    SmallVector<uint64_t, 1> Record;
+    ASSERT_EQ(RecordID, Stream.readRecord(Entry.ID, Record, &BlobOut));
+    EXPECT_TRUE(Record.empty());
+    EXPECT_EQ(BlobIn, BlobOut);
+  }
+}
+
 } // end anonymous namespace
diff --git a/unittests/Support/StreamingMemoryObjectTest.cpp b/unittests/Support/StreamingMemoryObjectTest.cpp
index 261f2144392..836dfa9084f 100644
--- a/unittests/Support/StreamingMemoryObjectTest.cpp
+++ b/unittests/Support/StreamingMemoryObjectTest.cpp
@@ -24,6 +24,21 @@ class NullDataStreamer : public DataStreamer {
   }
 };
 
+class BufferStreamer : public DataStreamer {
+  StringRef Buffer;
+
+public:
+  BufferStreamer(StringRef Buffer) : Buffer(Buffer) {}
+  size_t GetBytes(unsigned char *OutBuffer, size_t Length) override {
+    if (Length >= Buffer.size())
+      Length = Buffer.size();
+
+    std::copy(Buffer.begin(), Buffer.begin() + Length, OutBuffer);
+    Buffer = Buffer.drop_front(Length);
+    return Length;
+  }
+};
+
 TEST(StreamingMemoryObjectTest, isValidAddress) {
   auto DS = make_unique<NullDataStreamer>();
   StreamingMemoryObject O(std::move(DS));
@@ -39,4 +54,15 @@ TEST(StreamingMemoryObjectTest, setKnownObjectSize) {
   EXPECT_EQ(8u, O.readBytes(Buf, 16, 16));
 }
 
+TEST(StreamingMemoryObjectTest, getPointer) {
+  uint8_t InputBuffer[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07};
+  StreamingMemoryObject O(make_unique<BufferStreamer>(StringRef(
+      reinterpret_cast<const char *>(InputBuffer), sizeof(InputBuffer))));
+
+  EXPECT_TRUE(std::equal(InputBuffer + 1, InputBuffer + 2, O.getPointer(1, 2)));
+  EXPECT_TRUE(std::equal(InputBuffer + 3, InputBuffer + 7, O.getPointer(3, 4)));
+  EXPECT_TRUE(std::equal(InputBuffer + 4, InputBuffer + 8, O.getPointer(4, 5)));
+  EXPECT_TRUE(std::equal(InputBuffer, InputBuffer + 8, O.getPointer(0, 20)));
+}
+
 } // end namespace