[ADT] Guard our use of function pointers in a `PointerUnion` by a dynamic check that the pointer is actually sufficiently aligned.

[chandlerc]

Summary

We essentially use another bit to allow us to distinguish an embedded
function pointer from a double indirection through a well aligned
global. This makes it much easier to take advantage of the well aligned
function pointer when available, but degrade gracefully when unavailable
to an indirected function pointer.

This does come with a small cost. If folks are worried about this
runtime cost, we could just force every platform not explicitly
supported with the alignment attribute to use the indirected callback
and have no runtime overhead, but this feels a bit risky to me. Open to
suggestions.

Hopefully this fixes the reported segfault on some MacOS build bots and
with -Os and other build configurations where the alignment wasn't
actually reliable.

[bogner]

This fixes the segfault / asan error I was seeing in ADTTests with -Os on macOS, and the trade off with the runtime check seems pretty reasonable to me.

[chandlerc]

I'm gonna abandon this and submit a different patch.

It has really bad properties. In a bunch of cases the pointers still aren't aligned and with this patch when they *happen* to fall on 4-byte boundaries we'll end up embedding a pointer. This makes things like debugging awful due to the highly dynamic aspect of everything.

I'm going to add a layer of indirection until we have some much more reliable pattern here. It makes me really sad to have to do the layer of indirection, but this is breaking too many bots...

llvm/include/llvm/ADT/FunctionExtras.h

class unique_function<ReturnT(ParamTs...)> {
 
