Experiment with avoiding stack frames in swift_retain and swift_bridgeObjectRetain

Author: mikeash

stdlib/public/SwiftShims/swift/shims/RefCount.h

@@ -703,7 +703,7 @@ class RefCounts {
   // Out-of-line slow paths.
 
   SWIFT_NOINLINE
-  void incrementSlow(RefCountBits oldbits, uint32_t inc) SWIFT_CC(PreserveMost);
+  HeapObject *incrementSlow(RefCountBits oldbits, uint32_t inc) SWIFT_CC(PreserveMost);
 
   SWIFT_NOINLINE
   void incrementNonAtomicSlow(RefCountBits oldbits, uint32_t inc);
@@ -800,13 +800,13 @@ class RefCounts {
 
   // Increment the reference count.
   SWIFT_ALWAYS_INLINE
-  void increment(uint32_t inc = 1) {
+  HeapObject *increment(uint32_t inc = 1) {
     auto oldbits = refCounts.load(SWIFT_MEMORY_ORDER_CONSUME);
 
     // constant propagation will remove this in swift_retain, it should only
     // be present in swift_retain_n
     if (inc != 1 && oldbits.isImmortal(true)) {
-      return;
+      return getHeapObject();
     }
 
     RefCountBits newbits;
@@ -815,11 +815,12 @@ class RefCounts {
       bool fast = newbits.incrementStrongExtraRefCount(inc);
       if (SWIFT_UNLIKELY(!fast)) {
         if (oldbits.isImmortal(false))
-          return;
+          return getHeapObject();
         return incrementSlow(oldbits, inc);
       }
     } while (!refCounts.compare_exchange_weak(oldbits, newbits,
                                               std::memory_order_relaxed));
+    return getHeapObject();
   }
 
   SWIFT_ALWAYS_INLINE

stdlib/public/SwiftShims/swift/shims/Visibility.h

@@ -109,6 +109,12 @@
 #define SWIFT_WEAK_IMPORT
 #endif
 
+#if __has_attribute(musttail)
+#define SWIFT_MUSTTAIL __attribute__((musttail))
+#else
+#define SWIFT_MUSTTAIL
+#endif
+
 // Define the appropriate attributes for sharing symbols across
 // image (executable / shared-library) boundaries.
 //

stdlib/public/runtime/HeapObject.cpp

@@ -343,7 +343,7 @@ SWIFT_ALWAYS_INLINE
 static HeapObject *_swift_retain_(HeapObject *object) {
   SWIFT_RT_TRACK_INVOCATION(object, swift_retain);
   if (isValidPointerForNativeRetain(object))
-    object->refCounts.increment(1);
+    return object->refCounts.increment(1);
   return object;
 }
 

stdlib/public/runtime/RefCount.cpp

@@ -15,10 +15,10 @@
 namespace swift {
 
 template <typename RefCountBits>
-void RefCounts<RefCountBits>::incrementSlow(RefCountBits oldbits,
+HeapObject *RefCounts<RefCountBits>::incrementSlow(RefCountBits oldbits,
                                             uint32_t n) {
   if (oldbits.isImmortal(false)) {
-    return;
+    return getHeapObject();
   }
   else if (oldbits.hasSideTable()) {
     // Out-of-line slow path.
@@ -29,9 +29,10 @@ void RefCounts<RefCountBits>::incrementSlow(RefCountBits oldbits,
     // Retain count overflow.
     swift::swift_abortRetainOverflow();
   }
+  return getHeapObject();
 }
-template void RefCounts<InlineRefCountBits>::incrementSlow(InlineRefCountBits oldbits, uint32_t n);
-template void RefCounts<SideTableRefCountBits>::incrementSlow(SideTableRefCountBits oldbits, uint32_t n);
+template HeapObject *RefCounts<InlineRefCountBits>::incrementSlow(InlineRefCountBits oldbits, uint32_t n);
+template HeapObject *RefCounts<SideTableRefCountBits>::incrementSlow(SideTableRefCountBits oldbits, uint32_t n);
 
 template <typename RefCountBits>
 void RefCounts<RefCountBits>::incrementNonAtomicSlow(RefCountBits oldbits,

stdlib/public/runtime/SwiftObject.mm

@@ -574,6 +574,13 @@ static bool isBridgeObjectTaggedPointer(void *object) {
   return (void*)(uintptr_t(object) & ~unTaggedNonNativeBridgeObjectBits);
 }
 
+SWIFT_NOINLINE
+static void *objcRetainAndReturn(void *object) {
+  auto const objectRef = toPlainObject_unTagged_bridgeObject(object);
+  objc_retain(static_cast<id>(objectRef));
+  return object;
+}
+
 void *swift::swift_bridgeObjectRetain(void *object) {
 #if SWIFT_OBJC_INTEROP
   if (isObjCTaggedPointer(object) || isBridgeObjectTaggedPointer(object))
@@ -584,11 +591,9 @@ static bool isBridgeObjectTaggedPointer(void *object) {
 
 #if SWIFT_OBJC_INTEROP
   if (!isNonNative_unTagged_bridgeObject(object)) {
-    swift_retain(static_cast<HeapObject *>(objectRef));
-    return object;
+    return swift_retain(static_cast<HeapObject *>(objectRef));
   }
-  objc_retain(static_cast<id>(objectRef));
-  return object;
+  SWIFT_MUSTTAIL return objcRetainAndReturn(object);
 #else
   swift_retain(static_cast<HeapObject *>(objectRef));
   return object;