Adding RawSpan support to writeBytes in ByteBuffer (#3269)

Motivation:

With the new Swift version (6.2) we get access to spans for safely
interacting with the contiguous memory owned by other types. Allow users
write bytes to the ByteBuffer with RawSpans.

Modifications:

* ByteBuffer-core: Overload the setBytes, _setBytes, and
_setBytesAssumingUniqueBufferAccess to use RawSpan instead of
UnsafeRawBufferPointer
* ByteBuffer-aux: Overload the writeBytes to use RawSpan instead of
UnsafeRawBufferPointer

Result:

Users can now write the contiguous memory of other types into ByteBuffer
using the new RawSpan overloads, removing the need for unsafe pointer
handling.

Author: yigityazicilar

Sources/NIOCore/ByteBuffer-aux.swift

@@ -517,6 +517,22 @@ extension ByteBuffer {
         return written
     }
 
+    #if compiler(>=6.2)
+    /// Write `bytes` into this `ByteBuffer`. Moves the writer index forward by the number of bytes written.
+    ///
+    /// - Parameters:
+    ///   - bytes: A `RawSpan`
+    /// - Returns: The number of bytes written or `bytes.byteCount`.
+    @discardableResult
+    @inlinable
+    @available(macOS 26, iOS 26, tvOS 26, watchOS 26, visionOS 26, *)
+    public mutating func writeBytes(_ bytes: RawSpan) -> Int {
+        let written = self.setBytes(bytes, at: self.writerIndex)
+        self._moveWriterIndex(forwardBy: written)
+        return written
+    }
+    #endif
+
     /// Writes `byte` `count` times. Moves the writer index forward by the number of bytes written.
     ///
     /// - Parameters:

Sources/NIOCore/ByteBuffer-core.swift

@@ -508,6 +508,34 @@ public struct ByteBuffer {
         targetPtr.copyMemory(from: bytes)
     }
 
+    #if compiler(>=6.2)
+    @inlinable
+    @available(macOS 26, iOS 26, tvOS 26, watchOS 26, visionOS 26, *)
+    mutating func _setBytes(_ bytes: RawSpan, at index: _Index) -> _Capacity {
+        let bytesCount = bytes.byteCount
+        let newEndIndex: _Index = index + _toIndex(bytesCount)
+        if !isKnownUniquelyReferenced(&self._storage) {
+            let extraCapacity = newEndIndex > self._slice.upperBound ? newEndIndex - self._slice.upperBound : 0
+            self._copyStorageAndRebase(extraCapacity: extraCapacity)
+        }
+        self._ensureAvailableCapacity(_Capacity(bytesCount), at: index)
+        self._setBytesAssumingUniqueBufferAccess(bytes, at: index)
+        return _toCapacity(bytesCount)
+    }
+
+    @inlinable
+    @available(macOS 26, iOS 26, tvOS 26, watchOS 26, visionOS 26, *)
+    mutating func _setBytesAssumingUniqueBufferAccess(_ bytes: RawSpan, at index: _Index) {
+        let targetPtr = UnsafeMutableRawBufferPointer(
+            rebasing: self._slicedStorageBuffer.dropFirst(Int(index))
+        )
+
+        bytes.withUnsafeBytes {
+            targetPtr.copyMemory(from: $0)
+        }
+    }
+    #endif
+
     @inline(never)
     @inlinable
     @_specialize(where Bytes == CircularBuffer<UInt8>)
@@ -1035,6 +1063,16 @@ extension ByteBuffer {
         Int(self._setBytes(bytes, at: _toIndex(index)))
     }
 
+    #if compiler(>=6.2)
+    /// Copy `bytes` from a `RawSpan` into the `ByteBuffer` at `index`. Does not move the writer index.
+    @discardableResult
+    @inlinable
+    @available(macOS 26, iOS 26, tvOS 26, watchOS 26, visionOS 26, *)
+    public mutating func setBytes(_ bytes: RawSpan, at index: Int) -> Int {
+        Int(self._setBytes(bytes, at: _toIndex(index)))
+    }
+    #endif
+
     /// Move the reader index forward by `offset` bytes.
     ///
     /// - warning: By contract the bytes between (including) `readerIndex` and (excluding) `writerIndex` must be

Tests/NIOCoreTests/ByteBufferTest.swift

@@ -4420,3 +4420,46 @@ extension ByteBufferTest {
         XCTAssertNil(result, "peekUUIDBytes() should return nil when fewer than 16 bytes are readable.")
     }
 }
+
+#if compiler(>=6.2)
+extension ByteBufferTest {
+    @available(macOS 26, iOS 26, tvOS 26, watchOS 26, visionOS 26, *)
+    func testWriteBytesRawSpan() {
+        // Write 16 bytes into buffer using a RawSpan
+        let byteArray: [UInt8] = Array(0..<16)
+        let rawSpan = byteArray.span.bytes
+        let writeLength = self.buf.writeBytes(rawSpan)
+        XCTAssertEqual(writeLength, rawSpan.byteCount)
+
+        let read = self.buf.readBytes(length: 4)
+        XCTAssertEqual([0, 1, 2, 3], read)
+        XCTAssertEqual(buf.readerIndex, 4)
+
+        let peek = self.buf.peekBytes(length: 4)
+        XCTAssertEqual([4, 5, 6, 7], peek)
+        XCTAssertEqual(buf.readerIndex, 4)
+
+        let rest = self.buf.readBytes(length: 12)
+        XCTAssertEqual([4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], rest)
+        XCTAssertEqual(buf.readerIndex, 16)
+    }
+
+    @available(macOS 26, iOS 26, tvOS 26, watchOS 26, visionOS 26, *)
+    func testSetBytesRawSpan() {
+        // Write 4 bytes using setBytes
+        let byteArray: [UInt8] = Array(0..<4)
+        let rawSpan = byteArray.span.bytes
+        let writeLength = self.buf.setBytes(rawSpan, at: 0)
+        XCTAssertEqual(writeLength, rawSpan.byteCount)
+
+        // Should not be readable as writer index is not moved by setBytes
+        let shouldBeNil = self.buf.readBytes(length: 4)
+        XCTAssertNil(shouldBeNil)
+
+        // Move writer index
+        self.buf.moveWriterIndex(to: 4)
+        let result = self.buf.readBytes(length: 4)
+        XCTAssertEqual(Array(0..<4), result!)
+    }
+}
+#endif