Merge pull request #80517 from glessard/rdar138440979-MutableSpan-mutableBytes

[SE-0467] add mutableBytes to MutableSpan

Author: glessard

stdlib/public/core/Span/MutableRawSpan.swift

@@ -109,16 +109,17 @@ extension MutableRawSpan {
   }
 
   @_alwaysEmitIntoClient
-  @lifetime(copy elements)
+  @lifetime(&elements)
   public init<Element: BitwiseCopyable>(
-    _elements elements: consuming MutableSpan<Element>
+    _elements elements: inout MutableSpan<Element>
   ) {
-    let bytes = unsafe UnsafeMutableRawBufferPointer(
-      start: elements._pointer,
-      count: elements.count &* MemoryLayout<Element>.stride
+    let (start, count) = unsafe (elements._pointer, elements._count)
+    let span = unsafe MutableRawSpan(
+      _unchecked: start,
+      byteCount: count == 1 ? MemoryLayout<Element>.size
+                 : count &* MemoryLayout<Element>.stride
     )
-    let span = unsafe MutableRawSpan(_unsafeBytes: bytes)
-    self = unsafe _overrideLifetime(span, copying: elements)
+    self = unsafe _overrideLifetime(span, mutating: &elements)
   }
 }
 

stdlib/public/core/Span/MutableSpan.swift

@@ -256,6 +256,17 @@ extension MutableSpan where Element: BitwiseCopyable {
       RawSpan(_mutableSpan: self)
     }
   }
+
+  /// Construct a MutableRawSpan over the memory represented by this span
+  ///
+  /// - Returns: a MutableRawSpan over the memory represented by this span
+  @_alwaysEmitIntoClient
+  public var mutableBytes: MutableRawSpan {
+    @lifetime(&self)
+    mutating get {
+      MutableRawSpan(_elements: &self)
+    }
+  }
 }
 
 @available(SwiftCompatibilitySpan 5.0, *)

test/stdlib/Span/MutableRawSpanTests.swift

@@ -36,6 +36,38 @@ suite.test("Basic Initializer")
   }
 }
 
+private struct Padded: BitwiseCopyable {
+  var storage: (Int64, Int8)
+}
+
+suite.test("Initializer from MutableSpan")
+.require(.stdlib_6_2).code {
+  guard #available(SwiftStdlib 6.2, *) else { return }
+
+  var array = [0, 1, 2].map({ Padded(storage: (Int64($0), Int8($0))) })
+  array.withUnsafeMutableBufferPointer {
+    var span = MutableSpan(_unsafeElements: $0)
+    var rawSpan = MutableRawSpan(_elements: &span)
+
+    expectEqual(rawSpan.byteCount, $0.count * MemoryLayout<Padded>.stride)
+
+    rawSpan.storeBytes(of: 15, as: Int64.self)
+  }
+  expectEqual(array[0].storage.0, 15)
+
+  var slice = array.prefix(1)
+  slice.withUnsafeMutableBufferPointer {
+    expectEqual($0.count, 1)
+    var span = MutableSpan(_unsafeElements: $0)
+    var rawSpan = MutableRawSpan(_elements: &span)
+
+    expectEqual(rawSpan.byteCount, MemoryLayout<Padded>.size)
+
+    rawSpan.storeBytes(of: 3, as: Int64.self)
+  }
+  expectEqual(slice[0].storage.0, 3)
+}
+
 suite.test("isEmpty property")
 .skip(.custom(
   { if #available(SwiftStdlib 6.2, *) { false } else { true } },

test/stdlib/Span/MutableSpanTests.swift

@@ -122,7 +122,28 @@ suite.test("RawSpan from MutableSpan")
     let span = MutableSpan(_unsafeStart: p, count: c)
     let bytes  = span.bytes
     expectEqual(bytes.byteCount, count*MemoryLayout<Int>.stride)
+    let v = bytes.unsafeLoad(
+      fromByteOffset: MemoryLayout<Int>.stride, as: Int.self
+    )
+    expectEqual(v, 1)
+  }
+}
+
+suite.test("MutableRawSpan from MutableSpan")
+.require(.stdlib_6_2).code {
+  guard #available(SwiftStdlib 6.2, *) else { return }
+
+  let count = 4
+  var array = Array(0..<count)
+  expectEqual(array[0], 0)
+  array.withUnsafeMutableBufferPointer {
+    let (p, c) = ($0.baseAddress!, $0.count)
+    var span = MutableSpan(_unsafeStart: p, count: c)
+    var bytes  = span.mutableBytes
+    expectEqual(bytes.byteCount, count*MemoryLayout<Int>.stride)
+    bytes.storeBytes(of: 1, as: Int.self)
   }
+  expectEqual(array[0], 1)
 }
 
 suite.test("indices property")