Avoid mishandling retain_n of immortal objects where n >= 3

Author: Catfish-Man

stdlib/public/SwiftShims/RefCount.h

@@ -805,6 +805,13 @@ class RefCounts {
   // Increment the reference count.
   void 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;
+    }
+    
     RefCountBits newbits;
     do {
       newbits = oldbits;
@@ -820,6 +827,13 @@ class RefCounts {
 
   void incrementNonAtomic(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;
+    }
+    
     auto newbits = oldbits;
     bool fast = newbits.incrementStrongExtraRefCount(inc);
     if (SWIFT_UNLIKELY(!fast)) {
@@ -982,6 +996,12 @@ class RefCounts {
   bool doDecrementSlow(RefCountBits oldbits, uint32_t dec) {
     RefCountBits newbits;
 
+    // constant propagation will remove this in swift_release, it should only
+    // be present in swift_release_n
+    if (dec != 1 && oldbits.isImmortal(true)) {
+      return false;
+    }
+    
     bool deinitNow;
     do {
       newbits = oldbits;
@@ -1024,6 +1044,12 @@ class RefCounts {
   bool doDecrementNonAtomicSlow(RefCountBits oldbits, uint32_t dec) {
     bool deinitNow;
     auto newbits = oldbits;
+    
+    // constant propagation will remove this in swift_release, it should only
+    // be present in swift_release_n
+    if (dec != 1 && oldbits.isImmortal(true)) {
+      return false;
+    }
 
     bool fast =
       newbits.decrementStrongExtraRefCount(dec);
@@ -1066,6 +1092,12 @@ class RefCounts {
     auto oldbits = refCounts.load(SWIFT_MEMORY_ORDER_CONSUME);
     RefCountBits newbits;
 
+    // constant propagation will remove this in swift_release, it should only
+    // be present in swift_release_n
+    if (dec != 1 && oldbits.isImmortal(true)) {
+      return false;
+    }
+    
     do {
       newbits = oldbits;
       bool fast =
@@ -1248,7 +1280,15 @@ class RefCounts {
   // Return weak reference count.
   // Note that this is not equal to the number of outstanding weak pointers.
   uint32_t getWeakCount() const;
-  
+
+  // DO NOT TOUCH.
+  // This exists for the benefits of the Refcounting.cpp tests. Do not use it
+  // elsewhere.
+  auto getBitsValue()
+    -> decltype(auto) {
+    return refCounts.load(std::memory_order_relaxed).getBitsValue();
+  }
+
   private:
   HeapObject *getHeapObject();
 
@@ -1449,13 +1489,14 @@ inline bool RefCounts<InlineRefCountBits>::doDecrementNonAtomic(uint32_t dec) {
   // Use slow path if we can't guarantee atomicity.
   if (oldbits.hasSideTable() || oldbits.getUnownedRefCount() != 1)
     return doDecrementNonAtomicSlow<performDeinit>(oldbits, dec);
+    
+  if (oldbits.isImmortal(true)) {
+    return false;
+  }
 
   auto newbits = oldbits;
   bool fast = newbits.decrementStrongExtraRefCount(dec);
   if (!fast) {
-    if (oldbits.isImmortal(false)) {
-      return false;
-    }
     return doDecrementNonAtomicSlow<performDeinit>(oldbits, dec);
   }
 

unittests/runtime/Refcounting.cpp

@@ -17,6 +17,10 @@
 using namespace swift;
 
 struct TestObject : HeapObject {
+  constexpr TestObject(HeapMetadata const *newMetadata)
+    : HeapObject(newMetadata, InlineRefCounts::Immortal)
+    , Addr(NULL), Value(0) {}
+
   size_t *Addr;
   size_t Value;
 };
@@ -34,6 +38,8 @@ static const FullMetadata<ClassMetadata> TestClassObjectMetadata = {
   { { nullptr }, ClassFlags::UsesSwiftRefcounting, 0, 0, 0, 0, 0, 0 }
 };
 
+static TestObject ImmortalTestObject{&TestClassObjectMetadata};
+
 /// Create an object that, when deallocated, stores the given value to
 /// the given pointer.
 static TestObject *allocTestObject(size_t *addr, size_t value) {
@@ -220,3 +226,49 @@ TEST(RefcountingTest, nonatomic_unknown_retain_release_n) {
   EXPECT_EQ(0u, value);
   EXPECT_EQ(1u, swift_retainCount(object));
 }
+
+// Verify that refcounting operations on immortal objects never changes the
+// refcount field.
+TEST(RefcountingTest, immortal_retain_release) {
+  auto initialBitsValue = ImmortalTestObject.refCounts.getBitsValue();
+
+  swift_retain(&ImmortalTestObject);
+  EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+  swift_release(&ImmortalTestObject);
+  EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+  swift_nonatomic_retain(&ImmortalTestObject);
+  EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+  swift_nonatomic_release(&ImmortalTestObject);
+  EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+
+  swift_unownedRetain(&ImmortalTestObject);
+  EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+  swift_unownedRelease(&ImmortalTestObject);
+  EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+  swift_nonatomic_unownedRetain(&ImmortalTestObject);
+  EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+  swift_nonatomic_unownedRelease(&ImmortalTestObject);
+  EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+
+  for (unsigned i = 0; i < 32; i++) {
+    uint32_t amount = 1U << i;
+
+    swift_retain_n(&ImmortalTestObject, amount);
+    EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+    swift_release_n(&ImmortalTestObject, amount);
+    EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+    swift_nonatomic_retain_n(&ImmortalTestObject, amount);
+    EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+    swift_nonatomic_release_n(&ImmortalTestObject, amount);
+    EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+    
+    swift_unownedRetain_n(&ImmortalTestObject, amount);
+    EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+    swift_unownedRelease_n(&ImmortalTestObject, amount);
+    EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+    swift_nonatomic_unownedRetain_n(&ImmortalTestObject, amount);
+    EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+    swift_nonatomic_unownedRelease_n(&ImmortalTestObject, amount);
+    EXPECT_EQ(initialBitsValue, ImmortalTestObject.refCounts.getBitsValue());
+  }
+}