Skip to content

Commit d4bdb34

Browse files
authored
Add a no-separator variant of joined() to match the stdlib (apple#87)
* Add a no-separator variant of joined() to match the stdlib * @inlinable and @Frozen
1 parent 11d7cfe commit d4bdb34

File tree

3 files changed

+268
-89
lines changed

3 files changed

+268
-89
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the Swift Async Algorithms open source project
4+
//
5+
// Copyright (c) 2022 Apple Inc. and the Swift project authors
6+
// Licensed under Apache License v2.0 with Runtime Library Exception
7+
//
8+
// See https://swift.org/LICENSE.txt for license information
9+
//
10+
//===----------------------------------------------------------------------===//
11+
12+
extension AsyncSequence where Element: AsyncSequence {
13+
@inlinable
14+
public func joined<Separator: AsyncSequence>(separator: Separator) -> AsyncJoinedBySeparatorSequence<Self, Separator> {
15+
return AsyncJoinedBySeparatorSequence(self, separator: separator)
16+
}
17+
}
18+
19+
public struct AsyncJoinedBySeparatorSequence<Base: AsyncSequence, Separator: AsyncSequence>: AsyncSequence where Base.Element: AsyncSequence, Separator.Element == Base.Element.Element {
20+
public typealias Element = Base.Element.Element
21+
public typealias AsyncIterator = Iterator
22+
23+
public struct Iterator: AsyncIteratorProtocol {
24+
@usableFromInline
25+
enum State {
26+
@usableFromInline
27+
enum SeparatorState {
28+
case initial(Separator)
29+
case partialAsync(Separator.AsyncIterator, ContiguousArray<Element>)
30+
case cached(ContiguousArray<Element>)
31+
case partialCached(ContiguousArray<Element>.Iterator, ContiguousArray<Element>)
32+
33+
@usableFromInline
34+
func startSeparator() -> SeparatorState {
35+
switch self {
36+
case .initial(let separatorSequence):
37+
return .partialAsync(separatorSequence.makeAsyncIterator(), [])
38+
case .cached(let array):
39+
return .partialCached(array.makeIterator(), array)
40+
default:
41+
fatalError("Invalid separator sequence state")
42+
}
43+
}
44+
45+
@usableFromInline
46+
func next() async rethrows -> (Element?, SeparatorState) {
47+
switch self {
48+
case .partialAsync(var separatorIterator, var cache):
49+
guard let next = try await separatorIterator.next() else {
50+
return (nil, .cached(cache))
51+
}
52+
cache.append(next)
53+
return (next, .partialAsync(separatorIterator, cache))
54+
case .partialCached(var cacheIterator, let cache):
55+
guard let next = cacheIterator.next() else {
56+
return (nil, .cached(cache))
57+
}
58+
return (next, .partialCached(cacheIterator, cache))
59+
default:
60+
fatalError("Invalid separator sequence state")
61+
}
62+
}
63+
}
64+
65+
case initial(Base.AsyncIterator, Separator)
66+
case sequence(Base.AsyncIterator, Base.Element.AsyncIterator, SeparatorState)
67+
case separator(Base.AsyncIterator, SeparatorState, Base.Element)
68+
case terminal
69+
}
70+
71+
@usableFromInline
72+
var state: State
73+
74+
@inlinable
75+
init(_ iterator: Base.AsyncIterator, separator: Separator) {
76+
state = .initial(iterator, separator)
77+
}
78+
79+
@inlinable
80+
public mutating func next() async rethrows -> Base.Element.Element? {
81+
do {
82+
switch state {
83+
case .terminal:
84+
return nil
85+
case .initial(var outerIterator, let separator):
86+
guard let innerSequence = try await outerIterator.next() else {
87+
state = .terminal
88+
return nil
89+
}
90+
let innerIterator = innerSequence.makeAsyncIterator()
91+
state = .sequence(outerIterator, innerIterator, .initial(separator))
92+
return try await next()
93+
case .sequence(var outerIterator, var innerIterator, let separatorState):
94+
if let item = try await innerIterator.next() {
95+
state = .sequence(outerIterator, innerIterator, separatorState)
96+
return item
97+
}
98+
99+
guard let nextInner = try await outerIterator.next() else {
100+
state = .terminal
101+
return nil
102+
}
103+
104+
state = .separator(outerIterator, separatorState.startSeparator(), nextInner)
105+
return try await next()
106+
case .separator(let iterator, let separatorState, let nextBase):
107+
let (itemOpt, newSepState) = try await separatorState.next()
108+
guard let item = itemOpt else {
109+
state = .sequence(iterator, nextBase.makeAsyncIterator(), newSepState)
110+
return try await next()
111+
}
112+
state = .separator(iterator, newSepState, nextBase)
113+
return item
114+
}
115+
} catch {
116+
state = .terminal
117+
throw error
118+
}
119+
}
120+
}
121+
122+
@usableFromInline
123+
let base: Base
124+
125+
@usableFromInline
126+
let separator: Separator
127+
128+
@inlinable
129+
init(_ base: Base, separator: Separator) {
130+
self.base = base
131+
self.separator = separator
132+
}
133+
134+
@inlinable
135+
public func makeAsyncIterator() -> Iterator {
136+
return Iterator(base.makeAsyncIterator(), separator: separator)
137+
}
138+
}
139+
140+
extension AsyncJoinedBySeparatorSequence: Sendable
141+
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Separator: Sendable, Separator.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }
142+
extension AsyncJoinedBySeparatorSequence.Iterator: Sendable
143+
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Separator: Sendable, Separator.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }
144+
extension AsyncJoinedBySeparatorSequence.Iterator.State: Sendable
145+
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Separator: Sendable, Separator.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }
146+
extension AsyncJoinedBySeparatorSequence.Iterator.State.SeparatorState: Sendable
147+
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Separator: Sendable, Separator.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }

