diff --git a/Sources/ConnectionPoolModule/Max2Sequence.swift b/Sources/ConnectionPoolModule/Max2Sequence.swift new file mode 100644 index 00000000..6c330067 --- /dev/null +++ b/Sources/ConnectionPoolModule/Max2Sequence.swift @@ -0,0 +1,95 @@ +// A `Sequence` that can contain at most two elements. However it does not heap allocate. +@usableFromInline +struct Max2Sequence: Sequence { + @usableFromInline + private(set) var first: Element? + @usableFromInline + private(set) var second: Element? + + @inlinable + var count: Int { + if self.first == nil { return 0 } + if self.second == nil { return 1 } + return 2 + } + + @inlinable + var isEmpty: Bool { + self.first == nil + } + + @inlinable + init(_ first: Element?, _ second: Element? = nil) { + if let first = first { + self.first = first + self.second = second + } else { + self.first = second + self.second = nil + } + } + + @inlinable + init() { + self.first = nil + self.second = nil + } + + @inlinable + func makeIterator() -> Iterator { + Iterator(first: self.first, second: self.second) + } + + @usableFromInline + struct Iterator: IteratorProtocol { + @usableFromInline + let first: Element? + @usableFromInline + let second: Element? + + @usableFromInline + private(set) var index: UInt8 = 0 + + @inlinable + init(first: Element?, second: Element?) { + self.first = first + self.second = second + self.index = 0 + } + + @inlinable + mutating func next() -> Element? { + switch self.index { + case 0: + self.index += 1 + return self.first + case 1: + self.index += 1 + return self.second + default: + return nil + } + } + } + + @inlinable + mutating func append(_ element: Element) { + precondition(self.second == nil) + if self.first == nil { + self.first = element + } else if self.second == nil { + self.second = element + } else { + fatalError("Max2Sequence can only hold two Elements.") + } + } + + @inlinable + func map(_ transform: (Element) throws -> (NewElement)) rethrows -> Max2Sequence { + try Max2Sequence(self.first.flatMap(transform), self.second.flatMap(transform)) + } +} + +extension Max2Sequence: Equatable where Element: Equatable {} +extension Max2Sequence: Hashable where Element: Hashable {} +extension Max2Sequence: Sendable where Element: Sendable {} diff --git a/Tests/ConnectionPoolModuleTests/Max2SequenceTests.swift b/Tests/ConnectionPoolModuleTests/Max2SequenceTests.swift new file mode 100644 index 00000000..081e867b --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/Max2SequenceTests.swift @@ -0,0 +1,60 @@ +@testable import _ConnectionPoolModule +import XCTest + +final class Max2SequenceTests: XCTestCase { + func testCountAndIsEmpty() async { + var sequence = Max2Sequence() + XCTAssertEqual(sequence.count, 0) + XCTAssertEqual(sequence.isEmpty, true) + sequence.append(1) + XCTAssertEqual(sequence.count, 1) + XCTAssertEqual(sequence.isEmpty, false) + sequence.append(2) + XCTAssertEqual(sequence.count, 2) + XCTAssertEqual(sequence.isEmpty, false) + } + + func testOptionalInitializer() { + let emptySequence = Max2Sequence(nil, nil) + XCTAssertEqual(emptySequence.count, 0) + XCTAssertEqual(emptySequence.isEmpty, true) + var emptySequenceIterator = emptySequence.makeIterator() + XCTAssertNil(emptySequenceIterator.next()) + XCTAssertNil(emptySequenceIterator.next()) + XCTAssertNil(emptySequenceIterator.next()) + + let oneElemSequence1 = Max2Sequence(1, nil) + XCTAssertEqual(oneElemSequence1.count, 1) + XCTAssertEqual(oneElemSequence1.isEmpty, false) + var oneElemSequence1Iterator = oneElemSequence1.makeIterator() + XCTAssertEqual(oneElemSequence1Iterator.next(), 1) + XCTAssertNil(oneElemSequence1Iterator.next()) + XCTAssertNil(oneElemSequence1Iterator.next()) + + let oneElemSequence2 = Max2Sequence(nil, 2) + XCTAssertEqual(oneElemSequence2.count, 1) + XCTAssertEqual(oneElemSequence2.isEmpty, false) + var oneElemSequence2Iterator = oneElemSequence2.makeIterator() + XCTAssertEqual(oneElemSequence2Iterator.next(), 2) + XCTAssertNil(oneElemSequence2Iterator.next()) + XCTAssertNil(oneElemSequence2Iterator.next()) + + let twoElemSequence = Max2Sequence(1, 2) + XCTAssertEqual(twoElemSequence.count, 2) + XCTAssertEqual(twoElemSequence.isEmpty, false) + var twoElemSequenceIterator = twoElemSequence.makeIterator() + XCTAssertEqual(twoElemSequenceIterator.next(), 1) + XCTAssertEqual(twoElemSequenceIterator.next(), 2) + XCTAssertNil(twoElemSequenceIterator.next()) + } + + func testMap() { + let twoElemSequence = Max2Sequence(1, 2).map({ "\($0)" }) + XCTAssertEqual(twoElemSequence.count, 2) + XCTAssertEqual(twoElemSequence.isEmpty, false) + var twoElemSequenceIterator = twoElemSequence.makeIterator() + XCTAssertEqual(twoElemSequenceIterator.next(), "1") + XCTAssertEqual(twoElemSequenceIterator.next(), "2") + XCTAssertNil(twoElemSequenceIterator.next()) + } +}