[stdlib] _BridgingHashBuffer: Fix memory management

- Have the hash buffer include a reference to the original hash storage instance, along with a copy of its _HashTable, so that its lifetime can be independent from the deferred bridging object.
- Convert _BridgingHashBuffer to a ManagedBuffer so that we can easily put reference-counted properties in its header.

Author: lorentey

stdlib/public/core/Dictionary.swift

@@ -2831,7 +2831,7 @@ final internal class _SwiftDeferredNSDictionary<Key: Hashable, Value>
   @nonobjc
   private var _bridgedKeys_DoNotUse: AnyObject?
 
-  // This stored property must be stored at offset zero.  We perform atomic
+  // This stored property must be stored at offset one.  We perform atomic
   // operations on it.
   //
   // Do not access this property directly.
@@ -2858,11 +2858,6 @@ final internal class _SwiftDeferredNSDictionary<Key: Hashable, Value>
     _sanityCheckFailure("don't call this designated initializer")
   }
 
-  deinit {
-    _bridgedKeys?.deinitialize(elementsFrom: native.hashTable)
-    _bridgedValues?.deinitialize(elementsFrom: native.hashTable)
-  }
-
   @nonobjc
   private var _bridgedKeysPtr: UnsafeMutablePointer<AnyObject?> {
     return _getUnsafePointerToStoredProperties(self)
@@ -2894,20 +2889,14 @@ final internal class _SwiftDeferredNSDictionary<Key: Hashable, Value>
 
   /// Attach a buffer for bridged Dictionary keys.
   @nonobjc
-  private func _initializeBridgedKeys(_ storage: _BridgingHashBuffer) -> Bool {
-    return _stdlib_atomicInitializeARCRef(
-      object: _bridgedKeysPtr,
-      desired: storage)
+  private func _initializeBridgedKeys(_ storage: _BridgingHashBuffer) {
+    _stdlib_atomicInitializeARCRef(object: _bridgedKeysPtr, desired: storage)
   }
 
   /// Attach a buffer for bridged Dictionary values.
   @nonobjc
-  private func _initializeBridgedValues(
-    _ storage: _BridgingHashBuffer
-  ) -> Bool {
-    return _stdlib_atomicInitializeARCRef(
-      object: _bridgedValuesPtr,
-      desired: storage)
+  private func _initializeBridgedValues(_ storage: _BridgingHashBuffer) {
+    _stdlib_atomicInitializeARCRef(object: _bridgedValuesPtr, desired: storage)
   }
 
   @nonobjc
@@ -2917,19 +2906,16 @@ final internal class _SwiftDeferredNSDictionary<Key: Hashable, Value>
 
     // Allocate and initialize heap storage for bridged keys.
     let bridged = _BridgingHashBuffer.allocate(
-      bucketCount: native._storage._bucketCount)
+      owner: native._storage,
+      hashTable: native.hashTable)
     for index in native.hashTable {
       let object = _bridgeAnythingToObjectiveC(native.uncheckedKey(at: index))
       bridged.initialize(at: index, to: object)
     }
 
     // Atomically put the bridged keys in place.
-    if !_initializeBridgedKeys(bridged) {
-      // Lost the race.
-      bridged.deinitialize(elementsFrom: native.hashTable)
-      return _bridgedKeys!
-    }
-    return bridged
+    _initializeBridgedKeys(bridged)
+    return _bridgedKeys!
   }
 
   @nonobjc
@@ -2939,19 +2925,16 @@ final internal class _SwiftDeferredNSDictionary<Key: Hashable, Value>
 
     // Allocate and initialize heap storage for bridged values.
     let bridged = _BridgingHashBuffer.allocate(
-      bucketCount: native._storage._bucketCount)
+      owner: native._storage,
+      hashTable: native.hashTable)
     for index in native.hashTable {
       let object = _bridgeAnythingToObjectiveC(native.uncheckedValue(at: index))
       bridged.initialize(at: index, to: object)
     }
 
     // Atomically put the bridged values in place.
-    if !_initializeBridgedValues(bridged) {
-      // Lost the race.
-      bridged.deinitialize(elementsFrom: native.hashTable)
-      return _bridgedValues!
-    }
-    return bridged
+    _initializeBridgedValues(bridged)
+    return _bridgedValues!
   }
 
   @usableFromInline

stdlib/public/core/Hashing.swift

@@ -103,61 +103,52 @@ final internal class _SwiftEmptyNSEnumerator
 #if _runtime(_ObjC)
 /// This is a minimal class holding a single tail-allocated flat buffer,
 /// representing hash table storage for AnyObject elements. This is used to
-/// store bridged elements in deferred bridging scenarios. Lacking a _HashTable,
-/// instances of this class don't know which of their elements are initialized,
-/// so they can't be used on their own.
+/// store bridged elements in deferred bridging scenarios.
 ///
 /// Using a dedicated class for this rather than a _HeapBuffer makes it easy to
 /// recognize these in heap dumps etc.
