[llvm] [ADT] Add a missing call to a unique_function destructor after move (PR #98747)

[Dmitry Yanovsky via llvm-commits]

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

>From 2f4a7e7ec52d60c0f9cd57c0a9da5a5ccf9f2757 Mon Sep 17 00:00:00 2001
From: Dmitry Yanovsky <kerambyte at gmail.com>
Date: Sat, 13 Jul 2024 15:45:18 +0100
Subject: [PATCH 1/4] [ADT] Add a missing call to a unique_function destructor
 after move

Right now immediately after moving the captured state, we `disarm` the moved-from object's destructor by resetting the `RHS.CallbackAndInlineFlag`. This means that we never properly destroy the RHS object's captured state.
---
 llvm/include/llvm/ADT/FunctionExtras.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/llvm/include/llvm/ADT/FunctionExtras.h b/llvm/include/llvm/ADT/FunctionExtras.h
index 49e0e8ab0db40..38a3829fb5137 100644
--- a/llvm/include/llvm/ADT/FunctionExtras.h
+++ b/llvm/include/llvm/ADT/FunctionExtras.h
@@ -312,6 +312,7 @@ template <typename ReturnT, typename... ParamTs> class UniqueFunctionBase {
       // Non-trivial move, so dispatch to a type-erased implementation.
       getNonTrivialCallbacks()->MovePtr(getInlineStorage(),
                                         RHS.getInlineStorage());
+      getNonTrivialCallbacks()->DestroyPtr(RHS.getInlineStorage());
     }
 
     // Clear the old callback and inline flag to get back to as-if-null.

>From d3c36176002505dff00b013d26c572b934a995a1 Mon Sep 17 00:00:00 2001
From: kerambyte <kerambyte at gmail.com>
Date: Mon, 15 Jul 2024 13:03:51 +0100
Subject: [PATCH 2/4] add a unit test

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

diff --git a/llvm/unittests/ADT/FunctionExtrasTest.cpp b/llvm/unittests/ADT/FunctionExtrasTest.cpp
index fc856a976946b..80f7f3cf55665 100644
--- a/llvm/unittests/ADT/FunctionExtrasTest.cpp
+++ b/llvm/unittests/ADT/FunctionExtrasTest.cpp
@@ -310,4 +310,22 @@ class Incomplete {};
 Incomplete incompleteFunction() { return {}; }
 const Incomplete incompleteFunctionConst() { return {}; }
 
+// Check that the moved-from captured state is properly destroyed during
+// move construction/assignment.
+TEST(UniqueFunctionTest, MovedFromStateIsDestroyedCorrectly) {
+  static int NumOfDestructorsCalled = 0;
+  struct State {
+    State() = default;
+    State(State &&) = default;
+    State &operator=(State &&) = default;
+    ~State() { ++NumOfDestructorsCalled; }
+  };
+  {
+    unique_function<void()> CapturingFunction{[state = State{}] {}};
+    unique_function<void()> CapturingFunctionMoved{
+        std::move(CapturingFunction)};
+  }
+  EXPECT_EQ(NumOfDestructorsCalled, 4);
+}
+
 } // anonymous namespace

>From 1f59df8bee75577b50e45f9aca1033cc25707ae4 Mon Sep 17 00:00:00 2001
From: kerambyte <kerambyte at gmail.com>
Date: Mon, 15 Jul 2024 17:13:54 +0100
Subject: [PATCH 3/4] refactor the test

---
 llvm/unittests/ADT/FunctionExtrasTest.cpp | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/llvm/unittests/ADT/FunctionExtrasTest.cpp b/llvm/unittests/ADT/FunctionExtrasTest.cpp
index 80f7f3cf55665..01cdf499ac100 100644
--- a/llvm/unittests/ADT/FunctionExtrasTest.cpp
+++ b/llvm/unittests/ADT/FunctionExtrasTest.cpp
@@ -313,11 +313,14 @@ const Incomplete incompleteFunctionConst() { return {}; }
 // Check that the moved-from captured state is properly destroyed during
 // move construction/assignment.
 TEST(UniqueFunctionTest, MovedFromStateIsDestroyedCorrectly) {
+  static int NumOfMovesCalled = 0;
   static int NumOfDestructorsCalled = 0;
   struct State {
     State() = default;
-    State(State &&) = default;
-    State &operator=(State &&) = default;
+    State(const State &) = delete;
+    State(State &&) { ++NumOfMovesCalled; }
+    State &operator=(const State &) = delete;
+    State &operator=(State &&Rhs) { return *new (this) State{std::move(Rhs)}; }
     ~State() { ++NumOfDestructorsCalled; }
   };
   {
@@ -325,7 +328,7 @@ TEST(UniqueFunctionTest, MovedFromStateIsDestroyedCorrectly) {
     unique_function<void()> CapturingFunctionMoved{
         std::move(CapturingFunction)};
   }
-  EXPECT_EQ(NumOfDestructorsCalled, 4);
+  EXPECT_EQ(NumOfDestructorsCalled, 1 + NumOfMovesCalled);
 }
 
 } // anonymous namespace

>From e1c35030f6f0407f6668ba8175fc7bcb3cd0375d Mon Sep 17 00:00:00 2001
From: kerambyte <kerambyte at gmail.com>
Date: Mon, 15 Jul 2024 17:26:49 +0100
Subject: [PATCH 4/4] cleanup

---
 llvm/unittests/ADT/FunctionExtrasTest.cpp | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/llvm/unittests/ADT/FunctionExtrasTest.cpp b/llvm/unittests/ADT/FunctionExtrasTest.cpp
index 01cdf499ac100..23be6be037aa9 100644
--- a/llvm/unittests/ADT/FunctionExtrasTest.cpp
+++ b/llvm/unittests/ADT/FunctionExtrasTest.cpp
@@ -317,10 +317,7 @@ TEST(UniqueFunctionTest, MovedFromStateIsDestroyedCorrectly) {
   static int NumOfDestructorsCalled = 0;
   struct State {
     State() = default;
-    State(const State &) = delete;
     State(State &&) { ++NumOfMovesCalled; }
-    State &operator=(const State &) = delete;
-    State &operator=(State &&Rhs) { return *new (this) State{std::move(Rhs)}; }
     ~State() { ++NumOfDestructorsCalled; }
   };
   {
@@ -328,6 +325,7 @@ TEST(UniqueFunctionTest, MovedFromStateIsDestroyedCorrectly) {
     unique_function<void()> CapturingFunctionMoved{
         std::move(CapturingFunction)};
   }
+  printf("%i, %i\n", NumOfMovesCalled, NumOfDestructorsCalled);
   EXPECT_EQ(NumOfDestructorsCalled, 1 + NumOfMovesCalled);
 }