Optimize accessing Swift Arrays from ObjC

Author: Catfish-Man

stdlib/public/core/ContiguousArrayBuffer.swift

@@ -79,6 +79,97 @@ internal final class _ContiguousArrayStorage<
   }
 
 #if _runtime(_ObjC)
+  
+  internal final override func withUnsafeBufferOfObjects<R>(
+    _ body: (UnsafeBufferPointer<AnyObject>) throws -> R
+  ) rethrows -> R {
+    _internalInvariant(_isBridgedVerbatimToObjectiveC(Element.self))
+    let count = countAndCapacity.count
+    let elements = UnsafeRawPointer(_elementPointer)
+      .assumingMemoryBound(to: AnyObject.self)
+    defer { _fixLifetime(self) }
+    return try body(UnsafeBufferPointer(start: elements, count: count))
+  }
+  
+  @objc(countByEnumeratingWithState:objects:count:)
+  @_effects(releasenone)
+  internal final override func countByEnumerating(
+    with state: UnsafeMutablePointer<_SwiftNSFastEnumerationState>,
+    objects: UnsafeMutablePointer<AnyObject>?, count: Int
+  ) -> Int {
+    var enumerationState = state.pointee
+    
+    if enumerationState.state != 0 {
+      return 0
+    }
+    
+    return withUnsafeBufferOfObjects {
+      objects in
+      enumerationState.mutationsPtr = _fastEnumerationStorageMutationsPtr
+      enumerationState.itemsPtr =
+        AutoreleasingUnsafeMutablePointer(objects.baseAddress)
+      enumerationState.state = 1
+      state.pointee = enumerationState
+      return objects.count
+    }
+  }
+  
+  @inline(__always)
+  @_effects(readonly)
+  @nonobjc private func _objectAt(_ index: Int) -> Unmanaged<AnyObject> {
+    return withUnsafeBufferOfObjects {
+      objects in
+      _precondition(
+        _isValidArraySubscript(index, count: objects.count),
+        "Array index out of range")
+      return Unmanaged.passUnretained(objects[index])
+    }
+  }
+  
+  @objc(objectAtIndexedSubscript:)
+  @_effects(readonly)
+  final override internal func objectAtSubscript(_ index: Int) -> Unmanaged<AnyObject> {
+    return _objectAt(index)
+  }
+  
+  @objc(objectAtIndex:)
+  @_effects(readonly)
+  final override internal func objectAt(_ index: Int) -> Unmanaged<AnyObject> {
+    return _objectAt(index)
+  }
+  
+  @objc internal override final var count: Int {
+    @_effects(readonly) get {
+      return withUnsafeBufferOfObjects { $0.count }
+    }
+  }
+
+  @_effects(releasenone)
+  @objc internal override final func getObjects(
+    _ aBuffer: UnsafeMutablePointer<AnyObject>, range: _SwiftNSRange
+  ) {
+    return withUnsafeBufferOfObjects {
+      objects in
+      _precondition(
+        _isValidArrayIndex(range.location, count: objects.count),
+        "Array index out of range")
+
+      _precondition(
+        _isValidArrayIndex(
+          range.location + range.length, count: objects.count),
+        "Array index out of range")
+
+      if objects.isEmpty { return }
+
+      // These objects are "returned" at +0, so treat them as pointer values to
+      // avoid retains. Copy bytes via a raw pointer to circumvent reference
+      // counting while correctly aliasing with all other pointer types.
+      UnsafeMutableRawPointer(aBuffer).copyMemory(
+        from: objects.baseAddress! + range.location,
+        byteCount: range.length * MemoryLayout<AnyObject>.stride)
+    }
+  }
+  
   /// If the `Element` is bridged verbatim, invoke `body` on an
   /// `UnsafeBufferPointer` to the elements and return the result.
   /// Otherwise, return `nil`.

stdlib/public/core/SwiftNativeNSArray.swift

@@ -64,29 +64,32 @@ internal class __SwiftNativeNSArrayWithContiguousStorage
 private let NSNotFound: Int = .max
 
 // Implement the APIs required by NSArray 
