From 22b4496e86c32127e997f3e7385ef64d2d80cc4b Mon Sep 17 00:00:00 2001
From: Dmitry Yanovsky <kerambyte@gmail.com>
Date: Mon, 19 Aug 2024 19:56:45 +0100
Subject: [PATCH] [ADT] Fix alignment check in unique_function constructor
 (#99403)

Right now the check fails for any state-capturing lambda since this
expression - `alignof(decltype(StorageUnion.InlineStorage))` - returns 1
for the alignment value and not 4/8 as expected
([MSVC|Clang|GCC](https://godbolt.org/z/eTEdq4xjM)). So this check fails
for pretty much any state-capturing callable we try to store into a
`unique_function` and we take the out-of-line storage path:

\llvm-project\llvm\include\llvm\ADT\FunctionExtras.h,
`UniqueFunctionBase` constructor (line ~266):

```
if (sizeof(CallableT) > InlineStorageSize ||
    alignof(CallableT) > alignof(decltype(StorageUnion.InlineStorage))) {
  // ...
}
```

The fix is simply to use an explicit const variable to store the
alignment value.

There is no easy way to unit-test the fix since inline storage is
considered to be an implementation detail so we shouldn't assume how the
lambda ends up being stored.
---
 llvm/include/llvm/ADT/FunctionExtras.h    |  6 ++++--
 llvm/unittests/ADT/FunctionExtrasTest.cpp | 19 +++++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/ADT/FunctionExtras.h b/llvm/include/llvm/ADT/FunctionExtras.h
index 49e0e8ab0db400..ad5c89cbc8bda7 100644
--- a/llvm/include/llvm/ADT/FunctionExtras.h
+++ b/llvm/include/llvm/ADT/FunctionExtras.h
@@ -80,6 +80,7 @@ using EnableIfCallable = std::enable_if_t<std::disjunction<
 template <typename ReturnT, typename... ParamTs> class UniqueFunctionBase {
 protected:
   static constexpr size_t InlineStorageSize = sizeof(void *) * 3;
+  static constexpr size_t InlineStorageAlign = alignof(void *);
 
   template <typename T, class = void>
   struct IsSizeLessThanThresholdT : std::false_type {};
@@ -161,7 +162,8 @@ template <typename ReturnT, typename... ParamTs> class UniqueFunctionBase {
     // provide three pointers worth of storage here.
     // This is mutable as an inlined `const unique_function<void() const>` may
     // still modify its own mutable members.
-    alignas(void *) mutable std::byte InlineStorage[InlineStorageSize];
+    alignas(InlineStorageAlign) mutable std::byte
+        InlineStorage[InlineStorageSize];
   } StorageUnion;
 
   // A compressed pointer to either our dispatching callback or our table of
@@ -262,7 +264,7 @@ template <typename ReturnT, typename... ParamTs> class UniqueFunctionBase {
     bool IsInlineStorage = true;
     void *CallableAddr = getInlineStorage();
     if (sizeof(CallableT) > InlineStorageSize ||
-        alignof(CallableT) > alignof(decltype(StorageUnion.InlineStorage))) {
+        alignof(CallableT) > InlineStorageAlign) {
       IsInlineStorage = false;
       // Allocate out-of-line storage. FIXME: Use an explicit alignment
       // parameter in C++17 mode.
diff --git a/llvm/unittests/ADT/FunctionExtrasTest.cpp b/llvm/unittests/ADT/FunctionExtrasTest.cpp
index fc856a976946bf..7abdc8b77e0595 100644
--- a/llvm/unittests/ADT/FunctionExtrasTest.cpp
+++ b/llvm/unittests/ADT/FunctionExtrasTest.cpp
@@ -310,4 +310,23 @@ class Incomplete {};
 Incomplete incompleteFunction() { return {}; }
 const Incomplete incompleteFunctionConst() { return {}; }
 
+// Check that we can store a pointer-sized payload inline in the unique_function.
+TEST(UniqueFunctionTest, InlineStorageWorks) {
+  // We do assume a couple of implementation details of the unique_function here:
+  //  - It can store certain small-enough payload inline
+  //  - Inline storage size is at least >= sizeof(void*)
+  void *ptr;
+  unique_function<void(void *)> UniqueFunctionWithInlineStorage{
+      [ptr](void *self) {
+        auto mid = reinterpret_cast<uintptr_t>(&ptr);
+        auto beg = reinterpret_cast<uintptr_t>(self);
+        auto end = reinterpret_cast<uintptr_t>(self) +
+                   sizeof(unique_function<void(void *)>);
+        // Make sure the address of the captured pointer lies somewhere within
+        // the unique_function object.
+        EXPECT_TRUE(mid >= beg && mid < end);
+      }};
+  UniqueFunctionWithInlineStorage(&UniqueFunctionWithInlineStorage);
+}
+
 } // anonymous namespace