1 files changed, 50 insertions, 0 deletions

diff --git a/unittests/ExecutionEngine/JIT/JITTest.cpp b/unittests/ExecutionEngine/JIT/JITTest.cpp
index 84ee0e3..b85f724 100644
--- a/unittests/ExecutionEngine/JIT/JITTest.cpp
+++ b/unittests/ExecutionEngine/JIT/JITTest.cpp
@@ -65,6 +65,8 @@ public:
     stubsAllocated = 0;
   }
 
+  void setSizeRequired(bool Required) { SizeRequired = Required; }
+
   virtual void setMemoryWritable() { Base->setMemoryWritable(); }
   virtual void setMemoryExecutable() { Base->setMemoryExecutable(); }
   virtual void setPoisonMemory(bool poison) { Base->setPoisonMemory(poison); }
@@ -628,6 +630,54 @@ TEST_F(JITTest, AvailableExternallyFunctionIsntCompiled) {
                         << " not 7 from the IR version.";
 }
 
+TEST_F(JITTest, NeedsExactSizeWithManyGlobals) {
+  // PR5291: When the JMM needed the exact size of function bodies before
+  // starting to emit them, the JITEmitter would modify a set while iterating
+  // over it.
+  TheJIT->DisableLazyCompilation(true);
+  RJMM->setSizeRequired(true);
+
+  LoadAssembly("@A = global i32 42 "
+               "@B = global i32* @A "
+               "@C = global i32** @B "
+               "@D = global i32*** @C "
+               "@E = global i32**** @D "
+               "@F = global i32***** @E "
+               "@G = global i32****** @F "
+               "@H = global i32******* @G "
+               "@I = global i32******** @H "
+               "define i32********* @test() { "
+               "  ret i32********* @I "
+               "}");
+  Function *testIR = M->getFunction("test");
+  int32_t********* (*test)() = reinterpret_cast<int32_t*********(*)()>(
+    (intptr_t)TheJIT->getPointerToFunction(testIR));
+  EXPECT_EQ(42, *********test());
+}
+
+TEST_F(JITTest, EscapedLazyStubStillCallable) {
+  TheJIT->DisableLazyCompilation(false);
+  LoadAssembly("define internal i32 @stubbed() { "
+               "  ret i32 42 "
+               "} "
+               " "
+               "define i32()* @get_stub() { "
+               "  ret i32()* @stubbed "
+               "} ");
+  typedef int32_t(*StubTy)();
+
+  // Call get_stub() to get the address of @stubbed without actually JITting it.
+  Function *get_stubIR = M->getFunction("get_stub");
+  StubTy (*get_stub)() = reinterpret_cast<StubTy(*)()>(
+    (intptr_t)TheJIT->getPointerToFunction(get_stubIR));
+  StubTy stubbed = get_stub();
+  // Now get_stubIR is the only reference to stubbed's stub.
+  get_stubIR->eraseFromParent();
+  // Now there are no references inside the JIT, but we've got a pointer outside
+  // it.  The stub should be callable and return the right value.
+  EXPECT_EQ(42, stubbed());
+}
+
 // Converts the LLVM assembly to bitcode and returns it in a std::string.  An
 // empty string indicates an error.
 std::string AssembleToBitcode(LLVMContext &Context, const char *Assembly) {