[ORC] Add findSymbolIn() wrapper to C bindings.

[anarazel]

Summary

In many cases JIT users will know in which module a symbol
resides. Avoiding search other modules can be more efficient. It also
allows to handle duplicate symbol names between modules.

[anarazel]

Ping? This is a simple enough addition that seems to have been envisioned when designing the Orc C stack (given that the underlying stack isn't used anywhere else, and has findSymbolIn()).

[lhames]

Looks good to me. Do you have commit access? If not I can commit on your behalf.

[anarazel]

In D44889#1110390, @lhames wrote:

Looks good to me. Do you have commit access? If not I can commit on your behalf.

Thanks for the review. I recently got commit access, so I can take care of committing myself.

Thanks,

Andres

[lhames]

Great. Thanks for working on this Andres!

• Lang.

include/llvm-c/OrcBindings.h

 /**
  * Get symbol address from JIT instance.
  */
 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.
  */
 LLVMOrcErrorCode LLVMOrcDisposeInstance(LLVMOrcJITStackRef JITStack);
 
 #ifdef __cplusplus
 }
 #endif /* extern "C" */
 
 #endif /* LLVM_C_ORCBINDINGS_H */

lib/ExecutionEngine/Orc/OrcCBindings.cpp

 
 LLVMOrcErrorCode LLVMOrcGetSymbolAddress(LLVMOrcJITStackRef JITStack,
                                          LLVMOrcTargetAddress *RetAddr,
                                          const char *SymbolName) {
   OrcCBindingsStack &J = *unwrap(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();
   delete J;
   return Err;
 }

lib/ExecutionEngine/Orc/OrcCBindingsStack.h

JITSymbol findSymbol(const std::string &Name,
                                  bool ExportedSymbolsOnly) {
     if (auto Sym = IndirectStubsMgr->findStub(Name, ExportedSymbolsOnly))
       return Sym;
     return CODLayer.findSymbol(mangle(Name), ExportedSymbolsOnly);
   }
 
   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, Name, ExportedSymbolsOnly);
   }
 
   LLVMOrcErrorCode findSymbolAddress(JITTargetAddress &RetAddr,
                                      const std::string &Name,
                                      bool ExportedSymbolsOnly) {
     RetAddr = 0;
     if (auto Sym = findSymbol(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;
   }
 
+  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:
 
   LLVMOrcErrorCode mapError(Error Err) {
     LLVMOrcErrorCode Result = LLVMOrcErrSuccess;
     handleAllErrors(std::move(Err), [&](ErrorInfoBase &EIB) {
       // Handler of last resort.

unittests/ExecutionEngine/Orc/OrcCAPITest.cpp

LLVMOrcJITStackRef JIT =
     LLVMOrcCreateInstance(wrap(TM.get()));
 
   std::unique_ptr<Module> M = createTestModule(TM->getTargetTriple());
 
   LLVMOrcGetMangledSymbol(JIT, &testFuncName, "testFunc");
 
   LLVMOrcModuleHandle H;
   LLVMOrcAddEagerlyCompiledIR(JIT, &H, wrap(M.release()), myResolver, nullptr);
+
+  // 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);
 
   LLVMOrcDisposeMangledSymbol(testFuncName);
   LLVMOrcDisposeInstance(JIT);
 }
 
 TEST_F(OrcCAPIExecutionTest, TestLazyIRCompilation) {