-internal final class _BridgingHashBuffer {
-  internal var _bucketCount: Int
-
-  // This type is made with allocWithTailElems, so no init is ever called.
-  // But we still need to have an init to satisfy the compiler.
-  private init(_doNotUse: ()) {
-    _sanityCheckFailure("This class cannot be directly initialized")
+internal final class _BridgingHashBuffer
+  : ManagedBuffer<_BridgingHashBuffer.Header, AnyObject> {
+  struct Header {
+    internal var owner: AnyObject
+    internal var hashTable: _HashTable
+
+    init(owner: AnyObject, hashTable: _HashTable) {
+      self.owner = owner
+      self.hashTable = hashTable
+    }
   }
 
-  internal static func allocate(bucketCount: Int) -> _BridgingHashBuffer {
-    _sanityCheck(bucketCount & (bucketCount &- 1) == 0,
-      "Bucket count must be a power of two")
-    let object = Builtin.allocWithTailElems_1(
-      _BridgingHashBuffer.self,
-      bucketCount._builtinWordValue, AnyObject.self)
-    object._bucketCount = bucketCount
-    return object
+  internal static func allocate(
+    owner: AnyObject,
+    hashTable: _HashTable
+  ) -> _BridgingHashBuffer {
+    let buffer = self.create(minimumCapacity: hashTable.bucketCount) { _ in
+      Header(owner: owner, hashTable: hashTable)
+    }
+    return unsafeDowncast(buffer, to: _BridgingHashBuffer.self)
   }
 
   deinit {
-    _sanityCheck(_bucketCount == -1)
-  }
-
-  internal var _elements: UnsafeMutablePointer<AnyObject> {
-    @inline(__always) get {
-      let ptr = Builtin.projectTailElems(self, AnyObject.self)
-      return UnsafeMutablePointer(ptr)
+    for index in header.hashTable {
+      (firstElementAddress + index.bucket).deinitialize(count: 1)
     }
+    _fixLifetime(self)
   }
 
   internal subscript(index: _HashTable.Index) -> AnyObject {
     @inline(__always) get {
-      _sanityCheck(index.bucket >= 0 && index.bucket < _bucketCount)
-      return _elements[index.bucket]
+      _sanityCheck(header.hashTable.isOccupied(index))
+      defer { _fixLifetime(self) }
+      return firstElementAddress[index.bucket]
     }
   }
 
   @inline(__always)
   internal func initialize(at index: _HashTable.Index, to object: AnyObject) {
-    _sanityCheck(index.bucket >= 0 && index.bucket < _bucketCount)
-    (_elements + index.bucket).initialize(to: object)
-  }
-
-  @inline(__always)
-  internal func deinitialize(elementsFrom hashTable: _HashTable) {
-    for index in hashTable {
-      (_elements + index.bucket).deinitialize(count: 1)
-    }
-    _bucketCount = -1
+    _sanityCheck(header.hashTable.isOccupied(index))
+    (firstElementAddress + index.bucket).initialize(to: object)
+    _fixLifetime(self)
   }
 }
 #endif

stdlib/public/core/Set.swift

@@ -2156,6 +2156,7 @@ final internal class _SwiftSetNSEnumerator<Element: Hashable>
   }
 
   private func bridgedElement(at index: _HashTable.Index) -> AnyObject {
+    _sanityCheck(base.hashTable.isOccupied(index))
     if let bridgedElements = self.bridgedElements {
       return bridgedElements[index]
     }
@@ -2233,10 +2234,6 @@ final internal class _SwiftDeferredNSSet<Element: Hashable>
     super.init()
   }
 
-  deinit {
-    _bridgedElements?.deinitialize(elementsFrom: native.hashTable)
-  }
-
   /// Returns the pointer to the stored property, which contains bridged
   /// Set elements.
   @nonobjc
@@ -2256,10 +2253,8 @@ final internal class _SwiftDeferredNSSet<Element: Hashable>
 
   /// Attach a buffer for bridged Set elements.
   @nonobjc
-  private func _initializeBridgedElements(
-    _ storage: _BridgingHashBuffer
-  ) -> Bool {
-    return _stdlib_atomicInitializeARCRef(
+  private func _initializeBridgedElements(_ storage: _BridgingHashBuffer) {
+    _stdlib_atomicInitializeARCRef(
       object: _bridgedElementsPtr,
       desired: storage)
   }
@@ -2270,19 +2265,16 @@ final internal class _SwiftDeferredNSSet<Element: Hashable>
 
     // Allocate and initialize heap storage for bridged objects.
     let bridged = _BridgingHashBuffer.allocate(
-      bucketCount: native._storage._bucketCount)
+      owner: native._storage,
+      hashTable: native.hashTable)
     for index in native.hashTable {
       let object = _bridgeAnythingToObjectiveC(native.element(at: index))
       bridged.initialize(at: index, to: object)
     }
 
     // Atomically put the bridged elements in place.
-    if !_initializeBridgedElements(bridged) {
-      // Lost the race.
-      bridged.deinitialize(elementsFrom: native.hashTable)
-      return _bridgedElements!
-    }
-    return bridged
+    _initializeBridgedElements(bridged)
+    return _bridgedElements!
   }
 
   @objc