1 files changed, 97 insertions, 10 deletions

diff --git a/lib/ExecutionEngine/Orc/IndirectionUtils.cpp b/lib/ExecutionEngine/Orc/IndirectionUtils.cpp
index 68397beae63a..9ca2c5cb4a55 100644
--- a/lib/ExecutionEngine/Orc/IndirectionUtils.cpp
+++ b/lib/ExecutionEngine/Orc/IndirectionUtils.cpp
@@ -13,38 +13,123 @@
 #include "llvm/ExecutionEngine/Orc/OrcABISupport.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/Support/Format.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include <sstream>
 
+using namespace llvm;
+using namespace llvm::orc;
+
+namespace {
+
+class CompileCallbackMaterializationUnit : public orc::MaterializationUnit {
+public:
+  using CompileFunction = JITCompileCallbackManager::CompileFunction;
+
+  CompileCallbackMaterializationUnit(SymbolStringPtr Name,
+                                     CompileFunction Compile)
+      : MaterializationUnit(SymbolFlagsMap({{Name, JITSymbolFlags::Exported}})),
+        Name(std::move(Name)), Compile(std::move(Compile)) {}
+
+private:
+  void materialize(MaterializationResponsibility R) {
+    SymbolMap Result;
+    Result[Name] = JITEvaluatedSymbol(Compile(), JITSymbolFlags::Exported);
+    R.resolve(Result);
+    R.finalize();
+  }
+
+  void discard(const VSO &V, SymbolStringPtr Name) {
+    llvm_unreachable("Discard should never occur on a LMU?");
+  }
+
+  SymbolStringPtr Name;
+  CompileFunction Compile;
+};
+
+} // namespace
+
 namespace llvm {
 namespace orc {
 
 void JITCompileCallbackManager::anchor() {}
 void IndirectStubsManager::anchor() {}
 
+Expected<JITTargetAddress>
+JITCompileCallbackManager::getCompileCallback(CompileFunction Compile) {
+  if (auto TrampolineAddr = getAvailableTrampolineAddr()) {
+    auto CallbackName = ES.getSymbolStringPool().intern(
+        std::string("cc") + std::to_string(++NextCallbackId));
+
+    std::lock_guard<std::mutex> Lock(CCMgrMutex);
+    AddrToSymbol[*TrampolineAddr] = CallbackName;
+    cantFail(CallbacksVSO.define(
+        llvm::make_unique<CompileCallbackMaterializationUnit>(
+            std::move(CallbackName), std::move(Compile))));
+    return *TrampolineAddr;
+  } else
+    return TrampolineAddr.takeError();
+}
+
+JITTargetAddress JITCompileCallbackManager::executeCompileCallback(
+    JITTargetAddress TrampolineAddr) {
+  SymbolStringPtr Name;
+
+  {
+    std::unique_lock<std::mutex> Lock(CCMgrMutex);
+    auto I = AddrToSymbol.find(TrampolineAddr);
+
+    // If this address is not associated with a compile callback then report an
+    // error to the execution session and return ErrorHandlerAddress to the
+    // callee.
+    if (I == AddrToSymbol.end()) {
+      Lock.unlock();
+      std::string ErrMsg;
+      {
+        raw_string_ostream ErrMsgStream(ErrMsg);
+        ErrMsgStream << "No compile callback for trampoline at "
+                     << format("0x%016x", TrampolineAddr);
+      }
+      ES.reportError(
+          make_error<StringError>(std::move(ErrMsg), inconvertibleErrorCode()));
+      return ErrorHandlerAddress;
+    } else
+      Name = I->second;
+  }
+
+  if (auto Sym = lookup({&CallbacksVSO}, Name))
+    return Sym->getAddress();
+  else {
+    // If anything goes wrong materializing Sym then report it to the session
+    // and return the ErrorHandlerAddress;
+    ES.reportError(Sym.takeError());
+    return ErrorHandlerAddress;
+  }
+}
+
 std::unique_ptr<JITCompileCallbackManager>
-createLocalCompileCallbackManager(const Triple &T,
+createLocalCompileCallbackManager(const Triple &T, ExecutionSession &ES,
                                   JITTargetAddress ErrorHandlerAddress) {
   switch (T.getArch()) {
     default: return nullptr;
 
     case Triple::aarch64: {
       typedef orc::LocalJITCompileCallbackManager<orc::OrcAArch64> CCMgrT;
-      return llvm::make_unique<CCMgrT>(ErrorHandlerAddress);
+      return llvm::make_unique<CCMgrT>(ES, ErrorHandlerAddress);
     }
 
     case Triple::x86: {
       typedef orc::LocalJITCompileCallbackManager<orc::OrcI386> CCMgrT;
-      return llvm::make_unique<CCMgrT>(ErrorHandlerAddress);
+      return llvm::make_unique<CCMgrT>(ES, ErrorHandlerAddress);
     }
 
     case Triple::x86_64: {
       if ( T.getOS() == Triple::OSType::Win32 ) {
         typedef orc::LocalJITCompileCallbackManager<orc::OrcX86_64_Win32> CCMgrT;
-        return llvm::make_unique<CCMgrT>(ErrorHandlerAddress);
+        return llvm::make_unique<CCMgrT>(ES, ErrorHandlerAddress);
       } else {
         typedef orc::LocalJITCompileCallbackManager<orc::OrcX86_64_SysV> CCMgrT;
-        return llvm::make_unique<CCMgrT>(ErrorHandlerAddress);
+        return llvm::make_unique<CCMgrT>(ES, ErrorHandlerAddress);
       }
     }
 
@@ -54,7 +139,11 @@ createLocalCompileCallbackManager(const Triple &T,
 std::function<std::unique_ptr<IndirectStubsManager>()>
 createLocalIndirectStubsManagerBuilder(const Triple &T) {
   switch (T.getArch()) {
-    default: return nullptr;
+    default:
+      return [](){
+        return llvm::make_unique<
+                       orc::LocalIndirectStubsManager<orc::OrcGenericABI>>();
+      };
 
     case Triple::aarch64:
       return [](){
@@ -176,7 +265,6 @@ void makeAllSymbolsExternallyAccessible(Module &M) {
 
 Function* cloneFunctionDecl(Module &Dst, const Function &F,
                             ValueToValueMapTy *VMap) {
-  assert(F.getParent() != &Dst && "Can't copy decl over existing function.");
   Function *NewF =
     Function::Create(cast<FunctionType>(F.getValueType()),
                      F.getLinkage(), F.getName(), &Dst);
@@ -214,7 +302,6 @@ void moveFunctionBody(Function &OrigF, ValueToValueMapTy &VMap,
 
 GlobalVariable* cloneGlobalVariableDecl(Module &Dst, const GlobalVariable &GV,
                                         ValueToValueMapTy *VMap) {
-  assert(GV.getParent() != &Dst && "Can't copy decl over existing global var.");
   GlobalVariable *NewGV = new GlobalVariable(
       Dst, GV.getValueType(), GV.isConstant(),
       GV.getLinkage(), nullptr, GV.getName(), nullptr,
@@ -236,8 +323,8 @@ void moveGlobalVariableInitializer(GlobalVariable &OrigGV,
     assert(VMap[&OrigGV] == NewGV &&
            "Incorrect global variable mapping in VMap.");
   assert(NewGV->getParent() != OrigGV.getParent() &&
-         "moveGlobalVariable should only be used to move initializers between "
-         "modules");
+         "moveGlobalVariableInitializer should only be used to move "
+         "initializers between modules");
 
   NewGV->setInitializer(MapValue(OrigGV.getInitializer(), VMap, RF_None,
                                  nullptr, Materializer));