[ORC] Refactor executor symbol lookup to use ExecutorSymbolDef (NFC) (#76989)

This migrates the dylib manager lookup and related APIs to replace
ExecutorAddress with ExecutorSymbolDef so that in the future we can
model JITSymbolFlags for these symbols. The current change should be NFC
as we are only setting the Exported symbol flag.

Author: benlangmuir

compiler-rt/lib/orc/executor_symbol_def.h

@@ -0,0 +1,151 @@
+//===--------- ExecutorSymbolDef.h - (Addr, Flags) pair ---------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Represents a defining location for a symbol in the executing program.
+//
+// This file was derived from
+// llvm/include/llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_EXECUTOR_SYMBOL_DEF_H
+#define ORC_RT_EXECUTOR_SYMBOL_DEF_H
+
+#include "bitmask_enum.h"
+#include "executor_address.h"
+#include "simple_packed_serialization.h"
+
+namespace __orc_rt {
+
+/// Flags for symbols in the JIT.
+class JITSymbolFlags {
+public:
+  using UnderlyingType = uint8_t;
+  using TargetFlagsType = uint8_t;
+
+  /// These values must be kept in sync with \c JITSymbolFlags in the JIT.
+  enum FlagNames : UnderlyingType {
+    None = 0,
+    HasError = 1U << 0,
+    Weak = 1U << 1,
+    Common = 1U << 2,
+    Absolute = 1U << 3,
+    Exported = 1U << 4,
+    Callable = 1U << 5,
+    MaterializationSideEffectsOnly = 1U << 6,
+    ORC_RT_MARK_AS_BITMASK_ENUM( // LargestValue =
+        MaterializationSideEffectsOnly)
+  };
+
+  /// Default-construct a JITSymbolFlags instance.
+  JITSymbolFlags() = default;
+
+  /// Construct a JITSymbolFlags instance from the given flags and target
+  ///        flags.
+  JITSymbolFlags(FlagNames Flags, TargetFlagsType TargetFlags)
+      : TargetFlags(TargetFlags), Flags(Flags) {}
+
+  bool operator==(const JITSymbolFlags &RHS) const {
+    return Flags == RHS.Flags && TargetFlags == RHS.TargetFlags;
+  }
+
+  /// Get the underlying flags value as an integer.
+  UnderlyingType getRawFlagsValue() const {
+    return static_cast<UnderlyingType>(Flags);
+  }
+
+  /// Return a reference to the target-specific flags.
+  TargetFlagsType &getTargetFlags() { return TargetFlags; }
+
+  /// Return a reference to the target-specific flags.
+  const TargetFlagsType &getTargetFlags() const { return TargetFlags; }
+
+private:
+  TargetFlagsType TargetFlags = 0;
+  FlagNames Flags = None;
+};
+
+/// Represents a defining location for a JIT symbol.
+class ExecutorSymbolDef {
+public:
+  ExecutorSymbolDef() = default;
+  ExecutorSymbolDef(ExecutorAddr Addr, JITSymbolFlags Flags)
+      : Addr(Addr), Flags(Flags) {}
+
+  const ExecutorAddr &getAddress() const { return Addr; }
+
+  const JITSymbolFlags &getFlags() const { return Flags; }
+
+  friend bool operator==(const ExecutorSymbolDef &LHS,
+                         const ExecutorSymbolDef &RHS) {
+    return LHS.getAddress() == RHS.getAddress() &&
+           LHS.getFlags() == RHS.getFlags();
+  }
+
+private:
+  ExecutorAddr Addr;
+  JITSymbolFlags Flags;
+};
+
+using SPSJITSymbolFlags =
+    SPSTuple<JITSymbolFlags::UnderlyingType, JITSymbolFlags::TargetFlagsType>;
+
+/// SPS serializatior for JITSymbolFlags.
+template <> class SPSSerializationTraits<SPSJITSymbolFlags, JITSymbolFlags> {
+  using FlagsArgList = SPSJITSymbolFlags::AsArgList;
+
+public:
+  static size_t size(const JITSymbolFlags &F) {
+    return FlagsArgList::size(F.getRawFlagsValue(), F.getTargetFlags());
+  }
+
+  static bool serialize(SPSOutputBuffer &BOB, const JITSymbolFlags &F) {
+    return FlagsArgList::serialize(BOB, F.getRawFlagsValue(),
+                                   F.getTargetFlags());
+  }
+
+  static bool deserialize(SPSInputBuffer &BIB, JITSymbolFlags &F) {
+    JITSymbolFlags::UnderlyingType RawFlags;
+    JITSymbolFlags::TargetFlagsType TargetFlags;
+    if (!FlagsArgList::deserialize(BIB, RawFlags, TargetFlags))
+      return false;
+    F = JITSymbolFlags{static_cast<JITSymbolFlags::FlagNames>(RawFlags),
+                       TargetFlags};
+    return true;
+  }
+};
+
+using SPSExecutorSymbolDef = SPSTuple<SPSExecutorAddr, SPSJITSymbolFlags>;
+
+/// SPS serializatior for ExecutorSymbolDef.
+template <>
+class SPSSerializationTraits<SPSExecutorSymbolDef, ExecutorSymbolDef> {
+  using DefArgList = SPSExecutorSymbolDef::AsArgList;
+
+public:
+  static size_t size(const ExecutorSymbolDef &ESD) {
+    return DefArgList::size(ESD.getAddress(), ESD.getFlags());
+  }
+
+  static bool serialize(SPSOutputBuffer &BOB, const ExecutorSymbolDef &ESD) {
+    return DefArgList::serialize(BOB, ESD.getAddress(), ESD.getFlags());
+  }
+
+  static bool deserialize(SPSInputBuffer &BIB, ExecutorSymbolDef &ESD) {
+    ExecutorAddr Addr;
+    JITSymbolFlags Flags;
+    if (!DefArgList::deserialize(BIB, Addr, Flags))
+      return false;
+    ESD = ExecutorSymbolDef{Addr, Flags};
+    return true;
+  }
+};
+
+} // End namespace __orc_rt
+
+#endif // ORC_RT_EXECUTOR_SYMBOL_DEF_H

compiler-rt/lib/orc/tests/unit/CMakeLists.txt

@@ -5,6 +5,7 @@ set(UNITTEST_SOURCES
   endian_test.cpp
   error_test.cpp
   executor_address_test.cpp
+  executor_symbol_def_test.cpp
   extensible_rtti_test.cpp
   interval_map_test.cpp
   interval_set_test.cpp

compiler-rt/lib/orc/tests/unit/executor_symbol_def_test.cpp

@@ -0,0 +1,19 @@
+//===-- executor_symbol_def_test.cpp --------------------------------------===//
+//
+// 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 "executor_symbol_def.h"
+#include "simple_packed_serialization_utils.h"
+#include "gtest/gtest.h"
+
+using namespace __orc_rt;
+
+TEST(ExecutorSymbolDefTest, Serialization) {
+  blobSerializationRoundTrip<SPSExecutorSymbolDef>(ExecutorSymbolDef{});
+  blobSerializationRoundTrip<SPSExecutorSymbolDef>(
+      ExecutorSymbolDef{ExecutorAddr{0x70}, {JITSymbolFlags::Callable, 9}});
+}

compiler-rt/lib/orc/tests/unit/simple_packed_serialization_test.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "simple_packed_serialization.h"
+#include "simple_packed_serialization_utils.h"
 #include "gtest/gtest.h"
 
 using namespace __orc_rt;
@@ -48,25 +49,6 @@ TEST(SimplePackedSerializationTest, SPSInputBuffer) {
   EXPECT_FALSE(IB.read(&C, 1));
 }
 
-template <typename SPSTagT, typename T>
-static void blobSerializationRoundTrip(const T &Value) {
-  using BST = SPSSerializationTraits<SPSTagT, T>;
-
-  size_t Size = BST::size(Value);
-  auto Buffer = std::make_unique<char[]>(Size);
-  SPSOutputBuffer OB(Buffer.get(), Size);
-
-  EXPECT_TRUE(BST::serialize(OB, Value));
-
-  SPSInputBuffer IB(Buffer.get(), Size);
-
-  T DSValue;
-  EXPECT_TRUE(BST::deserialize(IB, DSValue));
-
-  EXPECT_EQ(Value, DSValue)
-      << "Incorrect value after serialization/deserialization round-trip";
-}
-
 template <typename T> static void testFixedIntegralTypeSerialization() {
   blobSerializationRoundTrip<T, T>(0);
   blobSerializationRoundTrip<T, T>(static_cast<T>(1));

compiler-rt/lib/orc/tests/unit/simple_packed_serialization_utils.h

@@ -0,0 +1,34 @@
+//===-- simple_packed_serialization_utils.h -------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_TEST_SIMPLE_PACKED_SERIALIZATION_UTILS_H
+#define ORC_RT_TEST_SIMPLE_PACKED_SERIALIZATION_UTILS_H
+
+#include "simple_packed_serialization.h"
+#include "gtest/gtest.h"
+
+template <typename SPSTagT, typename T>
+static void blobSerializationRoundTrip(const T &Value) {
+  using BST = __orc_rt::SPSSerializationTraits<SPSTagT, T>;
+
+  size_t Size = BST::size(Value);
+  auto Buffer = std::make_unique<char[]>(Size);
+  __orc_rt::SPSOutputBuffer OB(Buffer.get(), Size);
+
+  EXPECT_TRUE(BST::serialize(OB, Value));
+
+  __orc_rt::SPSInputBuffer IB(Buffer.get(), Size);
+
+  T DSValue;
+  EXPECT_TRUE(BST::deserialize(IB, DSValue));
+
+  EXPECT_EQ(Value, DSValue)
+      << "Incorrect value after serialization/deserialization round-trip";
+}
+
+#endif // ORC_RT_TEST_SIMPLE_PACKED_SERIALIZATION_UTILS_H

llvm/include/llvm/ExecutionEngine/Orc/EPCGenericDylibManager.h

@@ -19,6 +19,7 @@
 #define LLVM_EXECUTIONENGINE_ORC_EPCGENERICDYLIBMANAGER_H
 
 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
 #include "llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h"
 
 namespace llvm {
@@ -49,11 +50,11 @@ class EPCGenericDylibManager {
   Expected<tpctypes::DylibHandle> open(StringRef Path, uint64_t Mode);
 
   /// Looks up symbols within the given dylib.
-  Expected<std::vector<ExecutorAddr>> lookup(tpctypes::DylibHandle H,
-                                             const SymbolLookupSet &Lookup);
+  Expected<std::vector<ExecutorSymbolDef>>
+  lookup(tpctypes::DylibHandle H, const SymbolLookupSet &Lookup);
 
   /// Looks up symbols within the given dylib.
-  Expected<std::vector<ExecutorAddr>>
+  Expected<std::vector<ExecutorSymbolDef>>
   lookup(tpctypes::DylibHandle H, const RemoteSymbolLookupSet &Lookup);
 
 private:

llvm/include/llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h

@@ -15,6 +15,7 @@
 
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/ExecutionEngine/Orc/Shared/SimplePackedSerialization.h"
 
 namespace llvm {
 namespace orc {
@@ -48,6 +49,63 @@ class ExecutorSymbolDef {
   JITSymbolFlags Flags;
 };
 
+namespace shared {
+
+using SPSJITSymbolFlags =
+    SPSTuple<JITSymbolFlags::UnderlyingType, JITSymbolFlags::TargetFlagsType>;
+
+/// SPS serializatior for JITSymbolFlags.
+template <> class SPSSerializationTraits<SPSJITSymbolFlags, JITSymbolFlags> {
+  using FlagsArgList = SPSJITSymbolFlags::AsArgList;
+
+public:
+  static size_t size(const JITSymbolFlags &F) {
+    return FlagsArgList::size(F.getRawFlagsValue(), F.getTargetFlags());
+  }
+
+  static bool serialize(SPSOutputBuffer &BOB, const JITSymbolFlags &F) {
+    return FlagsArgList::serialize(BOB, F.getRawFlagsValue(),
+                                   F.getTargetFlags());
+  }
+
+  static bool deserialize(SPSInputBuffer &BIB, JITSymbolFlags &F) {
+    JITSymbolFlags::UnderlyingType RawFlags;
+    JITSymbolFlags::TargetFlagsType TargetFlags;
+    if (!FlagsArgList::deserialize(BIB, RawFlags, TargetFlags))
+      return false;
+    F = JITSymbolFlags{static_cast<JITSymbolFlags::FlagNames>(RawFlags),
+                       TargetFlags};
+    return true;
+  }
+};
+
+using SPSExecutorSymbolDef = SPSTuple<SPSExecutorAddr, SPSJITSymbolFlags>;
+
+/// SPS serializatior for ExecutorSymbolDef.
+template <>
+class SPSSerializationTraits<SPSExecutorSymbolDef, ExecutorSymbolDef> {
+  using DefArgList = SPSExecutorSymbolDef::AsArgList;
+
+public:
+  static size_t size(const ExecutorSymbolDef &ESD) {
+    return DefArgList::size(ESD.getAddress(), ESD.getFlags());
+  }
+
+  static bool serialize(SPSOutputBuffer &BOB, const ExecutorSymbolDef &ESD) {
+    return DefArgList::serialize(BOB, ESD.getAddress(), ESD.getFlags());
+  }
+
+  static bool deserialize(SPSInputBuffer &BIB, ExecutorSymbolDef &ESD) {
+    ExecutorAddr Addr;
+    JITSymbolFlags Flags;
+    if (!DefArgList::deserialize(BIB, Addr, Flags))
+      return false;
+    ESD = ExecutorSymbolDef{Addr, Flags};
+    return true;
+  }
+};
+
+} // End namespace shared.
 } // End namespace orc.
 } // End namespace llvm.
 

llvm/include/llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h

@@ -14,6 +14,7 @@
 #define LLVM_EXECUTIONENGINE_ORC_SHARED_ORCRTBRIDGE_H
 
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
 #include "llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h"
 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
 
@@ -54,7 +55,7 @@ using SPSSimpleExecutorDylibManagerOpenSignature =
                                                  shared::SPSString, uint64_t);
 
 using SPSSimpleExecutorDylibManagerLookupSignature =
-    shared::SPSExpected<shared::SPSSequence<shared::SPSExecutorAddr>>(
+    shared::SPSExpected<shared::SPSSequence<shared::SPSExecutorSymbolDef>>(
         shared::SPSExecutorAddr, shared::SPSExecutorAddr,
         shared::SPSRemoteSymbolLookupSet);
 

llvm/include/llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h

@@ -19,6 +19,7 @@
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/Orc/Shared/AllocationActions.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
 #include "llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h"
 #include "llvm/ExecutionEngine/Orc/Shared/SimplePackedSerialization.h"
 #include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
@@ -113,7 +114,7 @@ struct PointerWrite {
 /// A handle used to represent a loaded dylib in the target process.
 using DylibHandle = ExecutorAddr;
 
-using LookupResult = std::vector<ExecutorAddr>;
+using LookupResult = std::vector<ExecutorSymbolDef>;
 
 } // end namespace tpctypes
 

llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleExecutorDylibManager.h

@@ -18,6 +18,7 @@
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
 #include "llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h"
 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
 #include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
@@ -37,8 +38,8 @@ class SimpleExecutorDylibManager : public ExecutorBootstrapService {
   virtual ~SimpleExecutorDylibManager();
 
   Expected<tpctypes::DylibHandle> open(const std::string &Path, uint64_t Mode);
-  Expected<std::vector<ExecutorAddr>> lookup(tpctypes::DylibHandle H,
-                                             const RemoteSymbolLookupSet &L);
+  Expected<std::vector<ExecutorSymbolDef>>
+  lookup(tpctypes::DylibHandle H, const RemoteSymbolLookupSet &L);
 
   Error shutdown() override;
   void addBootstrapSymbols(StringMap<ExecutorAddr> &M) override;