Implement unknownRetain_n/Release_n entry points

rdar://21945003

Swift SVN r31902

Author: trentxintong

Diff for: include/swift/Runtime/HeapObject.h

@@ -430,13 +430,20 @@ extern "C" void *swift_bridgeObjectRetain(void *value);
 /// Increment the strong retain count of an object which might not be a native
 /// Swift object.
 extern "C" void *swift_unknownRetain(void *value);
+/// Increment the strong retain count of an object which might not be a native
+/// Swift object by n.
+extern "C" void *swift_unknownRetain_n(void *value, int n);
 
 #else
 
 static inline void swift_unknownRetain(void *value) {
   swift_retain(static_cast<HeapObject *>(value));
 }
 
+static inline void swift_unknownRetain_n(void *value, int n) {
+  swift_retain_n(static_cast<HeapObject *>(value), n);
+}
+
 #endif /* SWIFT_OBJC_INTEROP */
 
 extern "C" void swift_bridgeObjectRelease(void *value);
@@ -446,13 +453,20 @@ extern "C" void swift_bridgeObjectRelease(void *value);
 /// Decrement the strong retain count of an object which might not be a native
 /// Swift object.
 extern "C" void swift_unknownRelease(void *value);
+/// Decrement the strong retain count of an object which might not be a native
+/// Swift object by n.
+extern "C" void swift_unknownRelease_n(void *value, int n);
 
 #else
 
 static inline void swift_unknownRelease(void *value) {
   swift_release(static_cast<HeapObject *>(value));
 }
 
+static inline void swift_unknownRelease_n(void *value, int n) {
+  swift_release_n(static_cast<HeapObject *>(value), n);
+}
+
 #endif /* SWIFT_OBJC_INTEROP */
 
 #if SWIFT_OBJC_INTEROP

Diff for: stdlib/public/runtime/SwiftObject.mm

@@ -566,18 +566,36 @@ static bool usesNativeSwiftReferenceCounting_unowned(const void *object) {
   return usesNativeSwiftReferenceCounting_allocated(object);
 }
 
+void *swift::swift_unknownRetain_n(void *object, int n) {
+  void *objc_ret = nullptr;
+  if (isObjCTaggedPointerOrNull(object)) return object;
+  if (usesNativeSwiftReferenceCounting_allocated(object))
+    return swift_retain_n(static_cast<HeapObject *>(object), n);
+  for (int i = 0; i < n; ++i)
+    objc_ret = objc_retain(static_cast<id>(object));
+  return objc_ret;
+}
+
+void swift::swift_unknownRelease_n(void *object, int n) {
+  if (isObjCTaggedPointerOrNull(object)) return;
+  if (usesNativeSwiftReferenceCounting_allocated(object))
+    return swift_release_n(static_cast<HeapObject *>(object), n);
+  for (int i = 0; i < n; ++i)
+    objc_release(static_cast<id>(object));
+}
+
 void *swift::swift_unknownRetain(void *object) {
   if (isObjCTaggedPointerOrNull(object)) return object;
   if (usesNativeSwiftReferenceCounting_allocated(object))
-    return swift_retain((HeapObject*) object);
-  return objc_retain((id) object);
+    return swift_retain(static_cast<HeapObject *>(object));
+  return objc_retain(static_cast<id>(object));
 }
 
 void swift::swift_unknownRelease(void *object) {
   if (isObjCTaggedPointerOrNull(object)) return;
   if (usesNativeSwiftReferenceCounting_allocated(object))
-    return swift_release((HeapObject*) object);
-  return objc_release((id) object);
+    return swift_release(static_cast<HeapObject *>(object));
+  return objc_release(static_cast<id>(object));
 }
 
 /// Return true iff the given BridgeObject is not known to use native

Diff for: unittests/runtime/CMakeLists.txt

