Fix sys::Cache for ref-counted values which caused memory leaks in SourceKit

include/swift/Basic/Cache.h

@@ -21,13 +21,6 @@
 namespace swift {
 namespace sys {
 
-template <typename T>
-struct CacheTypeMgmtInfo {
-  static void *enterCache(const T &Val) { return new T(Val); }
-  static void exitCache(void *Ptr) { delete static_cast<T*>(Ptr); }
-  static const T &getFromCache(void *Ptr) { return *static_cast<T*>(Ptr); }
-};
-
 template <typename T>
 struct CacheKeyHashInfo {
   static uintptr_t getHashValue(const T &Val) {
@@ -40,19 +33,24 @@ struct CacheKeyHashInfo {
 };
 
 template <typename T>
-struct CacheValueCostInfo {
-  static size_t getCost(const T &Val) { return sizeof(Val); }
+struct CacheKeyInfo : public CacheKeyHashInfo<T> {
+  static void *enterCache(const T &Val) { return new T(Val); }
+  static void exitCache(void *Ptr) { delete static_cast<T*>(Ptr); }
+  static const void *getLookupKey(const T *Val) { return Val; }
+  static const T &getFromCache(void *Ptr) { return *static_cast<T*>(Ptr); }
 };
 
 template <typename T>
-struct CacheKeyInfo : public CacheKeyHashInfo<T>,
-                      public CacheTypeMgmtInfo<T> {
-  static const void *getLookupKey(const T *Val) { return Val; }
+struct CacheValueCostInfo {
+  static size_t getCost(const T &Val) { return sizeof(Val); }
 };
 
 template <typename T>
-struct CacheValueInfo : public CacheValueCostInfo<T>,
-                        public CacheTypeMgmtInfo<T> {
+struct CacheValueInfo : public CacheValueCostInfo<T> {
+  static void *enterCache(const T &Val) { return new T(Val); }
+  static void retain(void *Ptr) {}
+  static void release(void *Ptr) { delete static_cast<T *>(Ptr); }
+  static const T &getFromCache(void *Ptr) { return *static_cast<T *>(Ptr); }
 };
 
 /// The underlying implementation of the caching mechanism.
@@ -68,7 +66,8 @@ class CacheImpl {
     bool (*keyIsEqualCB)(void *Key1, void *Key2, void *UserData);
 
     void (*keyDestroyCB)(void *Key, void *UserData);
-    void (*valueDestroyCB)(void *Value, void *UserData);
+    void (*valueRetainCB)(void *Value, void *UserData);
+    void (*valueReleaseCB)(void *Value, void *UserData);
   };
 
 protected:
@@ -154,7 +153,8 @@ class Cache : CacheImpl {
       keyHash,
       keyIsEqual,
       keyDestroy,
-      valueDestroy
+      valueRetain,
+      valueRelease,
     };
     Impl = create(Name, CBs);
   }
@@ -203,19 +203,23 @@ class Cache : CacheImpl {
   static void keyDestroy(void *Key, void *UserData) {
     KeyInfoT::exitCache(Key);
   }
-  static void valueDestroy(void *Value, void *UserData) {
-    ValueInfoT::exitCache(Value);
+  static void valueRetain(void *Value, void *UserData) {
+    ValueInfoT::retain(Value);
+  }
+  static void valueRelease(void *Value, void *UserData) {
+    ValueInfoT::release(Value);
   }
 };
 
 template <typename T>
 struct CacheValueInfo<llvm::IntrusiveRefCntPtr<T>>{
   static void *enterCache(const llvm::IntrusiveRefCntPtr<T> &Val) {
-    T *Ptr = Val.get();
-    Ptr->Retain();
-    return Ptr;
+    return const_cast<T *>(Val.get());
+  }
+  static void retain(void *Ptr) {
+    static_cast<T*>(Ptr)->Retain();
   }
-  static void exitCache(void *Ptr) {
+  static void release(void *Ptr) {
     static_cast<T*>(Ptr)->Release();
   }
   static llvm::IntrusiveRefCntPtr<T> getFromCache(void *Ptr) {

lib/Basic/Cache.cpp

@@ -78,11 +78,14 @@ void CacheImpl::setAndRetain(void *Key, void *Value, size_t Cost) {
   DefaultCacheKey CKey(Key, &DCache.CBs);
   auto Entry = DCache.Entries.find(CKey);
   if (Entry != DCache.Entries.end()) {
+    if (Entry->second == Value)
+      return;
     DCache.CBs.keyDestroyCB(Entry->first.Key, nullptr);
-    DCache.CBs.valueDestroyCB(Entry->second, nullptr);
+    DCache.CBs.valueReleaseCB(Entry->second, nullptr);
     DCache.Entries.erase(Entry);
   }
 
+  DCache.CBs.valueRetainCB(Value, nullptr);
   DCache.Entries[CKey] = Value;
 
   // FIXME: Not thread-safe! It should avoid deleting the value until
@@ -116,7 +119,7 @@ bool CacheImpl::remove(const void *Key) {
   auto Entry = DCache.Entries.find(CKey);
   if (Entry != DCache.Entries.end()) {
     DCache.CBs.keyDestroyCB(Entry->first.Key, nullptr);
-    DCache.CBs.valueDestroyCB(Entry->second, nullptr);
+    DCache.CBs.valueReleaseCB(Entry->second, nullptr);
     DCache.Entries.erase(Entry);
     return true;
   }
@@ -129,7 +132,7 @@ void CacheImpl::removeAll() {
 
   for (auto Entry : DCache.Entries) {
     DCache.CBs.keyDestroyCB(Entry.first.Key, nullptr);
-    DCache.CBs.valueDestroyCB(Entry.second, nullptr);
+    DCache.CBs.valueReleaseCB(Entry.second, nullptr);
   }
   DCache.Entries.clear();
 }

lib/Basic/Darwin/Cache-Mac.cpp

@@ -29,11 +29,11 @@ CacheImpl::ImplTy CacheImpl::create(StringRef Name, const CallBacks &CBs) {
     CBs.keyIsEqualCB,
     nullptr,
     CBs.keyDestroyCB,
-    CBs.valueDestroyCB,
+    CBs.valueReleaseCB,
     nullptr,
     nullptr,
     CBs.UserData,
-    nullptr
+    CBs.valueRetainCB,
   };
 
   cache_t *cache_out = nullptr;

unittests/Basic/CMakeLists.txt

@@ -7,6 +7,7 @@ handle_gyb_sources(
 
 add_swift_unittest(SwiftBasicTests
   BlotMapVectorTest.cpp
+  CacheTest.cpp
   ClusteredBitVectorTest.cpp
   DemangleTest.cpp
   DiverseStackTest.cpp

unittests/Basic/CacheTest.cpp

@@ -0,0 +1,186 @@
+//===--- CacheTest.cpp ----------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "swift/Basic/Cache.h"
+#include "gtest/gtest.h"
+
+#if defined(__APPLE__)
+#define USES_LIBCACHE 1
+#else
+#define USES_LIBCACHE 0
+#endif
+
+namespace {
+struct Counter {
+  mutable int enter = 0;
+  mutable int exit = 0;
+};
+}
+
+namespace swift {
+namespace sys {
+template <>
+struct CacheValueInfo<Counter>{
+  static void *enterCache(const Counter &value) {
+    return const_cast<Counter *>(&value);
+  }
+  static void retain(void *ptr) {
+    static_cast<Counter*>(ptr)->enter+= 1;
+  }
+  static void release(void *ptr) {
+    static_cast<Counter*>(ptr)->exit += 1;
+  }
+  static const Counter &getFromCache(void *ptr) {
+    return *static_cast<Counter *>(ptr);
+  }
+  static size_t getCost(const Counter &value) {
+    return 0;
+  }
+};
+}
+}
+
+namespace {
+struct RefCntToken : llvm::RefCountedBase<RefCntToken> {
+  bool &freed;
+  RefCntToken(bool &freed) : freed(freed) {}
+  ~RefCntToken() { freed = true; }
+};
+}
+
+TEST(Cache, sameKey) {
+  Counter c1, c2;
+  swift::sys::Cache<const char *, Counter> cache(__func__);
+  cache.set("a", c1);
+  EXPECT_EQ(1, c1.enter);
+  EXPECT_EQ(0, c1.exit);
+
+  cache.set("a", c2);
+  EXPECT_EQ(1, c1.enter);
+  EXPECT_EQ(1, c1.exit);
+  EXPECT_EQ(1, c2.enter);
+  EXPECT_EQ(0, c2.exit);
+}
+
+TEST(Cache, sameValue) {
+  Counter c;
+  swift::sys::Cache<const char *, Counter> cache(__func__);
+  cache.set("a", c);
+  EXPECT_EQ(1, c.enter);
+  EXPECT_EQ(0, c.exit);
+
+  cache.set("b", c);
+#if USES_LIBCACHE
+  EXPECT_EQ(1, c.enter); // value is shared.
+#else
+  EXPECT_EQ(2, c.enter);
+#endif
+  EXPECT_EQ(0, c.exit);
+
+  cache.remove("a");
+#if USES_LIBCACHE
+  EXPECT_EQ(1, c.enter); // value is shared.
+  EXPECT_EQ(0, c.exit);
+#else
+  EXPECT_EQ(2, c.enter);
+  EXPECT_EQ(1, c.exit);
+#endif
+
+  cache.remove("b");
+  EXPECT_EQ(c.enter, c.exit);
+}
+
+TEST(Cache, sameKeyValue) {
+  Counter c;
+  swift::sys::Cache<const char *, Counter> cache(__func__);
+  cache.set("a", c);
+  EXPECT_EQ(1, c.enter);
+  EXPECT_EQ(0, c.exit);
+
+  cache.set("a", c);
+  EXPECT_EQ(1, c.enter);
+  EXPECT_EQ(0, c.exit);
+
+  cache.remove("a");
+  EXPECT_EQ(1, c.enter);
+  EXPECT_EQ(1, c.exit);
+}
+
+TEST(Cache, sameKeyIntrusiveRefCountPter) {
+  bool freed1 = false;
+  bool freed2 = false;
+  swift::sys::Cache<const char *, llvm::IntrusiveRefCntPtr<RefCntToken>> cache(__func__);
+  {
+    llvm::IntrusiveRefCntPtr<RefCntToken> c1(new RefCntToken(freed1));
+    llvm::IntrusiveRefCntPtr<RefCntToken> c2(new RefCntToken(freed2));
+    cache.set("a", c1);
+    cache.set("a", c2);
+  }
+  EXPECT_TRUE(freed1);
+  EXPECT_FALSE(freed2);
+  cache.remove("a");
+  EXPECT_TRUE(freed2);
+}
+
+TEST(Cache, sameValueIntrusiveRefCountPter) {
+  bool freed = false;
+  swift::sys::Cache<const char *, llvm::IntrusiveRefCntPtr<RefCntToken>> cache(__func__);
+  {
+    llvm::IntrusiveRefCntPtr<RefCntToken> c(new RefCntToken(freed));
+    cache.set("a", c);
+    EXPECT_FALSE(freed);
+
+    cache.set("b", c);
+    EXPECT_FALSE(freed);
+
+    cache.remove("a");
+    EXPECT_FALSE(freed);
+
+    cache.remove("b");
+    EXPECT_FALSE(freed);
+  }
+  EXPECT_TRUE(freed);
+}
+
+TEST(Cache, sameKeyValueIntrusiveRefCountPter) {
+  bool freed = false;
+  swift::sys::Cache<const char *, llvm::IntrusiveRefCntPtr<RefCntToken>> cache(__func__);
+  {
+    llvm::IntrusiveRefCntPtr<RefCntToken> c(new RefCntToken(freed));
+    cache.set("a", c);
+    EXPECT_FALSE(freed);
+    cache.set("a", c);
+    EXPECT_FALSE(freed);
+  }
+  EXPECT_FALSE(freed);
+  cache.remove("a");
+  EXPECT_TRUE(freed);
+}
+
+TEST(Cache, copyValue) {
+  struct S {
+    int ident, copy;
+    S(int ident) : ident(ident), copy(0) {}
+    S(const S &other) : ident(other.ident), copy(other.copy+1) {}
+    S(S &&other) : ident(other.ident), copy(other.copy) {}
+  };
+  swift::sys::Cache<const char *, struct S> cache(__func__);
+  S s{0};
+  EXPECT_EQ(0, s.ident);
+  EXPECT_EQ(0, s.copy);
+  cache.set("a", s);
+  EXPECT_EQ(0, cache.get("a")->ident);
+  EXPECT_EQ(2, cache.get("a")->copy); // return by value causes 2nd copy
+  cache.set("b", *cache.get("a"));
+  EXPECT_EQ(0, cache.get("b")->ident);
+  EXPECT_EQ(4, cache.get("b")->copy); // return by value causes 2nd copy
+}