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

[via llvm-commits]

Author: Dmitry Yanovsky
Date: 2024-08-19T12:38:06-07:00
New Revision: e8ca306fbefdb8385c53d6cc029251d1d5be7049

URL: https://github.com/llvm/llvm-project/commit/e8ca306fbefdb8385c53d6cc029251d1d5be7049
DIFF: https://github.com/llvm/llvm-project/commit/e8ca306fbefdb8385c53d6cc029251d1d5be7049.diff

LOG: [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.

Added: 
    

Modified: 
    llvm/include/llvm/ADT/FunctionExtras.h
    llvm/unittests/ADT/CountCopyAndMove.cpp
    llvm/unittests/ADT/CountCopyAndMove.h
    llvm/unittests/ADT/FunctionExtrasTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/FunctionExtras.h b/llvm/include/llvm/ADT/FunctionExtras.h
index ad5c89cbc8bda7..d92868e3715f40 100644
--- a/llvm/include/llvm/ADT/FunctionExtras.h
+++ b/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.

diff  --git a/llvm/unittests/ADT/CountCopyAndMove.cpp b/llvm/unittests/ADT/CountCopyAndMove.cpp
index fe1e2f4a5b89fd..022dbee40dc6f1 100644
--- a/llvm/unittests/ADT/CountCopyAndMove.cpp
+++ b/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;

diff  --git a/llvm/unittests/ADT/CountCopyAndMove.h b/llvm/unittests/ADT/CountCopyAndMove.h
index 126054427b81a3..c91d34bd8ebb50 100644
--- a/llvm/unittests/ADT/CountCopyAndMove.h
+++ b/llvm/unittests/ADT/CountCopyAndMove.h
@@ -12,6 +12,8 @@
 namespace llvm {
 
 struct CountCopyAndMove {
+  static int DefaultConstructions;
+  static int ValueConstructions;
   static int CopyConstructions;
   static int CopyAssignments;
   static int MoveConstructions;
@@ -19,8 +21,8 @@ struct CountCopyAndMove {
   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,6 +42,8 @@ struct CountCopyAndMove {
   ~CountCopyAndMove() { ++Destructions; }
 
   static void ResetCounts() {
+    DefaultConstructions = 0;
+    ValueConstructions = 0;
     CopyConstructions = 0;
     CopyAssignments = 0;
     MoveConstructions = 0;
@@ -47,6 +51,11 @@ struct CountCopyAndMove {
     Destructions = 0;
   }
 
+  static int TotalConstructions() {
+    return DefaultConstructions + ValueConstructions + MoveConstructions +
+           CopyConstructions;
+  }
+
   static int TotalCopies() { return CopyConstructions + CopyAssignments; }
 
   static int TotalMoves() { return MoveConstructions + MoveAssignments; }

diff  --git a/llvm/unittests/ADT/FunctionExtrasTest.cpp b/llvm/unittests/ADT/FunctionExtrasTest.cpp
index 7abdc8b77e0595..f64b888dbe1591 100644
--- a/llvm/unittests/ADT/FunctionExtrasTest.cpp
+++ b/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