Fix ThreadSafeRefCountedBase*::Release() (#3673)

C++ atomic's fetch_sub returns the previous value, where we want to
check the new value. This was causing massive memory leaks in SourceKit.

For ThreadSafeRefCountedBase, just switch to the one in LLVM that's
already correct.  We should move the VPTR one to LLVM as well and then
we can get rid of this header.

rdar://problem/27358273

Author: benlangmuir

include/swift/Basic/ThreadSafeRefCounted.h

@@ -15,37 +15,10 @@
 
 #include <atomic>
 #include <cassert>
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
 
 namespace swift {
 
-/// A thread-safe version of \c llvm::RefCountedBase.
-///
-/// A generic base class for objects that wish to have their lifetimes managed
-/// using reference counts. Classes subclass \c ThreadSafeRefCountedBase to
-/// obtain such functionality, and are typically handled with
-/// \c IntrusiveRefCntPtr "smart pointers" which automatically handle the
-/// management of reference counts.
-/// FIXME: This should eventually move to llvm.
-template <class Derived>
-class ThreadSafeRefCountedBase {
-  mutable std::atomic<unsigned> ref_cnt;
-
-protected:
-  ThreadSafeRefCountedBase() : ref_cnt(0) {}
-
-public:
-  void Retain() const {
-    ref_cnt.fetch_add(1, std::memory_order_acq_rel);
-  }
-
-  void Release() const {
-    int refCount =
-        static_cast<int>(ref_cnt.fetch_sub(1, std::memory_order_acq_rel));
-    assert(refCount >= 0 && "Reference count was already zero.");
-    if (refCount == 0) delete static_cast<const Derived*>(this);
-  }
-};
-
 /// A class that has the same function as \c ThreadSafeRefCountedBase, but with
 /// a virtual destructor.
 ///
@@ -62,12 +35,11 @@ class ThreadSafeRefCountedBaseVPTR {
 
 public:
   void Retain() const {
-    ref_cnt.fetch_add(1, std::memory_order_acq_rel);
+    ref_cnt += 1;
   }
 
   void Release() const {
-    int refCount =
-        static_cast<int>(ref_cnt.fetch_sub(1, std::memory_order_acq_rel));
+    int refCount = static_cast<int>(--ref_cnt);
     assert(refCount >= 0 && "Reference count was already zero.");
     if (refCount == 0) delete this;
   }

include/swift/IDE/CodeCompletionCache.h

@@ -49,7 +49,7 @@ class CodeCompletionCache {
     }
   };
 
-  struct Value : public ThreadSafeRefCountedBase<Value> {
+  struct Value : public llvm::ThreadSafeRefCountedBase<Value> {
     llvm::sys::TimeValue ModuleModificationTime;
     CodeCompletionResultSink Sink;
   };

tools/SourceKit/include/SourceKit/Core/LLVM.h

@@ -43,6 +43,7 @@ namespace llvm {
   // Reference counting.
   template <typename T> class IntrusiveRefCntPtr;
   template <typename T> struct IntrusiveRefCntPtrInfo;
+  template <class Derived> class ThreadSafeRefCountedBase;
 
   class raw_ostream;
   // TODO: DenseMap, ...
@@ -69,7 +70,6 @@ namespace llvm {
 }
 
 namespace swift {
-  template <class Derived> class ThreadSafeRefCountedBase;
   class ThreadSafeRefCountedBaseVPTR;
 }
 
@@ -95,7 +95,7 @@ namespace SourceKit {
   // Reference counting.
   using llvm::IntrusiveRefCntPtr;
   using llvm::IntrusiveRefCntPtrInfo;
-  using swift::ThreadSafeRefCountedBase;
+  using llvm::ThreadSafeRefCountedBase;
   using swift::ThreadSafeRefCountedBaseVPTR;
   template <typename T> class ThreadSafeRefCntPtr;
 

tools/SourceKit/lib/SwiftLang/SwiftInterfaceGenContext.h

@@ -31,7 +31,7 @@ namespace SourceKit {
   typedef IntrusiveRefCntPtr<ASTUnit> ASTUnitRef;
 
 class SwiftInterfaceGenContext :
-  public swift::ThreadSafeRefCountedBase<SwiftInterfaceGenContext> {
+  public llvm::ThreadSafeRefCountedBase<SwiftInterfaceGenContext> {
 public:
   static SwiftInterfaceGenContextRef create(StringRef DocumentName,
                                             bool IsModule,

tools/SourceKit/lib/SwiftLang/SwiftInvocation.h

@@ -26,7 +26,7 @@ namespace SourceKit {
 
 /// Encompasses an invocation for getting an AST. This is used to control AST
 /// sharing among different requests.
-class SwiftInvocation : public swift::ThreadSafeRefCountedBase<SwiftInvocation> {
+class SwiftInvocation : public llvm::ThreadSafeRefCountedBase<SwiftInvocation> {
 public:
   ~SwiftInvocation();
 

unittests/Basic/CMakeLists.txt

@@ -19,6 +19,7 @@ add_swift_unittest(SwiftBasicTests
   SourceManager.cpp
   StringExtrasTest.cpp
   SuccessorMapTest.cpp
+  ThreadSafeRefCntPointerTests.cpp
   TreeScopedHashTableTests.cpp
   Unicode.cpp
   ${generated_tests}

unittests/Basic/ThreadSafeRefCntPointerTests.cpp

@@ -0,0 +1,51 @@
+#include "swift/Basic/ThreadSafeRefCounted.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "gtest/gtest.h"
+
+using llvm::IntrusiveRefCntPtr;
+
+struct TestRelease : llvm::ThreadSafeRefCountedBase<TestRelease> {
+  bool &destroy;
+  TestRelease(bool &destroy) : destroy(destroy) {}
+  ~TestRelease() { destroy = true; }
+};
+
+TEST(ThreadSafeRefCountedBase, ReleaseSimple) {
+  bool destroyed = false;
+  {
+    IntrusiveRefCntPtr<TestRelease> ref = new TestRelease(destroyed);
+  }
+  EXPECT_TRUE(destroyed);
+}
+TEST(ThreadSafeRefCountedBase, Release) {
+  bool destroyed = false;
+  {
+    IntrusiveRefCntPtr<TestRelease> ref = new TestRelease(destroyed);
+    ref->Retain();
+    ref->Release();
+  }
+  EXPECT_TRUE(destroyed);
+}
+
+struct TestReleaseVPTR : swift::ThreadSafeRefCountedBaseVPTR {
+  bool &destroy;
+  TestReleaseVPTR(bool &destroy) : destroy(destroy) {}
+  virtual ~TestReleaseVPTR() { destroy = true; }
+};
+
+TEST(ThreadSafeRefCountedBaseVPTR, ReleaseSimple) {
+  bool destroyed = false;
+  {
+    IntrusiveRefCntPtr<TestReleaseVPTR> ref = new TestReleaseVPTR(destroyed);
+  }
+  EXPECT_TRUE(destroyed);
+}
+TEST(ThreadSafeRefCountedBaseVPTR, Release) {
+  bool destroyed = false;
+  {
+    IntrusiveRefCntPtr<TestReleaseVPTR> ref = new TestReleaseVPTR(destroyed);
+    ref->Retain();
+    ref->Release();
+  }
+  EXPECT_TRUE(destroyed);
+}