Use associated objects to attach contiguous array buffers to lazy ones (#75148)

rdar://132124808

Author: Catfish-Man

stdlib/public/SwiftShims/swift/shims/CMakeLists.txt

@@ -9,6 +9,7 @@ set(sources
   LibcOverlayShims.h
   LibcShims.h
   MetadataSections.h
+  ObjCShims.h
   Random.h
   RefCount.h
   Reflection.h

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

@@ -30,9 +30,11 @@ extern "C" {
 #if __LLP64__
 typedef unsigned long long _swift_shims_CFHashCode;
 typedef signed long long _swift_shims_CFIndex;
+
 #else
 typedef unsigned long _swift_shims_CFHashCode;
 typedef signed long _swift_shims_CFIndex;
+
 #endif
 
 // Consider creating SwiftMacTypes.h for these
@@ -74,7 +76,7 @@ _swift_stdlib_NSStringGetCStringTrampoline(id _Nonnull obj,
 SWIFT_RUNTIME_STDLIB_API
 __swift_uint8_t
 _swift_stdlib_dyld_is_objc_constant_string(const void * _Nonnull addr);
-  
+
 #endif // __OBJC2__
 
 #ifdef __cplusplus

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

@@ -0,0 +1,51 @@
+//===--- ObjCShims.h - Access to libobjc for the core stdlib ---------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+//  Using the ObjectiveC module in the core stdlib would create a
+//  circular dependency, so instead we import these declarations as
+//  part of SwiftShims.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_STDLIB_SHIMS_OBJCSHIMS_H
+#define SWIFT_STDLIB_SHIMS_OBJCSHIMS_H
+
+#include "SwiftStdint.h"
+#include "Visibility.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifdef __OBJC2__
+
+extern void objc_setAssociatedObject(void);
+extern void objc_getAssociatedObject(void);
+int objc_sync_enter(id _Nonnull object);
+int objc_sync_exit(id _Nonnull object);
+
+static void * _Nonnull getSetAssociatedObjectPtr() {
+  return (void *)&objc_setAssociatedObject;
+}
+
+static void * _Nonnull getGetAssociatedObjectPtr() {
+  return (void *)&objc_getAssociatedObject;
+}
+
+#endif // __OBJC2__
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#endif // SWIFT_STDLIB_SHIMS_COREFOUNDATIONSHIMS_H
+

stdlib/public/SwiftShims/swift/shims/module.modulemap

@@ -1,6 +1,7 @@
 module SwiftShims {
   header "AssertionReporting.h"
   header "CoreFoundationShims.h"
+  header "ObjCShims.h"
   header "EmbeddedShims.h"
   header "FoundationShims.h"
   header "GlobalObjects.h"

stdlib/public/core/ArrayBuffer.swift

@@ -567,6 +567,81 @@ extension _ArrayBuffer {
     }
   }
 
+  @inlinable @_alwaysEmitIntoClient
+  static var associationKey: UnsafeRawPointer {
+    //We never dereference this, we just need an address to use as a unique key
+    UnsafeRawPointer(Builtin.addressof(&_swiftEmptyArrayStorage))
+  }
+  
+  @inlinable @_alwaysEmitIntoClient
+  internal func getAssociatedBuffer() -> _ContiguousArrayBuffer<Element>? {
+    let getter = unsafeBitCast(
+      getGetAssociatedObjectPtr(),
+      to: (@convention(c)(
+        AnyObject,
+        UnsafeRawPointer
+      ) -> UnsafeRawPointer?).self
+    )
+    if let assocPtr = getter(
+      _storage.objCInstance,
+      _ArrayBuffer.associationKey
+    ) {
+      let buffer = assocPtr.loadUnaligned(
+        as: _ContiguousArrayStorage<Element>.self
+      )
+      return _ContiguousArrayBuffer(buffer)
+    }
+    return nil
+  }
+  
+  @inlinable @_alwaysEmitIntoClient
+  internal func setAssociatedBuffer(_ buffer: _ContiguousArrayBuffer<Element>) {
+    let setter = unsafeBitCast(getSetAssociatedObjectPtr(), to: (@convention(c)(
+      AnyObject,
+      UnsafeRawPointer,
+      AnyObject?,
+      UInt
+    ) -> Void).self)
+    setter(
+      _storage.objCInstance,
+      _ArrayBuffer.associationKey,
+      buffer._storage,
+      1 //OBJC_ASSOCIATION_RETAIN_NONATOMIC
+    )
+  }
+  
+  @_alwaysEmitIntoClient @inline(never)
+  internal func withUnsafeBufferPointer_nonNative<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R {
+    let unwrapped: _ContiguousArrayBuffer<Element>
+    // libobjc already provides the necessary memory barriers for
+    // double checked locking to be safe, per comments on
+    // https://github.com/swiftlang/swift/pull/75148
+    if let associatedBuffer = getAssociatedBuffer() {
+      unwrapped = associatedBuffer
+    } else {
+      let lock = _storage.objCInstance
+      objc_sync_enter(lock)
+      var associatedBuffer = getAssociatedBuffer()
+      if let associatedBuffer {
+        unwrapped = associatedBuffer
+      } else {
+        associatedBuffer = ContiguousArray(self)._buffer
+        unwrapped = associatedBuffer.unsafelyUnwrapped
+        setAssociatedBuffer(unwrapped)
+      }
+      defer { _fixLifetime(unwrapped) }
+      objc_sync_exit(lock)
+    }
+    return try body(
+      UnsafeBufferPointer(
+        start: unwrapped.firstElementAddress,
+        count: unwrapped.count
+      )
+    )
+  }
+  
   /// Call `body(p)`, where `p` is an `UnsafeBufferPointer` over the
   /// underlying contiguous storage.  If no such storage exists, it is
   /// created on-demand.
@@ -579,7 +654,7 @@ extension _ArrayBuffer {
       return try body(
         UnsafeBufferPointer(start: firstElementAddress, count: count))
     }
-    return try ContiguousArray(self).withUnsafeBufferPointer(body)
+    return try withUnsafeBufferPointer_nonNative(body)
   }
 
   /// Call `body(p)`, where `p` is an `UnsafeMutableBufferPointer`

stdlib/public/core/CocoaArray.swift

@@ -51,12 +51,16 @@ internal struct _CocoaArrayWrapper: RandomAccessCollection {
 
   @usableFromInline
   internal var endIndex: Int {
-    return core.count
+    @_effects(releasenone) get {
+      core.count
+    }
   }
 
   @usableFromInline
   internal subscript(i: Int) -> AnyObject {
-    return core.objectAt(i)
+    @_effects(releasenone) get {
+      core.objectAt(i)
+    }
   }
 
   @usableFromInline

stdlib/public/stubs/FoundationHelpers.mm

@@ -112,5 +112,6 @@ typedef __swift_uint8_t (*getCStringImplPtr)(id,
           && SWIFT_RUNTIME_WEAK_USE(_dyld_is_objc_constant(dyld_objc_string_kind, addr))) ? 1 : 0;
 }
 
+
 #endif