getStreamedBitcodeModule: take StreamingMemoryObject instead of DataStreamer

[jvoung]

Summary

Currently, getStreamedBitcodeModule takes a DataStreamer and then
creates a StreamingMemoryObject from that DataStreamer for the BitcodeReader.
That allows flexibility in the supplied DataStreamer implementation but does
not allow flexibility in the StreamingMemoryObject implementation.

For example, the StreamingMemoryObject might want the flexibilty of caching bytes
in a thread-local way.

Make StreamingMemoryObject the interface instead so that both DataStreamer
and StreamingMemoryObject implementations can be customized.

[jvoung]

include/llvm/Bitcode/ReaderWriter.h

 #define LLVM_BITCODE_READERWRITER_H
 
 #include "llvm/Support/ErrorOr.h"
 #include <string>
 
 namespace llvm {
   class BitstreamWriter;
   class MemoryBuffer;
-  class DataStreamer;
+  class StreamingMemoryObject;
   class LLVMContext;
   class Module;
   class ModulePass;
   class raw_ostream;
 
   /// Read the header of the specified bitcode buffer and prepare for lazy
   /// deserialization of function bodies.  If successful, this takes ownership
   /// of 'buffer. On error, this *does not* take ownership of Buffer.
   ErrorOr<Module *> getLazyBitcodeModule(MemoryBuffer *Buffer,
                                          LLVMContext &Context);
 
   /// getStreamedBitcodeModule - Read the header of the specified stream
   /// and prepare for lazy deserialization and streaming of function bodies.
+  /// This takes ownership of 'Streamer' on both success and error.
   /// On error, this returns null, and fills in *ErrMsg with an error
   /// description if ErrMsg is non-null.
   Module *getStreamedBitcodeModule(const std::string &name,
-                                   DataStreamer *streamer,
+                                   StreamingMemoryObject *Streamer,
                                    LLVMContext &Context,
                                    std::string *ErrMsg = nullptr);
 
   /// getBitcodeTargetTriple - Read the header of the specified bitcode
   /// buffer and extract just the triple information. If successful,
   /// this returns a string and *does not* take ownership
   /// of 'buffer'. On error, this returns "", and fills in *ErrMsg
   /// if ErrMsg is non-null.

include/llvm/Support/StreamableMemoryObject.h

public:
   /// isObjectEnd    - Returns true if the address is one past the end of the
   ///                  object (i.e. if it is equal to base + extent)
   ///                  May block until (address - base) bytes have been read
   /// @param address - address of the byte, in the same space as getBase()
   /// @result        - true if the address is equal to base + extent
   virtual bool isObjectEnd(uint64_t address) const = 0;
 };
 
