Add OneElementFastSequence to be used in ConnectionPool (#420)

Sources/ConnectionPoolModule/OneElementFastSequence.swift

@@ -0,0 +1,151 @@
+/// A `Sequence` that does not heap allocate, if it only carries a single element
+@usableFromInline
+struct OneElementFastSequence<Element>: Sequence {
+    @usableFromInline
+    enum Base {
+        case none(reserveCapacity: Int)
+        case one(Element, reserveCapacity: Int)
+        case n([Element])
+    }
+
+    @usableFromInline
+    private(set) var base: Base
+
+    @inlinable
+    init() {
+        self.base = .none(reserveCapacity: 0)
+    }
+
+    @inlinable
+    init(_ element: Element) {
+        self.base = .one(element, reserveCapacity: 1)
+    }
+
+    @inlinable
+    init(_ collection: some Collection<Element>) {
+        switch collection.count {
+        case 0:
+            self.base = .none(reserveCapacity: 0)
+        case 1:
+            self.base = .one(collection.first!, reserveCapacity: 0)
+        default:
+            if let collection = collection as? Array<Element> {
+                self.base = .n(collection)
+            } else {
+                self.base = .n(Array(collection))
+            }
+        }
+    }
+
+    @usableFromInline
+    var count: Int {
+        switch self.base {
+        case .none:
+            return 0
+        case .one:
+            return 1
+        case .n(let array):
+            return array.count
+        }
+    }
+
+    @inlinable
+    var first: Element? {
+        switch self.base {
+        case .none:
+            return nil
+        case .one(let element, _):
+            return element
+        case .n(let array):
+            return array.first
+        }
+    }
+
+    @usableFromInline
+    var isEmpty: Bool {
+        switch self.base {
+        case .none:
+            return true
+        case .one, .n:
+            return false
+        }
+    }
+
+    @inlinable
+    mutating func reserveCapacity(_ minimumCapacity: Int) {
+        switch self.base {
+        case .none(let reservedCapacity):
+            self.base = .none(reserveCapacity: Swift.max(reservedCapacity, minimumCapacity))
+        case .one(let element, let reservedCapacity):
+            self.base = .one(element, reserveCapacity: Swift.max(reservedCapacity, minimumCapacity))
+        case .n(var array):
+            self.base = .none(reserveCapacity: 0) // prevent CoW
+            array.reserveCapacity(minimumCapacity)
+            self.base = .n(array)
+        }
+    }
+
+    @inlinable
+    mutating func append(_ element: Element) {
+        switch self.base {
+        case .none(let reserveCapacity):
+            self.base = .one(element, reserveCapacity: reserveCapacity)
+        case .one(let existing, let reserveCapacity):
+            var new = [Element]()
+            new.reserveCapacity(reserveCapacity)
+            new.append(existing)
+            new.append(element)
+            self.base = .n(new)
+        case .n(var existing):
+            self.base = .none(reserveCapacity: 0) // prevent CoW
+            existing.append(element)
+            self.base = .n(existing)
+        }
+    }
+
+    @inlinable
+    func makeIterator() -> Iterator {
+        Iterator(self)
+    }
+
+    @usableFromInline
+    struct Iterator: IteratorProtocol {
+        @usableFromInline private(set) var index: Int = 0
+        @usableFromInline private(set) var backing: OneElementFastSequence<Element>
+
+        @inlinable
+        init(_ backing: OneElementFastSequence<Element>) {
+            self.backing = backing
+        }
+
+        @inlinable
+        mutating func next() -> Element? {
+            switch self.backing.base {
+            case .none:
+                return nil
+            case .one(let element, _):
+                if self.index == 0 {
+                    self.index += 1
+                    return element
+                }
+                return nil
+
+            case .n(let array):
+                if self.index < array.endIndex {
+                    defer { self.index += 1}
+                    return array[self.index]
+                }
+                return nil
+            }
+        }
+    }
+}
+
+extension OneElementFastSequence: Equatable where Element: Equatable {}
+extension OneElementFastSequence.Base: Equatable where Element: Equatable {}
+
+extension OneElementFastSequence: Hashable where Element: Hashable {}
+extension OneElementFastSequence.Base: Hashable where Element: Hashable {}
+
+extension OneElementFastSequence: Sendable where Element: Sendable {}
+extension OneElementFastSequence.Base: Sendable where Element: Sendable {}

Tests/ConnectionPoolModuleTests/OneElementFastSequence.swift

@@ -0,0 +1,70 @@
+@testable import _ConnectionPoolModule
+import XCTest
+
+final class OneElementFastSequenceTests: XCTestCase {
+    func testCountIsEmptyAndIterator() async {
+        var sequence = OneElementFastSequence<Int>()
+        XCTAssertEqual(sequence.count, 0)
+        XCTAssertEqual(sequence.isEmpty, true)
+        XCTAssertEqual(sequence.first, nil)
+        XCTAssertEqual(Array(sequence), [])
+        sequence.append(1)
+        XCTAssertEqual(sequence.count, 1)
+        XCTAssertEqual(sequence.isEmpty, false)
+        XCTAssertEqual(sequence.first, 1)
+        XCTAssertEqual(Array(sequence), [1])
+        sequence.append(2)
+        XCTAssertEqual(sequence.count, 2)
+        XCTAssertEqual(sequence.isEmpty, false)
+        XCTAssertEqual(sequence.first, 1)
+        XCTAssertEqual(Array(sequence), [1, 2])
+        sequence.append(3)
+        XCTAssertEqual(sequence.count, 3)
+        XCTAssertEqual(sequence.isEmpty, false)
+        XCTAssertEqual(sequence.first, 1)
+        XCTAssertEqual(Array(sequence), [1, 2, 3])
+    }
+
+    func testReserveCapacityIsForwarded() {
+        var emptySequence = OneElementFastSequence<Int>()
+        emptySequence.reserveCapacity(8)
+        emptySequence.append(1)
+        emptySequence.append(2)
+        guard case .n(let array) = emptySequence.base else {
+            return XCTFail("Expected sequence to be backed by an array")
+        }
+        XCTAssertEqual(array.capacity, 8)
+
+        var oneElemSequence = OneElementFastSequence<Int>(1)
+        oneElemSequence.reserveCapacity(8)
+        oneElemSequence.append(2)
+        guard case .n(let array) = oneElemSequence.base else {
+            return XCTFail("Expected sequence to be backed by an array")
+        }
+        XCTAssertEqual(array.capacity, 8)
+
+        var twoElemSequence = OneElementFastSequence<Int>([1, 2])
+        twoElemSequence.reserveCapacity(8)
+        guard case .n(let array) = twoElemSequence.base else {
+            return XCTFail("Expected sequence to be backed by an array")
+        }
+        XCTAssertEqual(array.capacity, 8)
+    }
+
+    func testNewSequenceSlowPath() {
+        let sequence = OneElementFastSequence<UInt8>("AB".utf8)
+        XCTAssertEqual(Array(sequence), [UInt8(ascii: "A"), UInt8(ascii: "B")])
+    }
+
+    func testSingleItem() {
+        let sequence = OneElementFastSequence<UInt8>("A".utf8)
+        XCTAssertEqual(Array(sequence), [UInt8(ascii: "A")])
+    }
+
+    func testEmptyCollection() {
+        let sequence = OneElementFastSequence<UInt8>("".utf8)
+        XCTAssertTrue(sequence.isEmpty)
+        XCTAssertEqual(sequence.count, 0)
+        XCTAssertEqual(Array(sequence), [])
+    }
+}