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

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.

Author: kerambyte

llvm/include/llvm/ADT/FunctionExtras.h

@@ -314,6 +314,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.

llvm/unittests/ADT/CountCopyAndMove.cpp

@@ -10,6 +10,8 @@
 
 using namespace llvm;
 
+int CountCopyAndMove::DefaultConstructions = 0;
+int CountCopyAndMove::ValueConstructions = 0;
 int CountCopyAndMove::CopyConstructions = 0;
 int CountCopyAndMove::CopyAssignments = 0;
 int CountCopyAndMove::MoveConstructions = 0;

llvm/unittests/ADT/CountCopyAndMove.h

@@ -12,15 +12,17 @@
 namespace llvm {
 
 struct CountCopyAndMove {
+  static int DefaultConstructions;
+  static int ValueConstructions;
   static int CopyConstructions;
   static int CopyAssignments;
   static int MoveConstructions;
   static int MoveAssignments;
   static int Destructions;
   int val;
 
-  CountCopyAndMove() = default;
-  explicit CountCopyAndMove(int val) : val(val) {}
+  CountCopyAndMove() { ++DefaultConstructions; }
+  explicit CountCopyAndMove(int val) : val(val) { ++ValueConstructions; }
   CountCopyAndMove(const CountCopyAndMove &other) : val(other.val) {
     ++CopyConstructions;
   }
@@ -40,13 +42,20 @@ struct CountCopyAndMove {
   ~CountCopyAndMove() { ++Destructions; }
 
   static void ResetCounts() {
+    DefaultConstructions = 0;
+    ValueConstructions = 0;
     CopyConstructions = 0;
     CopyAssignments = 0;
     MoveConstructions = 0;
     MoveAssignments = 0;
     Destructions = 0;
   }
 
+  static int TotalConstructions() {
+    return DefaultConstructions + ValueConstructions + MoveConstructions +
+           CopyConstructions;
+  }
+
   static int TotalCopies() { return CopyConstructions + CopyAssignments; }
 
   static int TotalMoves() { return MoveConstructions + MoveAssignments; }

llvm/unittests/ADT/FunctionExtrasTest.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/FunctionExtras.h"
+#include "CountCopyAndMove.h"
 #include "gtest/gtest.h"
 
 #include <memory>
@@ -329,4 +330,18 @@ TEST(UniqueFunctionTest, InlineStorageWorks) {
   UniqueFunctionWithInlineStorage(&UniqueFunctionWithInlineStorage);
 }
 
+// Check that the moved-from captured state is properly destroyed during
+// move construction/assignment.
+TEST(UniqueFunctionTest, MovedFromStateIsDestroyedCorrectly) {
+  CountCopyAndMove::ResetCounts();
+  {
+    unique_function<void()> CapturingFunction{
+        [Counter = CountCopyAndMove{}] {}};
+    unique_function<void()> CapturingFunctionMoved{
+        std::move(CapturingFunction)};
+  }
+  EXPECT_EQ(CountCopyAndMove::TotalConstructions(),
+            CountCopyAndMove::Destructions);
+}
+
 } // anonymous namespace