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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
#![no_std]
#![doc = include_str!("../README.md")]

use core::pin::Pin;
use core::task::{Context, Poll};

/// A stream that produces items that are ordered according to some token.
///
/// The main advantage of this trait over the standard `Stream` trait is the ability to implement a
/// [`join`](join()) function that does not either block until both source streams produce an item
/// or contain a race condition when rejoining streams that originated from a common well-ordered
/// source.
pub trait OrderedStream {
    /// The type ordered by this stream.
    ///
    /// Each stream must produce values that are in ascending order according to this function,
    /// although there is no requirement that the values be strictly ascending.
    type Ordering: Ord;

    /// The unordered data carried by this stream
    ///
    /// This is split from the `Ordering` type to allow specifying a smaller or cheaper-to-generate
    /// type as the ordering key.  This is especially useful if you generate values to pass in to
    /// `before`.
    type Data;

    /// Attempt to pull out the next value of this stream, registering the current task for wakeup
    /// if needed, and returning `NoneBefore` if it is known that the stream will not produce any
    /// more values ordered before the given point.
    ///
    /// # Return value
    ///
    /// There are several possible return values, each indicating a distinct stream state depending
    /// on the value passed in `before`:
    ///
    /// - If `before` was `None`, `Poll::Pending` means that this stream's next value is not ready
    /// yet. Implementations will ensure that the current task is notified when the next value may
    /// be ready.
    ///
    /// - If `before` was `Some`, `Poll::Pending` means that this stream's next value is not ready
    /// and that it is not yet known if the stream will produce a value ordered prior to the given
    /// ordering value.  Implementations will ensure that the current task is notified when either
    /// the next value is ready or once it is known that no such value will be produced.
    ///
    /// - `Poll::Ready(PollResult::Item)` means that the stream has successfully produced
    /// an item.  The stream may produce further values on subsequent `poll_next_before` calls.
    /// The returned ordering value must not be less than any prior ordering value returned by this
    /// stream.  The returned ordering value **may** be greater than the value passed to `before`.
    ///
    /// - `Poll::Ready(PollResult::Terminated)` means that the stream has terminated, and
    /// `poll_next_before` should not be invoked again.
    ///
    /// - `Poll::Ready(PollResult::NoneBefore)` means that the stream will not produce
    /// any further ordering tokens less than the given token.  Subsequent `poll_next_before` calls
    /// may still produce additional items, but their tokens will be greater than or equal to the
    /// given token.  It does not make sense to return this value if `before` was `None`.
    fn poll_next_before(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        before: Option<&Self::Ordering>,
    ) -> Poll<PollResult<Self::Ordering, Self::Data>>;

    /// The minimum value of the ordering for any future items.
    ///
    /// If this does not return `None`, the returned ordering must be less than or equal to the
    /// ordering of any future item returned from [`Self::poll_next_before`].  This value should
    /// (but is not required to) be greater than or equal to the ordering of the most recent item
    /// returned.
    fn position_hint(&self) -> Option<MaybeBorrowed<'_, Self::Ordering>> {
        None
    }

    /// Returns the bounds on the remaining length of the stream.
    fn size_hint(&self) -> (usize, Option<usize>) {
        (0, None)
    }
}

/// A value that is either borrowed or owned.
///
/// This is similar to `std::borrow::Cow`, but does not require the ability to convert from
/// borrowed to owned.
#[derive(Debug)]
pub enum MaybeBorrowed<'a, T> {
    Borrowed(&'a T),
    Owned(T),
}

impl<'a, T> AsRef<T> for MaybeBorrowed<'a, T> {
    fn as_ref(&self) -> &T {
        match self {
            Self::Borrowed(t) => t,
            Self::Owned(t) => t,
        }
    }
}

impl<'a, T> core::ops::Deref for MaybeBorrowed<'a, T> {
    type Target = T;

    fn deref(&self) -> &T {
        match self {
            Self::Borrowed(t) => t,
            Self::Owned(t) => t,
        }
    }
}

impl<P> OrderedStream for Pin<P>
where
    P: core::ops::DerefMut + Unpin,
    P::Target: OrderedStream,
{
    type Data = <P::Target as OrderedStream>::Data;
    type Ordering = <P::Target as OrderedStream>::Ordering;

    fn poll_next_before(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        before: Option<&Self::Ordering>,
    ) -> Poll<PollResult<Self::Ordering, Self::Data>> {
        self.get_mut().as_mut().poll_next_before(cx, before)
    }

    fn position_hint(&self) -> Option<MaybeBorrowed<'_, Self::Ordering>> {
        (**self).position_hint()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (**self).size_hint()
    }
}

impl<S> OrderedStream for Option<S>
where
    S: OrderedStream,
{
    type Data = S::Data;
    type Ordering = S::Ordering;

    fn poll_next_before(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        before: Option<&Self::Ordering>,
    ) -> Poll<PollResult<Self::Ordering, Self::Data>> {
        match self.as_pin_mut() {
            Some(s) => s.poll_next_before(cx, before),
            None => Poll::Ready(PollResult::Terminated),
        }
    }

    fn position_hint(&self) -> Option<MaybeBorrowed<'_, Self::Ordering>> {
        self.as_ref().and_then(|s| s.position_hint())
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.as_ref().map_or((0, Some(0)), |s| s.size_hint())
    }
}