Sources/AsyncAlgorithms/AsyncJoinedSequence.swift

+36-88
Original file line numberDiff line numberDiff line change
@@ -11,137 +11,85 @@
1111

1212
extension AsyncSequence where Element: AsyncSequence {
1313
@inlinable
14-
public func joined<Separator: AsyncSequence>(separator: Separator) -> AsyncJoinedSequence<Self, Separator> {
15-
return AsyncJoinedSequence(self, separator: separator)
14+
public func joined() -> AsyncJoinedSequence<Self> {
15+
return AsyncJoinedSequence(self)
1616
}
1717
}
1818

19-
public struct AsyncJoinedSequence<Base: AsyncSequence, Separator: AsyncSequence>: AsyncSequence where Base.Element: AsyncSequence, Separator.Element == Base.Element.Element {
19+
@frozen
20+
public struct AsyncJoinedSequence<Base: AsyncSequence>: AsyncSequence where Base.Element: AsyncSequence {
2021
public typealias Element = Base.Element.Element
2122
public typealias AsyncIterator = Iterator
2223

24+
@frozen
2325
public struct Iterator: AsyncIteratorProtocol {
2426
@usableFromInline
2527
enum State {
26-
@usableFromInline
27-
enum SeparatorState {
28-
case initial(Separator)
29-
case partialAsync(Separator.AsyncIterator, ContiguousArray<Element>)
30-
case cached(ContiguousArray<Element>)
31-
case partialCached(ContiguousArray<Element>.Iterator, ContiguousArray<Element>)
32-
33-
@usableFromInline
34-
func startSeparator() -> SeparatorState {
35-
switch self {
36-
case .initial(let separatorSequence):
37-
return .partialAsync(separatorSequence.makeAsyncIterator(), [])
38-
case .cached(let array):
39-
return .partialCached(array.makeIterator(), array)
40-
default:
41-
fatalError("Invalid separator sequence state")
42-
}
43-
}
44-
45-
@usableFromInline
46-
func next() async rethrows -> (Element?, SeparatorState) {
47-
switch self {
48-
case .partialAsync(var separatorIterator, var cache):
49-
guard let next = try await separatorIterator.next() else {
50-
return (nil, .cached(cache))
51-
}
52-
cache.append(next)
53-
return (next, .partialAsync(separatorIterator, cache))
54-
case .partialCached(var cacheIterator, let cache):
55-
guard let next = cacheIterator.next() else {
56-
return (nil, .cached(cache))
57-
}
58-
return (next, .partialCached(cacheIterator, cache))
59-
default:
60-
fatalError("Invalid separator sequence state")
61-
}
62-
}
63-
}
64-
65-
case initial(Base.AsyncIterator, Separator)
66-
case sequence(Base.AsyncIterator, Base.Element.AsyncIterator, SeparatorState)
67-
case separator(Base.AsyncIterator, SeparatorState, Base.Element)
28+
case initial(Base.AsyncIterator)
29+
case sequence(Base.AsyncIterator, Base.Element.AsyncIterator)
6830
case terminal
6931
}
70-
32+
7133
@usableFromInline
7234
var state: State
73-
74-
@usableFromInline
75-
init(_ iterator: Base.AsyncIterator, separator: Separator) {
76-
state = .initial(iterator, separator)
35+
36+
@inlinable
37+
init(_ iterator: Base.AsyncIterator) {
38+
state = .initial(iterator)
7739
}
78-
40+
7941
@inlinable
8042
public mutating func next() async rethrows -> Base.Element.Element? {
8143
do {
8244
switch state {
8345
case .terminal:
8446
return nil
85-
case .initial(var outerIterator, let separator):
47+
case .initial(var outerIterator):
8648
guard let innerSequence = try await outerIterator.next() else {
8749
state = .terminal
8850
return nil
8951
}
9052
let innerIterator = innerSequence.makeAsyncIterator()
91-
state = .sequence(outerIterator, innerIterator, .initial(separator))
53+
state = .sequence(outerIterator, innerIterator)
9254
return try await next()
93-
case .sequence(var outerIterator, var innerIterator, let separatorState):
94-
if let item = try await innerIterator.next() {
95-
state = .sequence(outerIterator, innerIterator, separatorState)
96-
return item
97-
}
98-
99-
guard let nextInner = try await outerIterator.next() else {
100-
state = .terminal
101-
return nil
102-
}
103-
104-
state = .separator(outerIterator, separatorState.startSeparator(), nextInner)
105-
return try await next()
106-
case .separator(let iterator, let separatorState, let nextBase):
107-
let (itemOpt, newSepState) = try await separatorState.next()
108-
guard let item = itemOpt else {
109-
state = .sequence(iterator, nextBase.makeAsyncIterator(), newSepState)
110-
return try await next()
55+
case .sequence(var outerIterator, var innerIterator):
56+
if let item = try await innerIterator.next() {
57+
state = .sequence(outerIterator, innerIterator)
58+
return item
11159
}
112-
state = .separator(iterator, newSepState, nextBase)
113-
return item
60+
61+
guard let nextInner = try await outerIterator.next() else {
62+
state = .terminal
63+
return nil
64+
}
65+
66+
state = .sequence(outerIterator, nextInner.makeAsyncIterator())
67+
return try await next()
11468
}
11569
} catch {
11670
state = .terminal
11771
throw error
11872
}
11973
}
12074
}
121-
75+
12276
@usableFromInline
12377
let base: Base
124-
125-
@usableFromInline
126-
let separator: Separator
127-
128-
@usableFromInline
129-
init(_ base: Base, separator: Separator) {
78+
79+
@inlinable
80+
init(_ base: Base) {
13081
self.base = base
131-
self.separator = separator
13282
}
133-
83+
13484
@inlinable
13585
public func makeAsyncIterator() -> Iterator {
136-
return Iterator(base.makeAsyncIterator(), separator: separator)
86+
return Iterator(base.makeAsyncIterator())
13787
}
13888
}
13989

14090
extension AsyncJoinedSequence: Sendable
141-
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Separator: Sendable, Separator.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }
91+
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }
14292
extension AsyncJoinedSequence.Iterator: Sendable
143-
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Separator: Sendable, Separator.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }
93+
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }
14494
extension AsyncJoinedSequence.Iterator.State: Sendable
145-
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Separator: Sendable, Separator.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }
146-
extension AsyncJoinedSequence.Iterator.State.SeparatorState: Sendable
147-
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Separator: Sendable, Separator.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }
95+
where Base: Sendable, Base.Element: Sendable, Base.Element.Element: Sendable, Base.AsyncIterator: Sendable, Base.Element.AsyncIterator: Sendable { }

0 commit comments

Comments
 (0)