+
 /// StreamingMemoryObject - interface to data which is actually streamed from
 /// a DataStreamer. In addition to inherited members, it has the
 /// dropLeadingBytes and setKnownObjectSize methods which are not applicable
 /// to non-streamed objects.
 class StreamingMemoryObject : public StreamableMemoryObject {
 public:
-  StreamingMemoryObject(DataStreamer *streamer);
+  /// Drop s bytes from the front of the stream, pushing the positions of the
+  /// remaining bytes down by s. This is used to skip past the bitcode header,
+  /// since we don't know a priori if it's present, and we can't put bytes
+  /// back into the stream once we've read them.
+  virtual bool dropLeadingBytes(size_t s) = 0;
+
+  /// If the data object size is known in advance, many of the operations can
+  /// be made more efficient, so this method should be called before reading
+  /// starts (although it can be called anytime).
+  virtual void setKnownObjectSize(size_t size) = 0;
+};
+
+
+/// StreamingMemoryObjectImpl - An implementation of a StreamingMemoryObject.
+class StreamingMemoryObjectImpl : public StreamingMemoryObject {
+public:
+  /// Constructs a StreamingMemoryObject and takes ownership of DataStreamer.
+  StreamingMemoryObjectImpl(DataStreamer *streamer);
   uint64_t getBase() const override { return 0; }
   uint64_t getExtent() const override;
   int readByte(uint64_t address, uint8_t *ptr) const override;
   int readBytes(uint64_t address, uint64_t size,
                 uint8_t *buf) const override;
   const uint8_t *getPointer(uint64_t address, uint64_t size) const override {
     // 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 don't stream.
     assert(0 && "getPointer in streaming memory objects not allowed");
     return nullptr;
   }
   bool isValidAddress(uint64_t address) const override;
   bool isObjectEnd(uint64_t address) const override;
 
   /// Drop s bytes from the front of the stream, pushing the positions of the
   /// remaining bytes down by s. This is used to skip past the bitcode header,
   /// since we don't know a priori if it's present, and we can't put bytes
   /// back into the stream once we've read them.
-  bool dropLeadingBytes(size_t s);
+  bool dropLeadingBytes(size_t s) override;
 
   /// If the data object size is known in advance, many of the operations can
   /// be made more efficient, so this method should be called before reading
   /// starts (although it can be called anytime).
-  void setKnownObjectSize(size_t size);
+  void setKnownObjectSize(size_t size) override;
 
 private:
   const static uint32_t kChunkSize = 4096 * 4;
   mutable std::vector<unsigned char> Bytes;
   std::unique_ptr<DataStreamer> Streamer;
   mutable size_t BytesRead;   // Bytes read from stream
   size_t BytesSkipped;// Bytes skipped at start of stream (e.g. wrapper/header)
   mutable size_t ObjectSize; // 0 if unknown, set if wrapper seen or EOF reached
while (Pos >= BytesRead) {
           EOFReached = true;
           return false;
         }
       }
     }
     return true;
   }
 
-  StreamingMemoryObject(const StreamingMemoryObject&) LLVM_DELETED_FUNCTION;
-  void operator=(const StreamingMemoryObject&) LLVM_DELETED_FUNCTION;
+  StreamingMemoryObjectImpl(
+      const StreamingMemoryObjectImpl&) LLVM_DELETED_FUNCTION;
+  void operator=(
+      const StreamingMemoryObjectImpl&) LLVM_DELETED_FUNCTION;
 };
 
 StreamableMemoryObject *getNonStreamedMemoryObject(
     const unsigned char *Start, const unsigned char *End);
 
 }
 #endif  // STREAMABLEMEMORYOBJECT_H_

lib/Bitcode/Reader/BitcodeReader.h

 
 class BitcodeReader : public GVMaterializer {
   LLVMContext &Context;
   Module *TheModule;
   MemoryBuffer *Buffer;
   bool BufferOwned;
   std::unique_ptr<BitstreamReader> StreamFile;
   BitstreamCursor Stream;
-  DataStreamer *LazyStreamer;
+  StreamingMemoryObject* LazyStreamer;  // owned by StreamFile.
   uint64_t NextUnreadBit;
   bool SeenValueSymbolTable;
 
   std::vector<Type*> TypeList;
   BitcodeReaderValueList ValueList;
   BitcodeReaderMDValueList MDValueList;
   SmallVector<Instruction *, 64> InstructionList;
   SmallVector<SmallVector<uint64_t, 64>, 64> UseListRecords;
public:
   }
 
   explicit BitcodeReader(MemoryBuffer *buffer, LLVMContext &C)
     : Context(C), TheModule(0), Buffer(buffer), BufferOwned(false),
       LazyStreamer(0), NextUnreadBit(0), SeenValueSymbolTable(false),
       ValueList(C), MDValueList(C),
       SeenFirstFunctionBody(false), UseRelativeIDs(false) {
   }
-  explicit BitcodeReader(DataStreamer *streamer, LLVMContext &C)
+  explicit BitcodeReader(StreamingMemoryObject *Streamer, LLVMContext &C)
     : Context(C), TheModule(0), Buffer(0), BufferOwned(false),
