Handle forward referenced function when streaming bitcode.

Without this the included unit test would assert in

  assert(BasicBlockFwdRefs.empty() && "Unresolved blockaddress fwd references");

llvm-svn: 239871

lib/Bitcode/Reader/BitcodeReader.cpp

@@ -4597,6 +4597,35 @@ const std::error_category &llvm::BitcodeErrorCategory() {
 // External interface
 //===----------------------------------------------------------------------===//
 
+static ErrorOr<std::unique_ptr<Module>>
+getBitcodeModuleImpl(std::unique_ptr<DataStreamer> Streamer, StringRef Name,
+                     BitcodeReader *R, LLVMContext &Context,
+                     bool MaterializeAll, bool ShouldLazyLoadMetadata) {
+  std::unique_ptr<Module> M = make_unique<Module>(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<MemoryBuffer> &&Buffer,
                          LLVMContext &Context, bool MaterializeAll,
                          DiagnosticHandlerFunction DiagnosticHandler,
                          bool ShouldLazyLoadMetadata = false) {
-  std::unique_ptr<Module> M =
-      make_unique<Module>(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<std::unique_ptr<Module>> 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<std::unique_ptr<Module>> llvm::getLazyBitcodeModule(
@@ -4652,10 +4664,9 @@ ErrorOr<std::unique_ptr<Module>> llvm::getStreamedBitcodeModule(
     LLVMContext &Context, DiagnosticHandlerFunction DiagnosticHandler) {
   std::unique_ptr<Module> M = make_unique<Module>(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<std::unique_ptr<Module>>

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<Module> getLazyModuleFromAssembly(LLVMContext &Context,
   return std::move(ModuleOrErr.get());
 }
 
+class BufferDataStreamer : public DataStreamer {
+  std::unique_ptr<MemoryBuffer> 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<MemoryBuffer> Buffer)
+      : Buffer(std::move(Buffer)) {}
+};
+
+static std::unique_ptr<Module>
+getStreamedModuleFromAssembly(LLVMContext &Context, SmallString<1024> &Mem,
+                              const char *Assembly) {
+  writeModuleToBuffer(parseAssembly(Assembly), Mem);
+  std::unique_ptr<MemoryBuffer> Buffer =
+      MemoryBuffer::getMemBuffer(Mem.str(), "test", false);
+  auto Streamer = make_unique<BufferDataStreamer>(std::move(Buffer));
+  ErrorOr<std::unique_ptr<Module>> ModuleOrErr =
+      getStreamedBitcodeModule("test", std::move(Streamer), Context);
+  return std::move(ModuleOrErr.get());
+}
+
+TEST(BitReaderTest, MateralizeForwardRefWithStream) {
+  SmallString<1024> Mem;
+
+  LLVMContext Context;
+  std::unique_ptr<Module> 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;