[Orc] Expose the compile callback API through the C bindings.

llvm-svn: 251683

2 files changed, 71 insertions, 22 deletions

diff --git a/llvm/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp b/llvm/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp
index bddb6a76bbf..095d42599bf 100644
--- a/llvm/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp
@@ -51,21 +51,41 @@ protected:
     return 0;
   }
 
+  struct CompileContext {
+    CompileContext() : Compiled(false) { }
+
+    OrcCAPIExecutionTest* APIExecTest;
+    std::unique_ptr<Module> M;
+    LLVMOrcModuleHandle H;
+    bool Compiled;
+  };
+
+  static LLVMOrcTargetAddress myCompileCallback(LLVMOrcJITStackRef JITStack,
+                                                void *Ctx) {
+    CompileContext *CCtx = static_cast<CompileContext*>(Ctx);
+    auto *ET = CCtx->APIExecTest;
+    CCtx->M = ET->createTestModule(ET->TM->getTargetTriple());
+    CCtx->H = LLVMOrcAddEagerlyCompiledIR(JITStack, wrap(CCtx->M.get()),
+                                          myResolver, 0);
+    CCtx->Compiled = true;
+    LLVMOrcTargetAddress MainAddr = LLVMOrcGetSymbolAddress(JITStack, "main");
+    LLVMOrcSetIndirectStubPointer(JITStack, "foo", MainAddr);
+    return MainAddr;
+  }
+
 };
 
 char *OrcCAPIExecutionTest::testFuncName = 0;
 
 TEST_F(OrcCAPIExecutionTest, TestEagerIRCompilation) {
-  auto TM = getHostTargetMachineIfSupported();
-
   if (!TM)
     return;
 
-  std::unique_ptr<Module> M = createTestModule(TM->getTargetTriple());
-
   LLVMOrcJITStackRef JIT =
     LLVMOrcCreateInstance(wrap(TM.get()), LLVMGetGlobalContext());
 
+  std::unique_ptr<Module> M = createTestModule(TM->getTargetTriple());
+
   LLVMOrcGetMangledSymbol(JIT, &testFuncName, "testFunc");
 
   LLVMOrcModuleHandle H =
@@ -82,17 +102,16 @@ TEST_F(OrcCAPIExecutionTest, TestEagerIRCompilation) {
 }
 
 TEST_F(OrcCAPIExecutionTest, TestLazyIRCompilation) {
-  auto TM = getHostTargetMachineIfSupported();
-
   if (!TM)
     return;
 
-  std::unique_ptr<Module> M = createTestModule(TM->getTargetTriple());
-
   LLVMOrcJITStackRef JIT =
     LLVMOrcCreateInstance(wrap(TM.get()), LLVMGetGlobalContext());
 
+  std::unique_ptr<Module> M = createTestModule(TM->getTargetTriple());
+
   LLVMOrcGetMangledSymbol(JIT, &testFuncName, "testFunc");
+
   LLVMOrcModuleHandle H =
     LLVMOrcAddLazilyCompiledIR(JIT, wrap(M.get()), myResolver, 0);
   MainFnTy MainFn = (MainFnTy)LLVMOrcGetSymbolAddress(JIT, "main");
@@ -106,4 +125,37 @@ TEST_F(OrcCAPIExecutionTest, TestLazyIRCompilation) {
   LLVMOrcDisposeInstance(JIT);
 }
 
+TEST_F(OrcCAPIExecutionTest, TestDirectCallbacksAPI) {
+  if (!TM)
+    return;
+
+  LLVMOrcJITStackRef JIT =
+    LLVMOrcCreateInstance(wrap(TM.get()), LLVMGetGlobalContext());
+
+  LLVMOrcGetMangledSymbol(JIT, &testFuncName, "testFunc");
+
+  CompileContext C;
+  C.APIExecTest = this;
+  LLVMOrcCreateIndirectStub(JIT, "foo",
+                            LLVMOrcCreateLazyCompileCallback(JIT,
+                                                             myCompileCallback,
+                                                             &C));
+  MainFnTy FooFn = (MainFnTy)LLVMOrcGetSymbolAddress(JIT, "foo");
+  int Result = FooFn();
+  EXPECT_TRUE(C.Compiled)
+    << "Function wasn't lazily compiled";
+  EXPECT_EQ(Result, 42)
+    << "Direct-callback JIT'd code did not return expected result";
+
+  C.Compiled = false;
+  FooFn();
+  EXPECT_FALSE(C.Compiled)
+    << "Direct-callback JIT'd code was JIT'd twice";
+
+  LLVMOrcRemoveModule(JIT, C.H);
+
+  LLVMOrcDisposeMangledSymbol(testFuncName);
+  LLVMOrcDisposeInstance(JIT);
+}
+
 }
diff --git a/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h b/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h
index bfdaced05b1..15d9f54a37f 100644
--- a/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h
+++ b/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h
@@ -38,23 +38,20 @@ public:
       InitializeNativeTargetAsmPrinter();
       NativeTargetInitialized = true;
     }
-  };
-
-  // Get a target machine for the host if it supports JIT execution.
-  std::unique_ptr<TargetMachine> getHostTargetMachineIfSupported() {
-    std::unique_ptr<TargetMachine> TM(EngineBuilder().selectTarget());
-
-    if (!TM)
-      return nullptr;
-
-    const Triple& TT = TM->getTargetTriple();
 
-    if (TT.getArch() != Triple::x86_64 || !TT.isOSDarwin())
-      return nullptr;
+    // Try to select a TargetMachine for the host.
+    TM.reset(EngineBuilder().selectTarget());
 
-    return TM;
-  }
+    if (TM) {
+      // If we found a TargetMachine, check that it's one that Orc supports.
+      const Triple& TT = TM->getTargetTriple();
+      if (TT.getArch() != Triple::x86_64 || !TT.isOSDarwin())
+        TM = nullptr;
+    }
+  };
 
+protected:
+  std::unique_ptr<TargetMachine> TM;
 private:
   static bool NativeTargetInitialized;
 };