-
Notifications
You must be signed in to change notification settings - Fork 10.3k
/
AsyncPrefixSequence.swift
148 lines (133 loc) · 5 KB
/
AsyncPrefixSequence.swift
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2021 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
import Swift
@available(SwiftStdlib 5.1, *)
extension AsyncSequence {
/// Returns an asynchronous sequence, up to the specified maximum length,
/// containing the initial elements of the base asynchronous sequence.
///
/// Use `prefix(_:)` to reduce the number of elements produced by the
/// asynchronous sequence.
///
/// In this example, an asynchronous sequence called `Counter` produces `Int`
/// values from `1` to `10`. The `prefix(_:)` method causes the modified
/// sequence to pass through the first six values, then end.
///
/// for await number in Counter(howHigh: 10).prefix(6) {
/// print(number, terminator: " ")
/// }
/// // Prints "1 2 3 4 5 6 "
///
/// If the count passed to `prefix(_:)` exceeds the number of elements in the
/// base sequence, the result contains all of the elements in the sequence.
///
/// - Parameter count: The maximum number of elements to return. The value of
/// `count` must be greater than or equal to zero.
/// - Returns: An asynchronous sequence starting at the beginning of the
/// base sequence with at most `count` elements.
@inlinable
public __consuming func prefix(
_ count: Int
) -> AsyncPrefixSequence<Self> {
precondition(count >= 0,
"Can't prefix a negative number of elements from an async sequence")
return AsyncPrefixSequence(self, count: count)
}
}
/// An asynchronous sequence, up to a specified maximum length,
/// containing the initial elements of a base asynchronous sequence.
@available(SwiftStdlib 5.1, *)
public struct AsyncPrefixSequence<Base: AsyncSequence> {
@usableFromInline
let base: Base
@usableFromInline
let count: Int
@usableFromInline
init(_ base: Base, count: Int) {
self.base = base
self.count = count
}
}
@available(SwiftStdlib 5.1, *)
extension AsyncPrefixSequence: AsyncSequence {
/// The type of element produced by this asynchronous sequence.
///
/// The prefix sequence produces whatever type of element its base iterator
/// produces.
public typealias Element = Base.Element
/// The type of the error that can be produced by the sequence.
///
/// The prefix sequence produces whatever type of error its
/// base sequence does.
@available(SwiftStdlib 6.0, *)
public typealias Failure = Base.Failure
/// The type of iterator that produces elements of the sequence.
public typealias AsyncIterator = Iterator
/// The iterator that produces elements of the prefix sequence.
public struct Iterator: AsyncIteratorProtocol {
// FIXME: Remove when $AssociatedTypeImplements is no longer needed
@available(SwiftStdlib 6.0, *)
public typealias Failure = Base.Failure
@usableFromInline
var baseIterator: Base.AsyncIterator
@usableFromInline
var remaining: Int
@usableFromInline
init(_ baseIterator: Base.AsyncIterator, count: Int) {
self.baseIterator = baseIterator
self.remaining = count
}
/// Produces the next element in the prefix sequence.
///
/// Until reaching the number of elements to include, this iterator calls
/// `next()` on its base iterator and passes through the result. After
/// reaching the maximum number of elements, subsequent calls to `next()`
/// return `nil`.
@inlinable
public mutating func next() async rethrows -> Base.Element? {
if remaining != 0 {
remaining &-= 1
return try await baseIterator.next()
} else {
return nil
}
}
/// Produces the next element in the prefix sequence.
///
/// Until reaching the number of elements to include, this iterator calls
/// `next(isolation:)` on its base iterator and passes through the
/// result. After reaching the maximum number of elements, subsequent calls
/// to `next(isolation:)` return `nil`.
@available(SwiftStdlib 6.0, *)
@inlinable
public mutating func next(isolation actor: isolated (any Actor)?) async throws(Failure) -> Base.Element? {
if remaining != 0 {
remaining &-= 1
return try await baseIterator.next(isolation: actor)
} else {
return nil
}
}
}
@inlinable
public __consuming func makeAsyncIterator() -> Iterator {
return Iterator(base.makeAsyncIterator(), count: count)
}
}
@available(SwiftStdlib 5.1, *)
extension AsyncPrefixSequence: Sendable
where Base: Sendable,
Base.Element: Sendable { }
@available(SwiftStdlib 5.1, *)
extension AsyncPrefixSequence.Iterator: Sendable
where Base.AsyncIterator: Sendable,
Base.Element: Sendable { }