1 files changed, 37 insertions, 38 deletions

diff --git a/unittests/Bitcode/BitReaderTest.cpp b/unittests/Bitcode/BitReaderTest.cpp
index 3c56ea04c744..420aca2443bb 100644
--- a/unittests/Bitcode/BitReaderTest.cpp
+++ b/unittests/Bitcode/BitReaderTest.cpp
@@ -10,6 +10,7 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/AsmParser/Parser.h"
+#include "llvm/Bitcode/BitstreamReader.h"
 #include "llvm/Bitcode/BitstreamWriter.h"
 #include "llvm/Bitcode/ReaderWriter.h"
 #include "llvm/IR/Constants.h"
@@ -21,6 +22,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/StreamingMemoryObject.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
@@ -89,6 +91,41 @@ getStreamedModuleFromAssembly(LLVMContext &Context, SmallString<1024> &Mem,
   return std::move(ModuleOrErr.get());
 }
 
+// Checks if we correctly detect eof if we try to read N bits when there are not
+// enough bits left on the input stream to read N bits, and we are using a data
+// streamer. In particular, it checks if we properly set the object size when
+// the eof is reached under such conditions.
+TEST(BitReaderTest, TestForEofAfterReadFailureOnDataStreamer) {
+  // Note: Because StreamingMemoryObject does a call to method GetBytes in it's
+  // constructor, using internal constant kChunkSize, we must fill the input
+  // with more characters than that amount.
+  static size_t InputSize = StreamingMemoryObject::kChunkSize + 5;
+  char *Text = new char[InputSize];
+  std::memset(Text, 'a', InputSize);
+  Text[InputSize - 1] = '\0';
+  StringRef Input(Text);
+
+  // Build bitsteam reader using data streamer.
+  auto MemoryBuf = MemoryBuffer::getMemBuffer(Input);
+  std::unique_ptr<DataStreamer> Streamer(
+      new BufferDataStreamer(std::move(MemoryBuf)));
+  auto OwnedBytes =
+      llvm::make_unique<StreamingMemoryObject>(std::move(Streamer));
+  auto Reader = llvm::make_unique<BitstreamReader>(std::move(OwnedBytes));
+  BitstreamCursor Cursor;
+  Cursor.init(Reader.get());
+
+  // Jump to two bytes before end of stream.
+  Cursor.JumpToBit((InputSize - 4) * CHAR_BIT);
+  // Try to read 4 bytes when only 2 are present, resulting in error value 0.
+  const size_t ReadErrorValue = 0;
+  EXPECT_EQ(ReadErrorValue, Cursor.Read(32));
+  // Should be at eof now.
+  EXPECT_TRUE(Cursor.AtEndOfStream());
+
+  delete[] Text;
+}
+
 TEST(BitReaderTest, MateralizeForwardRefWithStream) {
   SmallString<1024> Mem;
 
@@ -103,30 +140,6 @@ TEST(BitReaderTest, MateralizeForwardRefWithStream) {
   EXPECT_FALSE(M->getFunction("func")->empty());
 }
 
-TEST(BitReaderTest, DematerializeFunctionPreservesLinkageType) {
-  SmallString<1024> Mem;
-
-  LLVMContext Context;
-  std::unique_ptr<Module> M = getLazyModuleFromAssembly(
-      Context, Mem, "define internal i32 @func() {\n"
-                      "ret i32 0\n"
-                    "}\n");
-
-  EXPECT_FALSE(verifyModule(*M, &dbgs()));
-
-  M->getFunction("func")->materialize();
-  EXPECT_FALSE(M->getFunction("func")->empty());
-  EXPECT_TRUE(M->getFunction("func")->getLinkage() ==
-              GlobalValue::InternalLinkage);
-
-  // Check that the linkage type is preserved after dematerialization.
-  M->getFunction("func")->dematerialize();
-  EXPECT_TRUE(M->getFunction("func")->empty());
-  EXPECT_TRUE(M->getFunction("func")->getLinkage() ==
-              GlobalValue::InternalLinkage);
-  EXPECT_FALSE(verifyModule(*M, &dbgs()));
-}
-
 // Tests that lazy evaluation can parse functions out of order.
 TEST(BitReaderTest, MaterializeFunctionsOutOfOrder) {
   SmallString<1024> Mem;
@@ -203,10 +216,6 @@ TEST(BitReaderTest, MaterializeFunctionsForBlockAddr) { // PR11677
                     "  unreachable\n"
                     "}\n");
   EXPECT_FALSE(verifyModule(*M, &dbgs()));
-
-  // Try (and fail) to dematerialize @func.
-  M->getFunction("func")->dematerialize();
-  EXPECT_FALSE(M->getFunction("func")->empty());
 }
 
 TEST(BitReaderTest, MaterializeFunctionsForBlockAddrInFunctionBefore) {
@@ -234,11 +243,6 @@ TEST(BitReaderTest, MaterializeFunctionsForBlockAddrInFunctionBefore) {
   EXPECT_FALSE(M->getFunction("func")->empty());
   EXPECT_TRUE(M->getFunction("other")->empty());
   EXPECT_FALSE(verifyModule(*M, &dbgs()));
-
-  // Try (and fail) to dematerialize @func.
-  M->getFunction("func")->dematerialize();
-  EXPECT_FALSE(M->getFunction("func")->empty());
-  EXPECT_FALSE(verifyModule(*M, &dbgs()));
 }
 
 TEST(BitReaderTest, MaterializeFunctionsForBlockAddrInFunctionAfter) {
@@ -266,11 +270,6 @@ TEST(BitReaderTest, MaterializeFunctionsForBlockAddrInFunctionAfter) {
   EXPECT_FALSE(M->getFunction("func")->empty());
   EXPECT_TRUE(M->getFunction("other")->empty());
   EXPECT_FALSE(verifyModule(*M, &dbgs()));
-
-  // Try (and fail) to dematerialize @func.
-  M->getFunction("func")->dematerialize();
-  EXPECT_FALSE(M->getFunction("func")->empty());
-  EXPECT_FALSE(verifyModule(*M, &dbgs()));
 }
 
 } // end namespace