/// An [`OrderedStream`] that tracks if the underlying stream should be polled.
pub trait FusedOrderedStream: OrderedStream {
    /// Returns `true` if the stream should no longer be polled.
    fn is_terminated(&self) -> bool;
}

impl<P> FusedOrderedStream for Pin<P>
where
    P: core::ops::DerefMut + Unpin,
    P::Target: FusedOrderedStream,
{
    fn is_terminated(&self) -> bool {
        (**self).is_terminated()
    }
}

impl<S> FusedOrderedStream for Option<S>
where
    S: FusedOrderedStream,
{
    fn is_terminated(&self) -> bool {
        self.as_ref().map_or(true, |s| s.is_terminated())
    }
}

/// The result of a [`OrderedStream::poll_next_before`] operation.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum PollResult<Ordering, Data> {
    /// An item with a corresponding ordering token.
    Item { data: Data, ordering: Ordering },
    /// This stream will not return any items prior to the given point.
    NoneBefore,
    /// This stream is terminated and should not be polled again.
    Terminated,
}

impl<D, T> PollResult<T, D> {
    /// Extract the data from the result.
    pub fn into_data(self) -> Option<D> {
        match self {
            Self::Item { data, .. } => Some(data),
            _ => None,
        }
    }

    /// Extract the item from the result.
    pub fn into_tuple(self) -> Option<(T, D)> {
        match self {
            Self::Item { data, ordering } => Some((ordering, data)),
            _ => None,
        }
    }

    /// Apply a closure to the data.
    pub fn map_data<R>(self, f: impl FnOnce(D) -> R) -> PollResult<T, R> {
        match self {
            Self::Item { data, ordering } => PollResult::Item {
                data: f(data),
                ordering,
            },
            Self::NoneBefore => PollResult::NoneBefore,
            Self::Terminated => PollResult::Terminated,
        }
    }
}

impl<T, D, E> PollResult<T, Result<D, E>> {
    /// Extract the error of a [`Result`] item.
    pub fn transpose_result(self) -> Result<PollResult<T, D>, E> {
        self.transpose_result_item().map_err(|(_, e)| e)
    }

    /// Extract the error and ordering from a [`Result`] item.
    pub fn transpose_result_item(self) -> Result<PollResult<T, D>, (T, E)> {
        match self {
            Self::Item {
                data: Ok(data),
                ordering,
            } => Ok(PollResult::Item { data, ordering }),
            Self::Item {
                data: Err(data),
                ordering,
            } => Err((ordering, data)),
            Self::NoneBefore => Ok(PollResult::NoneBefore),
            Self::Terminated => Ok(PollResult::Terminated),
        }
    }
}

/// A [`Future`](core::future::Future) that produces an item with an associated ordering.
///
/// This is equivalent to an [`OrderedStream`] that always produces exactly one item.  This trait
/// is not very useful on its own; see [`FromFuture`] to convert it to a stream.
///
/// It is valid to implement both [`Future`](core::future::Future) and [`OrderedFuture`] on the
/// same type.  In this case, unless otherwise documented by the implementing type, neither poll
/// function should be invoked after either returns an output value.
pub trait OrderedFuture {
    /// See [`OrderedStream::Ordering`].
    type Ordering: Ord;

    /// See [`OrderedStream::Data`].
    type Output;

    /// Attempt to pull out the value of this future, registering the current task for wakeup if
    /// needed, and returning `None` if it is known that the future will not produce a value
    /// ordered before the given point.
    ///
    /// # Return value
    ///
    /// There are several possible return values, each indicating a distinct state depending on the
    /// value passed in `before`:
    ///
    /// - If `before` was `None`, `Poll::Pending` means that this future's value is not ready yet.
    /// Implementations will ensure that the current task is notified when the next value may be
    /// ready.
    ///
    /// - If `before` was `Some`, `Poll::Pending` means that this future's value is not ready and
    /// that it is not yet known if the value will be ordered prior to the given ordering value.
    /// Implementations will ensure that the current task is notified when either the next value is
    /// ready or once it is known that no such value will be produced.
    ///
    /// - `Poll::Ready(Some(Data))` means that the future has successfully terminated.  The
    /// returned ordering value **may** be greater than the value passed to `before`.  The
    /// `poll_before` function should not be invoked again.
    ///
    /// - `Poll::Ready(None)` means that this future will not produce an ordering token less than
    /// the given token.  It is an error to return `None` if `before` was `None`.
    fn poll_before(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        before: Option<&Self::Ordering>,
    ) -> Poll<Option<(Self::Ordering, Self::Output)>>;

    /// The minimum value of the ordering of the item.
    ///
    /// See [`OrderedStream::position_hint`] for details.
    fn position_hint(&self) -> Option<MaybeBorrowed<'_, Self::Ordering>> {
        None
    }
}

mod adapters;
pub use adapters::*;
mod join;
pub use join::*;
mod multi;
pub use multi::*;