diff --git a/Sources/ConnectionPoolModule/OneElementFastSequence.swift b/Sources/ConnectionPoolModule/OneElementFastSequence.swift new file mode 100644 index 00000000..1bb3b8e4 --- /dev/null +++ b/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: 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) { + 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 { + 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 + + @inlinable + init(_ backing: OneElementFastSequence) { + 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 {} diff --git a/Tests/ConnectionPoolModuleTests/OneElementFastSequence.swift b/Tests/ConnectionPoolModuleTests/OneElementFastSequence.swift new file mode 100644 index 00000000..8098438f --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/OneElementFastSequence.swift @@ -0,0 +1,70 @@ +@testable import _ConnectionPoolModule +import XCTest + +final class OneElementFastSequenceTests: XCTestCase { + func testCountIsEmptyAndIterator() async { + var sequence = OneElementFastSequence() + 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() + 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(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([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("AB".utf8) + XCTAssertEqual(Array(sequence), [UInt8(ascii: "A"), UInt8(ascii: "B")]) + } + + func testSingleItem() { + let sequence = OneElementFastSequence("A".utf8) + XCTAssertEqual(Array(sequence), [UInt8(ascii: "A")]) + } + + func testEmptyCollection() { + let sequence = OneElementFastSequence("".utf8) + XCTAssertTrue(sequence.isEmpty) + XCTAssertEqual(sequence.count, 0) + XCTAssertEqual(Array(sequence), []) + } +}