runtime: Fix overflow of swift_unownedRetain reference counts

include/swift/Runtime/Debug.h

@@ -124,6 +124,10 @@ void swift_abortRetainOverflow();
 LLVM_ATTRIBUTE_NORETURN LLVM_ATTRIBUTE_NOINLINE
 void swift_abortRetainUnowned(const void *object);
 
+// Halt due to an overflow in swift_unownedRetain().
+LLVM_ATTRIBUTE_NORETURN LLVM_ATTRIBUTE_NOINLINE
+void swift_abortUnownedRetainOverflow();
+
 /// This function dumps one line of a stack trace. It is assumed that \p framePC
 /// is the address of the stack frame at index \p index. If \p shortOutput is
 /// true, this functions prints only the name of the symbol and offset, ignores

stdlib/public/SwiftShims/RefCount.h

@@ -564,10 +564,12 @@ class RefCountBitsT {
     else
       return doDecrementStrongExtraRefCount<DontClearPinnedFlag>(dec);
   }
-
+  // Returns the old reference count before the increment.
   LLVM_ATTRIBUTE_ALWAYS_INLINE
-  void incrementUnownedRefCount(uint32_t inc) {
-    setUnownedRefCount(getUnownedRefCount() + inc);
+  uint32_t incrementUnownedRefCount(uint32_t inc) {
+    uint32_t old = getUnownedRefCount();
+    setUnownedRefCount(old + inc);
+    return old;
   }
 
   LLVM_ATTRIBUTE_ALWAYS_INLINE
@@ -757,6 +759,9 @@ class RefCounts {
   LLVM_ATTRIBUTE_NOINLINE
   bool tryIncrementNonAtomicSlow(RefCountBits oldbits);
 
+  LLVM_ATTRIBUTE_NOINLINE
+  void incrementUnownedSlow(uint32_t inc);
+
   public:
   enum Initialized_t { Initialized };
 
@@ -1139,8 +1144,12 @@ class RefCounts {
 
       newbits = oldbits;
       assert(newbits.getUnownedRefCount() != 0);
-      newbits.incrementUnownedRefCount(inc);
-      // FIXME: overflow check?
+      uint32_t oldValue = newbits.incrementUnownedRefCount(inc);
+
+      // Check overflow and use the side table on overflow.
+      if (newbits.getUnownedRefCount() != oldValue + inc)
+        return incrementUnownedSlow(inc);
+
     } while (!refCounts.compare_exchange_weak(oldbits, newbits,
                                               std::memory_order_relaxed));
   }
@@ -1152,8 +1161,12 @@ class RefCounts {
 
     auto newbits = oldbits;
     assert(newbits.getUnownedRefCount() != 0);
-    newbits.incrementUnownedRefCount(inc);
-    // FIXME: overflow check?
+    uint32_t oldValue = newbits.incrementUnownedRefCount(inc);
+
+    // Check overflow and use the side table on overflow.
+    if (newbits.getUnownedRefCount() != oldValue + inc)
+      return incrementUnownedSlow(inc);
+
     refCounts.store(newbits, std::memory_order_relaxed);
   }
 

stdlib/public/runtime/Errors.cpp

@@ -339,6 +339,13 @@ void swift::swift_abortRetainOverflow() {
                     "fatal error: object was retained too many times");
 }
 
+// Crash due to an unowned retain count overflow.
+// FIXME: can't pass the object's address from InlineRefCounts without hacks
+void swift::swift_abortUnownedRetainOverflow() {
+  swift::fatalError(FatalErrorFlags::ReportBacktrace,
+                    "fatal error: object's unowned reference was retained too many times");
+}
+
 // Crash due to retain of a dead unowned reference.
 // FIXME: can't pass the object's address from InlineRefCounts without hacks
 void swift::swift_abortRetainUnowned(const void *object) {

stdlib/public/runtime/RefCount.cpp

@@ -143,6 +143,22 @@ HeapObjectSideTableEntry* RefCounts<InlineRefCountBits>::formWeakReference()
     return nullptr;
 }
 
+template <typename RefCountBits>
+void RefCounts<RefCountBits>::incrementUnownedSlow(uint32_t n) {
+  auto side = allocateSideTable();
+  if (side)
+    return side->incrementUnowned(n);
+  // Overflow but side table allocation failed.
+  swift_abortUnownedRetainOverflow();
+}
+
+template void RefCounts<InlineRefCountBits>::incrementUnownedSlow(uint32_t n);
+template <>
+void RefCounts<SideTableRefCountBits>::incrementUnownedSlow(uint32_t n) {
+  // Overflow from side table to a new side table?!
+  swift_abortUnownedRetainOverflow();
+}
+
 // namespace swift
 } // namespace swift
 

test/Interpreter/unowned_overflow.swift

@@ -0,0 +1,35 @@
+// RUN: %target-run-simple-swift | %FileCheck %s
+// REQUIRES: executable_test
+
+class Owner {
+  var children: [Child] = []
+
+  func addChild(_ c: Child) {
+    children.append(c)
+  }
+
+  func removeChildren() {
+    children.removeAll()
+  }
+
+  func test() {
+    // Overflow of unowned ref count on 32bit.
+    for _ in 0 ..< 500 {
+      addChild(Child(self))
+    }
+    removeChildren()
+  }
+}
+
+class Child {
+  unowned var owner: Owner
+
+  init(_ o: Owner) {
+    owner = o
+  }
+}
+
+let o = Owner()
+o.test()
+print("success")
+// CHECK: success

unittests/runtime/Refcounting.cpp

@@ -153,6 +153,24 @@ TEST(RefcountingTest, unowned_retain_release_n) {
   EXPECT_EQ(1u, value);
 }
 
+TEST(RefcountingTest, unowned_retain_release_n_overflow) {
+  // This test would test overflow on 32bit platforms.
+  // These platforms have 7 unowned reference count bits.
+  size_t value = 0;
+  auto object = allocTestObject(&value, 1);
+  EXPECT_EQ(0u, value);
+  swift_unownedRetain_n(object, 128);
+  EXPECT_EQ(129u, swift_unownedRetainCount(object));
+  swift_unownedRetain(object);
+  EXPECT_EQ(130u, swift_unownedRetainCount(object));
+  swift_unownedRelease_n(object, 128);
+  EXPECT_EQ(2u, swift_unownedRetainCount(object));
+  swift_unownedRelease(object);
+  EXPECT_EQ(1u, swift_unownedRetainCount(object));
+  swift_release(object);
+  EXPECT_EQ(1u, value);
+}
+
 TEST(RefcountingTest, isUniquelyReferenced) {
   size_t value = 0;
   auto object = allocTestObject(&value, 1);