-      LazyStreamer(streamer), NextUnreadBit(0), SeenValueSymbolTable(false),
+      LazyStreamer(Streamer), NextUnreadBit(0), SeenValueSymbolTable(false),
       ValueList(C), MDValueList(C),
       SeenFirstFunctionBody(false), UseRelativeIDs(false) {
   }
   ~BitcodeReader() {
     FreeState();
   }
 
   void materializeForwardReferencedFunctions();

lib/Bitcode/Reader/BitcodeReader.cpp

error_code BitcodeReader::InitStreamFromBuffer() {
   Stream.init(*StreamFile);
 
   return error_code::success();
 }
 
 error_code BitcodeReader::InitLazyStream() {
   // Check and strip off the bitcode wrapper; BitstreamReader expects never to
   // see it.
-  StreamingMemoryObject *Bytes = new StreamingMemoryObject(LazyStreamer);
-  StreamFile.reset(new BitstreamReader(Bytes));
+  StreamFile.reset(new BitstreamReader(LazyStreamer));
   Stream.init(*StreamFile);
 
   unsigned char buf[16];
-  if (Bytes->readBytes(0, 16, buf) == -1)
+  if (LazyStreamer->readBytes(0, 16, buf) == -1)
     return Error(BitcodeStreamInvalidSize);
 
   if (!isBitcode(buf, buf + 16))
     return Error(InvalidBitcodeSignature);
 
   if (isBitcodeWrapper(buf, buf + 4)) {
     const unsigned char *bitcodeStart = buf;
     const unsigned char *bitcodeEnd = buf + 16;
     SkipBitcodeWrapperHeader(bitcodeStart, bitcodeEnd, false);
-    Bytes->dropLeadingBytes(bitcodeStart - buf);
-    Bytes->setKnownObjectSize(bitcodeEnd - bitcodeStart);
+    LazyStreamer->dropLeadingBytes(bitcodeStart - buf);
+    LazyStreamer->setKnownObjectSize(bitcodeEnd - bitcodeStart);
   }
   return error_code::success();
 }
 
 namespace {
 class BitcodeErrorCategoryType : public error_category {
   const char *name() const override {
     return "llvm.bitcode";
ErrorOr<Module *> llvm::getLazyBitcodeModule(MemoryBuffer *Buffer,
 
   R->materializeForwardReferencedFunctions();
 
   return M;
 }
 
 
 Module *llvm::getStreamedBitcodeModule(const std::string &name,
-                                       DataStreamer *streamer,
+                                       StreamingMemoryObject *Streamer,
                                        LLVMContext &Context,
                                        std::string *ErrMsg) {
   Module *M = new Module(name, Context);
-  BitcodeReader *R = new BitcodeReader(streamer, Context);
+  BitcodeReader *R = new BitcodeReader(Streamer, Context);
   M->setMaterializer(R);
   if (error_code EC = R->ParseBitcodeInto(M)) {
     if (ErrMsg)
       *ErrMsg = EC.message();
     delete M;  // Also deletes R.
     return 0;
   }
   R->setBufferOwned(false); // no buffer to delete

lib/Support/StreamableMemoryObject.cpp

const uint8_t *RawMemoryObject::getPointer(uint64_t address,
                                            uint64_t size) const {
   return FirstChar + address;
 }
 } // anonymous namespace
 
 namespace llvm {
 // If the bitcode has a header, then its size is known, and we don't have to
 // block until we actually want to read it.
-bool StreamingMemoryObject::isValidAddress(uint64_t address) const {
+bool StreamingMemoryObjectImpl::isValidAddress(uint64_t address) const {
   if (ObjectSize && address < ObjectSize) return true;
     return fetchToPos(address);
 }
 
-bool StreamingMemoryObject::isObjectEnd(uint64_t address) const {
+bool StreamingMemoryObjectImpl::isObjectEnd(uint64_t address) const {
   if (ObjectSize) return address == ObjectSize;
   fetchToPos(address);
   return address == ObjectSize && address != 0;
 }
 
-uint64_t StreamingMemoryObject::getExtent() const {
+uint64_t StreamingMemoryObjectImpl::getExtent() const {
   if (ObjectSize) return ObjectSize;
   size_t pos = BytesRead + kChunkSize;
   // keep fetching until we run out of bytes
   while (fetchToPos(pos)) pos += kChunkSize;
   return ObjectSize;
 }
 
-int StreamingMemoryObject::readByte(uint64_t address, uint8_t* ptr) const {
+int StreamingMemoryObjectImpl::readByte(uint64_t address, uint8_t* ptr) const {
   if (!fetchToPos(address)) return -1;
   *ptr = Bytes[address + BytesSkipped];
   return 0;
 }
 
-int StreamingMemoryObject::readBytes(uint64_t address,
+int StreamingMemoryObjectImpl::readBytes(uint64_t address,
                                      uint64_t size,
                                      uint8_t *buf) const {
   if (!fetchToPos(address + size - 1)) return -1;
   memcpy(buf, &Bytes[address + BytesSkipped], size);
   return 0;
 }
 
-bool StreamingMemoryObject::dropLeadingBytes(size_t s) {
+bool StreamingMemoryObjectImpl::dropLeadingBytes(size_t s) {
   if (BytesRead < s) return true;
   BytesSkipped = s;
   BytesRead -= s;
   return false;
 }
 
-void StreamingMemoryObject::setKnownObjectSize(size_t size) {
+void StreamingMemoryObjectImpl::setKnownObjectSize(size_t size) {
   ObjectSize = size;
   Bytes.reserve(size);
 }
 
 StreamableMemoryObject *getNonStreamedMemoryObject(
     const unsigned char *Start, const unsigned char *End) {
   return new RawMemoryObject(Start, End);
 }
 
 StreamableMemoryObject::~StreamableMemoryObject() { }
 
-StreamingMemoryObject::StreamingMemoryObject(DataStreamer *streamer) :
+StreamingMemoryObjectImpl::StreamingMemoryObjectImpl(
+    DataStreamer *streamer) :
   Bytes(kChunkSize), Streamer(streamer), BytesRead(0), BytesSkipped(0),
   ObjectSize(0), EOFReached(false) {
   BytesRead = streamer->GetBytes(&Bytes[0], kChunkSize);
 }
 }

tools/llvm-dis/llvm-dis.cpp

 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/DataStream.h"
 #include "llvm/Support/FormattedStream.h"
 #include "llvm/Support/ManagedStatic.h"
-#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/Signals.h"
+#include "llvm/Support/StreamableMemoryObject.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/system_error.h"
 using namespace llvm;
 
 static cl::opt<std::string>
 InputFilename(cl::Positional, cl::desc("<input bitcode>"), cl::init("-"));
 
 static cl::opt<std::string>
int main(int argc, char **argv) {
 
 
   cl::ParseCommandLineOptions(argc, argv, "llvm .bc -> .ll disassembler\n");
 
   std::string ErrorMessage;
   std::unique_ptr<Module> M;
 
   // Use the bitcode streaming interface
-  DataStreamer *streamer = getDataFileStreamer(InputFilename, &ErrorMessage);
-  if (streamer) {
+  DataStreamer *Streamer = getDataFileStreamer(InputFilename, &ErrorMessage);
+  if (Streamer) {
+    std::unique_ptr<StreamingMemoryObjectImpl> StreamBytes(
+        new StreamingMemoryObjectImpl(Streamer));
     std::string DisplayFilename;
     if (InputFilename == "-")
       DisplayFilename = "<stdin>";
     else
       DisplayFilename = InputFilename;
-    M.reset(getStreamedBitcodeModule(DisplayFilename, streamer, Context,
-                                     &ErrorMessage));
+    M.reset(getStreamedBitcodeModule(DisplayFilename, StreamBytes.release(),
+                                     Context, &ErrorMessage));
     if(M.get() != 0) {
       if (error_code EC = M->materializeAllPermanently()) {
         ErrorMessage = EC.message();
         M.reset();
       }
     }
   }