[5.0][stdlib] Add withContiguousStorageIfAvailable (#21138)

* [stdlib] Add withContiguous{Mutable}StorageIfAvailable (#21092)

* Add MutableCollection.withContiguousMutableStorageIfAvailable

* Add withContiguousMutableStorageIfAvailable impls

* Add tests on concrete types

* Add Sequence.withContiguousStorageIfAvailable

* Implement withContiguousStorageIfAvailable in concrete types

* Add tests for withContiguousStorageIfAvailable

Author: airspeedswift

stdlib/private/StdlibCollectionUnittest/CheckMutableCollectionType.swift

@@ -438,13 +438,13 @@ if resiliencyChecks.subscriptRangeOnOutOfBoundsRangesBehavior != .none {
 }
 
 //===----------------------------------------------------------------------===//
-// _withUnsafeMutableBufferPointerIfSupported()
+// withContiguousMutableStorageIfAvailable()
 //===----------------------------------------------------------------------===//
 
-self.test("\(testNamePrefix)._withUnsafeMutableBufferPointerIfSupported()/semantics") {
+self.test("\(testNamePrefix).withContiguousMutableStorageIfAvailable()/semantics") {
   for test in subscriptRangeTests {
     var c = makeWrappedCollection(test.collection)
-    var result = c._withUnsafeMutableBufferPointerIfSupported {
+    var result = c.withContiguousMutableStorageIfAvailable {
       (bufferPointer) -> OpaqueValue<Array<OpaqueValue<Int>>> in
       let value = OpaqueValue(bufferPointer.map(extractValue))
       return value

stdlib/private/StdlibCollectionUnittest/LoggingWrappers.swift

@@ -83,6 +83,7 @@ public class SequenceLog {
   public static var prefixWhile = TypeIndexed(0)
   public static var prefixMaxLength = TypeIndexed(0)
   public static var suffixMaxLength = TypeIndexed(0)
+  public static var withContiguousStorageIfAvailable = TypeIndexed(0)
   public static var _customContainsEquatableElement = TypeIndexed(0)
   public static var _copyToContiguousArray = TypeIndexed(0)
   public static var _copyContents = TypeIndexed(0)  
@@ -112,6 +113,9 @@ public class SequenceLog {
   public static var _withUnsafeMutableBufferPointerIfSupported = TypeIndexed(0)
   public static var _withUnsafeMutableBufferPointerIfSupportedNonNilReturns =
     TypeIndexed(0)
+  public static var withContiguousMutableStorageIfAvailable = TypeIndexed(0)
+  public static var withContiguousMutableStorageIfAvailableNonNilReturns =
+    TypeIndexed(0)
   // RangeReplaceableCollection
   public static var init_ = TypeIndexed(0)
   public static var initRepeating = TypeIndexed(0)
@@ -210,6 +214,13 @@ extension LoggingSequence: Sequence {
     SequenceLog.underestimatedCount[selfType] += 1
     return base.underestimatedCount
   }
+  
+  public func withContiguousStorageIfAvailable<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    SequenceLog.withContiguousStorageIfAvailable[selfType] += 1
+    return try base.withContiguousStorageIfAvailable(body)
+  }
 
   public func _customContainsEquatableElement(_ element: Element) -> Bool? {
     SequenceLog._customContainsEquatableElement[selfType] += 1
@@ -385,6 +396,17 @@ extension LoggingMutableCollection: MutableCollection {
     return result
   }
 
+  public mutating func withContiguousMutableStorageIfAvailable<R>(
+    _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    MutableCollectionLog.withContiguousMutableStorageIfAvailable[selfType] += 1
+    let result = try base.withContiguousMutableStorageIfAvailable(body)
+    if result != nil {
+      Log.withContiguousMutableStorageIfAvailable[selfType] += 1
+    }
+    return result
+  }
+
 }
 
 public typealias LoggingMutableBidirectionalCollection<
@@ -501,10 +523,10 @@ public typealias LoggingRangeReplaceableRandomAccessCollection<
 > = LoggingRangeReplaceableCollection<Base>
 
 //===----------------------------------------------------------------------===//
-// Collections that count calls to `_withUnsafeMutableBufferPointerIfSupported`
+// Collections that count calls to `withContiguousMutableStorageIfAvailable`
 //===----------------------------------------------------------------------===//
 
-/// Interposes between `_withUnsafeMutableBufferPointerIfSupported` method calls
+/// Interposes between `withContiguousMutableStorageIfAvailable` method calls
 /// to increment a counter. Calls to this method from within dispatched methods
 /// are uncounted by the standard logging collection wrapper.
 public struct BufferAccessLoggingMutableCollection<
@@ -587,6 +609,17 @@ extension BufferAccessLoggingMutableCollection: MutableCollection {
     }
     return result
   }
+  
+  public mutating func withContiguousMutableStorageIfAvailable<R>(
+    _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    Log.withContiguousMutableStorageIfAvailable[selfType] += 1
+    let result = try base.withContiguousMutableStorageIfAvailable(body)
+    if result != nil {
+      Log.withContiguousMutableStorageIfAvailable[selfType] += 1
+    }
+    return result
+  }
 }
 
 public typealias BufferAccessLoggingMutableBidirectionalCollection<

stdlib/public/core/Array.swift

@@ -1281,6 +1281,26 @@ extension Array: RangeReplaceableCollection {
     }
   }
 
+  @inlinable
+  public mutating func withContiguousMutableStorageIfAvailable<R>(
+    _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    return try withUnsafeMutableBufferPointer {
+      (bufferPointer) -> R in
+      return try body(&bufferPointer)
+    }
+  }
+
+  @inlinable
+  public func withContiguousStorageIfAvailable<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    return try withUnsafeBufferPointer {
+      (bufferPointer) -> R in
+      return try body(bufferPointer)
+    }
+  }
+
   @inlinable
   public __consuming func _copyToContiguousArray() -> ContiguousArray<Element> {
     if let n = _buffer.requestNativeBuffer() {

stdlib/public/core/ArraySlice.swift

@@ -1075,6 +1075,26 @@ extension ArraySlice: RangeReplaceableCollection {
     }
   }
 
+  @inlinable
+  public mutating func withContiguousMutableStorageIfAvailable<R>(
+    _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    return try withUnsafeMutableBufferPointer {
+      (bufferPointer) -> R in
+      return try body(&bufferPointer)
+    }
+  }
+
+  @inlinable
+  public func withContiguousStorageIfAvailable<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    return try withUnsafeBufferPointer {
+      (bufferPointer) -> R in
+      return try body(bufferPointer)
+    }
+  }
+
   @inlinable
   public __consuming func _copyToContiguousArray() -> ContiguousArray<Element> {
     if let n = _buffer.requestNativeBuffer() {

stdlib/public/core/ContiguousArray.swift

@@ -923,6 +923,26 @@ extension ContiguousArray: RangeReplaceableCollection {
     }
   }
 
+  @inlinable
+  public mutating func withContiguousMutableStorageIfAvailable<R>(
+    _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    return try withUnsafeMutableBufferPointer {
+      (bufferPointer) -> R in
+      return try body(&bufferPointer)
+    }
+  }
+  
+  @inlinable
+  public func withContiguousStorageIfAvailable<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    return try withUnsafeBufferPointer {
+      (bufferPointer) -> R in
+      return try body(bufferPointer)
+    }
+  }
+
   @inlinable
   public __consuming func _copyToContiguousArray() -> ContiguousArray<Element> {
     if let n = _buffer.requestNativeBuffer() {

stdlib/public/core/MutableCollection.swift

@@ -180,6 +180,20 @@ where SubSequence: MutableCollection
   mutating func _withUnsafeMutableBufferPointerIfSupported<R>(
     _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
   ) rethrows -> R?
+
+  /// Call `body(p)`, where `p` is a pointer to the collection's
+  /// mutable contiguous storage.  If no such storage exists, it is
+  /// first created.  If the collection does not support an internal
+  /// representation in a form of mutable contiguous storage, `body` is not
+  /// called and `nil` is returned.
+  ///
+  /// Often, the optimizer can eliminate bounds- and uniqueness-checks
+  /// within an algorithm, but when that fails, invoking the
+  /// same algorithm on `body`\ 's argument lets you trade safety for
+  /// speed.
+  mutating func withContiguousMutableStorageIfAvailable<R>(
+    _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
+  ) rethrows -> R?
 }
 
 // TODO: swift-3-indexing-model - review the following
@@ -191,6 +205,13 @@ extension MutableCollection {
     return nil
   }
 
+  @inlinable
+  public mutating func withContiguousMutableStorageIfAvailable<R>(
+    _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    return nil
+  }
+
   /// Accesses a contiguous subrange of the collection's elements.
   ///
   /// The accessed slice uses the same indices for the same elements as the

stdlib/public/core/Sequence.swift

@@ -362,6 +362,10 @@ public protocol Sequence {
   __consuming func _copyContents(
     initializing ptr: UnsafeMutableBufferPointer<Element>
   ) -> (Iterator,UnsafeMutableBufferPointer<Element>.Index)
+  
+  func withContiguousStorageIfAvailable<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R?  
 }
 
 // Provides a default associated type witness for Iterator when the
@@ -1092,17 +1096,24 @@ extension Sequence {
   public __consuming func _copyContents(
     initializing buffer: UnsafeMutableBufferPointer<Element>
   ) -> (Iterator,UnsafeMutableBufferPointer<Element>.Index) {
-      var it = self.makeIterator()
-      guard var ptr = buffer.baseAddress else { return (it,buffer.startIndex) }
-      for idx in buffer.startIndex..<buffer.count {
-        guard let x = it.next() else {
-          return (it, idx)
-        }
-        ptr.initialize(to: x)
-        ptr += 1
+    var it = self.makeIterator()
+    guard var ptr = buffer.baseAddress else { return (it,buffer.startIndex) }
+    for idx in buffer.startIndex..<buffer.count {
+      guard let x = it.next() else {
+        return (it, idx)
       }
-      return (it,buffer.endIndex)
+      ptr.initialize(to: x)
+      ptr += 1
     }
+    return (it,buffer.endIndex)
+  }
+    
+  @inlinable
+  public func withContiguousStorageIfAvailable<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    return nil
+  }  
 }
 
 // FIXME(ABI)#182

stdlib/public/core/UnsafeBufferPointer.swift.gyb

@@ -432,6 +432,25 @@ extension Unsafe${Mutable}BufferPointer {
     return try body(&self)
   }
 
+  @inlinable
+  public mutating func withContiguousMutableStorageIfAvailable<R>(
+    _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    let (oldBase, oldCount) = (self.baseAddress, self.count)
+    defer { 
+      _debugPrecondition((oldBase, oldCount) == (self.baseAddress, self.count),
+      "UnsafeMutableBufferPointer.withUnsafeMutableBufferPointer: replacing the buffer is not allowed")
+    } 
+    return try body(&self)
+  }
+
+  @inlinable
+  public func withContiguousStorageIfAvailable<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    return try body(UnsafeBufferPointer(self))
+  }
+
 %  else:
 
   /// Creates an immutable typed buffer pointer referencing the same memory as the 
@@ -444,6 +463,13 @@ extension Unsafe${Mutable}BufferPointer {
     count = other.count
   }
 
+  @inlinable
+  public func withContiguousStorageIfAvailable<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    return try body(self)
+  }
+
 %  end
 
 %  if not Mutable:

test/api-digester/Outputs/stability-stdlib-abi.swift.expected

@@ -541,3 +541,6 @@ Class _StringStorage has changed its super class from _AbstractStringStorage to
 Constructor _AbstractStringStorage.init() has been removed
 Func _AbstractStringStorage.copy(with:) has been removed
 Func _AbstractStringStorage.getOrComputeBreadcrumbs() has been removed
+
+Func MutableCollection.withContiguousMutableStorageIfAvailable(_:) has been added as a protocol requirement
+Func Sequence.withContiguousStorageIfAvailable(_:) has been added as a protocol requirement

test/api-digester/Outputs/stability-stdlib-source.swift.expected

@@ -213,3 +213,6 @@ Func LazyCollectionProtocol.reversed() has been removed
 Var LazyCollectionProtocol.lazy has declared type change from LazyCollection<Self.Elements> to LazySequence<Self.Elements>
 
 Func Sequence.reduce(into:_:) has parameter 0 changing from Default to Owned
+
+Func MutableCollection.withContiguousMutableStorageIfAvailable(_:) has been added as a protocol requirement
+Func Sequence.withContiguousStorageIfAvailable(_:) has been added as a protocol requirement