[llvm] [ADT] Fix alignment check in unique_function constructor (PR #99403)

[Dmitry Yanovsky via llvm-commits]

https://github.com/kerambyte updated https://github.com/llvm/llvm-project/pull/99403

>From 77dff91260bfefb4f2ab90e489185ce0579670ec Mon Sep 17 00:00:00 2001
From: kerambyte <kerambyte at gmail.com>
Date: Thu, 18 Jul 2024 00:14:54 +0100
Subject: [PATCH 1/2] initial commit

---
 llvm/include/llvm/ADT/FunctionExtras.h | 6 ++++--
 1 file changed, 4 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.

>From 2e6e114460d081a094d2336f89f5576847eb3e13 Mon Sep 17 00:00:00 2001
From: kerambyte <kerambyte at gmail.com>
Date: Sun, 11 Aug 2024 11:11:00 +0100
Subject: [PATCH 2/2] added a unit test

---
 llvm/unittests/ADT/FunctionExtrasTest.cpp | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

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