Merge pull request #41288 from glessard/sr15433

[stdlib] Implement `withContiguousStorageIfAvailable` for raw buffer types

Author: glessard

stdlib/public/core/UnsafeRawBufferPointer.swift.gyb

@@ -830,6 +830,35 @@ extension Unsafe${Mutable}RawBufferPointer {
     let n = c / MemoryLayout<T>.stride
     return .init(start: .init(s._rawValue), count: n)
   }
+
+%  if Mutable:
+  @inlinable
+  @_alwaysEmitIntoClient
+  public func withContiguousMutableStorageIfAvailable<R>(
+    _ body: (inout UnsafeMutableBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    try withMemoryRebound(to: Element.self) { b in
+      var buffer = b
+      defer {
+        _debugPrecondition(
+          (b.baseAddress, b.count) == (buffer.baseAddress, buffer.count),
+          "UnsafeMutableRawBufferPointer.withContiguousMutableStorageIfAvailable: replacing the buffer is not allowed"
+        )
+      }
+      return try body(&buffer)
+    }
+  }
+
+%  end
+  @inlinable
+  @_alwaysEmitIntoClient
+  public func withContiguousStorageIfAvailable<R>(
+    _ body: (UnsafeBufferPointer<Element>) throws -> R
+  ) rethrows -> R? {
+    try withMemoryRebound(to: Element.self) {
+      try body(${ 'UnsafeBufferPointer<Element>($0)' if Mutable else '$0' })
+    }
+  }
 }
 
 extension Unsafe${Mutable}RawBufferPointer: CustomDebugStringConvertible {

validation-test/stdlib/UnsafeBufferPointer.swift.gyb

@@ -361,6 +361,31 @@ UnsafeMutableBufferPointerTestSuite.test("withContiguous(Mutable)StorageIfAvaila
   expectEqual(result1, result2)
 }
 
+UnsafeMutableRawBufferPointerTestSuite.test("withContiguous(Mutable)StorageIfAvailable") {
+  for c in [0, 1, 4] {
+    let pointer = UnsafeMutableRawPointer.allocate(byteCount: c, alignment: c)
+    defer { pointer.deallocate() }
+
+    let mutableBuffer = UnsafeMutableRawBufferPointer(start: pointer, count: c)
+    mutableBuffer.copyBytes(from: 0..<UInt8(c))
+
+    let r: ()? = mutableBuffer.withContiguousMutableStorageIfAvailable {
+      for i in $0.indices { $0[i] = 2*($0[i]+1) }
+    }
+    expectNotNil(r)
+    let s = mutableBuffer.withContiguousStorageIfAvailable {
+      $0.reduce(0, +)
+    }
+    expectEqual(s.map(Int.init), c*(c+1))
+
+    let immutableBuffer = UnsafeRawBufferPointer(mutableBuffer)
+    let t = immutableBuffer.withContiguousStorageIfAvailable {
+      $0.reduce(0, +)
+    }
+    expectEqual(s, t)
+  }
+}
+
 UnsafeMutableBufferPointerTestSuite.test("initialize(from: Slice)") {
   let o = 91
   let c = 1_000
@@ -386,6 +411,7 @@ UnsafeMutableBufferPointerTestSuite.test("Slice.withContiguousStorageIfAvailable
     let c = test.collection.map { MinimalEquatableValue($0.value) }
     let buffer = UnsafeMutableBufferPointer<MinimalEquatableValue>.allocate(
       capacity: test.collection.count)
+    defer { buffer.deallocate() }
     var (it, idx) = buffer.initialize(from: c)
     expectEqual(buffer.endIndex, idx)
     expectNil(it.next())
@@ -402,6 +428,28 @@ UnsafeMutableBufferPointerTestSuite.test("Slice.withContiguousStorageIfAvailable
   }
 }
 
+UnsafeMutableRawBufferPointerTestSuite.test("Slice.withContiguousStorageIfAvailable") {
+  for test in subscriptRangeTests {
+    let c = test.collection.map({ UInt8($0.value/1000) })
+    let buffer = UnsafeMutableRawBufferPointer.allocate(
+      byteCount: test.collection.count, alignment: 8
+    )
+    defer { buffer.deallocate() }
+    buffer.copyBytes(from: c)
+
+    let expected = test.expected.map({ UInt8($0.value/1000) })
+    let r1: ()? = buffer[test.bounds].withContiguousStorageIfAvailable {
+      expectTrue(expected.elementsEqual($0))
+    }
+    expectNotNil(r1)
+    let r2: ()? = buffer[test.bounds].withContiguousMutableStorageIfAvailable {
+      for i in $0.indices { $0[i] += 1 }
+    }
+    expectNotNil(r2)
+    expectTrue(buffer[test.bounds].elementsEqual(expected.map({ $0+1 })))
+  }
+}
+
 UnsafeMutableBufferPointerTestSuite.test("sort") {
   var values = (0..<1000).map({ _ in Int.random(in: 0..<100) })
   let sortedValues = values.sorted()
@@ -948,6 +996,8 @@ UnsafeMutable${'Raw' if IsRaw else ''}BufferPointerTestSuite.test("subscript/${R
   if _isDebugAssertConfiguration() || isOutOfBounds() {
     if expectedValues == nil { expectCrashLater() }
   }
+%     else:
+  _ = isOutOfBounds
 %     end
   let range = test.rangeSelection.${RangeName}(in: buffer)