[stdlib] Slice: customize withContiguous[Mutable]StorageIfAvailable (swiftlang#28883)

* [stdlib] Slice: customize withContiguous[Mutable]StorageIfAvailable

We can easily make an UnsafeBufferPointer that slices another UnsafeBufferPointer, so let’s allow Slice to vend a slice of the base collection’s contiguous storage, if it provides access to one.

We need to do some index distance calculations to implement this, but those will be constant-time in the usual case where the base collection is a RAC.

https://bugs.swift.org/browse/SR-11957
rdar://58090587

* [test] UnsafeBufferPointer: fix some warnings

* [stdlib] Slice: don’t calculate index distances unless the base provides contiguous mutable storage

Author: lorentey

stdlib/public/core/Slice.swift

@@ -215,6 +215,18 @@ extension Slice: Collection {
   public func _failEarlyRangeCheck(_ range: Range<Index>, bounds: Range<Index>) {
     _base._failEarlyRangeCheck(range, bounds: bounds)
   }
+
+  @_alwaysEmitIntoClient @inlinable
+  public func withContiguousStorageIfAvailable<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    try _base.withContiguousStorageIfAvailable { buffer in
+      let start = _base.distance(from: _base.startIndex, to: _startIndex)
+      let count = _base.distance(from: _startIndex, to: _endIndex)
+      let slice = UnsafeBufferPointer(rebasing: buffer[start ..< start + count])
+      return try body(slice)
+    }
+  }
 }
 
 extension Slice: BidirectionalCollection where Base: BidirectionalCollection {
@@ -258,6 +270,34 @@ extension Slice: MutableCollection where Base: MutableCollection {
       _writeBackMutableSlice(&self, bounds: bounds, slice: newValue)
     }
   }
+
+  @_alwaysEmitIntoClient @inlinable
+  public mutating func withContiguousMutableStorageIfAvailable<R>(
+    _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    // We're calling `withContiguousMutableStorageIfAvailable` twice here so
+    // that we don't calculate index distances unless we know we'll use them.
+    // The expectation here is that the base collection will make itself
+    // contiguous on the first try and the second call will be relatively cheap.
+    guard _base.withContiguousMutableStorageIfAvailable({ _ in }) != nil
+    else {
+      return nil
+    }
+    let start = _base.distance(from: _base.startIndex, to: _startIndex)
+    let count = _base.distance(from: _startIndex, to: _endIndex)
+    return try _base.withContiguousMutableStorageIfAvailable { buffer in
+      var slice = UnsafeMutableBufferPointer(
+        rebasing: buffer[start ..< start + count])
+      let copy = slice
+      defer {
+        _precondition(
+          slice.baseAddress == copy.baseAddress &&
+          slice.count == copy.count,
+          "Slice.withUnsafeMutableBufferPointer: replacing the buffer is not allowed")
+      }
+      return try body(&slice)
+    }
+  }
 }
 
 

validation-test/stdlib/Slice.swift.gyb

