Add an option to ByteBuffer.clear() to specify minimumCapacity (#1204)

Sources/NIO/ByteBuffer-core.swift

@@ -235,7 +235,7 @@ public struct ByteBuffer {
             return self.allocateStorage(capacity: self.capacity)
         }
 
-        private func allocateStorage(capacity: _Capacity) -> _Storage {
+        fileprivate func allocateStorage(capacity: _Capacity) -> _Storage {
             let newCapacity = capacity == 0 ? 0 : capacity.nextPowerOf2ClampedToMax()
             return _Storage(bytesNoCopy: _Storage.allocateAndPrepareRawMemory(bytes: newCapacity, allocator: self.allocator),
                             capacity: newCapacity,
@@ -667,6 +667,27 @@ public struct ByteBuffer {
         self._moveWriterIndex(to: 0)
         self._moveReaderIndex(to: 0)
     }
+
+    /// Set both reader index and writer index to `0`. This will reset the state of this `ByteBuffer` to the state
+    /// of a freshly allocated one, if possible without allocations. This is the cheapest way to recycle a `ByteBuffer`
+    /// for a new use-case.
+    ///
+    /// - note: This method will allocate if the underlying storage is referenced by another `ByteBuffer`. Even if an
+    ///         allocation is necessary this will be cheaper as the copy of the storage is elided.
+    ///
+    /// - parameters:
+    ///     - minimumCapacity: The minimum capacity that will be (re)allocated for this buffer
+    public mutating func clear(minimumCapacity: _Capacity) {
+        if !isKnownUniquelyReferenced(&self._storage) {
+            self._storage = self._storage.allocateStorage(capacity: minimumCapacity)
+        } else if minimumCapacity > self._storage.capacity {
+            self._storage.reallocStorage(capacity: minimumCapacity)
+        }
+        self._slice = self._storage.fullSlice
+
+        self._moveWriterIndex(to: 0)
+        self._moveReaderIndex(to: 0)
+    }
 }
 
 extension ByteBuffer: CustomStringConvertible {

Tests/NIOTests/ByteBufferTest+XCTest.swift

@@ -107,6 +107,13 @@ extension ByteBufferTest {
                 ("testWritableBytesAccountsForSlicing", testWritableBytesAccountsForSlicing),
                 ("testClearDupesStorageIfTheresTwoBuffersSharingStorage", testClearDupesStorageIfTheresTwoBuffersSharingStorage),
                 ("testClearDoesNotDupeStorageIfTheresOnlyOneBuffer", testClearDoesNotDupeStorageIfTheresOnlyOneBuffer),
+                ("testClearWithBiggerMinimumCapacityDupesStorageIfTheresTwoBuffersSharingStorage", testClearWithBiggerMinimumCapacityDupesStorageIfTheresTwoBuffersSharingStorage),
+                ("testClearWithSmallerMinimumCapacityDupesStorageIfTheresTwoBuffersSharingStorage", testClearWithSmallerMinimumCapacityDupesStorageIfTheresTwoBuffersSharingStorage),
+                ("testClearWithBiggerMinimumCapacityDoesNotDupeStorageIfTheresOnlyOneBuffer", testClearWithBiggerMinimumCapacityDoesNotDupeStorageIfTheresOnlyOneBuffer),
+                ("testClearWithSmallerMinimumCapacityDoesNotDupeStorageIfTheresOnlyOneBuffer", testClearWithSmallerMinimumCapacityDoesNotDupeStorageIfTheresOnlyOneBuffer),
+                ("testClearWithBiggerCapacityDoesReallocateStorageCorrectlyIfTheresOnlyOneBuffer", testClearWithBiggerCapacityDoesReallocateStorageCorrectlyIfTheresOnlyOneBuffer),
+                ("testClearWithSmallerCapacityDoesReallocateStorageCorrectlyIfTheresOnlyOneBuffer", testClearWithSmallerCapacityDoesReallocateStorageCorrectlyIfTheresOnlyOneBuffer),
+                ("testClearDoesAllocateStorageCorrectlyIfTheresTwoBuffersSharingStorage", testClearDoesAllocateStorageCorrectlyIfTheresTwoBuffersSharingStorage),
                 ("testClearResetsTheSliceCapacityIfTheresOnlyOneBuffer", testClearResetsTheSliceCapacityIfTheresOnlyOneBuffer),
                 ("testClearResetsTheSliceCapacityIfTheresTwoSlicesSharingStorage", testClearResetsTheSliceCapacityIfTheresTwoSlicesSharingStorage),
                 ("testWeUseFastWriteForContiguousCollections", testWeUseFastWriteForContiguousCollections),

Tests/NIOTests/ByteBufferTest.swift

@@ -1227,6 +1227,120 @@ class ByteBufferTest: XCTestCase {
         }
         XCTAssertEqual(bufPtrValPre, bufPtrValPost)
     }
+
+    func testClearWithBiggerMinimumCapacityDupesStorageIfTheresTwoBuffersSharingStorage() throws {
+        let alloc = ByteBufferAllocator()
+        let buf1 = alloc.buffer(capacity: 16)
+        var buf2 = buf1
+
+        var buf1PtrVal: UInt = 1
+        var buf2PtrVal: UInt = 2
+
+        buf1PtrVal = buf1.storagePointerIntegerValue()
+        buf2PtrVal = buf2.storagePointerIntegerValue()
+
+        XCTAssertEqual(buf1PtrVal, buf2PtrVal)
+
+        buf2.clear(minimumCapacity: 32)
+
+        buf1PtrVal = buf1.storagePointerIntegerValue()
+        buf2PtrVal = buf2.storagePointerIntegerValue()
+
+        XCTAssertNotEqual(buf1PtrVal, buf2PtrVal)
+        XCTAssertLessThan(buf1.capacity, 32)
+        XCTAssertGreaterThanOrEqual(buf2.capacity, 32)
+    }
+
+    func testClearWithSmallerMinimumCapacityDupesStorageIfTheresTwoBuffersSharingStorage() throws {
+        let alloc = ByteBufferAllocator()
+        let buf1 = alloc.buffer(capacity: 16)
+        var buf2 = buf1
+
+        var buf1PtrVal: UInt = 1
+        var buf2PtrVal: UInt = 2
+
+        buf1PtrVal = buf1.storagePointerIntegerValue()
+        buf2PtrVal = buf2.storagePointerIntegerValue()
+
+        XCTAssertEqual(buf1PtrVal, buf2PtrVal)
+
+        buf2.clear(minimumCapacity: 4)
+
+        buf1PtrVal = buf1.storagePointerIntegerValue()
+        buf2PtrVal = buf2.storagePointerIntegerValue()
+
+        XCTAssertNotEqual(buf1PtrVal, buf2PtrVal)
+        XCTAssertGreaterThanOrEqual(buf1.capacity, 16)
+        XCTAssertLessThan(buf2.capacity, 16)
+    }
+
+    func testClearWithBiggerMinimumCapacityDoesNotDupeStorageIfTheresOnlyOneBuffer() throws {
+        let alloc = ByteBufferAllocator()
+        var buf = alloc.buffer(capacity: 16)
+
+        var bufPtrValPre: UInt = 1
+        var bufPtrValPost: UInt = 2
+
+        XCTAssertLessThan(buf.capacity, 32)
+        let preCapacity = buf.capacity
+
+        bufPtrValPre = buf.storagePointerIntegerValue()
+        buf.clear(minimumCapacity: 32)
+        bufPtrValPost = buf.storagePointerIntegerValue()
+        let postCapacity = buf.capacity
+
+        XCTAssertNotEqual(bufPtrValPre, bufPtrValPost)
+        XCTAssertGreaterThanOrEqual(buf.capacity, 32)
+        XCTAssertNotEqual(preCapacity, postCapacity)
+    }
+
+    func testClearWithSmallerMinimumCapacityDoesNotDupeStorageIfTheresOnlyOneBuffer() throws {
+        let alloc = ByteBufferAllocator()
+        var buf = alloc.buffer(capacity: 16)
+
+        var bufPtrValPre: UInt = 1
+        var bufPtrValPost: UInt = 2
+
+        let preCapacity = buf.capacity
+        XCTAssertGreaterThanOrEqual(buf.capacity, 16)
+
+        bufPtrValPre = buf.storagePointerIntegerValue()
+        buf.clear(minimumCapacity: 8)
+        bufPtrValPost = buf.storagePointerIntegerValue()
+        let postCapacity = buf.capacity
+
+        XCTAssertEqual(bufPtrValPre, bufPtrValPost)
+        XCTAssertEqual(preCapacity, postCapacity)
+    }
+
+    func testClearWithBiggerCapacityDoesReallocateStorageCorrectlyIfTheresOnlyOneBuffer() throws {
+        let alloc = ByteBufferAllocator()
+        var buf = alloc.buffer(capacity: 16)
+
+        buf.clear(minimumCapacity: 32)
+
+        XCTAssertEqual(buf._storage.capacity, 32)
+    }
+
+    func testClearWithSmallerCapacityDoesReallocateStorageCorrectlyIfTheresOnlyOneBuffer() throws {
+        let alloc = ByteBufferAllocator()
+        var buf = alloc.buffer(capacity: 16)
+
+        buf.clear(minimumCapacity: 8)
+
+        XCTAssertEqual(buf._storage.capacity, 16)
+    }
+
+    func testClearDoesAllocateStorageCorrectlyIfTheresTwoBuffersSharingStorage() throws {
+        let alloc = ByteBufferAllocator()
+        var buf1 = alloc.buffer(capacity: 16)
+        let buf2 = buf1
+
+        buf1.clear(minimumCapacity: 8)
+
+        XCTAssertEqual(buf1._storage.capacity, 8)
+        XCTAssertEqual(buf2._storage.capacity, 16)
+    }
 
     func testClearResetsTheSliceCapacityIfTheresOnlyOneBuffer() {
         let alloc = ByteBufferAllocator()
@@ -2160,3 +2274,13 @@ private func testReserveCapacityLarger_reallocHook(_ ptr: UnsafeMutableRawPointe
 private func testReserveCapacityLarger_memcpyHook(_ dst: UnsafeMutableRawPointer, _ src: UnsafeRawPointer, _ count: Int) -> Void {
     // No copying
 }
+
+extension ByteBuffer {
+    func storagePointerIntegerValue() -> UInt {
+        var pointer: UInt = 0
+        self.withUnsafeReadableBytes { ptr in
+            pointer = UInt(bitPattern: ptr.baseAddress!)
+        }
+        return pointer
+    }
+}