diff --git a/include/llvm-c/OrcBindings.h b/include/llvm-c/OrcBindings.h index d2e2e7de4d3..ddba6430834 100644 --- a/include/llvm-c/OrcBindings.h +++ b/include/llvm-c/OrcBindings.h @@ -140,6 +140,15 @@ LLVMOrcErrorCode LLVMOrcGetSymbolAddress(LLVMOrcJITStackRef JITStack, LLVMOrcTargetAddress *RetAddr, const char *SymbolName); +/** + * Get symbol address from JIT instance, searching only the specified + * handle. + */ +LLVMOrcErrorCode LLVMOrcGetSymbolAddressIn(LLVMOrcJITStackRef JITStack, + LLVMOrcTargetAddress *RetAddr, + LLVMOrcModuleHandle H, + const char *SymbolName); + /** * Dispose of an ORC JIT stack. */ diff --git a/lib/ExecutionEngine/Orc/OrcCBindings.cpp b/lib/ExecutionEngine/Orc/OrcCBindings.cpp index 0ea7a663989..1cf8a3a745a 100644 --- a/lib/ExecutionEngine/Orc/OrcCBindings.cpp +++ b/lib/ExecutionEngine/Orc/OrcCBindings.cpp @@ -112,6 +112,14 @@ LLVMOrcErrorCode LLVMOrcGetSymbolAddress(LLVMOrcJITStackRef JITStack, return J.findSymbolAddress(*RetAddr, SymbolName, true); } +LLVMOrcErrorCode LLVMOrcGetSymbolAddressIn(LLVMOrcJITStackRef JITStack, + LLVMOrcTargetAddress *RetAddr, + LLVMOrcModuleHandle H, + const char *SymbolName) { + OrcCBindingsStack &J = *unwrap(JITStack); + return J.findSymbolAddressIn(*RetAddr, H, SymbolName, true); +} + LLVMOrcErrorCode LLVMOrcDisposeInstance(LLVMOrcJITStackRef JITStack) { auto *J = unwrap(JITStack); auto Err = J->shutdown(); diff --git a/lib/ExecutionEngine/Orc/OrcCBindingsStack.h b/lib/ExecutionEngine/Orc/OrcCBindingsStack.h index 96b553f6fdd..0a18f6b731c 100644 --- a/lib/ExecutionEngine/Orc/OrcCBindingsStack.h +++ b/lib/ExecutionEngine/Orc/OrcCBindingsStack.h @@ -380,6 +380,7 @@ public: JITSymbol findSymbolIn(orc::VModuleKey K, const std::string &Name, bool ExportedSymbolsOnly) { + assert(KeyLayers.count(K) && "looking up symbol in unknown module"); return KeyLayers[K]->findSymbolIn(K, mangle(Name), ExportedSymbolsOnly); } @@ -403,6 +404,27 @@ public: return LLVMOrcErrSuccess; } + LLVMOrcErrorCode findSymbolAddressIn(JITTargetAddress &RetAddr, + orc::VModuleKey K, + const std::string &Name, + bool ExportedSymbolsOnly) { + RetAddr = 0; + if (auto Sym = findSymbolIn(K, Name, ExportedSymbolsOnly)) { + // Successful lookup, non-null symbol: + if (auto AddrOrErr = Sym.getAddress()) { + RetAddr = *AddrOrErr; + return LLVMOrcErrSuccess; + } else + return mapError(AddrOrErr.takeError()); + } else if (auto Err = Sym.takeError()) { + // Lookup failure - report error. + return mapError(std::move(Err)); + } + // Otherwise we had a successful lookup but got a null result. We already + // set RetAddr to '0' above, so just return success. + return LLVMOrcErrSuccess; + } + const std::string &getErrorMessage() const { return ErrMsg; } private: diff --git a/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp b/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp index fe3f694f560..b288b6bab2c 100644 --- a/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp +++ b/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp @@ -98,12 +98,26 @@ TEST_F(OrcCAPIExecutionTest, TestEagerIRCompilation) { LLVMOrcModuleHandle H; LLVMOrcAddEagerlyCompiledIR(JIT, &H, wrap(M.release()), myResolver, nullptr); - LLVMOrcTargetAddress MainAddr; - LLVMOrcGetSymbolAddress(JIT, &MainAddr, "main"); - MainFnTy MainFn = (MainFnTy)MainAddr; - int Result = MainFn(); - EXPECT_EQ(Result, 42) - << "Eagerly JIT'd code did not return expected result"; + + // get symbol address searching the entire stack + { + LLVMOrcTargetAddress MainAddr; + LLVMOrcGetSymbolAddress(JIT, &MainAddr, "main"); + MainFnTy MainFn = (MainFnTy)MainAddr; + int Result = MainFn(); + EXPECT_EQ(Result, 42) + << "Eagerly JIT'd code did not return expected result"; + } + + // and then just searching a single handle + { + LLVMOrcTargetAddress MainAddr; + LLVMOrcGetSymbolAddressIn(JIT, &MainAddr, H, "main"); + MainFnTy MainFn = (MainFnTy)MainAddr; + int Result = MainFn(); + EXPECT_EQ(Result, 42) + << "Eagerly JIT'd code did not return expected result"; + } LLVMOrcRemoveModule(JIT, H);