@@ -113,30 +113,35 @@ SliceTests.test("${Collection}.Slice.{startIndex,endIndex}") {
 %   end
 % end
 
+SliceTests.test("${Collection}.Slice.withContiguousStorageIfAvailable") {
+  for test in subscriptRangeTests {
+    let c = ${Collection}(elements: test.collection)
+    let bounds = test.bounds(in: c)
+    var slice = Slice(base: c, bounds: bounds)
+    let r = slice.withContiguousStorageIfAvailable { _ in }
+    expectNil(r) // None of the minimal collection types implement this.
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // MutableSlice<Base>
 //===----------------------------------------------------------------------===//
 
-/*
-FIXME: uncomment this test when the following bug is fixed:
-
-<rdar://problem/21935030> Recast Slice and MutableSlice in terms of Collection
-and MutableCollection
-
-extension MutableSlice {
+extension Slice {
   func _checkBaseSubSequenceElementIsElement() {
+    expectEqualType(
       Element.self,
       Iterator.Element.self)
     expectEqualType(
       Element.self,
-      Iterator.Element.self,
       Base.Iterator.Element.self)
     expectEqualType(
       Element.self,
-      Iterator.Element.self,
+      Iterator.Element.self)
+    expectEqualType(
+      Element.self,
       Base.SubSequence.Iterator.Element.self)
   }
 }
-*/
 
 runAllTests()

validation-test/stdlib/UnsafeBufferPointer.swift.gyb

@@ -63,7 +63,7 @@ ${SelfName}TestSuite.test("AssociatedTypes") {
 %   if IsRaw:
     iteratorType: ${SelfType}.Iterator.self,
 %   else:
-    iteratorType: UnsafeBufferPointerIterator<Float>.self,
+    iteratorType: UnsafeBufferPointer<Float>.Iterator.self,
 %   end
     subSequenceType: Slice<${SelfType}>.self,
     indexType: Int.self,
@@ -361,6 +361,31 @@ UnsafeMutableBufferPointerTestSuite.test("withContiguous(Mutable)StorageIfAvaila
   expectEqual(result1, result2)
 }
 
+UnsafeMutableBufferPointerTestSuite.test("Slice.withContiguousStorageIfAvailable") {
+  guard #available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *) else {
+    return
+  }
+
+  for test in subscriptRangeTests {
+    let c = test.collection.map { MinimalEquatableValue($0.value) }
+    let buffer = UnsafeMutableBufferPointer<MinimalEquatableValue>.allocate(
+      capacity: test.collection.count)
+    var (it, idx) = buffer.initialize(from: c)
+    expectEqual(buffer.endIndex, idx)
+    expectNil(it.next())
+
+    let expected = test.expected.map { MinimalEquatableValue($0.value) }
+    let r1: Void? = buffer[test.bounds].withContiguousStorageIfAvailable { b in
+      expectTrue(expected.elementsEqual(b))
+    }
+    expectNotNil(r1)
+    let r2: Void? = buffer[test.bounds].withContiguousMutableStorageIfAvailable { b in
+      expectTrue(expected.elementsEqual(b))
+    }
+    expectNotNil(r2)
+  }
+}
+
 UnsafeMutableBufferPointerTestSuite.test("sort") {
   var values = (0..<1000).map({ _ in Int.random(in: 0..<100) })
   let sortedValues = values.sorted()
@@ -389,7 +414,7 @@ for testIndex in (0..<bufCount) {
 % Type = 'Unsafe' + ('Mutable' if mutable else '') + ('Raw' if raw else '') + 'BufferPointer'
   ${Type}TestSuite.test("${action}Element\(testIndex)ViaSlice") {
 % if raw:
-    let allocated = UnsafeMutableRawPointer.allocate(bytes: bufCount, alignedTo: 8)
+    let allocated = UnsafeMutableRawPointer.allocate(byteCount: bufCount, alignment: 8)
     for i in 0..<bufCount {
       allocated.storeBytes(of: UInt8(i), toByteOffset: i, as: UInt8.self)
     }
@@ -400,7 +425,7 @@ for testIndex in (0..<bufCount) {
     defer { allocated.deallocate() }
 
     let buffer = ${Type}(start: allocated, count: bufCount)
-    ${'var' if mutable else 'let'} slice = buffer[sliceRange]
+    ${'var' if mutable and action <> 'read' else 'let'} slice = buffer[sliceRange]
 
     if _isDebugAssertConfiguration(),
        testIndex < sliceRange.lowerBound ||
@@ -430,7 +455,7 @@ for testRange in testRanges {
 % Type = 'Unsafe' + ('Mutable' if mutable else '') + ('Raw' if raw else '') + 'BufferPointer'
   ${Type}TestSuite.test("${action}Slice\(testRange)ViaSlice") {
 % if raw:
-    let allocated = UnsafeMutableRawPointer.allocate(bytes: bufCount+2, alignedTo: 8)
+    let allocated = UnsafeMutableRawPointer.allocate(byteCount: bufCount+2, alignment: 8)
     for i in 0..<bufCount+2 {
       allocated.storeBytes(of: UInt8(i), toByteOffset: i, as: UInt8.self)
     }
@@ -441,7 +466,7 @@ for testRange in testRanges {
     defer { allocated.deallocate() }
 
     let buffer = ${Type}(start: allocated, count: bufCount+2)
-    ${'var' if mutable else 'let'} slice = buffer[sliceRange]
+    ${'var' if mutable and action <> 'read' else 'let'} slice = buffer[sliceRange]
 
     if _isDebugAssertConfiguration(),
       testRange.lowerBound < sliceRange.lowerBound ||
@@ -937,7 +962,7 @@ UnsafeMutable${'Raw' if IsRaw else ''}BufferPointerTestSuite.test("subscript/${R
 
 UnsafeMutable${'Raw' if IsRaw else ''}BufferPointerTestSuite.test("subscript/set/overlaps") {
 % if IsRaw:
-  let buffer = UnsafeMutableRawBufferPointer.allocate(count: 4)
+  let buffer = UnsafeMutableRawBufferPointer.allocate(byteCount: 4, alignment: 1)
 % else:
   let buffer = UnsafeMutableBufferPointer<Int>.allocate(capacity: 4)
 % end
@@ -970,7 +995,7 @@ UnsafeMutable${'Raw' if IsRaw else ''}BufferPointerTestSuite.test("subscript/set
 
 UnsafeMutable${'Raw' if IsRaw else ''}BufferPointerTestSuite.test("nonmutating-swapAt") {
 % if IsRaw:
-  let allocated = UnsafeMutableRawPointer.allocate(bytes: 4, alignedTo: 8)
+  let allocated = UnsafeMutableRawPointer.allocate(byteCount: 4, alignment: 8)
   let buffer = UnsafeMutableRawBufferPointer(start: allocated, count: 3)
   allocated.storeBytes(of: UInt8.max, toByteOffset: 3, as: UInt8.self)
 % else: