[Foundation] Remove inlinable shim calls

We want to enable overlays to import their shims as @_implementationOnly, so that the shims disappear from the public interface.

However, this isn’t possible when a shim is called from an @inlinable func, because then the existence (and definition) of the shim needs to be available to all callers of it.

Unfortunately Foundation’s Data has three instances where it calls  _SwiftFoundationOverlayShims._withStackOrHeapBuffer within @inlinable code:

- Data.init<S: Sequence>(_: S)
- Data.append<S: Sequence>(contentsOf: S)
- Data.replaceSubrange<C: Collection>(_: Range<Int>, with: C)

Rewrite the first two to write sequence contents directly into the target Data instance (saving a memcpy and possibly a memory allocation).

In replaceSubrange, add fast paths for contiguous collection cases, falling back to a Swift version of _withStackOrHeapBuffer with a 32-byte inline buffer.

The expectation is that this will be an overall speedup in most cases, with the possible exception of replaceSubrange invocations with a large non-contiguous collection.

rdar://58132561

Author: lorentey

stdlib/public/Darwin/Foundation/Data.swift

@@ -62,6 +62,39 @@ internal func __NSDataIsCompact(_ data: NSData) -> Bool {
 
 #endif
 
+@_alwaysEmitIntoClient
+internal func _withStackOrHeapBuffer(capacity: Int, _ body: (UnsafeMutableBufferPointer<UInt8>) -> Void) {
+    guard capacity > 0 else {
+        body(UnsafeMutableBufferPointer(start: nil, count: 0))
+        return
+    }
+    typealias InlineBuffer = ( // 32 bytes
+        UInt8, UInt8, UInt8, UInt8, UInt8, UInt8, UInt8, UInt8,
+        UInt8, UInt8, UInt8, UInt8, UInt8, UInt8, UInt8, UInt8,
+        UInt8, UInt8, UInt8, UInt8, UInt8, UInt8, UInt8, UInt8,
+        UInt8, UInt8, UInt8, UInt8, UInt8, UInt8, UInt8, UInt8
+    )
+    let inlineCount = MemoryLayout<InlineBuffer>.size
+    if capacity <= inlineCount {
+        var buffer: InlineBuffer = (
+            0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0
+        )
+        withUnsafeMutableBytes(of: &buffer) { buffer in
+            assert(buffer.count == inlineCount)
+            let start = buffer.baseAddress!.assumingMemoryBound(to: UInt8.self)
+            body(UnsafeMutableBufferPointer(start: start, count: capacity))
+        }
+        return
+    }
+
+    let buffer = UnsafeMutableBufferPointer<UInt8>.allocate(capacity: capacity)
+    defer { buffer.deallocate() }
+    body(buffer)
+}
+
 // Underlying storage representation for medium and large data.
 // Inlinability strategy: methods from here should not inline into InlineSlice or LargeSlice unless trivial.
 // NOTE: older overlays called this class _DataStorage. The two must
@@ -2054,44 +2087,43 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
         // The sequence might still be able to provide direct access to typed memory.
         // NOTE: It's safe to do this because we're already guarding on S's element as `UInt8`. This would not be safe on arbitrary sequences.
         let representation = elements.withContiguousStorageIfAvailable {
-            return _Representation(UnsafeRawBufferPointer($0))
+            _Representation(UnsafeRawBufferPointer($0))
         }
-
         if let representation = representation {
             _representation = representation
-        } else {
-            // Dummy assignment so we can capture below.
-            _representation = _Representation(capacity: 0)
-
-            // Copy as much as we can in one shot from the sequence.
-            let underestimatedCount = Swift.max(elements.underestimatedCount, 1)
-            _withStackOrHeapBuffer(underestimatedCount) { (buffer) in
-                // In order to copy from the sequence, we have to bind the buffer to UInt8.
-                // This is safe since we'll copy out of this buffer as raw memory later.
-                let capacity = buffer.pointee.capacity
-                let base = buffer.pointee.memory.bindMemory(to: UInt8.self, capacity: capacity)
-                var (iter, endIndex) = elements._copyContents(initializing: UnsafeMutableBufferPointer(start: base, count: capacity))
-
-                // Copy the contents of buffer...
-                _representation = _Representation(UnsafeRawBufferPointer(start: base, count: endIndex))
-
-                // ... and append the rest byte-wise, buffering through an InlineData.
-                var buffer = InlineData()
-                while let element = iter.next() {
-                    buffer.append(byte: element)
-                    if buffer.count == buffer.capacity {
-                        buffer.withUnsafeBytes { _representation.append(contentsOf: $0) }
-                        buffer.count = 0
-                    }
-                }
+            return
+        }
 
-                // If we've still got bytes left in the buffer (i.e. the loop ended before we filled up the buffer and cleared it out), append them.
-                if buffer.count > 0 {
-                    buffer.withUnsafeBytes { _representation.append(contentsOf: $0) }
-                    buffer.count = 0
-                }
+        // Copy as much as we can in one shot from the sequence.
+        let underestimatedCount = elements.underestimatedCount
+        _representation = _Representation(count: underestimatedCount)
+        var (iter, endIndex): (S.Iterator, Int) = _representation.withUnsafeMutableBytes { buffer in
+            let start = buffer.baseAddress!.assumingMemoryBound(to: UInt8.self)
+            let b = UnsafeMutableBufferPointer(start: start, count: buffer.count)
+            return elements._copyContents(initializing: b)
+        }
+        guard endIndex == _representation.count else {
+            // We can't trap here. We have to allow an underfilled buffer
+            // to emulate the previous implementation.
+            _representation.replaceSubrange(endIndex ..< _representation.endIndex, with: nil, count: 0)
+            return
+        }
+
+        // Append the rest byte-wise, buffering through an InlineData.
+        var buffer = InlineData()
+        while let element = iter.next() {
+            buffer.append(byte: element)
+            if buffer.count == buffer.capacity {
+                buffer.withUnsafeBytes { _representation.append(contentsOf: $0) }
+                buffer.count = 0
             }
         }
+
+        // If we've still got bytes left in the buffer (i.e. the loop ended before we filled up the buffer and cleared it out), append them.
+        if buffer.count > 0 {
+            buffer.withUnsafeBytes { _representation.append(contentsOf: $0) }
+            buffer.count = 0
+        }
     }
 
     @available(swift, introduced: 4.2)
@@ -2347,43 +2379,43 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
 
         // The sequence might still be able to provide direct access to typed memory.
         // NOTE: It's safe to do this because we're already guarding on S's element as `UInt8`. This would not be safe on arbitrary sequences.
-        var appended = false
-        elements.withContiguousStorageIfAvailable {
+        let appended: Void? = elements.withContiguousStorageIfAvailable {
             _representation.append(contentsOf: UnsafeRawBufferPointer($0))
-            appended = true
         }
-
-        guard !appended else { return }
+        guard appended == nil else { return }
 
         // The sequence is really not contiguous.
         // Copy as much as we can in one shot.
-        let underestimatedCount = Swift.max(elements.underestimatedCount, 1)
-        _withStackOrHeapBuffer(underestimatedCount) { (buffer) in
-            // In order to copy from the sequence, we have to bind the temporary buffer to `UInt8`.
-            // This is safe since we're the only owners of the buffer and we copy out as raw memory below anyway.
-            let capacity = buffer.pointee.capacity
-            let base = buffer.pointee.memory.bindMemory(to: UInt8.self, capacity: capacity)
-            var (iter, endIndex) = elements._copyContents(initializing: UnsafeMutableBufferPointer(start: base, count: capacity))
-
-            // Copy the contents of the buffer...
-            _representation.append(contentsOf: UnsafeRawBufferPointer(start: base, count: endIndex))
-
-            // ... and append the rest byte-wise, buffering through an InlineData.
-            var buffer = InlineData()
-            while let element = iter.next() {
-                buffer.append(byte: element)
-                if buffer.count == buffer.capacity {
-                    buffer.withUnsafeBytes { _representation.append(contentsOf: $0) }
-                    buffer.count = 0
-                }
-            }
+        let underestimatedCount = elements.underestimatedCount
+        let originalCount = _representation.count
+        _representation.count += underestimatedCount
+        var (iter, endIndex): (S.Iterator, Int) = _representation.withUnsafeMutableBytes { buffer in
+            let start = buffer.baseAddress!.assumingMemoryBound(to: UInt8.self) + originalCount
+            let b = UnsafeMutableBufferPointer(start: start, count: buffer.count - originalCount)
+            return elements._copyContents(initializing: b)
+        }
+        guard endIndex == underestimatedCount else {
+            // We can't trap here. We have to allow an underfilled buffer
+            // to emulate the previous implementation.
+            _representation.replaceSubrange(originalCount + endIndex ..< _representation.endIndex, with: nil, count: 0)
+            return
+        }
 
-            // If we've still got bytes left in the buffer (i.e. the loop ended before we filled up the buffer and cleared it out), append them.
-            if buffer.count > 0 {
+        // Append the rest byte-wise, buffering through an InlineData.
+        var buffer = InlineData()
+        while let element = iter.next() {
+            buffer.append(byte: element)
+            if buffer.count == buffer.capacity {
                 buffer.withUnsafeBytes { _representation.append(contentsOf: $0) }
                 buffer.count = 0
             }
         }
+
+        // If we've still got bytes left in the buffer (i.e. the loop ended before we filled up the buffer and cleared it out), append them.
+        if buffer.count > 0 {
+            buffer.withUnsafeBytes { _representation.append(contentsOf: $0) }
+            buffer.count = 0
+        }
     }
 
     // MARK: -
@@ -2436,15 +2468,26 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
     /// - parameter newElements: The replacement bytes.
     @inlinable // This is @inlinable as generic and reasonably small.
     public mutating func replaceSubrange<ByteCollection : Collection>(_ subrange: Range<Index>, with newElements: ByteCollection) where ByteCollection.Iterator.Element == Data.Iterator.Element {
+        // If the collection is already contiguous, access the underlying raw memory directly.
+        if let contiguous = newElements as? ContiguousBytes {
+            contiguous.withUnsafeBytes { buffer in
+                _representation.replaceSubrange(subrange, with: buffer.baseAddress, count: buffer.count)
+            }
+            return
+        }
+        // The collection might still be able to provide direct access to typed memory.
+        // NOTE: It's safe to do this because we're already guarding on ByteCollection's element as `UInt8`. This would not be safe on arbitrary collections.
+        let replaced: Void? = newElements.withContiguousStorageIfAvailable { buffer in
+            _representation.replaceSubrange(subrange, with: buffer.baseAddress, count: buffer.count)
+        }
+        guard replaced == nil else { return }
+
         let totalCount = Int(newElements.count)
-        _withStackOrHeapBuffer(totalCount) { conditionalBuffer in
-            let buffer = UnsafeMutableBufferPointer(start: conditionalBuffer.pointee.memory.assumingMemoryBound(to: UInt8.self), count: totalCount)
+        _withStackOrHeapBuffer(capacity: totalCount) { buffer in
             var (iterator, index) = newElements._copyContents(initializing: buffer)
-            while let byte = iterator.next() {
-                buffer[index] = byte
-                index = buffer.index(after: index)
-            }
-            replaceSubrange(subrange, with: conditionalBuffer.pointee.memory, count: totalCount)
+            precondition(index == buffer.endIndex, "Collection has less elements than its count")
+            precondition(iterator.next() == nil, "Collection has more elements than its count")
+            _representation.replaceSubrange(subrange, with: buffer.baseAddress, count: totalCount)
         }
     }
 

stdlib/public/SwiftShims/FoundationOverlayShims.h

@@ -33,41 +33,3 @@
 #import "NSLocaleShims.h"
 #import "NSTimeZoneShims.h"
 #import "NSUndoManagerShims.h"
-
-typedef struct {
-    void *_Nonnull memory;
-    size_t capacity;
-    _Bool onStack;
-} _ConditionalAllocationBuffer;
-
-static inline _Bool _resizeConditionalAllocationBuffer(_ConditionalAllocationBuffer *_Nonnull buffer, size_t amt) {
-    size_t amount = malloc_good_size(amt);
-    if (amount <= buffer->capacity) { return true; }
-    void *newMemory;
-    if (buffer->onStack) {
-        newMemory = malloc(amount);
-        if (newMemory == NULL) { return false; }
-        memcpy(newMemory, buffer->memory, buffer->capacity);
-        buffer->onStack = false;
-    } else {
-        newMemory = realloc(buffer->memory, amount);
-        if (newMemory == NULL) { return false; }
-    }
-    if (newMemory == NULL) { return false; }
-    buffer->memory = newMemory;
-    buffer->capacity = amount;
-    return true;
-}
-
-static inline _Bool _withStackOrHeapBuffer(size_t amount, void (__attribute__((noescape)) ^ _Nonnull applier)(_ConditionalAllocationBuffer *_Nonnull)) {
-    _ConditionalAllocationBuffer buffer;
-    buffer.capacity = malloc_good_size(amount);
-    buffer.onStack = (pthread_main_np() != 0 ? buffer.capacity < 2048 : buffer.capacity < 512);
-    buffer.memory = buffer.onStack ? alloca(buffer.capacity) : malloc(buffer.capacity);
-    if (buffer.memory == NULL) { return false; }
-    applier(&buffer);
-    if (!buffer.onStack) {
-        free(buffer.memory);
-    }
-    return true;
-}