diff --git a/lib/Bitcode/Reader/BitcodeReader.cpp b/lib/Bitcode/Reader/BitcodeReader.cpp index 541b056c08f..10e0410c025 100644 --- a/lib/Bitcode/Reader/BitcodeReader.cpp +++ b/lib/Bitcode/Reader/BitcodeReader.cpp @@ -4597,6 +4597,35 @@ const std::error_category &llvm::BitcodeErrorCategory() { // External interface //===----------------------------------------------------------------------===// +static ErrorOr> +getBitcodeModuleImpl(std::unique_ptr Streamer, StringRef Name, + BitcodeReader *R, LLVMContext &Context, + bool MaterializeAll, bool ShouldLazyLoadMetadata) { + std::unique_ptr M = make_unique(Name, Context); + M->setMaterializer(R); + + auto cleanupOnError = [&](std::error_code EC) { + R->releaseBuffer(); // Never take ownership on error. + return EC; + }; + + // Delay parsing Metadata if ShouldLazyLoadMetadata is true. + if (std::error_code EC = R->parseBitcodeInto(std::move(Streamer), M.get(), + ShouldLazyLoadMetadata)) + return cleanupOnError(EC); + + if (MaterializeAll) { + // Read in the entire module, and destroy the BitcodeReader. + if (std::error_code EC = M->materializeAllPermanently()) + return cleanupOnError(EC); + } else { + // Resolve forward references from blockaddresses. + if (std::error_code EC = R->materializeForwardReferencedFunctions()) + return cleanupOnError(EC); + } + return std::move(M); +} + /// \brief Get a lazy one-at-time loading module from bitcode. /// /// This isn't always used in a lazy context. In particular, it's also used by @@ -4610,34 +4639,17 @@ getLazyBitcodeModuleImpl(std::unique_ptr &&Buffer, LLVMContext &Context, bool MaterializeAll, DiagnosticHandlerFunction DiagnosticHandler, bool ShouldLazyLoadMetadata = false) { - std::unique_ptr M = - make_unique(Buffer->getBufferIdentifier(), Context); BitcodeReader *R = new BitcodeReader(Buffer.get(), Context, DiagnosticHandler); - M->setMaterializer(R); - auto cleanupOnError = [&](std::error_code EC) { - R->releaseBuffer(); // Never take ownership on error. - return EC; - }; - - // Delay parsing Metadata if ShouldLazyLoadMetadata is true. - if (std::error_code EC = - R->parseBitcodeInto(nullptr, M.get(), ShouldLazyLoadMetadata)) - return cleanupOnError(EC); - - if (MaterializeAll) { - // Read in the entire module, and destroy the BitcodeReader. - if (std::error_code EC = M->materializeAllPermanently()) - return EC; - } else { - // Resolve forward references from blockaddresses. - if (std::error_code EC = R->materializeForwardReferencedFunctions()) - return cleanupOnError(EC); - } + ErrorOr> Ret = + getBitcodeModuleImpl(nullptr, Buffer->getBufferIdentifier(), R, Context, + MaterializeAll, ShouldLazyLoadMetadata); + if (!Ret) + return Ret; Buffer.release(); // The BitcodeReader owns it now. - return std::move(M); + return Ret; } ErrorOr> llvm::getLazyBitcodeModule( @@ -4652,10 +4664,9 @@ ErrorOr> llvm::getStreamedBitcodeModule( LLVMContext &Context, DiagnosticHandlerFunction DiagnosticHandler) { std::unique_ptr M = make_unique(Name, Context); BitcodeReader *R = new BitcodeReader(Context, DiagnosticHandler); - M->setMaterializer(R); - if (std::error_code EC = R->parseBitcodeInto(std::move(Streamer), M.get())) - return EC; - return std::move(M); + + return getBitcodeModuleImpl(std::move(Streamer), Name, R, Context, false, + false); } ErrorOr> diff --git a/unittests/Bitcode/BitReaderTest.cpp b/unittests/Bitcode/BitReaderTest.cpp index 691a217d20e..81d3a28a004 100644 --- a/unittests/Bitcode/BitReaderTest.cpp +++ b/unittests/Bitcode/BitReaderTest.cpp @@ -8,6 +8,7 @@ //===----------------------------------------------------------------------===// #include "llvm/ADT/SmallString.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/AsmParser/Parser.h" #include "llvm/Bitcode/BitstreamWriter.h" #include "llvm/Bitcode/ReaderWriter.h" @@ -16,6 +17,7 @@ #include "llvm/IR/LLVMContext.h" #include "llvm/IR/Module.h" #include "llvm/IR/Verifier.h" +#include "llvm/Support/DataStream.h" #include "llvm/Support/Debug.h" #include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/SourceMgr.h" @@ -58,6 +60,49 @@ static std::unique_ptr getLazyModuleFromAssembly(LLVMContext &Context, return std::move(ModuleOrErr.get()); } +class BufferDataStreamer : public DataStreamer { + std::unique_ptr Buffer; + unsigned Pos = 0; + size_t GetBytes(unsigned char *Out, size_t Len) override { + StringRef Buf = Buffer->getBuffer(); + size_t Left = Buf.size() - Pos; + Len = std::min(Left, Len); + memcpy(Out, Buffer->getBuffer().substr(Pos).data(), Len); + Pos += Len; + return Len; + } + +public: + BufferDataStreamer(std::unique_ptr Buffer) + : Buffer(std::move(Buffer)) {} +}; + +static std::unique_ptr +getStreamedModuleFromAssembly(LLVMContext &Context, SmallString<1024> &Mem, + const char *Assembly) { + writeModuleToBuffer(parseAssembly(Assembly), Mem); + std::unique_ptr Buffer = + MemoryBuffer::getMemBuffer(Mem.str(), "test", false); + auto Streamer = make_unique(std::move(Buffer)); + ErrorOr> ModuleOrErr = + getStreamedBitcodeModule("test", std::move(Streamer), Context); + return std::move(ModuleOrErr.get()); +} + +TEST(BitReaderTest, MateralizeForwardRefWithStream) { + SmallString<1024> Mem; + + LLVMContext Context; + std::unique_ptr M = getStreamedModuleFromAssembly( + Context, Mem, "@table = constant i8* blockaddress(@func, %bb)\n" + "define void @func() {\n" + " unreachable\n" + "bb:\n" + " unreachable\n" + "}\n"); + EXPECT_FALSE(M->getFunction("func")->empty()); +} + TEST(BitReaderTest, DematerializeFunctionPreservesLinkageType) { SmallString<1024> Mem;