   // The type of the erased function pointer we use as a callback to dispatch to
   // the stored callable when it is trivial to move and destroy.
   using CallPtrT = ReturnT (*)(void *CallableAddr,
                                AdjustedParamT<ParamTs>... Params);
   using MovePtrT = void (*)(void *LHSCallableAddr, void *RHSCallableAddr);
   using DestroyPtrT = void (*)(void *CallableAddr);
 
+  /// A struct to hold a single trivial callback with sufficient alignment for
+  /// our bitpacking.
+  struct alignas(8) TrivialCallback {
+    CallPtrT CallPtr;
+  };
+
   /// A struct we use to aggregate three callbacks when we need full set of
   /// operations.
-  struct NonTrivialCallbacks {
+  struct alignas(8) NonTrivialCallbacks {
     CallPtrT CallPtr;
     MovePtrT MovePtr;
     DestroyPtrT DestroyPtr;
   };
 
-  // Now we can create a pointer union between either a direct, trivial call
-  // pointer and a pointer to a static struct of the call, move, and destroy
-  // pointers. We do this to keep the footprint in this object a single pointer
-  // while supporting all the necessary type-erased operation.
-  using CallbackPointerUnionT = PointerUnion<CallPtrT, NonTrivialCallbacks *>;
+  // Now we can create a pointer union between either a direct trivial call
+  // pointer, a pointer to a static trivial call pointer, or a pointer to
+  // a static struct of the call, move, and destroy pointers. We do this to
+  // keep the footprint in this object a single pointer while supporting all
+  // the necessary type-erased operations with a minimum number of
+  // indirections.
+  using CallbackPointerUnionT =
+      PointerUnion3<TrivialCallback *, NonTrivialCallbacks *, CallPtrT>;
 
   // The main storage buffer. This will either have a pointer to out-of-line
   // storage or an inline buffer storing the callable.
   union StorageUnionT {
     // For out-of-line storage we keep a pointer to the underlying storage and
     // the size. This is enough to deallocate the memory.
     struct OutOfLineStorageT {
       void *StoragePtr;
class unique_function<ReturnT(ParamTs...)> {
   size_t getOutOfLineStorageAlignment() const {
     return StorageUnion.OutOfLineStorage.Alignment;
   }
 
   void setOutOfLineStorage(void *Ptr, size_t Size, size_t Alignment) {
     StorageUnion.OutOfLineStorage = {Ptr, Size, Alignment};
   }
 
+  // We really want to have a well aligned function pointer if possible. That
+  // allows us to use a much more efficient representation when dispatching.
+  //
+  // To do this, on platforms where we expect it to be useful, we add an
+  // alignment attribute to try and get the function pointer to be well aligned.
+  // This is *not* a correctness requirement though. We will check to see if we
+  // ended up with adequate alignment below and use a safe approach otherwise.
+  //
+  // We currently support this on Clang and GCC-compatible compilers when
+  // targeting x86, PowerPC, AArch64, and ARM (but not in thumb mode as it has
+  // no utility then). Other architectures can be added as needed.
+#if (defined(__clang__) || defined(__GNUC__)) &&                               \
+    (defined(__i386__) || defined(__x86_64__) || defined(__ppc__) ||           \
+     defined(__aarch64__) || (defined(__arm__) && !defined(__thumb__)))
+#define CALL_IMPL_ALIGNMENT __attribute__((aligned(16)))
+#else
+#define CALL_IMPL_ALIGNMENT
+#endif
   template <typename CallableT>
-  static ReturnT CallImpl(void *CallableAddr,
-                          AdjustedParamT<ParamTs>... Params) {
+  static ReturnT CallImpl(void *CallableAddr, AdjustedParamT<ParamTs>... Params)
+      CALL_IMPL_ALIGNMENT {
     return (*reinterpret_cast<CallableT *>(CallableAddr))(
         std::forward<ParamTs>(Params)...);
   }
+#undef CALL_IMPL_ALIGNMENT
 
   template <typename CallableT>
   static void MoveImpl(void *LHSCallableAddr, void *RHSCallableAddr) noexcept {
     new (LHSCallableAddr)
         CallableT(std::move(*reinterpret_cast<CallableT *>(RHSCallableAddr)));
   }
 
   template <typename CallableT>
if (sizeof(CallableT) > InlineStorageSize ||
       auto Alignment = alignof(CallableT);
       CallableAddr = allocate_buffer(Size, Alignment);
       setOutOfLineStorage(CallableAddr, Size, Alignment);
     }
 
     // Now move into the storage.
     new (CallableAddr) CallableT(std::move(Callable));
 
-#ifndef _MSC_VER
-    // See if we can create a trivial callback.
-    //
-    // This requires two things. First, we need to put a function pointer into
-    // a `PointerIntPair` and a `PointerUnion`. We assume that function
-    // pointers on almost every platform have at least 4-byte alignment which
-    // allows this to work, but on some platforms this appears to not hold, and
-    // so we need to disable this optimization on those platforms.
-    //
-    // FIXME: Currently, Windows appears to fail the asserts that check this.
-    // We should investigate why, and if fixed removed the #ifndef above.
-    //
-    // Second, we need the callable to be trivially moved and trivially
-    // destroyed so that we don't have to store type erased callbacks for those
-    // operations.
+    // See if we can create a trivial callback. We need the callable to be
+    // trivially moved and trivially destroyed so that we don't have to store
+    // type erased callbacks for those operations.
     //
     // FIXME: We should use constexpr if here and below to avoid instantiating
     // the non-trivial static objects when unnecessary. While the linker should
     // remove them, it is still wasteful.
     if (llvm::is_trivially_move_constructible<CallableT>::value &&
         std::is_trivially_destructible<CallableT>::value) {
-      CallbackAndInlineFlag = {&CallImpl<CallableT>, IsInlineStorage};
+      // Now see if we can embed the pointer to the trivial callback. We have to
+      // do this at runtime because there is no reliably way to statically
+      // determine (or ensure) that the function pointer we get will be
+      // sufficiently aligned. We take steps above to try and make this likely
+      // on common LLVM platforms.
+      CallPtrT CallImplPtr = &CallImpl<CallableT>;
+      uintptr_t CallImplAddr = reinterpret_cast<uintptr_t>(CallImplPtr);
+      using PtrTraits = PointerLikeTypeTraits<CallPtrT>;
+      uintptr_t LowBitMask = ~((uintptr_t)-1 << PtrTraits::NumLowBitsAvailable);
+      if ((CallImplAddr & LowBitMask) == 0) {
+        // The necessary bits are clear in the function pointer so we can embed
+        // it directly.
+        CallbackAndInlineFlag = {CallImplPtr, IsInlineStorage};
+        return;
+      }
+
+      // Otherwise, we need to create a nicely aligned object. We use a static
+      // variable for this because it is a trivial struct.
+      static TrivialCallback Callback = { CallImplPtr };
+
+      CallbackAndInlineFlag = {&Callback, IsInlineStorage};
       return;
     }
-#endif
 
-    // Otherwise, we need to point at an object with a vtable that contains all
-    // the different type erased behaviors needed. Create a static instance of
-    // the derived type here and then use a pointer to that.
+    // Otherwise, we need to point at an object that contains all the different
+    // type erased behaviors needed. Create a static instance of the struct type
+    // here and then use a pointer to that.
     static NonTrivialCallbacks Callbacks = {
         &CallImpl<CallableT>, &MoveImpl<CallableT>, &DestroyImpl<CallableT>};
 
     CallbackAndInlineFlag = {&Callbacks, IsInlineStorage};
   }
 
   ReturnT operator()(ParamTs... Params) {
     void *CallableAddr =

llvm/include/llvm/Support/PointerLikeTypeTraits.h

static inline FunctionPointerT getFromVoidPointer(void *P) {
     return reinterpret_cast<FunctionPointerT>(P);
   }
 };
 
 /// Provide a default specialization for function pointers that assumes 4-byte
 /// alignment.
 ///
 /// We assume here that functions used with this are always at least 4-byte
-/// aligned. This means that, for example, thumb functions won't work or systems
-/// with weird unaligned function pointers won't work. But all practical systems
-/// we support satisfy this requirement.
+/// aligned. This means that, for example, thumb functions won't work. There
+/// are many ways that a function may end up not satisfying this alignment.
+/// Users of such pointers with things like `PointerIntPair` and `PointerUnion`
+/// must take care themselves that the function pointers they use are valid
+/// here. This may have to rely on a runtime check in some cases due to the
+/// complexity of determining this statically.
 template <typename ReturnT, typename... ParamTs>
 struct PointerLikeTypeTraits<ReturnT (*)(ParamTs...)>
     : FunctionPointerLikeTypeTraits<4, ReturnT (*)(ParamTs...)> {};
 
 } // end namespace llvm
 
 #endif