[compiler-rt] ac0c8cb - [ORC-RT] Add "wrap" and "unwrap" steps to __orc_rt::ExecutorAddr toPtr/fromPtr.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2022-08-26T15:58:10-07:00
New Revision: ac0c8cb3ff7100d8a4480bed492e54ffc80814a1

URL: https://github.com/llvm/llvm-project/commit/ac0c8cb3ff7100d8a4480bed492e54ffc80814a1
DIFF: https://github.com/llvm/llvm-project/commit/ac0c8cb3ff7100d8a4480bed492e54ffc80814a1.diff

LOG: [ORC-RT] Add "wrap" and "unwrap" steps to __orc_rt::ExecutorAddr toPtr/fromPtr.

The wrap/unwrap operations are applied to pointers after/before conversion to/from
raw addresses. They can be used to tag, untag, sign, or strip signing from
pointers. They currently default to 'rawPtr' (identity) on all platforms, but it
is expected that the default will be set based on the host architecture, e.g.
they would default to signing/stripping for arm64e.

This is the ORC runtime counterpart to f14cb494a34:

Added: 
    

Modified: 
    compiler-rt/lib/orc/executor_address.h
    compiler-rt/lib/orc/unittests/executor_address_test.cpp

Removed: 
    


################################################################################
diff  --git a/compiler-rt/lib/orc/executor_address.h b/compiler-rt/lib/orc/executor_address.h
index f8f417cb636d2..e541d13627bb1 100644
--- a/compiler-rt/lib/orc/executor_address.h
+++ b/compiler-rt/lib/orc/executor_address.h
@@ -39,24 +39,71 @@ class ExecutorAddrDiff {
 /// Represents an address in the executor process.
 class ExecutorAddr {
 public:
+  /// A wrap/unwrap function that leaves pointers unmodified.
+  template <typename T> using rawPtr = __orc_rt::identity<T>;
+
+  /// Default wrap function to use on this host.
+  template <typename T> using defaultWrap = rawPtr<T>;
+
+  /// Default unwrap function to use on this host.
+  template <typename T> using defaultUnwrap = rawPtr<T>;
+
+  /// Merges a tag into the raw address value:
+  ///   P' = P | (TagValue << TagOffset).
+  class Tag {
+  public:
+    constexpr Tag(uintptr_t TagValue, uintptr_t TagOffset)
+        : TagMask(TagValue << TagOffset) {}
+
+    template <typename T> constexpr T *operator()(T *P) {
+      return reinterpret_cast<T *>(reinterpret_cast<uintptr_t>(P) | TagMask);
+    }
+
+  private:
+    uintptr_t TagMask;
+  };
+
+  /// Strips a tag of the given length from the given offset within the pointer:
+  /// P' = P & ~(((1 << TagLen) -1) << TagOffset)
+  class Untag {
+  public:
+    constexpr Untag(uintptr_t TagLen, uintptr_t TagOffset)
+        : UntagMask(~(((uintptr_t(1) << TagLen) - 1) << TagOffset)) {}
+
+    template <typename T> constexpr T *operator()(T *P) {
+      return reinterpret_cast<T *>(reinterpret_cast<uintptr_t>(P) & UntagMask);
+    }
+
+  private:
+    uintptr_t UntagMask;
+  };
+
   ExecutorAddr() = default;
   explicit ExecutorAddr(uint64_t Addr) : Addr(Addr) {}
 
   /// Create an ExecutorAddr from the given pointer.
-  /// Warning: This should only be used when JITing in-process.
-  template <typename T> static ExecutorAddr fromPtr(T *Value) {
+  template <typename T, typename UnwrapFn = defaultUnwrap<T *>>
+  static ExecutorAddr fromPtr(T *Ptr, UnwrapFn &&Unwrap = UnwrapFn()) {
     return ExecutorAddr(
-        static_cast<uint64_t>(reinterpret_cast<uintptr_t>(Value)));
+        static_cast<uint64_t>(reinterpret_cast<uintptr_t>(Unwrap(Ptr))));
   }
 
   /// Cast this ExecutorAddr to a pointer of the given type.
-  /// Warning: This should only be esude when JITing in-process.
-  template <typename T> T toPtr() const {
-    static_assert(std::is_pointer<T>::value, "T must be a pointer type");
+  template <typename T, typename WrapFn = defaultWrap<T>>
+  std::enable_if_t<std::is_pointer<T>::value, T>
+  toPtr(WrapFn &&Wrap = WrapFn()) const {
+    uintptr_t IntPtr = static_cast<uintptr_t>(Addr);
+    assert(IntPtr == Addr && "ExecutorAddr value out of range for uintptr_t");
+    return Wrap(reinterpret_cast<T>(IntPtr));
+  }
+
+  /// Cast this ExecutorAddr to a pointer of the given function type.
+  template <typename T, typename WrapFn = defaultWrap<T *>>
+  std::enable_if_t<std::is_function<T>::value, T *>
+  toPtr(WrapFn &&Wrap = WrapFn()) const {
     uintptr_t IntPtr = static_cast<uintptr_t>(Addr);
-    assert(IntPtr == Addr &&
-           "JITTargetAddress value out of range for uintptr_t");
-    return reinterpret_cast<T>(IntPtr);
+    assert(IntPtr == Addr && "ExecutorAddr value out of range for uintptr_t");
+    return Wrap(reinterpret_cast<T *>(IntPtr));
   }
 
   uint64_t getValue() const { return Addr; }

diff  --git a/compiler-rt/lib/orc/unittests/executor_address_test.cpp b/compiler-rt/lib/orc/unittests/executor_address_test.cpp
index 7712ef9a773f1..e0466b65392cc 100644
--- a/compiler-rt/lib/orc/unittests/executor_address_test.cpp
+++ b/compiler-rt/lib/orc/unittests/executor_address_test.cpp
@@ -52,6 +52,37 @@ TEST(ExecutorAddrTest, PtrConversion) {
   EXPECT_EQ(XPtr, &X);
 }
 
+static void F() {}
+
+TEST(ExecutorAddrTest, PtrConversionWithFunctionType) {
+  // Test that function types (as opposed to function pointer types) can be
+  // used with toPtr.
+  auto FAddr = ExecutorAddr::fromPtr(F);
+  void (*FPtr)() = FAddr.toPtr<void()>();
+
+  EXPECT_EQ(FPtr, &F);
+}
+
+TEST(ExecutorAddrTest, WrappingAndUnwrapping) {
+  constexpr uintptr_t RawAddr = 0x123456;
+  int *RawPtr = (int *)RawAddr;
+
+  constexpr uintptr_t TagOffset = 8 * (sizeof(uintptr_t) - 1);
+  uintptr_t TagVal = 0xA5;
+  uintptr_t TagBits = TagVal << TagOffset;
+  void *TaggedPtr = (void *)((uintptr_t)RawPtr | TagBits);
+
+  ExecutorAddr EA =
+      ExecutorAddr::fromPtr(TaggedPtr, ExecutorAddr::Untag(8, TagOffset));
+
+  EXPECT_EQ(EA.getValue(), RawAddr);
+
+  void *ReconstitutedTaggedPtr =
+      EA.toPtr<void *>(ExecutorAddr::Tag(TagVal, TagOffset));
+
+  EXPECT_EQ(TaggedPtr, ReconstitutedTaggedPtr);
+}
+
 TEST(ExecutorAddrTest, AddrRanges) {
   ExecutorAddr A0(0), A1(1), A2(2), A3(3);
   ExecutorAddrRange R0(A0, A1), R1(A1, A2), R2(A2, A3), R3(A0, A2), R4(A1, A3);