Index: include/llvm/Support/StreamingMemoryObject.h =================================================================== --- include/llvm/Support/StreamingMemoryObject.h +++ include/llvm/Support/StreamingMemoryObject.h @@ -58,26 +58,27 @@ mutable size_t ObjectSize; // 0 if unknown, set if wrapper seen or EOF reached mutable bool EOFReached; - // Fetch enough bytes such that Pos can be read or EOF is reached - // (i.e. BytesRead > Pos). Return true if Pos can be read. - // Unlike most of the functions in BitcodeReader, returns true on success. - // Most of the requests will be small, but we fetch at kChunkSize bytes - // at a time to avoid making too many potentially expensive GetBytes calls + // Fetch enough bytes such that Pos can be read (i.e. BytesRead > + // Pos). Returns true if Pos can be read. Unlike most of the + // functions in BitcodeReader, returns true on success. Most of the + // requests will be small, but we fetch at kChunkSize bytes at a + // time to avoid making too many potentially expensive GetBytes + // calls. bool fetchToPos(size_t Pos) const { - if (EOFReached) - return Pos < ObjectSize; while (Pos >= BytesRead) { + if (EOFReached) + return false; Bytes.resize(BytesRead + BytesSkipped + kChunkSize); size_t bytes = Streamer->GetBytes(&Bytes[BytesRead + BytesSkipped], kChunkSize); BytesRead += bytes; - if (bytes != kChunkSize) { // reached EOF/ran out of bytes - ObjectSize = BytesRead; + if (bytes == 0) { // reached EOF/ran out of bytes + if (ObjectSize == 0) + ObjectSize = BytesRead; EOFReached = true; - break; } } - return Pos < BytesRead; + return !ObjectSize || Pos < ObjectSize; } StreamingMemoryObject(const StreamingMemoryObject&) = delete; Index: lib/Support/StreamingMemoryObject.cpp =================================================================== --- lib/Support/StreamingMemoryObject.cpp +++ lib/Support/StreamingMemoryObject.cpp @@ -73,7 +73,7 @@ // block until we actually want to read it. bool StreamingMemoryObject::isValidAddress(uint64_t address) const { if (ObjectSize && address < ObjectSize) return true; - return fetchToPos(address); + return fetchToPos(address); } uint64_t StreamingMemoryObject::getExtent() const { @@ -87,18 +87,30 @@ uint64_t StreamingMemoryObject::readBytes(uint8_t *Buf, uint64_t Size, uint64_t Address) const { fetchToPos(Address + Size - 1); - if (Address >= BytesRead) + // Note: For wrapped bitcode files will set ObjectSize after the + // first call to fetchToPos. In such cases, ObjectSize can be + // smaller than BytesRead. + size_t MaxAddress = + (ObjectSize && ObjectSize < BytesRead) ? ObjectSize : BytesRead; + if (Address >= MaxAddress) return 0; uint64_t End = Address + Size; - if (End > BytesRead) - End = BytesRead; - assert(static_cast(End - Address) >= 0); + if (End > MaxAddress) + End = MaxAddress; + assert(End >= Address); Size = End - Address; memcpy(Buf, &Bytes[Address + BytesSkipped], Size); return Size; } +void StreamingMemoryObject::setKnownObjectSize(size_t size) { + ObjectSize = size; + Bytes.reserve(size); + if (ObjectSize <= BytesRead) + EOFReached = true; +} + bool StreamingMemoryObject::dropLeadingBytes(size_t s) { if (BytesRead < s) return true; BytesSkipped = s; @@ -106,11 +118,6 @@ return false; } -void StreamingMemoryObject::setKnownObjectSize(size_t size) { - ObjectSize = size; - Bytes.reserve(size); -} - MemoryObject *getNonStreamedMemoryObject(const unsigned char *Start, const unsigned char *End) { return new RawMemoryObject(Start, End); Index: unittests/Support/StreamingMemoryObject.cpp =================================================================== --- unittests/Support/StreamingMemoryObject.cpp +++ unittests/Support/StreamingMemoryObject.cpp @@ -27,3 +27,12 @@ StreamingMemoryObject O(DS); EXPECT_TRUE(O.isValidAddress(32 * 1024)); } + +TEST(StreamingMemoryObject, TestSetKnownObjectSize) { + auto *DS = new NullDataStreamer(); + StreamingMemoryObject O(DS); + uint8_t Buf[32]; + EXPECT_EQ((uint64_t) 16, O.readBytes(Buf, 16, 0)); + O.setKnownObjectSize(24); + EXPECT_EQ((uint64_t) 8, O.readBytes(Buf, 16, 16)); +}