-extension __SwiftNativeNSArrayWithContiguousStorage: _NSArrayCore {
+extension __SwiftNativeNSArrayWithContiguousStorage {
   @objc internal var count: Int {
     return withUnsafeBufferOfObjects { $0.count }
   }
 
   @inline(__always)
-  @nonobjc private func _objectAt(_ index: Int) -> AnyObject {
+  @_effects(readonly)
+  @nonobjc private func _objectAt(_ index: Int) -> Unmanaged<AnyObject> {
     return withUnsafeBufferOfObjects {
       objects in
       _precondition(
         _isValidArraySubscript(index, count: objects.count),
         "Array index out of range")
-      return objects[index]
+      return Unmanaged.passUnretained(objects[index])
     }
   }
 
   @objc(objectAtIndexedSubscript:)
-  dynamic internal func objectAtSubscript(_ index: Int) -> AnyObject {
+  @_effects(readonly)
+  dynamic internal func objectAtSubscript(_ index: Int) -> Unmanaged<AnyObject> {
     return _objectAt(index)
   }
 
   @objc(objectAtIndex:)
-  dynamic internal func objectAt(_ index: Int) -> AnyObject {
+  @_effects(readonly)
+  dynamic internal func objectAt(_ index: Int) -> Unmanaged<AnyObject> {
     return _objectAt(index)
   }
 
@@ -146,7 +149,7 @@ extension __SwiftNativeNSArrayWithContiguousStorage: _NSArrayCore {
 @_fixed_layout
 @usableFromInline
 @objc internal final class _SwiftNSMutableArray :
-  _SwiftNativeNSMutableArray, _NSArrayCore
+  _SwiftNativeNSMutableArray
 {
   internal var contents: [AnyObject]
 
@@ -160,15 +163,17 @@ extension __SwiftNativeNSArrayWithContiguousStorage: _NSArrayCore {
   }
 
   @objc(objectAtIndexedSubscript:)
-  dynamic internal func objectAtSubscript(_ index: Int) -> AnyObject {
+  @_effects(readonly)
+  dynamic internal func objectAtSubscript(_ index: Int) -> Unmanaged<AnyObject> {
     //TODO: exception instead of precondition, once that's possible
-    return contents[index]
+    return Unmanaged.passUnretained(contents[index])
   }
 
   @objc(objectAtIndex:)
-  dynamic internal func objectAt(_ index: Int) -> AnyObject {
+  @_effects(readonly)
+  dynamic internal func objectAt(_ index: Int) -> Unmanaged<AnyObject> {
     //TODO: exception instead of precondition, once that's possible
-    return contents[index]
+    return Unmanaged.passUnretained(contents[index])
   }
 
   @objc internal func getObjects(

validation-test/stdlib/ArrayNew.swift.gyb

@@ -541,8 +541,6 @@ ArrayTestSuite.test("BridgedToObjC/Verbatim/objectAtIndex") {
   expectEqual(idValue0, unsafeBitCast(a.object(at: 0) as AnyObject, to: UInt.self))
   expectEqual(idValue1, unsafeBitCast(a.object(at: 1) as AnyObject, to: UInt.self))
   expectEqual(idValue2, unsafeBitCast(a.object(at: 2) as AnyObject, to: UInt.self))
-
-  expectAutoreleasedKeysAndValues(unopt: (0, 3))
 }
 
 for indexRange in [
@@ -602,8 +600,6 @@ ArrayTestSuite.test("BridgedToObjC/Verbatim/getObjects") {
 
   buffer.deallocate()
   _fixLifetime(a)
-
-  expectAutoreleasedKeysAndValues(unopt: (0, 3))
 }
 
 ArrayTestSuite.test("BridgedToObjC/Verbatim/copyWithZone") {
@@ -721,8 +717,6 @@ ArrayTestSuite.test("BridgedToObjC/Verbatim/BridgeBack/Reallocate") {
   expectEqual(idValue0, unsafeBitCast(a.object(at: 0) as AnyObject, to: UInt.self))
   expectEqual(idValue1, unsafeBitCast(a.object(at: 1) as AnyObject, to: UInt.self))
   expectEqual(idValue2, unsafeBitCast(a.object(at: 2) as AnyObject, to: UInt.self))
-
-  expectAutoreleasedKeysAndValues(unopt: (0, 3))
 }
 
 ArrayTestSuite.test("BridgedToObjC/Verbatim/BridgeBack/Adopt") {
@@ -812,7 +806,6 @@ ArrayTestSuite.test("BridgedToObjC/Custom/objectAtIndex") {
   expectEqual(idValue2, unsafeBitCast(a.object(at: 2) as AnyObject, to: UInt.self))
 
   expectEqual(3, TestBridgedValueTy.bridgeOperations)
-  expectAutoreleasedKeysAndValues(unopt: (0, 3))
 }
 
 for indexRange in [
@@ -874,7 +867,6 @@ ArrayTestSuite.test("BridgedToObjC/Custom/getObjects") {
   _fixLifetime(a)
 
   expectEqual(3, TestBridgedValueTy.bridgeOperations)
-  expectAutoreleasedKeysAndValues(unopt: (0, 3))
 }
 
 ArrayTestSuite.test("BridgedToObjC/Custom/copyWithZone") {
@@ -1011,8 +1003,6 @@ ArrayTestSuite.test("BridgedToObjC/Custom/BridgeBack/Cast") {
   expectEqual(idValue0, unsafeBitCast(a.object(at: 0) as AnyObject, to: UInt.self))
   expectEqual(idValue1, unsafeBitCast(a.object(at: 1) as AnyObject, to: UInt.self))
   expectEqual(idValue2, unsafeBitCast(a.object(at: 2) as AnyObject, to: UInt.self))
-
-  expectAutoreleasedKeysAndValues(unopt: (0, 3))
 }
 
 ArrayTestSuite.test("BridgedToObjC/Custom/BridgeBack/Reallocate") {
@@ -1045,8 +1035,6 @@ ArrayTestSuite.test("BridgedToObjC/Custom/BridgeBack/Reallocate") {
   expectEqual(idValue0, unsafeBitCast(a.object(at: 0) as AnyObject, to: UInt.self))
   expectEqual(idValue1, unsafeBitCast(a.object(at: 1) as AnyObject, to: UInt.self))
   expectEqual(idValue2, unsafeBitCast(a.object(at: 2) as AnyObject, to: UInt.self))
-
-  expectAutoreleasedKeysAndValues(unopt: (0, 3))
 }
 
 ArrayTestSuite.test("BridgedToObjC/Custom/BridgeBack/Adopt") {