[ORC] Introduce lookupAndRecordAddrs utility.

Accepts a vector of (SymbolStringPtr, ExecutorAddress*) pairs, looks up all the
symbols, then writes their address to each of the corresponding
ExecutorAddresses.

This idiom (looking up and recording addresses into a specific set of variables)
is used in MachOPlatform and the (temporarily reverted) ELFNixPlatform, and is
likely to be used in other places in the near future, so wrapping it in a
utility function should save us some boilerplate.

llvm/include/llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h

@@ -0,0 +1,69 @@
+//===-- LookupAndRecordAddrs.h - Symbol lookup support utility --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Record the addresses of a set of symbols into ExecutorAddress objects.
+//
+// This can be used to avoid repeated lookup (via ExecutionSession::lookup) of
+// the given symbols.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_LOOKUPANDRECORDADDRS_H
+#define LLVM_EXECUTIONENGINE_ORC_LOOKUPANDRECORDADDRS_H
+
+#include "llvm/ADT/FunctionExtras.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+
+#include <vector>
+
+namespace llvm {
+namespace orc {
+
+/// Record addresses of the given symbols in the given ExecutorAddresses.
+///
+/// Useful for making permanent records of symbol addreses to call or
+/// access in the executor (e.g. runtime support functions in Platform
+/// subclasses).
+///
+/// By default the symbols are looked up using
+/// SymbolLookupFlags::RequiredSymbol, and an error will be generated if any of
+/// the requested symbols are not defined.
+///
+/// If SymbolLookupFlags::WeaklyReferencedSymbol is used then any missing
+/// symbols will have their corresponding address objects set to zero, and
+/// this function will never generate an error (the caller will need to check
+/// addresses before using them).
+///
+/// Asynchronous version.
+void lookupAndRecordAddrs(
+    unique_function<void(Error)> OnRecorded, ExecutionSession &ES, LookupKind K,
+    const JITDylibSearchOrder &SearchOrder,
+    std::vector<std::pair<SymbolStringPtr, ExecutorAddress *>> Pairs,
+    SymbolLookupFlags LookupFlags = SymbolLookupFlags::RequiredSymbol);
+
+/// Record addresses of the given symbols in the given ExecutorAddresses.
+///
+/// Blocking version.
+Error lookupAndRecordAddrs(
+    ExecutionSession &ES, LookupKind K, const JITDylibSearchOrder &SearchOrder,
+    std::vector<std::pair<SymbolStringPtr, ExecutorAddress *>> Pairs,
+    SymbolLookupFlags LookupFlags = SymbolLookupFlags::RequiredSymbol);
+
+/// Record addresses of given symbols in the given ExecutorAddresses.
+///
+/// ExecutorProcessControl lookup version. Lookups are always implicitly
+/// weak.
+Error lookupAndRecordAddrs(
+    ExecutorProcessControl &EPC, tpctypes::DylibHandle H,
+    std::vector<std::pair<SymbolStringPtr, ExecutorAddress *>> Pairs);
+
+} // End namespace orc
+} // End namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_LOOKUPANDRECORDADDRS_H

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -15,6 +15,7 @@ add_llvm_component_library(LLVMOrcJIT
   JITTargetMachineBuilder.cpp
   LazyReexports.cpp
   Layer.cpp
+  LookupAndRecordAddrs.cpp
   LLJIT.cpp
   MachOPlatform.cpp
   Mangling.cpp

llvm/lib/ExecutionEngine/Orc/LookupAndRecordAddrs.cpp

@@ -0,0 +1,82 @@
+//===------- LookupAndRecordAddrs.h - Symbol lookup support utility -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
+
+#include <future>
+
+namespace llvm {
+namespace orc {
+
+void lookupAndRecordAddrs(
+    unique_function<void(Error)> OnRecorded, ExecutionSession &ES, LookupKind K,
+    const JITDylibSearchOrder &SearchOrder,
+    std::vector<std::pair<SymbolStringPtr, ExecutorAddress *>> Pairs,
+    SymbolLookupFlags LookupFlags) {
+
+  SymbolLookupSet Symbols;
+  for (auto &KV : Pairs)
+    Symbols.add(KV.first, LookupFlags);
+
+  ES.lookup(
+      K, SearchOrder, Symbols, SymbolState::Ready,
+      [Pairs = std::move(Pairs),
+       OnRec = std::move(OnRecorded)](Expected<SymbolMap> Result) mutable {
+        if (!Result)
+          return OnRec(Result.takeError());
+        for (auto &KV : Pairs) {
+          auto I = Result->find(KV.first);
+          KV.second->setValue((I != Result->end()) ? I->second.getAddress()
+                                                   : 0);
+        }
+        OnRec(Error::success());
+      },
+      NoDependenciesToRegister);
+}
+
+Error lookupAndRecordAddrs(
+    ExecutionSession &ES, LookupKind K, const JITDylibSearchOrder &SearchOrder,
+    std::vector<std::pair<SymbolStringPtr, ExecutorAddress *>> Pairs,
+    SymbolLookupFlags LookupFlags) {
+
+  std::promise<MSVCPError> ResultP;
+  auto ResultF = ResultP.get_future();
+  lookupAndRecordAddrs([&](Error Err) { ResultP.set_value(std::move(Err)); },
+                       ES, K, SearchOrder, Pairs, LookupFlags);
+  return ResultF.get();
+}
+
+Error lookupAndRecordAddrs(
+    ExecutorProcessControl &EPC, tpctypes::DylibHandle H,
+    std::vector<std::pair<SymbolStringPtr, ExecutorAddress *>> Pairs,
+    SymbolLookupFlags LookupFlags) {
+
+  SymbolLookupSet Symbols;
+  for (auto &KV : Pairs)
+    Symbols.add(KV.first, LookupFlags);
+
+  ExecutorProcessControl::LookupRequest LR(H, Symbols);
+  auto Result = EPC.lookupSymbols(LR);
+  if (!Result)
+    return Result.takeError();
+
+  if (Result->size() != 1)
+    return make_error<StringError>("Error in lookup result",
+                                   inconvertibleErrorCode());
+  if (Result->front().size() != Pairs.size())
+    return make_error<StringError>("Error in lookup result elements",
+                                   inconvertibleErrorCode());
+
+  for (unsigned I = 0; I != Pairs.size(); ++I)
+    Pairs[I].second->setValue(Result->front()[I]);
+
+  return Error::success();
+}
+
+} // End namespace orc.
+} // End namespace llvm.

llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp

@@ -12,6 +12,7 @@
 #include "llvm/ExecutionEngine/JITLink/x86_64.h"
 #include "llvm/ExecutionEngine/Orc/DebugUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
 #include "llvm/Support/BinaryByteStream.h"
 #include "llvm/Support/Debug.h"
 
@@ -496,31 +497,17 @@ void MachOPlatform::rt_lookupSymbol(SendSymbolAddressFn SendResult,
 
 Error MachOPlatform::bootstrapMachORuntime(JITDylib &PlatformJD) {
 
-  std::pair<const char *, ExecutorAddress *> Symbols[] = {
-      {"___orc_rt_macho_platform_bootstrap", &orc_rt_macho_platform_bootstrap},
-      {"___orc_rt_macho_platform_shutdown", &orc_rt_macho_platform_shutdown},
-      {"___orc_rt_macho_register_object_sections",
-       &orc_rt_macho_register_object_sections},
-      {"___orc_rt_macho_create_pthread_key", &orc_rt_macho_create_pthread_key}};
-
-  SymbolLookupSet RuntimeSymbols;
-  std::vector<std::pair<SymbolStringPtr, ExecutorAddress *>> AddrsToRecord;
-  for (const auto &KV : Symbols) {
-    auto Name = ES.intern(KV.first);
-    RuntimeSymbols.add(Name);
-    AddrsToRecord.push_back({std::move(Name), KV.second});
-  }
-
-  auto RuntimeSymbolAddrs = ES.lookup(
-      {{&PlatformJD, JITDylibLookupFlags::MatchAllSymbols}}, RuntimeSymbols);
-  if (!RuntimeSymbolAddrs)
-    return RuntimeSymbolAddrs.takeError();
-
-  for (const auto &KV : AddrsToRecord) {
-    auto &Name = KV.first;
-    assert(RuntimeSymbolAddrs->count(Name) && "Missing runtime symbol?");
-    KV.second->setValue((*RuntimeSymbolAddrs)[Name].getAddress());
-  }
+  if (auto Err = lookupAndRecordAddrs(
+          ES, LookupKind::Static, makeJITDylibSearchOrder(&PlatformJD),
+          {{ES.intern("___orc_rt_macho_platform_bootstrap"),
+            &orc_rt_macho_platform_bootstrap},
+           {ES.intern("___orc_rt_macho_platform_shutdown"),
+            &orc_rt_macho_platform_shutdown},
+           {ES.intern("___orc_rt_macho_register_object_sections"),
+            &orc_rt_macho_register_object_sections},
+           {ES.intern("___orc_rt_macho_create_pthread_key"),
+            &orc_rt_macho_create_pthread_key}}))
+    return Err;
 
   if (auto Err =
           ES.callSPSWrapper<void()>(orc_rt_macho_platform_bootstrap.getValue()))

llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt

@@ -20,6 +20,7 @@ add_llvm_unittest(OrcJITTests
   IndirectionUtilsTest.cpp
   JITTargetMachineBuilderTest.cpp
   LazyCallThroughAndReexportsTest.cpp
+  LookupAndRecordAddrsTest.cpp
   ObjectLinkingLayerTest.cpp
   OrcCAPITest.cpp
   OrcTestCommon.cpp

llvm/unittests/ExecutionEngine/Orc/LookupAndRecordAddrsTest.cpp

@@ -0,0 +1,108 @@
+//===- LookupAndRecordAddrsTest.cpp - Unit tests for LookupAndRecordAddrs -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "OrcTestCommon.h"
+
+#include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
+#include "llvm/Support/MSVCErrorWorkarounds.h"
+#include "llvm/Testing/Support/Error.h"
+
+#include <future>
+
+using namespace llvm;
+using namespace llvm::orc;
+
+class LookupAndRecordAddrsTest : public CoreAPIsBasedStandardTest {};
+
+namespace {
+
+TEST_F(LookupAndRecordAddrsTest, AsyncRequiredSuccess) {
+  cantFail(JD.define(absoluteSymbols({{Foo, FooSym}, {Bar, BarSym}})));
+
+  ExecutorAddress FooAddress, BarAddress;
+  std::promise<MSVCPError> ErrP;
+
+  lookupAndRecordAddrs([&](Error Err) { ErrP.set_value(std::move(Err)); }, ES,
+                       LookupKind::Static, makeJITDylibSearchOrder(&JD),
+                       {{Foo, &FooAddress}, {Bar, &BarAddress}});
+
+  Error Err = ErrP.get_future().get();
+
+  EXPECT_THAT_ERROR(std::move(Err), Succeeded());
+  EXPECT_EQ(FooAddress.getValue(), FooAddr);
+  EXPECT_EQ(BarAddress.getValue(), BarAddr);
+}
+
+TEST_F(LookupAndRecordAddrsTest, AsyncRequiredFailure) {
+  ExecutorAddress FooAddress, BarAddress;
+  std::promise<MSVCPError> ErrP;
+
+  lookupAndRecordAddrs([&](Error Err) { ErrP.set_value(std::move(Err)); }, ES,
+                       LookupKind::Static, makeJITDylibSearchOrder(&JD),
+                       {{Foo, &FooAddress}, {Bar, &BarAddress}});
+
+  Error Err = ErrP.get_future().get();
+
+  EXPECT_THAT_ERROR(std::move(Err), Failed());
+}
+
+TEST_F(LookupAndRecordAddrsTest, AsyncWeakReference) {
+  cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
+
+  ExecutorAddress FooAddress, BarAddress;
+  std::promise<MSVCPError> ErrP;
+
+  lookupAndRecordAddrs([&](Error Err) { ErrP.set_value(std::move(Err)); }, ES,
+                       LookupKind::Static, makeJITDylibSearchOrder(&JD),
+                       {{Foo, &FooAddress}, {Bar, &BarAddress}},
+                       SymbolLookupFlags::WeaklyReferencedSymbol);
+
+  Error Err = ErrP.get_future().get();
+
+  EXPECT_THAT_ERROR(std::move(Err), Succeeded());
+  EXPECT_EQ(FooAddress.getValue(), FooAddr);
+  EXPECT_EQ(BarAddress.getValue(), 0U);
+}
+
+TEST_F(LookupAndRecordAddrsTest, BlockingRequiredSuccess) {
+  cantFail(JD.define(absoluteSymbols({{Foo, FooSym}, {Bar, BarSym}})));
+
+  ExecutorAddress FooAddress, BarAddress;
+  auto Err =
+      lookupAndRecordAddrs(ES, LookupKind::Static, makeJITDylibSearchOrder(&JD),
+                           {{Foo, &FooAddress}, {Bar, &BarAddress}});
+
+  EXPECT_THAT_ERROR(std::move(Err), Succeeded());
+  EXPECT_EQ(FooAddress.getValue(), FooAddr);
+  EXPECT_EQ(BarAddress.getValue(), BarAddr);
+}
+
+TEST_F(LookupAndRecordAddrsTest, BlockingRequiredFailure) {
+  ExecutorAddress FooAddress, BarAddress;
+  auto Err =
+      lookupAndRecordAddrs(ES, LookupKind::Static, makeJITDylibSearchOrder(&JD),
+                           {{Foo, &FooAddress}, {Bar, &BarAddress}});
+
+  EXPECT_THAT_ERROR(std::move(Err), Failed());
+}
+
+TEST_F(LookupAndRecordAddrsTest, BlockingWeakReference) {
+  cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
+
+  ExecutorAddress FooAddress, BarAddress;
+  auto Err =
+      lookupAndRecordAddrs(ES, LookupKind::Static, makeJITDylibSearchOrder(&JD),
+                           {{Foo, &FooAddress}, {Bar, &BarAddress}},
+                           SymbolLookupFlags::WeaklyReferencedSymbol);
+
+  EXPECT_THAT_ERROR(std::move(Err), Succeeded());
+  EXPECT_EQ(FooAddress.getValue(), FooAddr);
+  EXPECT_EQ(BarAddress.getValue(), 0U);
+}
+
+} // namespace