@@ -7,6 +7,7 @@ if(("${SWIFT_HOST_VARIANT_SDK}" STREQUAL "${SWIFT_PRIMARY_VARIANT_SDK}") AND
     find_library(FOUNDATION_LIBRARY Foundation)
     list(APPEND PLATFORM_SOURCES
       weak.mm
+      Refcounting.mm
       )
     list(APPEND PLATFORM_TARGET_LINK_LIBRARIES
       ${FOUNDATION_LIBRARY}
@@ -17,7 +18,6 @@ if(("${SWIFT_HOST_VARIANT_SDK}" STREQUAL "${SWIFT_PRIMARY_VARIANT_SDK}") AND
   add_swift_unittest(SwiftRuntimeTests
     Metadata.cpp
     Enum.cpp
-    Refcounting.cpp
     ${PLATFORM_SOURCES}
     )
 

Diff for: unittests/runtime/Refcounting.cpp renamed to unittests/runtime/Refcounting.mm

@@ -10,12 +10,27 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <Foundation/NSObject.h>
+#include <objc/runtime.h>
 #include "swift/Runtime/HeapObject.h"
 #include "swift/Runtime/Metadata.h"
 #include "gtest/gtest.h"
 
 using namespace swift;
 
+static unsigned DestroyedObjCCount = 0;
+/// A trivial class that increments DestroyedObjCCount when deallocated.
+@interface ObjCTestClass : NSObject @end
+@implementation ObjCTestClass
+- (void) dealloc {
+  DestroyedObjCCount++;
+  [super dealloc];
+}
+@end
+static HeapObject *make_objc_object() {
+  return static_cast<HeapObject *>([ObjCTestClass new]);
+}
+
 struct TestObject : HeapObject {
   size_t *Addr;
   size_t Value;
@@ -29,18 +44,19 @@ static void destroyTestObject(HeapObject *_object) {
   swift_deallocObject(object, sizeof(TestObject), alignof(TestObject) - 1);
 }
 
-static const FullMetadata<HeapMetadata> TestObjectMetadata{
-  HeapMetadataHeader{{destroyTestObject}, {nullptr}},
-  HeapMetadata{Metadata{MetadataKind::HeapLocalVariable}}
+static const FullMetadata<ClassMetadata> TestClassObjectMetadata = {
+  { { &destroyTestObject }, { &_TWVBo } },
+  { { { MetadataKind::Class } }, 0, /*rodata*/ 1,
+  ClassFlags::UsesSwift1Refcounting, nullptr, 0, 0, 0, 0, 0 }
 };
 
 /// Create an object that, when deallocated, stores the given value to
 /// the given pointer.
 static TestObject *allocTestObject(size_t *addr, size_t value) {
   auto result =
-    static_cast<TestObject*>(swift_allocObject(&TestObjectMetadata,
-                                               sizeof(TestObject),
-                                               alignof(TestObject) - 1));
+    static_cast<TestObject *>(swift_allocObject(&TestClassObjectMetadata,
+                                                sizeof(TestObject),
+                                                alignof(TestObject) - 1));
   result->Addr = addr;
   result->Value = value;
   return result;
@@ -120,3 +136,40 @@ TEST(RefcountingTest, retain_release_n) {
   swift_release(object);
   EXPECT_EQ(1u, value);
 }
+
+TEST(RefcountingTest, unknown_retain_release_n) {
+  size_t value = 0;
+  auto object = allocTestObject(&value, 1);
+  EXPECT_EQ(0u, value);
+  auto retainResult = swift_unknownRetain_n(object, 32);
+  EXPECT_EQ(object, retainResult);
+  retainResult = swift_unknownRetain(object);
+  EXPECT_EQ(object, retainResult);
+  EXPECT_EQ(0u, value);
+  EXPECT_EQ(34u, swift_retainCount(object));
+  swift_unknownRelease_n(object, 31);
+  EXPECT_EQ(0u, value);
+  EXPECT_EQ(3u, swift_retainCount(object));
+  swift_unknownRelease(object);
+  EXPECT_EQ(0u, value);
+  EXPECT_EQ(2u, swift_retainCount(object));
+  swift_unknownRelease_n(object, 1);
+  EXPECT_EQ(0u, value);
+  EXPECT_EQ(1u, swift_retainCount(object));
+  swift_unknownRelease(object);
+  EXPECT_EQ(1u, value);
+}
+
+TEST(RefcountingTest, objc_unknown_retain_release_n) {
+  auto object = make_objc_object();
+  auto retainResult = swift_unknownRetain_n(object, 32);
+  EXPECT_EQ(object, retainResult);
+  retainResult = swift_unknownRetain(object);
+  EXPECT_EQ(object, retainResult);
+  swift_unknownRelease_n(object, 31);
+  swift_unknownRelease(object);
+  swift_unknownRelease_n(object, 1);
+  swift_unknownRelease(object);
+  // The object should be destroyed by now.
+  EXPECT_EQ(1u, DestroyedObjCCount);
+}