1 files changed, 139 insertions, 193 deletions

diff --git a/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp
index 62c6b7dfa31..75ccfc9ab0d 100644
--- a/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp
@@ -1,4 +1,4 @@
-//===- RTDyldObjectLinkingLayerTest.cpp - RTDyld linking layer unit tests -===//
+//===--- RTDyldObjectLinkingLayerTest.cpp - RTDyld linking layer tests ---===//
 //
 //                     The LLVM Compiler Infrastructure
 //
@@ -7,13 +7,14 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "OrcTestCommon.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
+#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
 #include "llvm/ExecutionEngine/Orc/Legacy.h"
 #include "llvm/ExecutionEngine/Orc/NullResolver.h"
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/LLVMContext.h"
@@ -25,258 +26,203 @@ using namespace llvm::orc;
 namespace {
 
 class RTDyldObjectLinkingLayerExecutionTest : public testing::Test,
-                                              public OrcExecutionTest {
-
-};
-
-class SectionMemoryManagerWrapper : public SectionMemoryManager {
-public:
-  int FinalizationCount = 0;
-  int NeedsToReserveAllocationSpaceCount = 0;
-
-  bool needsToReserveAllocationSpace() override {
-    ++NeedsToReserveAllocationSpaceCount;
-    return SectionMemoryManager::needsToReserveAllocationSpace();
-  }
-
-  bool finalizeMemory(std::string *ErrMsg = nullptr) override {
-    ++FinalizationCount;
-    return SectionMemoryManager::finalizeMemory(ErrMsg);
-  }
-};
+                                               public OrcExecutionTest {};
 
-TEST(RTDyldObjectLinkingLayerTest, TestSetProcessAllSections) {
+// Adds an object with a debug section to RuntimeDyld and then returns whether
+// the debug section was passed to the memory manager.
+static bool testSetProcessAllSections(std::unique_ptr<MemoryBuffer> Obj,
+                                      bool ProcessAllSections) {
   class MemoryManagerWrapper : public SectionMemoryManager {
   public:
     MemoryManagerWrapper(bool &DebugSeen) : DebugSeen(DebugSeen) {}
     uint8_t *allocateDataSection(uintptr_t Size, unsigned Alignment,
-                                 unsigned SectionID,
-                                 StringRef SectionName,
+                                 unsigned SectionID, StringRef SectionName,
                                  bool IsReadOnly) override {
       if (SectionName == ".debug_str")
         DebugSeen = true;
-      return SectionMemoryManager::allocateDataSection(Size, Alignment,
-                                                         SectionID,
-                                                         SectionName,
-                                                         IsReadOnly);
+      return SectionMemoryManager::allocateDataSection(
+          Size, Alignment, SectionID, SectionName, IsReadOnly);
     }
+
   private:
     bool &DebugSeen;
   };
 
   bool DebugSectionSeen = false;
-  auto MM = std::make_shared<MemoryManagerWrapper>(DebugSectionSeen);
 
   ExecutionSession ES;
+  auto &JD = ES.createJITDylib("main");
+  auto Foo = ES.intern("foo");
 
-  RTDyldObjectLinkingLayer ObjLayer(ES, [&MM](VModuleKey) {
-    return RTDyldObjectLinkingLayer::Resources{
-        MM, std::make_shared<NullResolver>()};
+  RTDyldObjectLinkingLayer ObjLayer(ES, [&DebugSectionSeen](VModuleKey) {
+    return llvm::make_unique<MemoryManagerWrapper>(DebugSectionSeen);
   });
 
+  auto OnResolveDoNothing = [](Expected<SymbolMap> R) {
+    cantFail(std::move(R));
+  };
+
+  auto OnReadyDoNothing = [](Error Err) { cantFail(std::move(Err)); };
+
+  ObjLayer.setProcessAllSections(ProcessAllSections);
+  auto K = ES.allocateVModule();
+  cantFail(ObjLayer.add(JD, K, std::move(Obj)));
+  ES.lookup({&JD}, {Foo}, OnResolveDoNothing, OnReadyDoNothing,
+            NoDependenciesToRegister);
+  return DebugSectionSeen;
+}
+
+TEST(RTDyldObjectLinkingLayerTest, TestSetProcessAllSections) {
   LLVMContext Context;
   auto M = llvm::make_unique<Module>("", Context);
   M->setTargetTriple("x86_64-unknown-linux-gnu");
   Type *Int32Ty = IntegerType::get(Context, 32);
   GlobalVariable *GV =
-    new GlobalVariable(*M, Int32Ty, false, GlobalValue::ExternalLinkage,
+      new GlobalVariable(*M, Int32Ty, false, GlobalValue::ExternalLinkage,
                          ConstantInt::get(Int32Ty, 42), "foo");
 
   GV->setSection(".debug_str");
 
-
   // Initialize the native target in case this is the first unit test
   // to try to build a TM.
   OrcNativeTarget::initialize();
-  std::unique_ptr<TargetMachine> TM(
-    EngineBuilder().selectTarget(Triple(M->getTargetTriple()), "", "",
-                                 SmallVector<std::string, 1>()));
+  std::unique_ptr<TargetMachine> TM(EngineBuilder().selectTarget(
+      Triple(M->getTargetTriple()), "", "", SmallVector<std::string, 1>()));
   if (!TM)
     return;
 
   auto Obj = SimpleCompiler(*TM)(*M);
 
-  {
-    // Test with ProcessAllSections = false (the default).
-    auto K = ES.allocateVModule();
-    cantFail(ObjLayer.addObject(
-        K, MemoryBuffer::getMemBufferCopy(Obj->getBuffer())));
-    cantFail(ObjLayer.emitAndFinalize(K));
-    EXPECT_EQ(DebugSectionSeen, false)
-      << "Unexpected debug info section";
-    cantFail(ObjLayer.removeObject(K));
-  }
-
-  {
-    // Test with ProcessAllSections = true.
-    ObjLayer.setProcessAllSections(true);
-    auto K = ES.allocateVModule();
-    cantFail(ObjLayer.addObject(K, std::move(Obj)));
-    cantFail(ObjLayer.emitAndFinalize(K));
-    EXPECT_EQ(DebugSectionSeen, true)
-      << "Expected debug info section not seen";
-    cantFail(ObjLayer.removeObject(K));
-  }
+  EXPECT_FALSE(testSetProcessAllSections(
+      MemoryBuffer::getMemBufferCopy(Obj->getBuffer()), false))
+      << "Debug section seen despite ProcessAllSections being false";
+  EXPECT_TRUE(testSetProcessAllSections(std::move(Obj), true))
+      << "Expected to see debug section when ProcessAllSections is true";
 }
 
-TEST_F(RTDyldObjectLinkingLayerExecutionTest, NoDuplicateFinalization) {
-  if (!SupportsJIT)
-    return;
+TEST(RTDyldObjectLinkingLayerTest, TestOverrideObjectFlags) {
 
-  ExecutionSession ES;
+  OrcNativeTarget::initialize();
 
-  auto MM = std::make_shared<SectionMemoryManagerWrapper>();
+  std::unique_ptr<TargetMachine> TM(
+      EngineBuilder().selectTarget(Triple("x86_64-unknown-linux-gnu"), "", "",
+                                   SmallVector<std::string, 1>()));
 
-  std::map<orc::VModuleKey, std::shared_ptr<orc::SymbolResolver>> Resolvers;
+  if (!TM)
+    return;
 
-  RTDyldObjectLinkingLayer ObjLayer(ES, [&](VModuleKey K) {
-    auto I = Resolvers.find(K);
-    assert(I != Resolvers.end() && "Missing resolver");
-    auto R = std::move(I->second);
-    Resolvers.erase(I);
-    return RTDyldObjectLinkingLayer::Resources{MM, std::move(R)};
-  });
-  SimpleCompiler Compile(*TM);
-
-  // Create a pair of modules that will trigger recursive finalization:
-  // Module 1:
-  //   int bar() { return 42; }
-  // Module 2:
-  //   int bar();
-  //   int foo() { return bar(); }
-  //
-  // Verify that the memory manager is only finalized once (for Module 2).
-  // Failure suggests that finalize is being called on the inner RTDyld
-  // instance (for Module 1) which is unsafe, as it will prevent relocation of
-  // Module 2.
-
-  ModuleBuilder MB1(Context, "", "dummy");
-  {
-    MB1.getModule()->setDataLayout(TM->createDataLayout());
-    Function *BarImpl = MB1.createFunctionDecl<int32_t(void)>("bar");
-    BasicBlock *BarEntry = BasicBlock::Create(Context, "entry", BarImpl);
-    IRBuilder<> Builder(BarEntry);
-    IntegerType *Int32Ty = IntegerType::get(Context, 32);
-    Value *FourtyTwo = ConstantInt::getSigned(Int32Ty, 42);
-    Builder.CreateRet(FourtyTwo);
-  }
+  // Our compiler is going to modify symbol visibility settings without telling
+  // ORC. This will test our ability to override the flags later.
+  class FunkySimpleCompiler : public SimpleCompiler {
+  public:
+    FunkySimpleCompiler(TargetMachine &TM) : SimpleCompiler(TM) {}
 
-  auto Obj1 = Compile(*MB1.getModule());
+    CompileResult operator()(Module &M) {
+      auto *Foo = M.getFunction("foo");
+      assert(Foo && "Expected function Foo not found");
+      Foo->setVisibility(GlobalValue::HiddenVisibility);
+      return SimpleCompiler::operator()(M);
+    }
+  };
 
-  ModuleBuilder MB2(Context, "", "dummy");
+  // Create a module with two void() functions: foo and bar.
+  ThreadSafeContext TSCtx(llvm::make_unique<LLVMContext>());
+  ThreadSafeModule M;
   {
-    MB2.getModule()->setDataLayout(TM->createDataLayout());
-    Function *BarDecl = MB2.createFunctionDecl<int32_t(void)>("bar");
-    Function *FooImpl = MB2.createFunctionDecl<int32_t(void)>("foo");
-    BasicBlock *FooEntry = BasicBlock::Create(Context, "entry", FooImpl);
-    IRBuilder<> Builder(FooEntry);
-    Builder.CreateRet(Builder.CreateCall(BarDecl));
+    ModuleBuilder MB(*TSCtx.getContext(), TM->getTargetTriple().str(), "dummy");
+    MB.getModule()->setDataLayout(TM->createDataLayout());
+
+    Function *FooImpl = MB.createFunctionDecl<void()>("foo");
+    BasicBlock *FooEntry =
+        BasicBlock::Create(*TSCtx.getContext(), "entry", FooImpl);
+    IRBuilder<> B1(FooEntry);
+    B1.CreateRetVoid();
+
+    Function *BarImpl = MB.createFunctionDecl<void()>("bar");
+    BasicBlock *BarEntry =
+        BasicBlock::Create(*TSCtx.getContext(), "entry", BarImpl);
+    IRBuilder<> B2(BarEntry);
+    B2.CreateRetVoid();
+
+    M = ThreadSafeModule(MB.takeModule(), std::move(TSCtx));
   }
-  auto Obj2 = Compile(*MB2.getModule());
 
-  auto K1 = ES.allocateVModule();
-  Resolvers[K1] = std::make_shared<NullResolver>();
-  cantFail(ObjLayer.addObject(K1, std::move(Obj1)));
+  // Create a simple stack and set the override flags option.
+  ExecutionSession ES;
+  auto &JD = ES.createJITDylib("main");
+  auto Foo = ES.intern("foo");
+  RTDyldObjectLinkingLayer ObjLayer(
+      ES, [](VModuleKey) { return llvm::make_unique<SectionMemoryManager>(); });
+  IRCompileLayer CompileLayer(ES, ObjLayer, FunkySimpleCompiler(*TM));
 
-  auto K2 = ES.allocateVModule();
-  auto LegacyLookup = [&](const std::string &Name) {
-    return ObjLayer.findSymbol(Name, true);
-  };
+  ObjLayer.setOverrideObjectFlagsWithResponsibilityFlags(true);
 
-  Resolvers[K2] = createSymbolResolver(
-      [&](const SymbolNameSet &Symbols) {
-        return cantFail(
-            getResponsibilitySetWithLegacyFn(Symbols, LegacyLookup));
-      },
-      [&](std::shared_ptr<AsynchronousSymbolQuery> Query,
-          const SymbolNameSet &Symbols) {
-        return lookupWithLegacyFn(ES, *Query, Symbols, LegacyLookup);
-      });
-
-  cantFail(ObjLayer.addObject(K2, std::move(Obj2)));
-  cantFail(ObjLayer.emitAndFinalize(K2));
-  cantFail(ObjLayer.removeObject(K2));
-
-  // Finalization of module 2 should trigger finalization of module 1.
-  // Verify that finalize on SMMW is only called once.
-  EXPECT_EQ(MM->FinalizationCount, 1)
-      << "Extra call to finalize";
+  cantFail(CompileLayer.add(JD, ES.allocateVModule(), std::move(M)));
+  ES.lookup({&JD}, {Foo}, [](Expected<SymbolMap> R) { cantFail(std::move(R)); },
+            [](Error Err) { cantFail(std::move(Err)); },
+            NoDependenciesToRegister);
 }
 
-TEST_F(RTDyldObjectLinkingLayerExecutionTest, NoPrematureAllocation) {
-  if (!SupportsJIT)
-    return;
+TEST(RTDyldObjectLinkingLayerTest, TestAutoClaimResponsibilityForSymbols) {
 
-  ExecutionSession ES;
+  OrcNativeTarget::initialize();
 
-  auto MM = std::make_shared<SectionMemoryManagerWrapper>();
+  std::unique_ptr<TargetMachine> TM(
+      EngineBuilder().selectTarget(Triple("x86_64-unknown-linux-gnu"), "", "",
+                                   SmallVector<std::string, 1>()));
 
-  RTDyldObjectLinkingLayer ObjLayer(ES, [&MM](VModuleKey K) {
-    return RTDyldObjectLinkingLayer::Resources{
-        MM, std::make_shared<NullResolver>()};
-  });
-  SimpleCompiler Compile(*TM);
-
-  // Create a pair of unrelated modules:
-  //
-  // Module 1:
-  //   int foo() { return 42; }
-  // Module 2:
-  //   int bar() { return 7; }
-  //
-  // Both modules will share a memory manager. We want to verify that the
-  // second object is not loaded before the first one is finalized. To do this
-  // in a portable way, we abuse the
-  // RuntimeDyld::MemoryManager::needsToReserveAllocationSpace hook, which is
-  // called once per object before any sections are allocated.
-
-  ModuleBuilder MB1(Context, "", "dummy");
-  {
-    MB1.getModule()->setDataLayout(TM->createDataLayout());
-    Function *BarImpl = MB1.createFunctionDecl<int32_t(void)>("foo");
-    BasicBlock *BarEntry = BasicBlock::Create(Context, "entry", BarImpl);
-    IRBuilder<> Builder(BarEntry);
-    IntegerType *Int32Ty = IntegerType::get(Context, 32);
-    Value *FourtyTwo = ConstantInt::getSigned(Int32Ty, 42);
-    Builder.CreateRet(FourtyTwo);
-  }
+  if (!TM)
+    return;
 
-  auto Obj1 = Compile(*MB1.getModule());
+  // Our compiler is going to add a new symbol without telling ORC.
+  // This will test our ability to auto-claim responsibility later.
+  class FunkySimpleCompiler : public SimpleCompiler {
+  public:
+    FunkySimpleCompiler(TargetMachine &TM) : SimpleCompiler(TM) {}
+
+    CompileResult operator()(Module &M) {
+      Function *BarImpl =
+          Function::Create(TypeBuilder<void(), false>::get(M.getContext()),
+                           GlobalValue::ExternalLinkage, "bar", &M);
+      BasicBlock *BarEntry =
+          BasicBlock::Create(M.getContext(), "entry", BarImpl);
+      IRBuilder<> B(BarEntry);
+      B.CreateRetVoid();
+
+      return SimpleCompiler::operator()(M);
+    }
+  };
 
-  ModuleBuilder MB2(Context, "", "dummy");
+  // Create a module with two void() functions: foo and bar.
+  ThreadSafeContext TSCtx(llvm::make_unique<LLVMContext>());
+  ThreadSafeModule M;
   {
-    MB2.getModule()->setDataLayout(TM->createDataLayout());
-    Function *BarImpl = MB2.createFunctionDecl<int32_t(void)>("bar");
-    BasicBlock *BarEntry = BasicBlock::Create(Context, "entry", BarImpl);
-    IRBuilder<> Builder(BarEntry);
-    IntegerType *Int32Ty = IntegerType::get(Context, 32);
-    Value *Seven = ConstantInt::getSigned(Int32Ty, 7);
-    Builder.CreateRet(Seven);
-  }
-  auto Obj2 = Compile(*MB2.getModule());
+    ModuleBuilder MB(*TSCtx.getContext(), TM->getTargetTriple().str(), "dummy");
+    MB.getModule()->setDataLayout(TM->createDataLayout());
 
-  auto K = ES.allocateVModule();
-  cantFail(ObjLayer.addObject(K, std::move(Obj1)));
-  cantFail(ObjLayer.addObject(ES.allocateVModule(), std::move(Obj2)));
-  cantFail(ObjLayer.emitAndFinalize(K));
-  cantFail(ObjLayer.removeObject(K));
-
-  // Only one call to needsToReserveAllocationSpace should have been made.
-  EXPECT_EQ(MM->NeedsToReserveAllocationSpaceCount, 1)
-      << "More than one call to needsToReserveAllocationSpace "
-         "(multiple unrelated objects loaded prior to finalization)";
-}
+    Function *FooImpl = MB.createFunctionDecl<void()>("foo");
+    BasicBlock *FooEntry =
+        BasicBlock::Create(*TSCtx.getContext(), "entry", FooImpl);
+    IRBuilder<> B(FooEntry);
+    B.CreateRetVoid();
 
-TEST_F(RTDyldObjectLinkingLayerExecutionTest, TestNotifyLoadedSignature) {
+    M = ThreadSafeModule(MB.takeModule(), std::move(TSCtx));
+  }
+
+  // Create a simple stack and set the override flags option.
   ExecutionSession ES;
+  auto &JD = ES.createJITDylib("main");
+  auto Foo = ES.intern("foo");
   RTDyldObjectLinkingLayer ObjLayer(
-      ES,
-      [](VModuleKey) {
-        return RTDyldObjectLinkingLayer::Resources{
-            nullptr, std::make_shared<NullResolver>()};
-      },
-      [](VModuleKey, const object::ObjectFile &obj,
-         const RuntimeDyld::LoadedObjectInfo &info) {});
+      ES, [](VModuleKey) { return llvm::make_unique<SectionMemoryManager>(); });
+  IRCompileLayer CompileLayer(ES, ObjLayer, FunkySimpleCompiler(*TM));
+
+  ObjLayer.setAutoClaimResponsibilityForObjectSymbols(true);
+
+  cantFail(CompileLayer.add(JD, ES.allocateVModule(), std::move(M)));
+  ES.lookup({&JD}, {Foo}, [](Expected<SymbolMap> R) { cantFail(std::move(R)); },
+            [](Error Err) { cantFail(std::move(Err)); },
+            NoDependenciesToRegister);
 }
 
 } // end anonymous namespace