diff options
Diffstat (limited to 'third_party/rust/futures-util/src/stream')
78 files changed, 11491 insertions, 0 deletions
diff --git a/third_party/rust/futures-util/src/stream/abortable.rs b/third_party/rust/futures-util/src/stream/abortable.rs new file mode 100644 index 0000000000..1fea895822 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/abortable.rs @@ -0,0 +1,19 @@ +use super::assert_stream; +use crate::stream::{AbortHandle, Abortable}; +use crate::Stream; + +/// Creates a new `Abortable` stream and an `AbortHandle` which can be used to stop it. +/// +/// This function is a convenient (but less flexible) alternative to calling +/// `AbortHandle::new` and `Abortable::new` manually. +/// +/// This function is only available when the `std` or `alloc` feature of this +/// library is activated, and it is activated by default. +pub fn abortable<St>(stream: St) -> (Abortable<St>, AbortHandle) +where + St: Stream, +{ + let (handle, reg) = AbortHandle::new_pair(); + let abortable = assert_stream::<St::Item, _>(Abortable::new(stream, reg)); + (abortable, handle) +} diff --git a/third_party/rust/futures-util/src/stream/empty.rs b/third_party/rust/futures-util/src/stream/empty.rs new file mode 100644 index 0000000000..e4fd87326b --- /dev/null +++ b/third_party/rust/futures-util/src/stream/empty.rs @@ -0,0 +1,45 @@ +use super::assert_stream; +use core::marker::PhantomData; +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; + +/// Stream for the [`empty`] function. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Empty<T> { + _phantom: PhantomData<T>, +} + +/// Creates a stream which contains no elements. +/// +/// The returned stream will always return `Ready(None)` when polled. +pub fn empty<T>() -> Empty<T> { + assert_stream::<T, _>(Empty { _phantom: PhantomData }) +} + +impl<T> Unpin for Empty<T> {} + +impl<T> FusedStream for Empty<T> { + fn is_terminated(&self) -> bool { + true + } +} + +impl<T> Stream for Empty<T> { + type Item = T; + + fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> { + Poll::Ready(None) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (0, Some(0)) + } +} + +impl<T> Clone for Empty<T> { + fn clone(&self) -> Self { + empty() + } +} diff --git a/third_party/rust/futures-util/src/stream/futures_ordered.rs b/third_party/rust/futures-util/src/stream/futures_ordered.rs new file mode 100644 index 0000000000..618bf1b7bd --- /dev/null +++ b/third_party/rust/futures-util/src/stream/futures_ordered.rs @@ -0,0 +1,244 @@ +use crate::stream::{FuturesUnordered, StreamExt}; +use alloc::collections::binary_heap::{BinaryHeap, PeekMut}; +use core::cmp::Ordering; +use core::fmt::{self, Debug}; +use core::iter::FromIterator; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::Stream; +use futures_core::{ + task::{Context, Poll}, + FusedStream, +}; +use pin_project_lite::pin_project; + +pin_project! { + #[must_use = "futures do nothing unless you `.await` or poll them"] + #[derive(Debug)] + struct OrderWrapper<T> { + #[pin] + data: T, // A future or a future's output + index: isize, + } +} + +impl<T> PartialEq for OrderWrapper<T> { + fn eq(&self, other: &Self) -> bool { + self.index == other.index + } +} + +impl<T> Eq for OrderWrapper<T> {} + +impl<T> PartialOrd for OrderWrapper<T> { + fn partial_cmp(&self, other: &Self) -> Option<Ordering> { + Some(self.cmp(other)) + } +} + +impl<T> Ord for OrderWrapper<T> { + fn cmp(&self, other: &Self) -> Ordering { + // BinaryHeap is a max heap, so compare backwards here. + other.index.cmp(&self.index) + } +} + +impl<T> Future for OrderWrapper<T> +where + T: Future, +{ + type Output = OrderWrapper<T::Output>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let index = self.index; + self.project().data.poll(cx).map(|output| OrderWrapper { data: output, index }) + } +} + +/// An unbounded queue of futures. +/// +/// This "combinator" is similar to [`FuturesUnordered`], but it imposes a FIFO order +/// on top of the set of futures. While futures in the set will race to +/// completion in parallel, results will only be returned in the order their +/// originating futures were added to the queue. +/// +/// Futures are pushed into this queue and their realized values are yielded in +/// order. This structure is optimized to manage a large number of futures. +/// Futures managed by `FuturesOrdered` will only be polled when they generate +/// notifications. This reduces the required amount of work needed to coordinate +/// large numbers of futures. +/// +/// When a `FuturesOrdered` is first created, it does not contain any futures. +/// Calling `poll` in this state will result in `Poll::Ready(None))` to be +/// returned. Futures are submitted to the queue using `push`; however, the +/// future will **not** be polled at this point. `FuturesOrdered` will only +/// poll managed futures when `FuturesOrdered::poll` is called. As such, it +/// is important to call `poll` after pushing new futures. +/// +/// If `FuturesOrdered::poll` returns `Poll::Ready(None)` this means that +/// the queue is currently not managing any futures. A future may be submitted +/// to the queue at a later time. At that point, a call to +/// `FuturesOrdered::poll` will either return the future's resolved value +/// **or** `Poll::Pending` if the future has not yet completed. When +/// multiple futures are submitted to the queue, `FuturesOrdered::poll` will +/// return `Poll::Pending` until the first future completes, even if +/// some of the later futures have already completed. +/// +/// Note that you can create a ready-made `FuturesOrdered` via the +/// [`collect`](Iterator::collect) method, or you can start with an empty queue +/// with the `FuturesOrdered::new` constructor. +/// +/// This type is only available when the `std` or `alloc` feature of this +/// library is activated, and it is activated by default. +#[must_use = "streams do nothing unless polled"] +pub struct FuturesOrdered<T: Future> { + in_progress_queue: FuturesUnordered<OrderWrapper<T>>, + queued_outputs: BinaryHeap<OrderWrapper<T::Output>>, + next_incoming_index: isize, + next_outgoing_index: isize, +} + +impl<T: Future> Unpin for FuturesOrdered<T> {} + +impl<Fut: Future> FuturesOrdered<Fut> { + /// Constructs a new, empty `FuturesOrdered` + /// + /// The returned `FuturesOrdered` does not contain any futures and, in this + /// state, `FuturesOrdered::poll_next` will return `Poll::Ready(None)`. + pub fn new() -> Self { + Self { + in_progress_queue: FuturesUnordered::new(), + queued_outputs: BinaryHeap::new(), + next_incoming_index: 0, + next_outgoing_index: 0, + } + } + + /// Returns the number of futures contained in the queue. + /// + /// This represents the total number of in-flight futures, both + /// those currently processing and those that have completed but + /// which are waiting for earlier futures to complete. + pub fn len(&self) -> usize { + self.in_progress_queue.len() + self.queued_outputs.len() + } + + /// Returns `true` if the queue contains no futures + pub fn is_empty(&self) -> bool { + self.in_progress_queue.is_empty() && self.queued_outputs.is_empty() + } + + /// Push a future into the queue. + /// + /// This function submits the given future to the internal set for managing. + /// This function will not call `poll` on the submitted future. The caller + /// must ensure that `FuturesOrdered::poll` is called in order to receive + /// task notifications. + #[deprecated(note = "use `push_back` instead")] + pub fn push(&mut self, future: Fut) { + self.push_back(future); + } + + /// Pushes a future to the back of the queue. + /// + /// This function submits the given future to the internal set for managing. + /// This function will not call `poll` on the submitted future. The caller + /// must ensure that `FuturesOrdered::poll` is called in order to receive + /// task notifications. + pub fn push_back(&mut self, future: Fut) { + let wrapped = OrderWrapper { data: future, index: self.next_incoming_index }; + self.next_incoming_index += 1; + self.in_progress_queue.push(wrapped); + } + + /// Pushes a future to the front of the queue. + /// + /// This function submits the given future to the internal set for managing. + /// This function will not call `poll` on the submitted future. The caller + /// must ensure that `FuturesOrdered::poll` is called in order to receive + /// task notifications. This future will be the next future to be returned + /// complete. + pub fn push_front(&mut self, future: Fut) { + let wrapped = OrderWrapper { data: future, index: self.next_outgoing_index - 1 }; + self.next_outgoing_index -= 1; + self.in_progress_queue.push(wrapped); + } +} + +impl<Fut: Future> Default for FuturesOrdered<Fut> { + fn default() -> Self { + Self::new() + } +} + +impl<Fut: Future> Stream for FuturesOrdered<Fut> { + type Item = Fut::Output; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let this = &mut *self; + + // Check to see if we've already received the next value + if let Some(next_output) = this.queued_outputs.peek_mut() { + if next_output.index == this.next_outgoing_index { + this.next_outgoing_index += 1; + return Poll::Ready(Some(PeekMut::pop(next_output).data)); + } + } + + loop { + match ready!(this.in_progress_queue.poll_next_unpin(cx)) { + Some(output) => { + if output.index == this.next_outgoing_index { + this.next_outgoing_index += 1; + return Poll::Ready(Some(output.data)); + } else { + this.queued_outputs.push(output) + } + } + None => return Poll::Ready(None), + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let len = self.len(); + (len, Some(len)) + } +} + +impl<Fut: Future> Debug for FuturesOrdered<Fut> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "FuturesOrdered {{ ... }}") + } +} + +impl<Fut: Future> FromIterator<Fut> for FuturesOrdered<Fut> { + fn from_iter<T>(iter: T) -> Self + where + T: IntoIterator<Item = Fut>, + { + let acc = Self::new(); + iter.into_iter().fold(acc, |mut acc, item| { + acc.push_back(item); + acc + }) + } +} + +impl<Fut: Future> FusedStream for FuturesOrdered<Fut> { + fn is_terminated(&self) -> bool { + self.in_progress_queue.is_terminated() && self.queued_outputs.is_empty() + } +} + +impl<Fut: Future> Extend<Fut> for FuturesOrdered<Fut> { + fn extend<I>(&mut self, iter: I) + where + I: IntoIterator<Item = Fut>, + { + for item in iter { + self.push_back(item); + } + } +} diff --git a/third_party/rust/futures-util/src/stream/futures_unordered/abort.rs b/third_party/rust/futures-util/src/stream/futures_unordered/abort.rs new file mode 100644 index 0000000000..1a42d24369 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/futures_unordered/abort.rs @@ -0,0 +1,12 @@ +pub(super) fn abort(s: &str) -> ! { + struct DoublePanic; + + impl Drop for DoublePanic { + fn drop(&mut self) { + panic!("panicking twice to abort the program"); + } + } + + let _bomb = DoublePanic; + panic!("{}", s); +} diff --git a/third_party/rust/futures-util/src/stream/futures_unordered/iter.rs b/third_party/rust/futures-util/src/stream/futures_unordered/iter.rs new file mode 100644 index 0000000000..20248c70fe --- /dev/null +++ b/third_party/rust/futures-util/src/stream/futures_unordered/iter.rs @@ -0,0 +1,172 @@ +use super::task::Task; +use super::FuturesUnordered; +use core::marker::PhantomData; +use core::pin::Pin; +use core::ptr; +use core::sync::atomic::Ordering::Relaxed; + +/// Mutable iterator over all futures in the unordered set. +#[derive(Debug)] +pub struct IterPinMut<'a, Fut> { + pub(super) task: *const Task<Fut>, + pub(super) len: usize, + pub(super) _marker: PhantomData<&'a mut FuturesUnordered<Fut>>, +} + +/// Mutable iterator over all futures in the unordered set. +#[derive(Debug)] +pub struct IterMut<'a, Fut: Unpin>(pub(super) IterPinMut<'a, Fut>); + +/// Immutable iterator over all futures in the unordered set. +#[derive(Debug)] +pub struct IterPinRef<'a, Fut> { + pub(super) task: *const Task<Fut>, + pub(super) len: usize, + pub(super) pending_next_all: *mut Task<Fut>, + pub(super) _marker: PhantomData<&'a FuturesUnordered<Fut>>, +} + +/// Immutable iterator over all the futures in the unordered set. +#[derive(Debug)] +pub struct Iter<'a, Fut: Unpin>(pub(super) IterPinRef<'a, Fut>); + +/// Owned iterator over all futures in the unordered set. +#[derive(Debug)] +pub struct IntoIter<Fut: Unpin> { + pub(super) len: usize, + pub(super) inner: FuturesUnordered<Fut>, +} + +impl<Fut: Unpin> Iterator for IntoIter<Fut> { + type Item = Fut; + + fn next(&mut self) -> Option<Self::Item> { + // `head_all` can be accessed directly and we don't need to spin on + // `Task::next_all` since we have exclusive access to the set. + let task = self.inner.head_all.get_mut(); + + if (*task).is_null() { + return None; + } + + unsafe { + // Moving out of the future is safe because it is `Unpin` + let future = (*(**task).future.get()).take().unwrap(); + + // Mutable access to a previously shared `FuturesUnordered` implies + // that the other threads already released the object before the + // current thread acquired it, so relaxed ordering can be used and + // valid `next_all` checks can be skipped. + let next = (**task).next_all.load(Relaxed); + *task = next; + if !task.is_null() { + *(**task).prev_all.get() = ptr::null_mut(); + } + self.len -= 1; + Some(future) + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (self.len, Some(self.len)) + } +} + +impl<Fut: Unpin> ExactSizeIterator for IntoIter<Fut> {} + +impl<'a, Fut> Iterator for IterPinMut<'a, Fut> { + type Item = Pin<&'a mut Fut>; + + fn next(&mut self) -> Option<Self::Item> { + if self.task.is_null() { + return None; + } + + unsafe { + let future = (*(*self.task).future.get()).as_mut().unwrap(); + + // Mutable access to a previously shared `FuturesUnordered` implies + // that the other threads already released the object before the + // current thread acquired it, so relaxed ordering can be used and + // valid `next_all` checks can be skipped. + let next = (*self.task).next_all.load(Relaxed); + self.task = next; + self.len -= 1; + Some(Pin::new_unchecked(future)) + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (self.len, Some(self.len)) + } +} + +impl<Fut> ExactSizeIterator for IterPinMut<'_, Fut> {} + +impl<'a, Fut: Unpin> Iterator for IterMut<'a, Fut> { + type Item = &'a mut Fut; + + fn next(&mut self) -> Option<Self::Item> { + self.0.next().map(Pin::get_mut) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.0.size_hint() + } +} + +impl<Fut: Unpin> ExactSizeIterator for IterMut<'_, Fut> {} + +impl<'a, Fut> Iterator for IterPinRef<'a, Fut> { + type Item = Pin<&'a Fut>; + + fn next(&mut self) -> Option<Self::Item> { + if self.task.is_null() { + return None; + } + + unsafe { + let future = (*(*self.task).future.get()).as_ref().unwrap(); + + // Relaxed ordering can be used since acquire ordering when + // `head_all` was initially read for this iterator implies acquire + // ordering for all previously inserted nodes (and we don't need to + // read `len_all` again for any other nodes). + let next = (*self.task).spin_next_all(self.pending_next_all, Relaxed); + self.task = next; + self.len -= 1; + Some(Pin::new_unchecked(future)) + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (self.len, Some(self.len)) + } +} + +impl<Fut> ExactSizeIterator for IterPinRef<'_, Fut> {} + +impl<'a, Fut: Unpin> Iterator for Iter<'a, Fut> { + type Item = &'a Fut; + + fn next(&mut self) -> Option<Self::Item> { + self.0.next().map(Pin::get_ref) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.0.size_hint() + } +} + +impl<Fut: Unpin> ExactSizeIterator for Iter<'_, Fut> {} + +// SAFETY: we do nothing thread-local and there is no interior mutability, +// so the usual structural `Send`/`Sync` apply. +unsafe impl<Fut: Send> Send for IterPinRef<'_, Fut> {} +unsafe impl<Fut: Sync> Sync for IterPinRef<'_, Fut> {} + +unsafe impl<Fut: Send> Send for IterPinMut<'_, Fut> {} +unsafe impl<Fut: Sync> Sync for IterPinMut<'_, Fut> {} + +unsafe impl<Fut: Send + Unpin> Send for IntoIter<Fut> {} +unsafe impl<Fut: Sync + Unpin> Sync for IntoIter<Fut> {} diff --git a/third_party/rust/futures-util/src/stream/futures_unordered/mod.rs b/third_party/rust/futures-util/src/stream/futures_unordered/mod.rs new file mode 100644 index 0000000000..6b5804dc41 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/futures_unordered/mod.rs @@ -0,0 +1,661 @@ +//! An unbounded set of futures. +//! +//! This module is only available when the `std` or `alloc` feature of this +//! library is activated, and it is activated by default. + +use crate::task::AtomicWaker; +use alloc::sync::{Arc, Weak}; +use core::cell::UnsafeCell; +use core::fmt::{self, Debug}; +use core::iter::FromIterator; +use core::marker::PhantomData; +use core::mem; +use core::pin::Pin; +use core::ptr; +use core::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release, SeqCst}; +use core::sync::atomic::{AtomicBool, AtomicPtr}; +use futures_core::future::Future; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use futures_task::{FutureObj, LocalFutureObj, LocalSpawn, Spawn, SpawnError}; + +mod abort; + +mod iter; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/102352 +pub use self::iter::{IntoIter, Iter, IterMut, IterPinMut, IterPinRef}; + +mod task; +use self::task::Task; + +mod ready_to_run_queue; +use self::ready_to_run_queue::{Dequeue, ReadyToRunQueue}; + +/// A set of futures which may complete in any order. +/// +/// See [`FuturesOrdered`](crate::stream::FuturesOrdered) for a version of this +/// type that preserves a FIFO order. +/// +/// This structure is optimized to manage a large number of futures. +/// Futures managed by [`FuturesUnordered`] will only be polled when they +/// generate wake-up notifications. This reduces the required amount of work +/// needed to poll large numbers of futures. +/// +/// [`FuturesUnordered`] can be filled by [`collect`](Iterator::collect)ing an +/// iterator of futures into a [`FuturesUnordered`], or by +/// [`push`](FuturesUnordered::push)ing futures onto an existing +/// [`FuturesUnordered`]. When new futures are added, +/// [`poll_next`](Stream::poll_next) must be called in order to begin receiving +/// wake-ups for new futures. +/// +/// Note that you can create a ready-made [`FuturesUnordered`] via the +/// [`collect`](Iterator::collect) method, or you can start with an empty set +/// with the [`FuturesUnordered::new`] constructor. +/// +/// This type is only available when the `std` or `alloc` feature of this +/// library is activated, and it is activated by default. +#[must_use = "streams do nothing unless polled"] +pub struct FuturesUnordered<Fut> { + ready_to_run_queue: Arc<ReadyToRunQueue<Fut>>, + head_all: AtomicPtr<Task<Fut>>, + is_terminated: AtomicBool, +} + +unsafe impl<Fut: Send> Send for FuturesUnordered<Fut> {} +unsafe impl<Fut: Sync> Sync for FuturesUnordered<Fut> {} +impl<Fut> Unpin for FuturesUnordered<Fut> {} + +impl Spawn for FuturesUnordered<FutureObj<'_, ()>> { + fn spawn_obj(&self, future_obj: FutureObj<'static, ()>) -> Result<(), SpawnError> { + self.push(future_obj); + Ok(()) + } +} + +impl LocalSpawn for FuturesUnordered<LocalFutureObj<'_, ()>> { + fn spawn_local_obj(&self, future_obj: LocalFutureObj<'static, ()>) -> Result<(), SpawnError> { + self.push(future_obj); + Ok(()) + } +} + +// FuturesUnordered is implemented using two linked lists. One which links all +// futures managed by a `FuturesUnordered` and one that tracks futures that have +// been scheduled for polling. The first linked list allows for thread safe +// insertion of nodes at the head as well as forward iteration, but is otherwise +// not thread safe and is only accessed by the thread that owns the +// `FuturesUnordered` value for any other operations. The second linked list is +// an implementation of the intrusive MPSC queue algorithm described by +// 1024cores.net. +// +// When a future is submitted to the set, a task is allocated and inserted in +// both linked lists. The next call to `poll_next` will (eventually) see this +// task and call `poll` on the future. +// +// Before a managed future is polled, the current context's waker is replaced +// with one that is aware of the specific future being run. This ensures that +// wake-up notifications generated by that specific future are visible to +// `FuturesUnordered`. When a wake-up notification is received, the task is +// inserted into the ready to run queue, so that its future can be polled later. +// +// Each task is wrapped in an `Arc` and thereby atomically reference counted. +// Also, each task contains an `AtomicBool` which acts as a flag that indicates +// whether the task is currently inserted in the atomic queue. When a wake-up +// notification is received, the task will only be inserted into the ready to +// run queue if it isn't inserted already. + +impl<Fut> Default for FuturesUnordered<Fut> { + fn default() -> Self { + Self::new() + } +} + +impl<Fut> FuturesUnordered<Fut> { + /// Constructs a new, empty [`FuturesUnordered`]. + /// + /// The returned [`FuturesUnordered`] does not contain any futures. + /// In this state, [`FuturesUnordered::poll_next`](Stream::poll_next) will + /// return [`Poll::Ready(None)`](Poll::Ready). + pub fn new() -> Self { + let stub = Arc::new(Task { + future: UnsafeCell::new(None), + next_all: AtomicPtr::new(ptr::null_mut()), + prev_all: UnsafeCell::new(ptr::null()), + len_all: UnsafeCell::new(0), + next_ready_to_run: AtomicPtr::new(ptr::null_mut()), + queued: AtomicBool::new(true), + ready_to_run_queue: Weak::new(), + woken: AtomicBool::new(false), + }); + let stub_ptr = Arc::as_ptr(&stub); + let ready_to_run_queue = Arc::new(ReadyToRunQueue { + waker: AtomicWaker::new(), + head: AtomicPtr::new(stub_ptr as *mut _), + tail: UnsafeCell::new(stub_ptr), + stub, + }); + + Self { + head_all: AtomicPtr::new(ptr::null_mut()), + ready_to_run_queue, + is_terminated: AtomicBool::new(false), + } + } + + /// Returns the number of futures contained in the set. + /// + /// This represents the total number of in-flight futures. + pub fn len(&self) -> usize { + let (_, len) = self.atomic_load_head_and_len_all(); + len + } + + /// Returns `true` if the set contains no futures. + pub fn is_empty(&self) -> bool { + // Relaxed ordering can be used here since we don't need to read from + // the head pointer, only check whether it is null. + self.head_all.load(Relaxed).is_null() + } + + /// Push a future into the set. + /// + /// This method adds the given future to the set. This method will not + /// call [`poll`](core::future::Future::poll) on the submitted future. The caller must + /// ensure that [`FuturesUnordered::poll_next`](Stream::poll_next) is called + /// in order to receive wake-up notifications for the given future. + pub fn push(&self, future: Fut) { + let task = Arc::new(Task { + future: UnsafeCell::new(Some(future)), + next_all: AtomicPtr::new(self.pending_next_all()), + prev_all: UnsafeCell::new(ptr::null_mut()), + len_all: UnsafeCell::new(0), + next_ready_to_run: AtomicPtr::new(ptr::null_mut()), + queued: AtomicBool::new(true), + ready_to_run_queue: Arc::downgrade(&self.ready_to_run_queue), + woken: AtomicBool::new(false), + }); + + // Reset the `is_terminated` flag if we've previously marked ourselves + // as terminated. + self.is_terminated.store(false, Relaxed); + + // Right now our task has a strong reference count of 1. We transfer + // ownership of this reference count to our internal linked list + // and we'll reclaim ownership through the `unlink` method below. + let ptr = self.link(task); + + // We'll need to get the future "into the system" to start tracking it, + // e.g. getting its wake-up notifications going to us tracking which + // futures are ready. To do that we unconditionally enqueue it for + // polling here. + self.ready_to_run_queue.enqueue(ptr); + } + + /// Returns an iterator that allows inspecting each future in the set. + pub fn iter(&self) -> Iter<'_, Fut> + where + Fut: Unpin, + { + Iter(Pin::new(self).iter_pin_ref()) + } + + /// Returns an iterator that allows inspecting each future in the set. + pub fn iter_pin_ref(self: Pin<&Self>) -> IterPinRef<'_, Fut> { + let (task, len) = self.atomic_load_head_and_len_all(); + let pending_next_all = self.pending_next_all(); + + IterPinRef { task, len, pending_next_all, _marker: PhantomData } + } + + /// Returns an iterator that allows modifying each future in the set. + pub fn iter_mut(&mut self) -> IterMut<'_, Fut> + where + Fut: Unpin, + { + IterMut(Pin::new(self).iter_pin_mut()) + } + + /// Returns an iterator that allows modifying each future in the set. + pub fn iter_pin_mut(mut self: Pin<&mut Self>) -> IterPinMut<'_, Fut> { + // `head_all` can be accessed directly and we don't need to spin on + // `Task::next_all` since we have exclusive access to the set. + let task = *self.head_all.get_mut(); + let len = if task.is_null() { 0 } else { unsafe { *(*task).len_all.get() } }; + + IterPinMut { task, len, _marker: PhantomData } + } + + /// Returns the current head node and number of futures in the list of all + /// futures within a context where access is shared with other threads + /// (mostly for use with the `len` and `iter_pin_ref` methods). + fn atomic_load_head_and_len_all(&self) -> (*const Task<Fut>, usize) { + let task = self.head_all.load(Acquire); + let len = if task.is_null() { + 0 + } else { + unsafe { + (*task).spin_next_all(self.pending_next_all(), Acquire); + *(*task).len_all.get() + } + }; + + (task, len) + } + + /// Releases the task. It destroys the future inside and either drops + /// the `Arc<Task>` or transfers ownership to the ready to run queue. + /// The task this method is called on must have been unlinked before. + fn release_task(&mut self, task: Arc<Task<Fut>>) { + // `release_task` must only be called on unlinked tasks + debug_assert_eq!(task.next_all.load(Relaxed), self.pending_next_all()); + unsafe { + debug_assert!((*task.prev_all.get()).is_null()); + } + + // The future is done, try to reset the queued flag. This will prevent + // `wake` from doing any work in the future + let prev = task.queued.swap(true, SeqCst); + + // Drop the future, even if it hasn't finished yet. This is safe + // because we're dropping the future on the thread that owns + // `FuturesUnordered`, which correctly tracks `Fut`'s lifetimes and + // such. + unsafe { + // Set to `None` rather than `take()`ing to prevent moving the + // future. + *task.future.get() = None; + } + + // If the queued flag was previously set, then it means that this task + // is still in our internal ready to run queue. We then transfer + // ownership of our reference count to the ready to run queue, and it'll + // come along and free it later, noticing that the future is `None`. + // + // If, however, the queued flag was *not* set then we're safe to + // release our reference count on the task. The queued flag was set + // above so all future `enqueue` operations will not actually + // enqueue the task, so our task will never see the ready to run queue + // again. The task itself will be deallocated once all reference counts + // have been dropped elsewhere by the various wakers that contain it. + if prev { + mem::forget(task); + } + } + + /// Insert a new task into the internal linked list. + fn link(&self, task: Arc<Task<Fut>>) -> *const Task<Fut> { + // `next_all` should already be reset to the pending state before this + // function is called. + debug_assert_eq!(task.next_all.load(Relaxed), self.pending_next_all()); + let ptr = Arc::into_raw(task); + + // Atomically swap out the old head node to get the node that should be + // assigned to `next_all`. + let next = self.head_all.swap(ptr as *mut _, AcqRel); + + unsafe { + // Store the new list length in the new node. + let new_len = if next.is_null() { + 1 + } else { + // Make sure `next_all` has been written to signal that it is + // safe to read `len_all`. + (*next).spin_next_all(self.pending_next_all(), Acquire); + *(*next).len_all.get() + 1 + }; + *(*ptr).len_all.get() = new_len; + + // Write the old head as the next node pointer, signaling to other + // threads that `len_all` and `next_all` are ready to read. + (*ptr).next_all.store(next, Release); + + // `prev_all` updates don't need to be synchronized, as the field is + // only ever used after exclusive access has been acquired. + if !next.is_null() { + *(*next).prev_all.get() = ptr; + } + } + + ptr + } + + /// Remove the task from the linked list tracking all tasks currently + /// managed by `FuturesUnordered`. + /// This method is unsafe because it has be guaranteed that `task` is a + /// valid pointer. + unsafe fn unlink(&mut self, task: *const Task<Fut>) -> Arc<Task<Fut>> { + // Compute the new list length now in case we're removing the head node + // and won't be able to retrieve the correct length later. + let head = *self.head_all.get_mut(); + debug_assert!(!head.is_null()); + let new_len = *(*head).len_all.get() - 1; + + let task = Arc::from_raw(task); + let next = task.next_all.load(Relaxed); + let prev = *task.prev_all.get(); + task.next_all.store(self.pending_next_all(), Relaxed); + *task.prev_all.get() = ptr::null_mut(); + + if !next.is_null() { + *(*next).prev_all.get() = prev; + } + + if !prev.is_null() { + (*prev).next_all.store(next, Relaxed); + } else { + *self.head_all.get_mut() = next; + } + + // Store the new list length in the head node. + let head = *self.head_all.get_mut(); + if !head.is_null() { + *(*head).len_all.get() = new_len; + } + + task + } + + /// Returns the reserved value for `Task::next_all` to indicate a pending + /// assignment from the thread that inserted the task. + /// + /// `FuturesUnordered::link` needs to update `Task` pointers in an order + /// that ensures any iterators created on other threads can correctly + /// traverse the entire `Task` list using the chain of `next_all` pointers. + /// This could be solved with a compare-exchange loop that stores the + /// current `head_all` in `next_all` and swaps out `head_all` with the new + /// `Task` pointer if the head hasn't already changed. Under heavy thread + /// contention, this compare-exchange loop could become costly. + /// + /// An alternative is to initialize `next_all` to a reserved pending state + /// first, perform an atomic swap on `head_all`, and finally update + /// `next_all` with the old head node. Iterators will then either see the + /// pending state value or the correct next node pointer, and can reload + /// `next_all` as needed until the correct value is loaded. The number of + /// retries needed (if any) would be small and will always be finite, so + /// this should generally perform better than the compare-exchange loop. + /// + /// A valid `Task` pointer in the `head_all` list is guaranteed to never be + /// this value, so it is safe to use as a reserved value until the correct + /// value can be written. + fn pending_next_all(&self) -> *mut Task<Fut> { + // The `ReadyToRunQueue` stub is never inserted into the `head_all` + // list, and its pointer value will remain valid for the lifetime of + // this `FuturesUnordered`, so we can make use of its value here. + Arc::as_ptr(&self.ready_to_run_queue.stub) as *mut _ + } +} + +impl<Fut: Future> Stream for FuturesUnordered<Fut> { + type Item = Fut::Output; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let len = self.len(); + + // Keep track of how many child futures we have polled, + // in case we want to forcibly yield. + let mut polled = 0; + let mut yielded = 0; + + // Ensure `parent` is correctly set. + self.ready_to_run_queue.waker.register(cx.waker()); + + loop { + // Safety: &mut self guarantees the mutual exclusion `dequeue` + // expects + let task = match unsafe { self.ready_to_run_queue.dequeue() } { + Dequeue::Empty => { + if self.is_empty() { + // We can only consider ourselves terminated once we + // have yielded a `None` + *self.is_terminated.get_mut() = true; + return Poll::Ready(None); + } else { + return Poll::Pending; + } + } + Dequeue::Inconsistent => { + // At this point, it may be worth yielding the thread & + // spinning a few times... but for now, just yield using the + // task system. + cx.waker().wake_by_ref(); + return Poll::Pending; + } + Dequeue::Data(task) => task, + }; + + debug_assert!(task != self.ready_to_run_queue.stub()); + + // Safety: + // - `task` is a valid pointer. + // - We are the only thread that accesses the `UnsafeCell` that + // contains the future + let future = match unsafe { &mut *(*task).future.get() } { + Some(future) => future, + + // If the future has already gone away then we're just + // cleaning out this task. See the comment in + // `release_task` for more information, but we're basically + // just taking ownership of our reference count here. + None => { + // This case only happens when `release_task` was called + // for this task before and couldn't drop the task + // because it was already enqueued in the ready to run + // queue. + + // Safety: `task` is a valid pointer + let task = unsafe { Arc::from_raw(task) }; + + // Double check that the call to `release_task` really + // happened. Calling it required the task to be unlinked. + debug_assert_eq!(task.next_all.load(Relaxed), self.pending_next_all()); + unsafe { + debug_assert!((*task.prev_all.get()).is_null()); + } + continue; + } + }; + + // Safety: `task` is a valid pointer + let task = unsafe { self.unlink(task) }; + + // Unset queued flag: This must be done before polling to ensure + // that the future's task gets rescheduled if it sends a wake-up + // notification **during** the call to `poll`. + let prev = task.queued.swap(false, SeqCst); + assert!(prev); + + // We're going to need to be very careful if the `poll` + // method below panics. We need to (a) not leak memory and + // (b) ensure that we still don't have any use-after-frees. To + // manage this we do a few things: + // + // * A "bomb" is created which if dropped abnormally will call + // `release_task`. That way we'll be sure the memory management + // of the `task` is managed correctly. In particular + // `release_task` will drop the future. This ensures that it is + // dropped on this thread and not accidentally on a different + // thread (bad). + // * We unlink the task from our internal queue to preemptively + // assume it'll panic, in which case we'll want to discard it + // regardless. + struct Bomb<'a, Fut> { + queue: &'a mut FuturesUnordered<Fut>, + task: Option<Arc<Task<Fut>>>, + } + + impl<Fut> Drop for Bomb<'_, Fut> { + fn drop(&mut self) { + if let Some(task) = self.task.take() { + self.queue.release_task(task); + } + } + } + + let mut bomb = Bomb { task: Some(task), queue: &mut *self }; + + // Poll the underlying future with the appropriate waker + // implementation. This is where a large bit of the unsafety + // starts to stem from internally. The waker is basically just + // our `Arc<Task<Fut>>` and can schedule the future for polling by + // enqueuing itself in the ready to run queue. + // + // Critically though `Task<Fut>` won't actually access `Fut`, the + // future, while it's floating around inside of wakers. + // These structs will basically just use `Fut` to size + // the internal allocation, appropriately accessing fields and + // deallocating the task if need be. + let res = { + let task = bomb.task.as_ref().unwrap(); + // We are only interested in whether the future is awoken before it + // finishes polling, so reset the flag here. + task.woken.store(false, Relaxed); + let waker = Task::waker_ref(task); + let mut cx = Context::from_waker(&waker); + + // Safety: We won't move the future ever again + let future = unsafe { Pin::new_unchecked(future) }; + + future.poll(&mut cx) + }; + polled += 1; + + match res { + Poll::Pending => { + let task = bomb.task.take().unwrap(); + // If the future was awoken during polling, we assume + // the future wanted to explicitly yield. + yielded += task.woken.load(Relaxed) as usize; + bomb.queue.link(task); + + // If a future yields, we respect it and yield here. + // If all futures have been polled, we also yield here to + // avoid starving other tasks waiting on the executor. + // (polling the same future twice per iteration may cause + // the problem: https://github.com/rust-lang/futures-rs/pull/2333) + if yielded >= 2 || polled == len { + cx.waker().wake_by_ref(); + return Poll::Pending; + } + continue; + } + Poll::Ready(output) => return Poll::Ready(Some(output)), + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let len = self.len(); + (len, Some(len)) + } +} + +impl<Fut> Debug for FuturesUnordered<Fut> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "FuturesUnordered {{ ... }}") + } +} + +impl<Fut> FuturesUnordered<Fut> { + /// Clears the set, removing all futures. + pub fn clear(&mut self) { + self.clear_head_all(); + + // we just cleared all the tasks, and we have &mut self, so this is safe. + unsafe { self.ready_to_run_queue.clear() }; + + self.is_terminated.store(false, Relaxed); + } + + fn clear_head_all(&mut self) { + while !self.head_all.get_mut().is_null() { + let head = *self.head_all.get_mut(); + let task = unsafe { self.unlink(head) }; + self.release_task(task); + } + } +} + +impl<Fut> Drop for FuturesUnordered<Fut> { + fn drop(&mut self) { + // When a `FuturesUnordered` is dropped we want to drop all futures + // associated with it. At the same time though there may be tons of + // wakers flying around which contain `Task<Fut>` references + // inside them. We'll let those naturally get deallocated. + self.clear_head_all(); + + // Note that at this point we could still have a bunch of tasks in the + // ready to run queue. None of those tasks, however, have futures + // associated with them so they're safe to destroy on any thread. At + // this point the `FuturesUnordered` struct, the owner of the one strong + // reference to the ready to run queue will drop the strong reference. + // At that point whichever thread releases the strong refcount last (be + // it this thread or some other thread as part of an `upgrade`) will + // clear out the ready to run queue and free all remaining tasks. + // + // While that freeing operation isn't guaranteed to happen here, it's + // guaranteed to happen "promptly" as no more "blocking work" will + // happen while there's a strong refcount held. + } +} + +impl<'a, Fut: Unpin> IntoIterator for &'a FuturesUnordered<Fut> { + type Item = &'a Fut; + type IntoIter = Iter<'a, Fut>; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl<'a, Fut: Unpin> IntoIterator for &'a mut FuturesUnordered<Fut> { + type Item = &'a mut Fut; + type IntoIter = IterMut<'a, Fut>; + + fn into_iter(self) -> Self::IntoIter { + self.iter_mut() + } +} + +impl<Fut: Unpin> IntoIterator for FuturesUnordered<Fut> { + type Item = Fut; + type IntoIter = IntoIter<Fut>; + + fn into_iter(mut self) -> Self::IntoIter { + // `head_all` can be accessed directly and we don't need to spin on + // `Task::next_all` since we have exclusive access to the set. + let task = *self.head_all.get_mut(); + let len = if task.is_null() { 0 } else { unsafe { *(*task).len_all.get() } }; + + IntoIter { len, inner: self } + } +} + +impl<Fut> FromIterator<Fut> for FuturesUnordered<Fut> { + fn from_iter<I>(iter: I) -> Self + where + I: IntoIterator<Item = Fut>, + { + let acc = Self::new(); + iter.into_iter().fold(acc, |acc, item| { + acc.push(item); + acc + }) + } +} + +impl<Fut: Future> FusedStream for FuturesUnordered<Fut> { + fn is_terminated(&self) -> bool { + self.is_terminated.load(Relaxed) + } +} + +impl<Fut> Extend<Fut> for FuturesUnordered<Fut> { + fn extend<I>(&mut self, iter: I) + where + I: IntoIterator<Item = Fut>, + { + for item in iter { + self.push(item); + } + } +} diff --git a/third_party/rust/futures-util/src/stream/futures_unordered/ready_to_run_queue.rs b/third_party/rust/futures-util/src/stream/futures_unordered/ready_to_run_queue.rs new file mode 100644 index 0000000000..4518705320 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/futures_unordered/ready_to_run_queue.rs @@ -0,0 +1,122 @@ +use crate::task::AtomicWaker; +use alloc::sync::Arc; +use core::cell::UnsafeCell; +use core::ptr; +use core::sync::atomic::AtomicPtr; +use core::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; + +use super::abort::abort; +use super::task::Task; + +pub(super) enum Dequeue<Fut> { + Data(*const Task<Fut>), + Empty, + Inconsistent, +} + +pub(super) struct ReadyToRunQueue<Fut> { + // The waker of the task using `FuturesUnordered`. + pub(super) waker: AtomicWaker, + + // Head/tail of the readiness queue + pub(super) head: AtomicPtr<Task<Fut>>, + pub(super) tail: UnsafeCell<*const Task<Fut>>, + pub(super) stub: Arc<Task<Fut>>, +} + +/// An MPSC queue into which the tasks containing the futures are inserted +/// whenever the future inside is scheduled for polling. +impl<Fut> ReadyToRunQueue<Fut> { + /// The enqueue function from the 1024cores intrusive MPSC queue algorithm. + pub(super) fn enqueue(&self, task: *const Task<Fut>) { + unsafe { + debug_assert!((*task).queued.load(Relaxed)); + + // This action does not require any coordination + (*task).next_ready_to_run.store(ptr::null_mut(), Relaxed); + + // Note that these atomic orderings come from 1024cores + let task = task as *mut _; + let prev = self.head.swap(task, AcqRel); + (*prev).next_ready_to_run.store(task, Release); + } + } + + /// The dequeue function from the 1024cores intrusive MPSC queue algorithm + /// + /// Note that this is unsafe as it required mutual exclusion (only one + /// thread can call this) to be guaranteed elsewhere. + pub(super) unsafe fn dequeue(&self) -> Dequeue<Fut> { + let mut tail = *self.tail.get(); + let mut next = (*tail).next_ready_to_run.load(Acquire); + + if tail == self.stub() { + if next.is_null() { + return Dequeue::Empty; + } + + *self.tail.get() = next; + tail = next; + next = (*next).next_ready_to_run.load(Acquire); + } + + if !next.is_null() { + *self.tail.get() = next; + debug_assert!(tail != self.stub()); + return Dequeue::Data(tail); + } + + if self.head.load(Acquire) as *const _ != tail { + return Dequeue::Inconsistent; + } + + self.enqueue(self.stub()); + + next = (*tail).next_ready_to_run.load(Acquire); + + if !next.is_null() { + *self.tail.get() = next; + return Dequeue::Data(tail); + } + + Dequeue::Inconsistent + } + + pub(super) fn stub(&self) -> *const Task<Fut> { + Arc::as_ptr(&self.stub) + } + + // Clear the queue of tasks. + // + // Note that each task has a strong reference count associated with it + // which is owned by the ready to run queue. This method just pulls out + // tasks and drops their refcounts. + // + // # Safety + // + // - All tasks **must** have had their futures dropped already (by FuturesUnordered::clear) + // - The caller **must** guarantee unique access to `self` + pub(crate) unsafe fn clear(&self) { + loop { + // SAFETY: We have the guarantee of mutual exclusion required by `dequeue`. + match self.dequeue() { + Dequeue::Empty => break, + Dequeue::Inconsistent => abort("inconsistent in drop"), + Dequeue::Data(ptr) => drop(Arc::from_raw(ptr)), + } + } + } +} + +impl<Fut> Drop for ReadyToRunQueue<Fut> { + fn drop(&mut self) { + // Once we're in the destructor for `Inner<Fut>` we need to clear out + // the ready to run queue of tasks if there's anything left in there. + + // All tasks have had their futures dropped already by the `FuturesUnordered` + // destructor above, and we have &mut self, so this is safe. + unsafe { + self.clear(); + } + } +} diff --git a/third_party/rust/futures-util/src/stream/futures_unordered/task.rs b/third_party/rust/futures-util/src/stream/futures_unordered/task.rs new file mode 100644 index 0000000000..ec2114effa --- /dev/null +++ b/third_party/rust/futures-util/src/stream/futures_unordered/task.rs @@ -0,0 +1,125 @@ +use alloc::sync::{Arc, Weak}; +use core::cell::UnsafeCell; +use core::sync::atomic::Ordering::{self, Relaxed, SeqCst}; +use core::sync::atomic::{AtomicBool, AtomicPtr}; + +use super::abort::abort; +use super::ReadyToRunQueue; +use crate::task::{waker_ref, ArcWake, WakerRef}; + +pub(super) struct Task<Fut> { + // The future + pub(super) future: UnsafeCell<Option<Fut>>, + + // Next pointer for linked list tracking all active tasks (use + // `spin_next_all` to read when access is shared across threads) + pub(super) next_all: AtomicPtr<Task<Fut>>, + + // Previous task in linked list tracking all active tasks + pub(super) prev_all: UnsafeCell<*const Task<Fut>>, + + // Length of the linked list tracking all active tasks when this node was + // inserted (use `spin_next_all` to synchronize before reading when access + // is shared across threads) + pub(super) len_all: UnsafeCell<usize>, + + // Next pointer in ready to run queue + pub(super) next_ready_to_run: AtomicPtr<Task<Fut>>, + + // Queue that we'll be enqueued to when woken + pub(super) ready_to_run_queue: Weak<ReadyToRunQueue<Fut>>, + + // Whether or not this task is currently in the ready to run queue + pub(super) queued: AtomicBool, + + // Whether the future was awoken during polling + // It is possible for this flag to be set to true after the polling, + // but it will be ignored. + pub(super) woken: AtomicBool, +} + +// `Task` can be sent across threads safely because it ensures that +// the underlying `Fut` type isn't touched from any of its methods. +// +// The parent (`super`) module is trusted not to access `future` +// across different threads. +unsafe impl<Fut> Send for Task<Fut> {} +unsafe impl<Fut> Sync for Task<Fut> {} + +impl<Fut> ArcWake for Task<Fut> { + fn wake_by_ref(arc_self: &Arc<Self>) { + let inner = match arc_self.ready_to_run_queue.upgrade() { + Some(inner) => inner, + None => return, + }; + + arc_self.woken.store(true, Relaxed); + + // It's our job to enqueue this task it into the ready to run queue. To + // do this we set the `queued` flag, and if successful we then do the + // actual queueing operation, ensuring that we're only queued once. + // + // Once the task is inserted call `wake` to notify the parent task, + // as it'll want to come along and run our task later. + // + // Note that we don't change the reference count of the task here, + // we merely enqueue the raw pointer. The `FuturesUnordered` + // implementation guarantees that if we set the `queued` flag that + // there's a reference count held by the main `FuturesUnordered` queue + // still. + let prev = arc_self.queued.swap(true, SeqCst); + if !prev { + inner.enqueue(Arc::as_ptr(arc_self)); + inner.waker.wake(); + } + } +} + +impl<Fut> Task<Fut> { + /// Returns a waker reference for this task without cloning the Arc. + pub(super) fn waker_ref(this: &Arc<Self>) -> WakerRef<'_> { + waker_ref(this) + } + + /// Spins until `next_all` is no longer set to `pending_next_all`. + /// + /// The temporary `pending_next_all` value is typically overwritten fairly + /// quickly after a node is inserted into the list of all futures, so this + /// should rarely spin much. + /// + /// When it returns, the correct `next_all` value is returned. + /// + /// `Relaxed` or `Acquire` ordering can be used. `Acquire` ordering must be + /// used before `len_all` can be safely read. + #[inline] + pub(super) fn spin_next_all( + &self, + pending_next_all: *mut Self, + ordering: Ordering, + ) -> *const Self { + loop { + let next = self.next_all.load(ordering); + if next != pending_next_all { + return next; + } + } + } +} + +impl<Fut> Drop for Task<Fut> { + fn drop(&mut self) { + // Since `Task<Fut>` is sent across all threads for any lifetime, + // regardless of `Fut`, we, to guarantee memory safety, can't actually + // touch `Fut` at any time except when we have a reference to the + // `FuturesUnordered` itself . + // + // Consequently it *should* be the case that we always drop futures from + // the `FuturesUnordered` instance. This is a bomb, just in case there's + // a bug in that logic. + unsafe { + if (*self.future.get()).is_some() { + abort("future still here when dropping"); + } + } + } +} diff --git a/third_party/rust/futures-util/src/stream/iter.rs b/third_party/rust/futures-util/src/stream/iter.rs new file mode 100644 index 0000000000..20471c2ed0 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/iter.rs @@ -0,0 +1,49 @@ +use super::assert_stream; +use core::pin::Pin; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; + +/// Stream for the [`iter`] function. +#[derive(Debug, Clone)] +#[must_use = "streams do nothing unless polled"] +pub struct Iter<I> { + iter: I, +} + +impl<I> Unpin for Iter<I> {} + +/// Converts an `Iterator` into a `Stream` which is always ready +/// to yield the next value. +/// +/// Iterators in Rust don't express the ability to block, so this adapter +/// simply always calls `iter.next()` and returns that. +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::stream::{self, StreamExt}; +/// +/// let stream = stream::iter(vec![17, 19]); +/// assert_eq!(vec![17, 19], stream.collect::<Vec<i32>>().await); +/// # }); +/// ``` +pub fn iter<I>(i: I) -> Iter<I::IntoIter> +where + I: IntoIterator, +{ + assert_stream::<I::Item, _>(Iter { iter: i.into_iter() }) +} + +impl<I> Stream for Iter<I> +where + I: Iterator, +{ + type Item = I::Item; + + fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<I::Item>> { + Poll::Ready(self.iter.next()) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.iter.size_hint() + } +} diff --git a/third_party/rust/futures-util/src/stream/mod.rs b/third_party/rust/futures-util/src/stream/mod.rs new file mode 100644 index 0000000000..ec685b9848 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/mod.rs @@ -0,0 +1,143 @@ +//! Asynchronous streams. +//! +//! This module contains: +//! +//! - The [`Stream`] trait, for objects that can asynchronously produce a +//! sequence of values. +//! - The [`StreamExt`] and [`TryStreamExt`] trait, which provides adapters for +//! chaining and composing streams. +//! - Top-level stream constructors like [`iter`](iter()) which creates a +//! stream from an iterator. + +#[cfg(feature = "alloc")] +pub use futures_core::stream::{BoxStream, LocalBoxStream}; +pub use futures_core::stream::{FusedStream, Stream, TryStream}; + +// Extension traits and combinators + +#[allow(clippy::module_inception)] +mod stream; +pub use self::stream::{ + Chain, Collect, Concat, Cycle, Enumerate, Filter, FilterMap, FlatMap, Flatten, Fold, ForEach, + Fuse, Inspect, Map, Next, NextIf, NextIfEq, Peek, PeekMut, Peekable, Scan, SelectNextSome, + Skip, SkipWhile, StreamExt, StreamFuture, Take, TakeUntil, TakeWhile, Then, Unzip, Zip, +}; + +#[cfg(feature = "std")] +pub use self::stream::CatchUnwind; + +#[cfg(feature = "alloc")] +pub use self::stream::Chunks; + +#[cfg(feature = "alloc")] +pub use self::stream::ReadyChunks; + +#[cfg(feature = "sink")] +#[cfg_attr(docsrs, doc(cfg(feature = "sink")))] +pub use self::stream::Forward; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +pub use self::stream::{BufferUnordered, Buffered, ForEachConcurrent}; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "sink")] +#[cfg_attr(docsrs, doc(cfg(feature = "sink")))] +#[cfg(feature = "alloc")] +pub use self::stream::{ReuniteError, SplitSink, SplitStream}; + +mod try_stream; +pub use self::try_stream::{ + try_unfold, AndThen, ErrInto, InspectErr, InspectOk, IntoStream, MapErr, MapOk, OrElse, + TryCollect, TryConcat, TryFilter, TryFilterMap, TryFlatten, TryFold, TryForEach, TryNext, + TrySkipWhile, TryStreamExt, TryTakeWhile, TryUnfold, +}; + +#[cfg(feature = "io")] +#[cfg_attr(docsrs, doc(cfg(feature = "io")))] +#[cfg(feature = "std")] +pub use self::try_stream::IntoAsyncRead; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +pub use self::try_stream::{TryBufferUnordered, TryBuffered, TryForEachConcurrent}; + +#[cfg(feature = "alloc")] +pub use self::try_stream::{TryChunks, TryChunksError}; + +// Primitive streams + +mod iter; +pub use self::iter::{iter, Iter}; + +mod repeat; +pub use self::repeat::{repeat, Repeat}; + +mod repeat_with; +pub use self::repeat_with::{repeat_with, RepeatWith}; + +mod empty; +pub use self::empty::{empty, Empty}; + +mod once; +pub use self::once::{once, Once}; + +mod pending; +pub use self::pending::{pending, Pending}; + +mod poll_fn; +pub use self::poll_fn::{poll_fn, PollFn}; + +mod poll_immediate; +pub use self::poll_immediate::{poll_immediate, PollImmediate}; + +mod select; +pub use self::select::{select, Select}; + +mod select_with_strategy; +pub use self::select_with_strategy::{select_with_strategy, PollNext, SelectWithStrategy}; + +mod unfold; +pub use self::unfold::{unfold, Unfold}; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +mod futures_ordered; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +pub use self::futures_ordered::FuturesOrdered; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +pub mod futures_unordered; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +#[doc(inline)] +pub use self::futures_unordered::FuturesUnordered; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +pub mod select_all; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +#[doc(inline)] +pub use self::select_all::{select_all, SelectAll}; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +mod abortable; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +pub use crate::abortable::{AbortHandle, AbortRegistration, Abortable, Aborted}; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +pub use abortable::abortable; + +// Just a helper function to ensure the streams we're returning all have the +// right implementations. +pub(crate) fn assert_stream<T, S>(stream: S) -> S +where + S: Stream<Item = T>, +{ + stream +} diff --git a/third_party/rust/futures-util/src/stream/once.rs b/third_party/rust/futures-util/src/stream/once.rs new file mode 100644 index 0000000000..ee21c8b594 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/once.rs @@ -0,0 +1,67 @@ +use super::assert_stream; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +/// Creates a stream of a single element. +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::stream::{self, StreamExt}; +/// +/// let stream = stream::once(async { 17 }); +/// let collected = stream.collect::<Vec<i32>>().await; +/// assert_eq!(collected, vec![17]); +/// # }); +/// ``` +pub fn once<Fut: Future>(future: Fut) -> Once<Fut> { + assert_stream::<Fut::Output, _>(Once::new(future)) +} + +pin_project! { + /// A stream which emits single element and then EOF. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Once<Fut> { + #[pin] + future: Option<Fut> + } +} + +impl<Fut> Once<Fut> { + pub(crate) fn new(future: Fut) -> Self { + Self { future: Some(future) } + } +} + +impl<Fut: Future> Stream for Once<Fut> { + type Item = Fut::Output; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + let v = match this.future.as_mut().as_pin_mut() { + Some(fut) => ready!(fut.poll(cx)), + None => return Poll::Ready(None), + }; + + this.future.set(None); + Poll::Ready(Some(v)) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.future.is_some() { + (1, Some(1)) + } else { + (0, Some(0)) + } + } +} + +impl<Fut: Future> FusedStream for Once<Fut> { + fn is_terminated(&self) -> bool { + self.future.is_none() + } +} diff --git a/third_party/rust/futures-util/src/stream/pending.rs b/third_party/rust/futures-util/src/stream/pending.rs new file mode 100644 index 0000000000..d7030ff3cc --- /dev/null +++ b/third_party/rust/futures-util/src/stream/pending.rs @@ -0,0 +1,45 @@ +use super::assert_stream; +use core::marker; +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; + +/// Stream for the [`pending()`] function. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Pending<T> { + _data: marker::PhantomData<T>, +} + +/// Creates a stream which never returns any elements. +/// +/// The returned stream will always return `Pending` when polled. +pub fn pending<T>() -> Pending<T> { + assert_stream::<T, _>(Pending { _data: marker::PhantomData }) +} + +impl<T> Unpin for Pending<T> {} + +impl<T> FusedStream for Pending<T> { + fn is_terminated(&self) -> bool { + true + } +} + +impl<T> Stream for Pending<T> { + type Item = T; + + fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> { + Poll::Pending + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (0, Some(0)) + } +} + +impl<T> Clone for Pending<T> { + fn clone(&self) -> Self { + pending() + } +} diff --git a/third_party/rust/futures-util/src/stream/poll_fn.rs b/third_party/rust/futures-util/src/stream/poll_fn.rs new file mode 100644 index 0000000000..b9bd7d1664 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/poll_fn.rs @@ -0,0 +1,57 @@ +//! Definition of the `PollFn` combinator + +use super::assert_stream; +use core::fmt; +use core::pin::Pin; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; + +/// Stream for the [`poll_fn`] function. +#[must_use = "streams do nothing unless polled"] +pub struct PollFn<F> { + f: F, +} + +impl<F> Unpin for PollFn<F> {} + +impl<F> fmt::Debug for PollFn<F> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PollFn").finish() + } +} + +/// Creates a new stream wrapping a function returning `Poll<Option<T>>`. +/// +/// Polling the returned stream calls the wrapped function. +/// +/// # Examples +/// +/// ``` +/// use futures::stream::poll_fn; +/// use futures::task::Poll; +/// +/// let mut counter = 1usize; +/// +/// let read_stream = poll_fn(move |_| -> Poll<Option<String>> { +/// if counter == 0 { return Poll::Ready(None); } +/// counter -= 1; +/// Poll::Ready(Some("Hello, World!".to_owned())) +/// }); +/// ``` +pub fn poll_fn<T, F>(f: F) -> PollFn<F> +where + F: FnMut(&mut Context<'_>) -> Poll<Option<T>>, +{ + assert_stream::<T, _>(PollFn { f }) +} + +impl<T, F> Stream for PollFn<F> +where + F: FnMut(&mut Context<'_>) -> Poll<Option<T>>, +{ + type Item = T; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> { + (&mut self.f)(cx) + } +} diff --git a/third_party/rust/futures-util/src/stream/poll_immediate.rs b/third_party/rust/futures-util/src/stream/poll_immediate.rs new file mode 100644 index 0000000000..c7e8a5b3c6 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/poll_immediate.rs @@ -0,0 +1,80 @@ +use core::pin::Pin; +use futures_core::task::{Context, Poll}; +use futures_core::Stream; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [poll_immediate](poll_immediate()) function. + /// + /// It will never return [Poll::Pending](core::task::Poll::Pending) + #[derive(Debug, Clone)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct PollImmediate<S> { + #[pin] + stream: Option<S> + } +} + +impl<T, S> Stream for PollImmediate<S> +where + S: Stream<Item = T>, +{ + type Item = Poll<T>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + let stream = match this.stream.as_mut().as_pin_mut() { + // inner is gone, so we can continue to signal that the stream is closed. + None => return Poll::Ready(None), + Some(inner) => inner, + }; + + match stream.poll_next(cx) { + Poll::Ready(Some(t)) => Poll::Ready(Some(Poll::Ready(t))), + Poll::Ready(None) => { + this.stream.set(None); + Poll::Ready(None) + } + Poll::Pending => Poll::Ready(Some(Poll::Pending)), + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.as_ref().map_or((0, Some(0)), Stream::size_hint) + } +} + +impl<S: Stream> super::FusedStream for PollImmediate<S> { + fn is_terminated(&self) -> bool { + self.stream.is_none() + } +} + +/// Creates a new stream that always immediately returns [Poll::Ready](core::task::Poll::Ready) when awaiting it. +/// +/// This is useful when immediacy is more important than waiting for the next item to be ready. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::stream::{self, StreamExt}; +/// use futures::task::Poll; +/// +/// let mut r = stream::poll_immediate(Box::pin(stream::iter(1_u32..3))); +/// assert_eq!(r.next().await, Some(Poll::Ready(1))); +/// assert_eq!(r.next().await, Some(Poll::Ready(2))); +/// assert_eq!(r.next().await, None); +/// +/// let mut p = stream::poll_immediate(Box::pin(stream::once(async { +/// futures::pending!(); +/// 42_u8 +/// }))); +/// assert_eq!(p.next().await, Some(Poll::Pending)); +/// assert_eq!(p.next().await, Some(Poll::Ready(42))); +/// assert_eq!(p.next().await, None); +/// # }); +/// ``` +pub fn poll_immediate<S: Stream>(s: S) -> PollImmediate<S> { + super::assert_stream::<Poll<S::Item>, PollImmediate<S>>(PollImmediate { stream: Some(s) }) +} diff --git a/third_party/rust/futures-util/src/stream/repeat.rs b/third_party/rust/futures-util/src/stream/repeat.rs new file mode 100644 index 0000000000..3f9aa87d5c --- /dev/null +++ b/third_party/rust/futures-util/src/stream/repeat.rs @@ -0,0 +1,58 @@ +use super::assert_stream; +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; + +/// Stream for the [`repeat`] function. +#[derive(Debug, Clone)] +#[must_use = "streams do nothing unless polled"] +pub struct Repeat<T> { + item: T, +} + +/// Create a stream which produces the same item repeatedly. +/// +/// The stream never terminates. Note that you likely want to avoid +/// usage of `collect` or such on the returned stream as it will exhaust +/// available memory as it tries to just fill up all RAM. +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::stream::{self, StreamExt}; +/// +/// let stream = stream::repeat(9); +/// assert_eq!(vec![9, 9, 9], stream.take(3).collect::<Vec<i32>>().await); +/// # }); +/// ``` +pub fn repeat<T>(item: T) -> Repeat<T> +where + T: Clone, +{ + assert_stream::<T, _>(Repeat { item }) +} + +impl<T> Unpin for Repeat<T> {} + +impl<T> Stream for Repeat<T> +where + T: Clone, +{ + type Item = T; + + fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> { + Poll::Ready(Some(self.item.clone())) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (usize::max_value(), None) + } +} + +impl<T> FusedStream for Repeat<T> +where + T: Clone, +{ + fn is_terminated(&self) -> bool { + false + } +} diff --git a/third_party/rust/futures-util/src/stream/repeat_with.rs b/third_party/rust/futures-util/src/stream/repeat_with.rs new file mode 100644 index 0000000000..f5a81b4ed4 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/repeat_with.rs @@ -0,0 +1,93 @@ +use super::assert_stream; +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; + +/// An stream that repeats elements of type `A` endlessly by +/// applying the provided closure `F: FnMut() -> A`. +/// +/// This `struct` is created by the [`repeat_with()`] function. +/// See its documentation for more. +#[derive(Debug, Clone)] +#[must_use = "streams do nothing unless polled"] +pub struct RepeatWith<F> { + repeater: F, +} + +impl<A, F: FnMut() -> A> Unpin for RepeatWith<F> {} + +impl<A, F: FnMut() -> A> Stream for RepeatWith<F> { + type Item = A; + + fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> { + Poll::Ready(Some((&mut self.repeater)())) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (usize::max_value(), None) + } +} + +impl<A, F: FnMut() -> A> FusedStream for RepeatWith<F> { + fn is_terminated(&self) -> bool { + false + } +} + +/// Creates a new stream that repeats elements of type `A` endlessly by +/// applying the provided closure, the repeater, `F: FnMut() -> A`. +/// +/// The `repeat_with()` function calls the repeater over and over again. +/// +/// Infinite stream like `repeat_with()` are often used with adapters like +/// [`stream.take()`], in order to make them finite. +/// +/// If the element type of the stream you need implements [`Clone`], and +/// it is OK to keep the source element in memory, you should instead use +/// the [`stream.repeat()`] function. +/// +/// # Examples +/// +/// Basic usage: +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::stream::{self, StreamExt}; +/// +/// // let's assume we have some value of a type that is not `Clone` +/// // or which don't want to have in memory just yet because it is expensive: +/// #[derive(PartialEq, Debug)] +/// struct Expensive; +/// +/// // a particular value forever: +/// let mut things = stream::repeat_with(|| Expensive); +/// +/// assert_eq!(Some(Expensive), things.next().await); +/// assert_eq!(Some(Expensive), things.next().await); +/// assert_eq!(Some(Expensive), things.next().await); +/// # }); +/// ``` +/// +/// Using mutation and going finite: +/// +/// ```rust +/// # futures::executor::block_on(async { +/// use futures::stream::{self, StreamExt}; +/// +/// // From the zeroth to the third power of two: +/// let mut curr = 1; +/// let mut pow2 = stream::repeat_with(|| { let tmp = curr; curr *= 2; tmp }) +/// .take(4); +/// +/// assert_eq!(Some(1), pow2.next().await); +/// assert_eq!(Some(2), pow2.next().await); +/// assert_eq!(Some(4), pow2.next().await); +/// assert_eq!(Some(8), pow2.next().await); +/// +/// // ... and now we're done +/// assert_eq!(None, pow2.next().await); +/// # }); +/// ``` +pub fn repeat_with<A, F: FnMut() -> A>(repeater: F) -> RepeatWith<F> { + assert_stream::<A, _>(RepeatWith { repeater }) +} diff --git a/third_party/rust/futures-util/src/stream/select.rs b/third_party/rust/futures-util/src/stream/select.rs new file mode 100644 index 0000000000..0c1e3af782 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/select.rs @@ -0,0 +1,117 @@ +use super::assert_stream; +use crate::stream::{select_with_strategy, PollNext, SelectWithStrategy}; +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`select()`] function. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Select<St1, St2> { + #[pin] + inner: SelectWithStrategy<St1, St2, fn(&mut PollNext)-> PollNext, PollNext>, + } +} + +/// This function will attempt to pull items from both streams. Each +/// stream will be polled in a round-robin fashion, and whenever a stream is +/// ready to yield an item that item is yielded. +/// +/// After one of the two input streams completes, the remaining one will be +/// polled exclusively. The returned stream completes when both input +/// streams have completed. +/// +/// Note that this function consumes both streams and returns a wrapped +/// version of them. +/// +/// ## Examples +/// +/// ```rust +/// # futures::executor::block_on(async { +/// use futures::stream::{ repeat, select, StreamExt }; +/// +/// let left = repeat(1); +/// let right = repeat(2); +/// +/// let mut out = select(left, right); +/// +/// for _ in 0..100 { +/// // We should be alternating. +/// assert_eq!(1, out.select_next_some().await); +/// assert_eq!(2, out.select_next_some().await); +/// } +/// # }); +/// ``` +pub fn select<St1, St2>(stream1: St1, stream2: St2) -> Select<St1, St2> +where + St1: Stream, + St2: Stream<Item = St1::Item>, +{ + fn round_robin(last: &mut PollNext) -> PollNext { + last.toggle() + } + + assert_stream::<St1::Item, _>(Select { + inner: select_with_strategy(stream1, stream2, round_robin), + }) +} + +impl<St1, St2> Select<St1, St2> { + /// Acquires a reference to the underlying streams that this combinator is + /// pulling from. + pub fn get_ref(&self) -> (&St1, &St2) { + self.inner.get_ref() + } + + /// Acquires a mutable reference to the underlying streams that this + /// combinator is pulling from. + /// + /// Note that care must be taken to avoid tampering with the state of the + /// stream which may otherwise confuse this combinator. + pub fn get_mut(&mut self) -> (&mut St1, &mut St2) { + self.inner.get_mut() + } + + /// Acquires a pinned mutable reference to the underlying streams that this + /// combinator is pulling from. + /// + /// Note that care must be taken to avoid tampering with the state of the + /// stream which may otherwise confuse this combinator. + pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut St1>, Pin<&mut St2>) { + let this = self.project(); + this.inner.get_pin_mut() + } + + /// Consumes this combinator, returning the underlying streams. + /// + /// Note that this may discard intermediate state of this combinator, so + /// care should be taken to avoid losing resources when this is called. + pub fn into_inner(self) -> (St1, St2) { + self.inner.into_inner() + } +} + +impl<St1, St2> FusedStream for Select<St1, St2> +where + St1: Stream, + St2: Stream<Item = St1::Item>, +{ + fn is_terminated(&self) -> bool { + self.inner.is_terminated() + } +} + +impl<St1, St2> Stream for Select<St1, St2> +where + St1: Stream, + St2: Stream<Item = St1::Item>, +{ + type Item = St1::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St1::Item>> { + let this = self.project(); + this.inner.poll_next(cx) + } +} diff --git a/third_party/rust/futures-util/src/stream/select_all.rs b/third_party/rust/futures-util/src/stream/select_all.rs new file mode 100644 index 0000000000..3474331adc --- /dev/null +++ b/third_party/rust/futures-util/src/stream/select_all.rs @@ -0,0 +1,254 @@ +//! An unbounded set of streams + +use core::fmt::{self, Debug}; +use core::iter::FromIterator; +use core::pin::Pin; + +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; + +use pin_project_lite::pin_project; + +use super::assert_stream; +use crate::stream::{futures_unordered, FuturesUnordered, StreamExt, StreamFuture}; + +pin_project! { + /// An unbounded set of streams + /// + /// This "combinator" provides the ability to maintain a set of streams + /// and drive them all to completion. + /// + /// Streams are pushed into this set and their realized values are + /// yielded as they become ready. Streams will only be polled when they + /// generate notifications. This allows to coordinate a large number of streams. + /// + /// Note that you can create a ready-made `SelectAll` via the + /// `select_all` function in the `stream` module, or you can start with an + /// empty set with the `SelectAll::new` constructor. + #[must_use = "streams do nothing unless polled"] + pub struct SelectAll<St> { + #[pin] + inner: FuturesUnordered<StreamFuture<St>>, + } +} + +impl<St: Debug> Debug for SelectAll<St> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "SelectAll {{ ... }}") + } +} + +impl<St: Stream + Unpin> SelectAll<St> { + /// Constructs a new, empty `SelectAll` + /// + /// The returned `SelectAll` does not contain any streams and, in this + /// state, `SelectAll::poll` will return `Poll::Ready(None)`. + pub fn new() -> Self { + Self { inner: FuturesUnordered::new() } + } + + /// Returns the number of streams contained in the set. + /// + /// This represents the total number of in-flight streams. + pub fn len(&self) -> usize { + self.inner.len() + } + + /// Returns `true` if the set contains no streams + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + /// Push a stream into the set. + /// + /// This function submits the given stream to the set for managing. This + /// function will not call `poll` on the submitted stream. The caller must + /// ensure that `SelectAll::poll` is called in order to receive task + /// notifications. + pub fn push(&mut self, stream: St) { + self.inner.push(stream.into_future()); + } + + /// Returns an iterator that allows inspecting each stream in the set. + pub fn iter(&self) -> Iter<'_, St> { + Iter(self.inner.iter()) + } + + /// Returns an iterator that allows modifying each stream in the set. + pub fn iter_mut(&mut self) -> IterMut<'_, St> { + IterMut(self.inner.iter_mut()) + } + + /// Clears the set, removing all streams. + pub fn clear(&mut self) { + self.inner.clear() + } +} + +impl<St: Stream + Unpin> Default for SelectAll<St> { + fn default() -> Self { + Self::new() + } +} + +impl<St: Stream + Unpin> Stream for SelectAll<St> { + type Item = St::Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + loop { + match ready!(self.inner.poll_next_unpin(cx)) { + Some((Some(item), remaining)) => { + self.push(remaining); + return Poll::Ready(Some(item)); + } + Some((None, _)) => { + // `FuturesUnordered` thinks it isn't terminated + // because it yielded a Some. + // We do not return, but poll `FuturesUnordered` + // in the next loop iteration. + } + None => return Poll::Ready(None), + } + } + } +} + +impl<St: Stream + Unpin> FusedStream for SelectAll<St> { + fn is_terminated(&self) -> bool { + self.inner.is_terminated() + } +} + +/// Convert a list of streams into a `Stream` of results from the streams. +/// +/// This essentially takes a list of streams (e.g. a vector, an iterator, etc.) +/// and bundles them together into a single stream. +/// The stream will yield items as they become available on the underlying +/// streams internally, in the order they become available. +/// +/// Note that the returned set can also be used to dynamically push more +/// streams into the set as they become available. +/// +/// This function is only available when the `std` or `alloc` feature of this +/// library is activated, and it is activated by default. +pub fn select_all<I>(streams: I) -> SelectAll<I::Item> +where + I: IntoIterator, + I::Item: Stream + Unpin, +{ + let mut set = SelectAll::new(); + + for stream in streams { + set.push(stream); + } + + assert_stream::<<I::Item as Stream>::Item, _>(set) +} + +impl<St: Stream + Unpin> FromIterator<St> for SelectAll<St> { + fn from_iter<T: IntoIterator<Item = St>>(iter: T) -> Self { + select_all(iter) + } +} + +impl<St: Stream + Unpin> Extend<St> for SelectAll<St> { + fn extend<T: IntoIterator<Item = St>>(&mut self, iter: T) { + for st in iter { + self.push(st) + } + } +} + +impl<St: Stream + Unpin> IntoIterator for SelectAll<St> { + type Item = St; + type IntoIter = IntoIter<St>; + + fn into_iter(self) -> Self::IntoIter { + IntoIter(self.inner.into_iter()) + } +} + +impl<'a, St: Stream + Unpin> IntoIterator for &'a SelectAll<St> { + type Item = &'a St; + type IntoIter = Iter<'a, St>; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl<'a, St: Stream + Unpin> IntoIterator for &'a mut SelectAll<St> { + type Item = &'a mut St; + type IntoIter = IterMut<'a, St>; + + fn into_iter(self) -> Self::IntoIter { + self.iter_mut() + } +} + +/// Immutable iterator over all streams in the unordered set. +#[derive(Debug)] +pub struct Iter<'a, St: Unpin>(futures_unordered::Iter<'a, StreamFuture<St>>); + +/// Mutable iterator over all streams in the unordered set. +#[derive(Debug)] +pub struct IterMut<'a, St: Unpin>(futures_unordered::IterMut<'a, StreamFuture<St>>); + +/// Owned iterator over all streams in the unordered set. +#[derive(Debug)] +pub struct IntoIter<St: Unpin>(futures_unordered::IntoIter<StreamFuture<St>>); + +impl<'a, St: Stream + Unpin> Iterator for Iter<'a, St> { + type Item = &'a St; + + fn next(&mut self) -> Option<Self::Item> { + let st = self.0.next()?; + let next = st.get_ref(); + // This should always be true because FuturesUnordered removes completed futures. + debug_assert!(next.is_some()); + next + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.0.size_hint() + } +} + +impl<St: Stream + Unpin> ExactSizeIterator for Iter<'_, St> {} + +impl<'a, St: Stream + Unpin> Iterator for IterMut<'a, St> { + type Item = &'a mut St; + + fn next(&mut self) -> Option<Self::Item> { + let st = self.0.next()?; + let next = st.get_mut(); + // This should always be true because FuturesUnordered removes completed futures. + debug_assert!(next.is_some()); + next + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.0.size_hint() + } +} + +impl<St: Stream + Unpin> ExactSizeIterator for IterMut<'_, St> {} + +impl<St: Stream + Unpin> Iterator for IntoIter<St> { + type Item = St; + + fn next(&mut self) -> Option<Self::Item> { + let st = self.0.next()?; + let next = st.into_inner(); + // This should always be true because FuturesUnordered removes completed futures. + debug_assert!(next.is_some()); + next + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.0.size_hint() + } +} + +impl<St: Stream + Unpin> ExactSizeIterator for IntoIter<St> {} diff --git a/third_party/rust/futures-util/src/stream/select_with_strategy.rs b/third_party/rust/futures-util/src/stream/select_with_strategy.rs new file mode 100644 index 0000000000..224d5f821c --- /dev/null +++ b/third_party/rust/futures-util/src/stream/select_with_strategy.rs @@ -0,0 +1,304 @@ +use super::assert_stream; +use core::{fmt, pin::Pin}; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +/// Type to tell [`SelectWithStrategy`] which stream to poll next. +#[derive(Debug, PartialEq, Eq, Copy, Clone, Hash)] +pub enum PollNext { + /// Poll the first stream. + Left, + /// Poll the second stream. + Right, +} + +impl PollNext { + /// Toggle the value and return the old one. + pub fn toggle(&mut self) -> Self { + let old = *self; + *self = self.other(); + old + } + + fn other(&self) -> PollNext { + match self { + PollNext::Left => PollNext::Right, + PollNext::Right => PollNext::Left, + } + } +} + +impl Default for PollNext { + fn default() -> Self { + PollNext::Left + } +} + +enum InternalState { + Start, + LeftFinished, + RightFinished, + BothFinished, +} + +impl InternalState { + fn finish(&mut self, ps: PollNext) { + match (&self, ps) { + (InternalState::Start, PollNext::Left) => { + *self = InternalState::LeftFinished; + } + (InternalState::Start, PollNext::Right) => { + *self = InternalState::RightFinished; + } + (InternalState::LeftFinished, PollNext::Right) + | (InternalState::RightFinished, PollNext::Left) => { + *self = InternalState::BothFinished; + } + _ => {} + } + } +} + +pin_project! { + /// Stream for the [`select_with_strategy()`] function. See function docs for details. + #[must_use = "streams do nothing unless polled"] + #[project = SelectWithStrategyProj] + pub struct SelectWithStrategy<St1, St2, Clos, State> { + #[pin] + stream1: St1, + #[pin] + stream2: St2, + internal_state: InternalState, + state: State, + clos: Clos, + } +} + +/// This function will attempt to pull items from both streams. You provide a +/// closure to tell [`SelectWithStrategy`] which stream to poll. The closure can +/// store state on `SelectWithStrategy` to which it will receive a `&mut` on every +/// invocation. This allows basing the strategy on prior choices. +/// +/// After one of the two input streams completes, the remaining one will be +/// polled exclusively. The returned stream completes when both input +/// streams have completed. +/// +/// Note that this function consumes both streams and returns a wrapped +/// version of them. +/// +/// ## Examples +/// +/// ### Priority +/// This example shows how to always prioritize the left stream. +/// +/// ```rust +/// # futures::executor::block_on(async { +/// use futures::stream::{ repeat, select_with_strategy, PollNext, StreamExt }; +/// +/// let left = repeat(1); +/// let right = repeat(2); +/// +/// // We don't need any state, so let's make it an empty tuple. +/// // We must provide some type here, as there is no way for the compiler +/// // to infer it. As we don't need to capture variables, we can just +/// // use a function pointer instead of a closure. +/// fn prio_left(_: &mut ()) -> PollNext { PollNext::Left } +/// +/// let mut out = select_with_strategy(left, right, prio_left); +/// +/// for _ in 0..100 { +/// // Whenever we poll out, we will alwas get `1`. +/// assert_eq!(1, out.select_next_some().await); +/// } +/// # }); +/// ``` +/// +/// ### Round Robin +/// This example shows how to select from both streams round robin. +/// Note: this special case is provided by [`futures-util::stream::select`]. +/// +/// ```rust +/// # futures::executor::block_on(async { +/// use futures::stream::{ repeat, select_with_strategy, PollNext, StreamExt }; +/// +/// let left = repeat(1); +/// let right = repeat(2); +/// +/// let rrobin = |last: &mut PollNext| last.toggle(); +/// +/// let mut out = select_with_strategy(left, right, rrobin); +/// +/// for _ in 0..100 { +/// // We should be alternating now. +/// assert_eq!(1, out.select_next_some().await); +/// assert_eq!(2, out.select_next_some().await); +/// } +/// # }); +/// ``` +pub fn select_with_strategy<St1, St2, Clos, State>( + stream1: St1, + stream2: St2, + which: Clos, +) -> SelectWithStrategy<St1, St2, Clos, State> +where + St1: Stream, + St2: Stream<Item = St1::Item>, + Clos: FnMut(&mut State) -> PollNext, + State: Default, +{ + assert_stream::<St1::Item, _>(SelectWithStrategy { + stream1, + stream2, + state: Default::default(), + internal_state: InternalState::Start, + clos: which, + }) +} + +impl<St1, St2, Clos, State> SelectWithStrategy<St1, St2, Clos, State> { + /// Acquires a reference to the underlying streams that this combinator is + /// pulling from. + pub fn get_ref(&self) -> (&St1, &St2) { + (&self.stream1, &self.stream2) + } + + /// Acquires a mutable reference to the underlying streams that this + /// combinator is pulling from. + /// + /// Note that care must be taken to avoid tampering with the state of the + /// stream which may otherwise confuse this combinator. + pub fn get_mut(&mut self) -> (&mut St1, &mut St2) { + (&mut self.stream1, &mut self.stream2) + } + + /// Acquires a pinned mutable reference to the underlying streams that this + /// combinator is pulling from. + /// + /// Note that care must be taken to avoid tampering with the state of the + /// stream which may otherwise confuse this combinator. + pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut St1>, Pin<&mut St2>) { + let this = self.project(); + (this.stream1, this.stream2) + } + + /// Consumes this combinator, returning the underlying streams. + /// + /// Note that this may discard intermediate state of this combinator, so + /// care should be taken to avoid losing resources when this is called. + pub fn into_inner(self) -> (St1, St2) { + (self.stream1, self.stream2) + } +} + +impl<St1, St2, Clos, State> FusedStream for SelectWithStrategy<St1, St2, Clos, State> +where + St1: Stream, + St2: Stream<Item = St1::Item>, + Clos: FnMut(&mut State) -> PollNext, +{ + fn is_terminated(&self) -> bool { + match self.internal_state { + InternalState::BothFinished => true, + _ => false, + } + } +} + +#[inline] +fn poll_side<St1, St2, Clos, State>( + select: &mut SelectWithStrategyProj<'_, St1, St2, Clos, State>, + side: PollNext, + cx: &mut Context<'_>, +) -> Poll<Option<St1::Item>> +where + St1: Stream, + St2: Stream<Item = St1::Item>, +{ + match side { + PollNext::Left => select.stream1.as_mut().poll_next(cx), + PollNext::Right => select.stream2.as_mut().poll_next(cx), + } +} + +#[inline] +fn poll_inner<St1, St2, Clos, State>( + select: &mut SelectWithStrategyProj<'_, St1, St2, Clos, State>, + side: PollNext, + cx: &mut Context<'_>, +) -> Poll<Option<St1::Item>> +where + St1: Stream, + St2: Stream<Item = St1::Item>, +{ + let first_done = match poll_side(select, side, cx) { + Poll::Ready(Some(item)) => return Poll::Ready(Some(item)), + Poll::Ready(None) => { + select.internal_state.finish(side); + true + } + Poll::Pending => false, + }; + let other = side.other(); + match poll_side(select, other, cx) { + Poll::Ready(None) => { + select.internal_state.finish(other); + if first_done { + Poll::Ready(None) + } else { + Poll::Pending + } + } + a => a, + } +} + +impl<St1, St2, Clos, State> Stream for SelectWithStrategy<St1, St2, Clos, State> +where + St1: Stream, + St2: Stream<Item = St1::Item>, + Clos: FnMut(&mut State) -> PollNext, +{ + type Item = St1::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St1::Item>> { + let mut this = self.project(); + + match this.internal_state { + InternalState::Start => { + let next_side = (this.clos)(this.state); + poll_inner(&mut this, next_side, cx) + } + InternalState::LeftFinished => match this.stream2.poll_next(cx) { + Poll::Ready(None) => { + *this.internal_state = InternalState::BothFinished; + Poll::Ready(None) + } + a => a, + }, + InternalState::RightFinished => match this.stream1.poll_next(cx) { + Poll::Ready(None) => { + *this.internal_state = InternalState::BothFinished; + Poll::Ready(None) + } + a => a, + }, + InternalState::BothFinished => Poll::Ready(None), + } + } +} + +impl<St1, St2, Clos, State> fmt::Debug for SelectWithStrategy<St1, St2, Clos, State> +where + St1: fmt::Debug, + St2: fmt::Debug, + State: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SelectWithStrategy") + .field("stream1", &self.stream1) + .field("stream2", &self.stream2) + .field("state", &self.state) + .finish() + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/all.rs b/third_party/rust/futures-util/src/stream/stream/all.rs new file mode 100644 index 0000000000..ba2baa5cf1 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/all.rs @@ -0,0 +1,92 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`all`](super::StreamExt::all) method. + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct All<St, Fut, F> { + #[pin] + stream: St, + f: F, + accum: Option<bool>, + #[pin] + future: Option<Fut>, + } +} + +impl<St, Fut, F> fmt::Debug for All<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("All") + .field("stream", &self.stream) + .field("accum", &self.accum) + .field("future", &self.future) + .finish() + } +} + +impl<St, Fut, F> All<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, accum: Some(true), future: None } + } +} + +impl<St, Fut, F> FusedFuture for All<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + fn is_terminated(&self) -> bool { + self.accum.is_none() && self.future.is_none() + } +} + +impl<St, Fut, F> Future for All<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + type Output = bool; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<bool> { + let mut this = self.project(); + Poll::Ready(loop { + if let Some(fut) = this.future.as_mut().as_pin_mut() { + // we're currently processing a future to produce a new accum value + let acc = this.accum.unwrap() && ready!(fut.poll(cx)); + if !acc { + break false; + } // early exit + *this.accum = Some(acc); + this.future.set(None); + } else if this.accum.is_some() { + // we're waiting on a new item from the stream + match ready!(this.stream.as_mut().poll_next(cx)) { + Some(item) => { + this.future.set(Some((this.f)(item))); + } + None => { + break this.accum.take().unwrap(); + } + } + } else { + panic!("All polled after completion") + } + }) + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/any.rs b/third_party/rust/futures-util/src/stream/stream/any.rs new file mode 100644 index 0000000000..f023125c70 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/any.rs @@ -0,0 +1,92 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`any`](super::StreamExt::any) method. + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Any<St, Fut, F> { + #[pin] + stream: St, + f: F, + accum: Option<bool>, + #[pin] + future: Option<Fut>, + } +} + +impl<St, Fut, F> fmt::Debug for Any<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Any") + .field("stream", &self.stream) + .field("accum", &self.accum) + .field("future", &self.future) + .finish() + } +} + +impl<St, Fut, F> Any<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, accum: Some(false), future: None } + } +} + +impl<St, Fut, F> FusedFuture for Any<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + fn is_terminated(&self) -> bool { + self.accum.is_none() && self.future.is_none() + } +} + +impl<St, Fut, F> Future for Any<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + type Output = bool; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<bool> { + let mut this = self.project(); + Poll::Ready(loop { + if let Some(fut) = this.future.as_mut().as_pin_mut() { + // we're currently processing a future to produce a new accum value + let acc = this.accum.unwrap() || ready!(fut.poll(cx)); + if acc { + break true; + } // early exit + *this.accum = Some(acc); + this.future.set(None); + } else if this.accum.is_some() { + // we're waiting on a new item from the stream + match ready!(this.stream.as_mut().poll_next(cx)) { + Some(item) => { + this.future.set(Some((this.f)(item))); + } + None => { + break this.accum.take().unwrap(); + } + } + } else { + panic!("Any polled after completion") + } + }) + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/buffer_unordered.rs b/third_party/rust/futures-util/src/stream/stream/buffer_unordered.rs new file mode 100644 index 0000000000..91b0f6bcce --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/buffer_unordered.rs @@ -0,0 +1,120 @@ +use crate::stream::{Fuse, FuturesUnordered, StreamExt}; +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`buffer_unordered`](super::StreamExt::buffer_unordered) + /// method. + #[must_use = "streams do nothing unless polled"] + pub struct BufferUnordered<St> + where + St: Stream, + { + #[pin] + stream: Fuse<St>, + in_progress_queue: FuturesUnordered<St::Item>, + max: usize, + } +} + +impl<St> fmt::Debug for BufferUnordered<St> +where + St: Stream + fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BufferUnordered") + .field("stream", &self.stream) + .field("in_progress_queue", &self.in_progress_queue) + .field("max", &self.max) + .finish() + } +} + +impl<St> BufferUnordered<St> +where + St: Stream, + St::Item: Future, +{ + pub(super) fn new(stream: St, n: usize) -> Self { + Self { + stream: super::Fuse::new(stream), + in_progress_queue: FuturesUnordered::new(), + max: n, + } + } + + delegate_access_inner!(stream, St, (.)); +} + +impl<St> Stream for BufferUnordered<St> +where + St: Stream, + St::Item: Future, +{ + type Item = <St::Item as Future>::Output; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + // First up, try to spawn off as many futures as possible by filling up + // our queue of futures. + while this.in_progress_queue.len() < *this.max { + match this.stream.as_mut().poll_next(cx) { + Poll::Ready(Some(fut)) => this.in_progress_queue.push(fut), + Poll::Ready(None) | Poll::Pending => break, + } + } + + // Attempt to pull the next value from the in_progress_queue + match this.in_progress_queue.poll_next_unpin(cx) { + x @ Poll::Pending | x @ Poll::Ready(Some(_)) => return x, + Poll::Ready(None) => {} + } + + // If more values are still coming from the stream, we're not done yet + if this.stream.is_done() { + Poll::Ready(None) + } else { + Poll::Pending + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let queue_len = self.in_progress_queue.len(); + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(queue_len); + let upper = match upper { + Some(x) => x.checked_add(queue_len), + None => None, + }; + (lower, upper) + } +} + +impl<St> FusedStream for BufferUnordered<St> +where + St: Stream, + St::Item: Future, +{ + fn is_terminated(&self) -> bool { + self.in_progress_queue.is_terminated() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for BufferUnordered<S> +where + S: Stream + Sink<Item>, + S::Item: Future, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/buffered.rs b/third_party/rust/futures-util/src/stream/stream/buffered.rs new file mode 100644 index 0000000000..5854eb7ea5 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/buffered.rs @@ -0,0 +1,118 @@ +use crate::stream::{Fuse, FusedStream, FuturesOrdered, StreamExt}; +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`buffered`](super::StreamExt::buffered) method. + #[must_use = "streams do nothing unless polled"] + pub struct Buffered<St> + where + St: Stream, + St::Item: Future, + { + #[pin] + stream: Fuse<St>, + in_progress_queue: FuturesOrdered<St::Item>, + max: usize, + } +} + +impl<St> fmt::Debug for Buffered<St> +where + St: Stream + fmt::Debug, + St::Item: Future, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Buffered") + .field("stream", &self.stream) + .field("in_progress_queue", &self.in_progress_queue) + .field("max", &self.max) + .finish() + } +} + +impl<St> Buffered<St> +where + St: Stream, + St::Item: Future, +{ + pub(super) fn new(stream: St, n: usize) -> Self { + Self { stream: super::Fuse::new(stream), in_progress_queue: FuturesOrdered::new(), max: n } + } + + delegate_access_inner!(stream, St, (.)); +} + +impl<St> Stream for Buffered<St> +where + St: Stream, + St::Item: Future, +{ + type Item = <St::Item as Future>::Output; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + // First up, try to spawn off as many futures as possible by filling up + // our queue of futures. + while this.in_progress_queue.len() < *this.max { + match this.stream.as_mut().poll_next(cx) { + Poll::Ready(Some(fut)) => this.in_progress_queue.push_back(fut), + Poll::Ready(None) | Poll::Pending => break, + } + } + + // Attempt to pull the next value from the in_progress_queue + let res = this.in_progress_queue.poll_next_unpin(cx); + if let Some(val) = ready!(res) { + return Poll::Ready(Some(val)); + } + + // If more values are still coming from the stream, we're not done yet + if this.stream.is_done() { + Poll::Ready(None) + } else { + Poll::Pending + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let queue_len = self.in_progress_queue.len(); + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(queue_len); + let upper = match upper { + Some(x) => x.checked_add(queue_len), + None => None, + }; + (lower, upper) + } +} + +impl<St> FusedStream for Buffered<St> +where + St: Stream, + St::Item: Future, +{ + fn is_terminated(&self) -> bool { + self.stream.is_done() && self.in_progress_queue.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Buffered<S> +where + S: Stream + Sink<Item>, + S::Item: Future, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/catch_unwind.rs b/third_party/rust/futures-util/src/stream/stream/catch_unwind.rs new file mode 100644 index 0000000000..09a6dc1b76 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/catch_unwind.rs @@ -0,0 +1,61 @@ +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; +use std::any::Any; +use std::panic::{catch_unwind, AssertUnwindSafe, UnwindSafe}; +use std::pin::Pin; + +pin_project! { + /// Stream for the [`catch_unwind`](super::StreamExt::catch_unwind) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct CatchUnwind<St> { + #[pin] + stream: St, + caught_unwind: bool, + } +} + +impl<St: Stream + UnwindSafe> CatchUnwind<St> { + pub(super) fn new(stream: St) -> Self { + Self { stream, caught_unwind: false } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St: Stream + UnwindSafe> Stream for CatchUnwind<St> { + type Item = Result<St::Item, Box<dyn Any + Send>>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + if *this.caught_unwind { + Poll::Ready(None) + } else { + let res = catch_unwind(AssertUnwindSafe(|| this.stream.as_mut().poll_next(cx))); + + match res { + Ok(poll) => poll.map(|opt| opt.map(Ok)), + Err(e) => { + *this.caught_unwind = true; + Poll::Ready(Some(Err(e))) + } + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.caught_unwind { + (0, Some(0)) + } else { + self.stream.size_hint() + } + } +} + +impl<St: FusedStream + UnwindSafe> FusedStream for CatchUnwind<St> { + fn is_terminated(&self) -> bool { + self.caught_unwind || self.stream.is_terminated() + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/chain.rs b/third_party/rust/futures-util/src/stream/stream/chain.rs new file mode 100644 index 0000000000..36ff1e533d --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/chain.rs @@ -0,0 +1,76 @@ +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`chain`](super::StreamExt::chain) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Chain<St1, St2> { + #[pin] + first: Option<St1>, + #[pin] + second: St2, + } +} + +// All interactions with `Pin<&mut Chain<..>>` happen through these methods +impl<St1, St2> Chain<St1, St2> +where + St1: Stream, + St2: Stream<Item = St1::Item>, +{ + pub(super) fn new(stream1: St1, stream2: St2) -> Self { + Self { first: Some(stream1), second: stream2 } + } +} + +impl<St1, St2> FusedStream for Chain<St1, St2> +where + St1: Stream, + St2: FusedStream<Item = St1::Item>, +{ + fn is_terminated(&self) -> bool { + self.first.is_none() && self.second.is_terminated() + } +} + +impl<St1, St2> Stream for Chain<St1, St2> +where + St1: Stream, + St2: Stream<Item = St1::Item>, +{ + type Item = St1::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + if let Some(first) = this.first.as_mut().as_pin_mut() { + if let Some(item) = ready!(first.poll_next(cx)) { + return Poll::Ready(Some(item)); + } + + this.first.set(None); + } + this.second.poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if let Some(first) = &self.first { + let (first_lower, first_upper) = first.size_hint(); + let (second_lower, second_upper) = self.second.size_hint(); + + let lower = first_lower.saturating_add(second_lower); + + let upper = match (first_upper, second_upper) { + (Some(x), Some(y)) => x.checked_add(y), + _ => None, + }; + + (lower, upper) + } else { + self.second.size_hint() + } + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/chunks.rs b/third_party/rust/futures-util/src/stream/stream/chunks.rs new file mode 100644 index 0000000000..2a71ebc6cc --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/chunks.rs @@ -0,0 +1,103 @@ +use crate::stream::Fuse; +use alloc::vec::Vec; +use core::mem; +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`chunks`](super::StreamExt::chunks) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Chunks<St: Stream> { + #[pin] + stream: Fuse<St>, + items: Vec<St::Item>, + cap: usize, // https://github.com/rust-lang/futures-rs/issues/1475 + } +} + +impl<St: Stream> Chunks<St> { + pub(super) fn new(stream: St, capacity: usize) -> Self { + assert!(capacity > 0); + + Self { + stream: super::Fuse::new(stream), + items: Vec::with_capacity(capacity), + cap: capacity, + } + } + + fn take(self: Pin<&mut Self>) -> Vec<St::Item> { + let cap = self.cap; + mem::replace(self.project().items, Vec::with_capacity(cap)) + } + + delegate_access_inner!(stream, St, (.)); +} + +impl<St: Stream> Stream for Chunks<St> { + type Item = Vec<St::Item>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.as_mut().project(); + loop { + match ready!(this.stream.as_mut().poll_next(cx)) { + // Push the item into the buffer and check whether it is full. + // If so, replace our buffer with a new and empty one and return + // the full one. + Some(item) => { + this.items.push(item); + if this.items.len() >= *this.cap { + return Poll::Ready(Some(self.take())); + } + } + + // Since the underlying stream ran out of values, return what we + // have buffered, if we have anything. + None => { + let last = if this.items.is_empty() { + None + } else { + let full_buf = mem::take(this.items); + Some(full_buf) + }; + + return Poll::Ready(last); + } + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let chunk_len = usize::from(!self.items.is_empty()); + let (lower, upper) = self.stream.size_hint(); + let lower = (lower / self.cap).saturating_add(chunk_len); + let upper = match upper { + Some(x) => x.checked_add(chunk_len), + None => None, + }; + (lower, upper) + } +} + +impl<St: FusedStream> FusedStream for Chunks<St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() && self.items.is_empty() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Chunks<S> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/collect.rs b/third_party/rust/futures-util/src/stream/stream/collect.rs new file mode 100644 index 0000000000..970ac26dbf --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/collect.rs @@ -0,0 +1,56 @@ +use core::mem; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`collect`](super::StreamExt::collect) method. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Collect<St, C> { + #[pin] + stream: St, + collection: C, + } +} + +impl<St: Stream, C: Default> Collect<St, C> { + fn finish(self: Pin<&mut Self>) -> C { + mem::take(self.project().collection) + } + + pub(super) fn new(stream: St) -> Self { + Self { stream, collection: Default::default() } + } +} + +impl<St, C> FusedFuture for Collect<St, C> +where + St: FusedStream, + C: Default + Extend<St::Item>, +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St, C> Future for Collect<St, C> +where + St: Stream, + C: Default + Extend<St::Item>, +{ + type Output = C; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<C> { + let mut this = self.as_mut().project(); + loop { + match ready!(this.stream.as_mut().poll_next(cx)) { + Some(e) => this.collection.extend(Some(e)), + None => return Poll::Ready(self.finish()), + } + } + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/concat.rs b/third_party/rust/futures-util/src/stream/stream/concat.rs new file mode 100644 index 0000000000..7e058b2315 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/concat.rs @@ -0,0 +1,62 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`concat`](super::StreamExt::concat) method. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Concat<St: Stream> { + #[pin] + stream: St, + accum: Option<St::Item>, + } +} + +impl<St> Concat<St> +where + St: Stream, + St::Item: Extend<<St::Item as IntoIterator>::Item> + IntoIterator + Default, +{ + pub(super) fn new(stream: St) -> Self { + Self { stream, accum: None } + } +} + +impl<St> Future for Concat<St> +where + St: Stream, + St::Item: Extend<<St::Item as IntoIterator>::Item> + IntoIterator + Default, +{ + type Output = St::Item; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut this = self.project(); + + loop { + match ready!(this.stream.as_mut().poll_next(cx)) { + None => return Poll::Ready(this.accum.take().unwrap_or_default()), + Some(e) => { + if let Some(a) = this.accum { + a.extend(e) + } else { + *this.accum = Some(e) + } + } + } + } + } +} + +impl<St> FusedFuture for Concat<St> +where + St: FusedStream, + St::Item: Extend<<St::Item as IntoIterator>::Item> + IntoIterator + Default, +{ + fn is_terminated(&self) -> bool { + self.accum.is_none() && self.stream.is_terminated() + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/count.rs b/third_party/rust/futures-util/src/stream/stream/count.rs new file mode 100644 index 0000000000..513cab7b6a --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/count.rs @@ -0,0 +1,53 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`count`](super::StreamExt::count) method. + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Count<St> { + #[pin] + stream: St, + count: usize + } +} + +impl<St> fmt::Debug for Count<St> +where + St: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Count").field("stream", &self.stream).field("count", &self.count).finish() + } +} + +impl<St: Stream> Count<St> { + pub(super) fn new(stream: St) -> Self { + Self { stream, count: 0 } + } +} + +impl<St: FusedStream> FusedFuture for Count<St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St: Stream> Future for Count<St> { + type Output = usize; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut this = self.project(); + + Poll::Ready(loop { + match ready!(this.stream.as_mut().poll_next(cx)) { + Some(_) => *this.count += 1, + None => break *this.count, + } + }) + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/cycle.rs b/third_party/rust/futures-util/src/stream/stream/cycle.rs new file mode 100644 index 0000000000..507431d24f --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/cycle.rs @@ -0,0 +1,68 @@ +use core::pin::Pin; +use core::usize; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`cycle`](super::StreamExt::cycle) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Cycle<St> { + orig: St, + #[pin] + stream: St, + } +} + +impl<St> Cycle<St> +where + St: Clone + Stream, +{ + pub(super) fn new(stream: St) -> Self { + Self { orig: stream.clone(), stream } + } +} + +impl<St> Stream for Cycle<St> +where + St: Clone + Stream, +{ + type Item = St::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + match ready!(this.stream.as_mut().poll_next(cx)) { + None => { + this.stream.set(this.orig.clone()); + this.stream.poll_next(cx) + } + item => Poll::Ready(item), + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + // the cycle stream is either empty or infinite + match self.orig.size_hint() { + size @ (0, Some(0)) => size, + (0, _) => (0, None), + _ => (usize::max_value(), None), + } + } +} + +impl<St> FusedStream for Cycle<St> +where + St: Clone + Stream, +{ + fn is_terminated(&self) -> bool { + // the cycle stream is either empty or infinite + if let (0, Some(0)) = self.size_hint() { + true + } else { + false + } + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/enumerate.rs b/third_party/rust/futures-util/src/stream/stream/enumerate.rs new file mode 100644 index 0000000000..1cf9d49aaa --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/enumerate.rs @@ -0,0 +1,64 @@ +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`enumerate`](super::StreamExt::enumerate) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Enumerate<St> { + #[pin] + stream: St, + count: usize, + } +} + +impl<St: Stream> Enumerate<St> { + pub(super) fn new(stream: St) -> Self { + Self { stream, count: 0 } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St: Stream + FusedStream> FusedStream for Enumerate<St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St: Stream> Stream for Enumerate<St> { + type Item = (usize, St::Item); + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let this = self.project(); + + match ready!(this.stream.poll_next(cx)) { + Some(item) => { + let prev_count = *this.count; + *this.count += 1; + Poll::Ready(Some((prev_count, item))) + } + None => Poll::Ready(None), + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Enumerate<S> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/filter.rs b/third_party/rust/futures-util/src/stream/stream/filter.rs new file mode 100644 index 0000000000..997fe9977e --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/filter.rs @@ -0,0 +1,117 @@ +use crate::fns::FnMut1; +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`filter`](super::StreamExt::filter) method. + #[must_use = "streams do nothing unless polled"] + pub struct Filter<St, Fut, F> + where St: Stream, + { + #[pin] + stream: St, + f: F, + #[pin] + pending_fut: Option<Fut>, + pending_item: Option<St::Item>, + } +} + +impl<St, Fut, F> fmt::Debug for Filter<St, Fut, F> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Filter") + .field("stream", &self.stream) + .field("pending_fut", &self.pending_fut) + .field("pending_item", &self.pending_item) + .finish() + } +} + +#[allow(single_use_lifetimes)] // https://github.com/rust-lang/rust/issues/55058 +impl<St, Fut, F> Filter<St, Fut, F> +where + St: Stream, + F: for<'a> FnMut1<&'a St::Item, Output = Fut>, + Fut: Future<Output = bool>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, pending_fut: None, pending_item: None } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F> FusedStream for Filter<St, Fut, F> +where + St: Stream + FusedStream, + F: FnMut(&St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + fn is_terminated(&self) -> bool { + self.pending_fut.is_none() && self.stream.is_terminated() + } +} + +#[allow(single_use_lifetimes)] // https://github.com/rust-lang/rust/issues/55058 +impl<St, Fut, F> Stream for Filter<St, Fut, F> +where + St: Stream, + F: for<'a> FnMut1<&'a St::Item, Output = Fut>, + Fut: Future<Output = bool>, +{ + type Item = St::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St::Item>> { + let mut this = self.project(); + Poll::Ready(loop { + if let Some(fut) = this.pending_fut.as_mut().as_pin_mut() { + let res = ready!(fut.poll(cx)); + this.pending_fut.set(None); + if res { + break this.pending_item.take(); + } + *this.pending_item = None; + } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) { + this.pending_fut.set(Some(this.f.call_mut(&item))); + *this.pending_item = Some(item); + } else { + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = usize::from(self.pending_item.is_some()); + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for Filter<S, Fut, F> +where + S: Stream + Sink<Item>, + F: FnMut(&S::Item) -> Fut, + Fut: Future<Output = bool>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/filter_map.rs b/third_party/rust/futures-util/src/stream/stream/filter_map.rs new file mode 100644 index 0000000000..6b7d0070df --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/filter_map.rs @@ -0,0 +1,111 @@ +use crate::fns::FnMut1; +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`filter_map`](super::StreamExt::filter_map) method. + #[must_use = "streams do nothing unless polled"] + pub struct FilterMap<St, Fut, F> { + #[pin] + stream: St, + f: F, + #[pin] + pending: Option<Fut>, + } +} + +impl<St, Fut, F> fmt::Debug for FilterMap<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FilterMap") + .field("stream", &self.stream) + .field("pending", &self.pending) + .finish() + } +} + +impl<St, Fut, F> FilterMap<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, pending: None } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F, T> FusedStream for FilterMap<St, Fut, F> +where + St: Stream + FusedStream, + F: FnMut1<St::Item, Output = Fut>, + Fut: Future<Output = Option<T>>, +{ + fn is_terminated(&self) -> bool { + self.pending.is_none() && self.stream.is_terminated() + } +} + +impl<St, Fut, F, T> Stream for FilterMap<St, Fut, F> +where + St: Stream, + F: FnMut1<St::Item, Output = Fut>, + Fut: Future<Output = Option<T>>, +{ + type Item = T; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> { + let mut this = self.project(); + Poll::Ready(loop { + if let Some(p) = this.pending.as_mut().as_pin_mut() { + // We have an item in progress, poll that until it's done + let item = ready!(p.poll(cx)); + this.pending.set(None); + if item.is_some() { + break item; + } + } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) { + // No item in progress, but the stream is still going + this.pending.set(Some(this.f.call_mut(item))); + } else { + // The stream is done + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = usize::from(self.pending.is_some()); + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for FilterMap<S, Fut, F> +where + S: Stream + Sink<Item>, + F: FnMut1<S::Item, Output = Fut>, + Fut: Future, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/flatten.rs b/third_party/rust/futures-util/src/stream/stream/flatten.rs new file mode 100644 index 0000000000..9f6b7a472d --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/flatten.rs @@ -0,0 +1,73 @@ +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`flatten`](super::StreamExt::flatten) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Flatten<St, U> { + #[pin] + stream: St, + #[pin] + next: Option<U>, + } +} + +impl<St, U> Flatten<St, U> { + pub(super) fn new(stream: St) -> Self { + Self { stream, next: None } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St> FusedStream for Flatten<St, St::Item> +where + St: FusedStream, + St::Item: Stream, +{ + fn is_terminated(&self) -> bool { + self.next.is_none() && self.stream.is_terminated() + } +} + +impl<St> Stream for Flatten<St, St::Item> +where + St: Stream, + St::Item: Stream, +{ + type Item = <St::Item as Stream>::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + Poll::Ready(loop { + if let Some(s) = this.next.as_mut().as_pin_mut() { + if let Some(item) = ready!(s.poll_next(cx)) { + break Some(item); + } else { + this.next.set(None); + } + } else if let Some(s) = ready!(this.stream.as_mut().poll_next(cx)) { + this.next.set(Some(s)); + } else { + break None; + } + }) + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Flatten<S, S::Item> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/flatten_unordered.rs b/third_party/rust/futures-util/src/stream/stream/flatten_unordered.rs new file mode 100644 index 0000000000..07f971c55a --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/flatten_unordered.rs @@ -0,0 +1,509 @@ +use alloc::sync::Arc; +use core::{ + cell::UnsafeCell, + convert::identity, + fmt, + num::NonZeroUsize, + pin::Pin, + sync::atomic::{AtomicU8, Ordering}, +}; + +use pin_project_lite::pin_project; + +use futures_core::{ + future::Future, + ready, + stream::{FusedStream, Stream}, + task::{Context, Poll, Waker}, +}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use futures_task::{waker, ArcWake}; + +use crate::stream::FuturesUnordered; + +/// There is nothing to poll and stream isn't being +/// polled or waking at the moment. +const NONE: u8 = 0; + +/// Inner streams need to be polled. +const NEED_TO_POLL_INNER_STREAMS: u8 = 1; + +/// The base stream needs to be polled. +const NEED_TO_POLL_STREAM: u8 = 0b10; + +/// It needs to poll base stream and inner streams. +const NEED_TO_POLL_ALL: u8 = NEED_TO_POLL_INNER_STREAMS | NEED_TO_POLL_STREAM; + +/// The current stream is being polled at the moment. +const POLLING: u8 = 0b100; + +/// Inner streams are being woken at the moment. +const WAKING_INNER_STREAMS: u8 = 0b1000; + +/// The base stream is being woken at the moment. +const WAKING_STREAM: u8 = 0b10000; + +/// The base stream and inner streams are being woken at the moment. +const WAKING_ALL: u8 = WAKING_STREAM | WAKING_INNER_STREAMS; + +/// The stream was waked and will be polled. +const WOKEN: u8 = 0b100000; + +/// Determines what needs to be polled, and is stream being polled at the +/// moment or not. +#[derive(Clone, Debug)] +struct SharedPollState { + state: Arc<AtomicU8>, +} + +impl SharedPollState { + /// Constructs new `SharedPollState` with the given state. + fn new(value: u8) -> SharedPollState { + SharedPollState { state: Arc::new(AtomicU8::new(value)) } + } + + /// Attempts to start polling, returning stored state in case of success. + /// Returns `None` if some waker is waking at the moment. + fn start_polling( + &self, + ) -> Option<(u8, PollStateBomb<'_, impl FnOnce(&SharedPollState) -> u8>)> { + let value = self + .state + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |value| { + if value & WAKING_ALL == NONE { + Some(POLLING) + } else { + None + } + }) + .ok()?; + let bomb = PollStateBomb::new(self, SharedPollState::reset); + + Some((value, bomb)) + } + + /// Starts the waking process and performs bitwise or with the given value. + fn start_waking( + &self, + to_poll: u8, + waking: u8, + ) -> Option<(u8, PollStateBomb<'_, impl FnOnce(&SharedPollState) -> u8>)> { + let value = self + .state + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |value| { + // Waking process for this waker already started + if value & waking != NONE { + return None; + } + let mut next_value = value | to_poll; + // Only start the waking process if we're not in the polling phase and the stream isn't woken already + if value & (WOKEN | POLLING) == NONE { + next_value |= waking; + } + + if next_value != value { + Some(next_value) + } else { + None + } + }) + .ok()?; + + if value & (WOKEN | POLLING) == NONE { + let bomb = PollStateBomb::new(self, move |state| state.stop_waking(waking)); + + Some((value, bomb)) + } else { + None + } + } + + /// Sets current state to + /// - `!POLLING` allowing to use wakers + /// - `WOKEN` if the state was changed during `POLLING` phase as waker will be called, + /// or `will_be_woken` flag supplied + /// - `!WAKING_ALL` as + /// * Wakers called during the `POLLING` phase won't propagate their calls + /// * `POLLING` phase can't start if some of the wakers are active + /// So no wrapped waker can touch the inner waker's cell, it's safe to poll again. + fn stop_polling(&self, to_poll: u8, will_be_woken: bool) -> u8 { + self.state + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |mut value| { + let mut next_value = to_poll; + + value &= NEED_TO_POLL_ALL; + if value != NONE || will_be_woken { + next_value |= WOKEN; + } + next_value |= value; + + Some(next_value & !POLLING & !WAKING_ALL) + }) + .unwrap() + } + + /// Toggles state to non-waking, allowing to start polling. + fn stop_waking(&self, waking: u8) -> u8 { + self.state + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |value| { + let mut next_value = value & !waking; + // Waker will be called only if the current waking state is the same as the specified waker state + if value & WAKING_ALL == waking { + next_value |= WOKEN; + } + + if next_value != value { + Some(next_value) + } else { + None + } + }) + .unwrap_or_else(identity) + } + + /// Resets current state allowing to poll the stream and wake up wakers. + fn reset(&self) -> u8 { + self.state.swap(NEED_TO_POLL_ALL, Ordering::AcqRel) + } +} + +/// Used to execute some function on the given state when dropped. +struct PollStateBomb<'a, F: FnOnce(&SharedPollState) -> u8> { + state: &'a SharedPollState, + drop: Option<F>, +} + +impl<'a, F: FnOnce(&SharedPollState) -> u8> PollStateBomb<'a, F> { + /// Constructs new bomb with the given state. + fn new(state: &'a SharedPollState, drop: F) -> Self { + Self { state, drop: Some(drop) } + } + + /// Deactivates bomb, forces it to not call provided function when dropped. + fn deactivate(mut self) { + self.drop.take(); + } + + /// Manually fires the bomb, returning supplied state. + fn fire(mut self) -> Option<u8> { + self.drop.take().map(|drop| (drop)(self.state)) + } +} + +impl<F: FnOnce(&SharedPollState) -> u8> Drop for PollStateBomb<'_, F> { + fn drop(&mut self) { + if let Some(drop) = self.drop.take() { + (drop)(self.state); + } + } +} + +/// Will update state with the provided value on `wake_by_ref` call +/// and then, if there is a need, call `inner_waker`. +struct InnerWaker { + inner_waker: UnsafeCell<Option<Waker>>, + poll_state: SharedPollState, + need_to_poll: u8, +} + +unsafe impl Send for InnerWaker {} +unsafe impl Sync for InnerWaker {} + +impl InnerWaker { + /// Replaces given waker's inner_waker for polling stream/futures which will + /// update poll state on `wake_by_ref` call. Use only if you need several + /// contexts. + /// + /// ## Safety + /// + /// This function will modify waker's `inner_waker` via `UnsafeCell`, so + /// it should be used only during `POLLING` phase. + unsafe fn replace_waker(self_arc: &mut Arc<Self>, cx: &Context<'_>) -> Waker { + *self_arc.inner_waker.get() = cx.waker().clone().into(); + waker(self_arc.clone()) + } + + /// Attempts to start the waking process for the waker with the given value. + /// If succeeded, then the stream isn't yet woken and not being polled at the moment. + fn start_waking(&self) -> Option<(u8, PollStateBomb<'_, impl FnOnce(&SharedPollState) -> u8>)> { + self.poll_state.start_waking(self.need_to_poll, self.waking_state()) + } + + /// Returns the corresponding waking state toggled by this waker. + fn waking_state(&self) -> u8 { + self.need_to_poll << 3 + } +} + +impl ArcWake for InnerWaker { + fn wake_by_ref(self_arc: &Arc<Self>) { + if let Some((_, state_bomb)) = self_arc.start_waking() { + // Safety: now state is not `POLLING` + let waker_opt = unsafe { self_arc.inner_waker.get().as_ref().unwrap() }; + + if let Some(inner_waker) = waker_opt.clone() { + // Stop waking to allow polling stream + let poll_state_value = state_bomb.fire().unwrap(); + + // Here we want to call waker only if stream isn't woken yet and + // also to optimize the case when two wakers are called at the same time. + // + // In this case the best strategy will be to propagate only the latest waker's awake, + // and then poll both entities in a single `poll_next` call + if poll_state_value & (WOKEN | WAKING_ALL) == self_arc.waking_state() { + // Wake up inner waker + inner_waker.wake(); + } + } + } + } +} + +pin_project! { + /// Future which contains optional stream. + /// + /// If it's `Some`, it will attempt to call `poll_next` on it, + /// returning `Some((item, next_item_fut))` in case of `Poll::Ready(Some(...))` + /// or `None` in case of `Poll::Ready(None)`. + /// + /// If `poll_next` will return `Poll::Pending`, it will be forwarded to + /// the future and current task will be notified by waker. + #[must_use = "futures do nothing unless you `.await` or poll them"] + struct PollStreamFut<St> { + #[pin] + stream: Option<St>, + } +} + +impl<St> PollStreamFut<St> { + /// Constructs new `PollStreamFut` using given `stream`. + fn new(stream: impl Into<Option<St>>) -> Self { + Self { stream: stream.into() } + } +} + +impl<St: Stream + Unpin> Future for PollStreamFut<St> { + type Output = Option<(St::Item, PollStreamFut<St>)>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut stream = self.project().stream; + + let item = if let Some(stream) = stream.as_mut().as_pin_mut() { + ready!(stream.poll_next(cx)) + } else { + None + }; + let next_item_fut = PollStreamFut::new(stream.get_mut().take()); + let out = item.map(|item| (item, next_item_fut)); + + Poll::Ready(out) + } +} + +pin_project! { + /// Stream for the [`flatten_unordered`](super::StreamExt::flatten_unordered) + /// method. + #[project = FlattenUnorderedProj] + #[must_use = "streams do nothing unless polled"] + pub struct FlattenUnordered<St> where St: Stream { + #[pin] + inner_streams: FuturesUnordered<PollStreamFut<St::Item>>, + #[pin] + stream: St, + poll_state: SharedPollState, + limit: Option<NonZeroUsize>, + is_stream_done: bool, + inner_streams_waker: Arc<InnerWaker>, + stream_waker: Arc<InnerWaker>, + } +} + +impl<St> fmt::Debug for FlattenUnordered<St> +where + St: Stream + fmt::Debug, + St::Item: Stream + fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FlattenUnordered") + .field("poll_state", &self.poll_state) + .field("inner_streams", &self.inner_streams) + .field("limit", &self.limit) + .field("stream", &self.stream) + .field("is_stream_done", &self.is_stream_done) + .finish() + } +} + +impl<St> FlattenUnordered<St> +where + St: Stream, + St::Item: Stream + Unpin, +{ + pub(super) fn new(stream: St, limit: Option<usize>) -> FlattenUnordered<St> { + let poll_state = SharedPollState::new(NEED_TO_POLL_STREAM); + + FlattenUnordered { + inner_streams: FuturesUnordered::new(), + stream, + is_stream_done: false, + limit: limit.and_then(NonZeroUsize::new), + inner_streams_waker: Arc::new(InnerWaker { + inner_waker: UnsafeCell::new(None), + poll_state: poll_state.clone(), + need_to_poll: NEED_TO_POLL_INNER_STREAMS, + }), + stream_waker: Arc::new(InnerWaker { + inner_waker: UnsafeCell::new(None), + poll_state: poll_state.clone(), + need_to_poll: NEED_TO_POLL_STREAM, + }), + poll_state, + } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St> FlattenUnorderedProj<'_, St> +where + St: Stream, +{ + /// Checks if current `inner_streams` size is less than optional limit. + fn is_exceeded_limit(&self) -> bool { + self.limit.map_or(false, |limit| self.inner_streams.len() >= limit.get()) + } +} + +impl<St> FusedStream for FlattenUnordered<St> +where + St: FusedStream, + St::Item: FusedStream + Unpin, +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() && self.inner_streams.is_empty() + } +} + +impl<St> Stream for FlattenUnordered<St> +where + St: Stream, + St::Item: Stream + Unpin, +{ + type Item = <St::Item as Stream>::Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut next_item = None; + let mut need_to_poll_next = NONE; + + let mut this = self.as_mut().project(); + + let (mut poll_state_value, state_bomb) = match this.poll_state.start_polling() { + Some(value) => value, + _ => { + // Waker was called, just wait for the next poll + return Poll::Pending; + } + }; + + if poll_state_value & NEED_TO_POLL_STREAM != NONE { + // Safety: now state is `POLLING`. + let stream_waker = unsafe { InnerWaker::replace_waker(this.stream_waker, cx) }; + + // Here we need to poll the base stream. + // + // To improve performance, we will attempt to place as many items as we can + // to the `FuturesUnordered` bucket before polling inner streams + loop { + if this.is_exceeded_limit() || *this.is_stream_done { + // We either exceeded the limit or the stream is exhausted + if !*this.is_stream_done { + // The stream needs to be polled in the next iteration + need_to_poll_next |= NEED_TO_POLL_STREAM; + } + + break; + } else { + match this.stream.as_mut().poll_next(&mut Context::from_waker(&stream_waker)) { + Poll::Ready(Some(inner_stream)) => { + // Add new stream to the inner streams bucket + this.inner_streams.as_mut().push(PollStreamFut::new(inner_stream)); + // Inner streams must be polled afterward + poll_state_value |= NEED_TO_POLL_INNER_STREAMS; + } + Poll::Ready(None) => { + // Mark the stream as done + *this.is_stream_done = true; + } + Poll::Pending => { + break; + } + } + } + } + } + + if poll_state_value & NEED_TO_POLL_INNER_STREAMS != NONE { + // Safety: now state is `POLLING`. + let inner_streams_waker = + unsafe { InnerWaker::replace_waker(this.inner_streams_waker, cx) }; + + match this + .inner_streams + .as_mut() + .poll_next(&mut Context::from_waker(&inner_streams_waker)) + { + Poll::Ready(Some(Some((item, next_item_fut)))) => { + // Push next inner stream item future to the list of inner streams futures + this.inner_streams.as_mut().push(next_item_fut); + // Take the received item + next_item = Some(item); + // On the next iteration, inner streams must be polled again + need_to_poll_next |= NEED_TO_POLL_INNER_STREAMS; + } + Poll::Ready(Some(None)) => { + // On the next iteration, inner streams must be polled again + need_to_poll_next |= NEED_TO_POLL_INNER_STREAMS; + } + _ => {} + } + } + + // We didn't have any `poll_next` panic, so it's time to deactivate the bomb + state_bomb.deactivate(); + + let mut force_wake = + // we need to poll the stream and didn't reach the limit yet + need_to_poll_next & NEED_TO_POLL_STREAM != NONE && !this.is_exceeded_limit() + // or we need to poll inner streams again + || need_to_poll_next & NEED_TO_POLL_INNER_STREAMS != NONE; + + // Stop polling and swap the latest state + poll_state_value = this.poll_state.stop_polling(need_to_poll_next, force_wake); + // If state was changed during `POLLING` phase, need to manually call a waker + force_wake |= poll_state_value & NEED_TO_POLL_ALL != NONE; + + let is_done = *this.is_stream_done && this.inner_streams.is_empty(); + + if next_item.is_some() || is_done { + Poll::Ready(next_item) + } else { + if force_wake { + cx.waker().wake_by_ref(); + } + + Poll::Pending + } + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<St, Item> Sink<Item> for FlattenUnordered<St> +where + St: Stream + Sink<Item>, +{ + type Error = St::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/fold.rs b/third_party/rust/futures-util/src/stream/stream/fold.rs new file mode 100644 index 0000000000..b8b55ecb67 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/fold.rs @@ -0,0 +1,88 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`fold`](super::StreamExt::fold) method. + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Fold<St, Fut, T, F> { + #[pin] + stream: St, + f: F, + accum: Option<T>, + #[pin] + future: Option<Fut>, + } +} + +impl<St, Fut, T, F> fmt::Debug for Fold<St, Fut, T, F> +where + St: fmt::Debug, + Fut: fmt::Debug, + T: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Fold") + .field("stream", &self.stream) + .field("accum", &self.accum) + .field("future", &self.future) + .finish() + } +} + +impl<St, Fut, T, F> Fold<St, Fut, T, F> +where + St: Stream, + F: FnMut(T, St::Item) -> Fut, + Fut: Future<Output = T>, +{ + pub(super) fn new(stream: St, f: F, t: T) -> Self { + Self { stream, f, accum: Some(t), future: None } + } +} + +impl<St, Fut, T, F> FusedFuture for Fold<St, Fut, T, F> +where + St: Stream, + F: FnMut(T, St::Item) -> Fut, + Fut: Future<Output = T>, +{ + fn is_terminated(&self) -> bool { + self.accum.is_none() && self.future.is_none() + } +} + +impl<St, Fut, T, F> Future for Fold<St, Fut, T, F> +where + St: Stream, + F: FnMut(T, St::Item) -> Fut, + Fut: Future<Output = T>, +{ + type Output = T; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { + let mut this = self.project(); + Poll::Ready(loop { + if let Some(fut) = this.future.as_mut().as_pin_mut() { + // we're currently processing a future to produce a new accum value + *this.accum = Some(ready!(fut.poll(cx))); + this.future.set(None); + } else if this.accum.is_some() { + // we're waiting on a new item from the stream + let res = ready!(this.stream.as_mut().poll_next(cx)); + let a = this.accum.take().unwrap(); + if let Some(item) = res { + this.future.set(Some((this.f)(a, item))); + } else { + break a; + } + } else { + panic!("Fold polled after completion") + } + }) + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/for_each.rs b/third_party/rust/futures-util/src/stream/stream/for_each.rs new file mode 100644 index 0000000000..5302b0e034 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/for_each.rs @@ -0,0 +1,78 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`for_each`](super::StreamExt::for_each) method. + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct ForEach<St, Fut, F> { + #[pin] + stream: St, + f: F, + #[pin] + future: Option<Fut>, + } +} + +impl<St, Fut, F> fmt::Debug for ForEach<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ForEach") + .field("stream", &self.stream) + .field("future", &self.future) + .finish() + } +} + +impl<St, Fut, F> ForEach<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = ()>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, future: None } + } +} + +impl<St, Fut, F> FusedFuture for ForEach<St, Fut, F> +where + St: FusedStream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = ()>, +{ + fn is_terminated(&self) -> bool { + self.future.is_none() && self.stream.is_terminated() + } +} + +impl<St, Fut, F> Future for ForEach<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = ()>, +{ + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + let mut this = self.project(); + loop { + if let Some(fut) = this.future.as_mut().as_pin_mut() { + ready!(fut.poll(cx)); + this.future.set(None); + } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) { + this.future.set(Some((this.f)(item))); + } else { + break; + } + } + Poll::Ready(()) + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/for_each_concurrent.rs b/third_party/rust/futures-util/src/stream/stream/for_each_concurrent.rs new file mode 100644 index 0000000000..6c18753eb9 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/for_each_concurrent.rs @@ -0,0 +1,119 @@ +use crate::stream::{FuturesUnordered, StreamExt}; +use core::fmt; +use core::num::NonZeroUsize; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`for_each_concurrent`](super::StreamExt::for_each_concurrent) + /// method. + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct ForEachConcurrent<St, Fut, F> { + #[pin] + stream: Option<St>, + f: F, + futures: FuturesUnordered<Fut>, + limit: Option<NonZeroUsize>, + } +} + +impl<St, Fut, F> fmt::Debug for ForEachConcurrent<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ForEachConcurrent") + .field("stream", &self.stream) + .field("futures", &self.futures) + .field("limit", &self.limit) + .finish() + } +} + +impl<St, Fut, F> ForEachConcurrent<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = ()>, +{ + pub(super) fn new(stream: St, limit: Option<usize>, f: F) -> Self { + Self { + stream: Some(stream), + // Note: `limit` = 0 gets ignored. + limit: limit.and_then(NonZeroUsize::new), + f, + futures: FuturesUnordered::new(), + } + } +} + +impl<St, Fut, F> FusedFuture for ForEachConcurrent<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = ()>, +{ + fn is_terminated(&self) -> bool { + self.stream.is_none() && self.futures.is_empty() + } +} + +impl<St, Fut, F> Future for ForEachConcurrent<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future<Output = ()>, +{ + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + let mut this = self.project(); + loop { + let mut made_progress_this_iter = false; + + // Check if we've already created a number of futures greater than `limit` + if this.limit.map(|limit| limit.get() > this.futures.len()).unwrap_or(true) { + let mut stream_completed = false; + let elem = if let Some(stream) = this.stream.as_mut().as_pin_mut() { + match stream.poll_next(cx) { + Poll::Ready(Some(elem)) => { + made_progress_this_iter = true; + Some(elem) + } + Poll::Ready(None) => { + stream_completed = true; + None + } + Poll::Pending => None, + } + } else { + None + }; + if stream_completed { + this.stream.set(None); + } + if let Some(elem) = elem { + this.futures.push((this.f)(elem)); + } + } + + match this.futures.poll_next_unpin(cx) { + Poll::Ready(Some(())) => made_progress_this_iter = true, + Poll::Ready(None) => { + if this.stream.is_none() { + return Poll::Ready(()); + } + } + Poll::Pending => {} + } + + if !made_progress_this_iter { + return Poll::Pending; + } + } + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/forward.rs b/third_party/rust/futures-util/src/stream/stream/forward.rs new file mode 100644 index 0000000000..1fe24273aa --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/forward.rs @@ -0,0 +1,75 @@ +use crate::stream::Fuse; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`forward`](super::StreamExt::forward) method. + #[project = ForwardProj] + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Forward<St, Si, Item> { + #[pin] + sink: Option<Si>, + #[pin] + stream: Fuse<St>, + buffered_item: Option<Item>, + } +} + +impl<St, Si, Item> Forward<St, Si, Item> { + pub(crate) fn new(stream: St, sink: Si) -> Self { + Self { sink: Some(sink), stream: Fuse::new(stream), buffered_item: None } + } +} + +impl<St, Si, Item, E> FusedFuture for Forward<St, Si, Item> +where + Si: Sink<Item, Error = E>, + St: Stream<Item = Result<Item, E>>, +{ + fn is_terminated(&self) -> bool { + self.sink.is_none() + } +} + +impl<St, Si, Item, E> Future for Forward<St, Si, Item> +where + Si: Sink<Item, Error = E>, + St: Stream<Item = Result<Item, E>>, +{ + type Output = Result<(), E>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let ForwardProj { mut sink, mut stream, buffered_item } = self.project(); + let mut si = sink.as_mut().as_pin_mut().expect("polled `Forward` after completion"); + + loop { + // If we've got an item buffered already, we need to write it to the + // sink before we can do anything else + if buffered_item.is_some() { + ready!(si.as_mut().poll_ready(cx))?; + si.as_mut().start_send(buffered_item.take().unwrap())?; + } + + match stream.as_mut().poll_next(cx)? { + Poll::Ready(Some(item)) => { + *buffered_item = Some(item); + } + Poll::Ready(None) => { + ready!(si.poll_close(cx))?; + sink.set(None); + return Poll::Ready(Ok(())); + } + Poll::Pending => { + ready!(si.poll_flush(cx))?; + return Poll::Pending; + } + } + } + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/fuse.rs b/third_party/rust/futures-util/src/stream/stream/fuse.rs new file mode 100644 index 0000000000..fe67813e81 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/fuse.rs @@ -0,0 +1,75 @@ +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`fuse`](super::StreamExt::fuse) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Fuse<St> { + #[pin] + stream: St, + done: bool, + } +} + +impl<St> Fuse<St> { + pub(super) fn new(stream: St) -> Self { + Self { stream, done: false } + } + + /// Returns whether the underlying stream has finished or not. + /// + /// If this method returns `true`, then all future calls to poll are + /// guaranteed to return `None`. If this returns `false`, then the + /// underlying stream is still in use. + pub fn is_done(&self) -> bool { + self.done + } + + delegate_access_inner!(stream, St, ()); +} + +impl<S: Stream> FusedStream for Fuse<S> { + fn is_terminated(&self) -> bool { + self.done + } +} + +impl<S: Stream> Stream for Fuse<S> { + type Item = S::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<S::Item>> { + let this = self.project(); + + if *this.done { + return Poll::Ready(None); + } + + let item = ready!(this.stream.poll_next(cx)); + if item.is_none() { + *this.done = true; + } + Poll::Ready(item) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.done { + (0, Some(0)) + } else { + self.stream.size_hint() + } + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S: Stream + Sink<Item>, Item> Sink<Item> for Fuse<S> { + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/into_future.rs b/third_party/rust/futures-util/src/stream/stream/into_future.rs new file mode 100644 index 0000000000..8abfddcccd --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/into_future.rs @@ -0,0 +1,90 @@ +use crate::stream::StreamExt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; + +/// Future for the [`into_future`](super::StreamExt::into_future) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct StreamFuture<St> { + stream: Option<St>, +} + +impl<St: Stream + Unpin> StreamFuture<St> { + pub(super) fn new(stream: St) -> Self { + Self { stream: Some(stream) } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + /// + /// This method returns an `Option` to account for the fact that `StreamFuture`'s + /// implementation of `Future::poll` consumes the underlying stream during polling + /// in order to return it to the caller of `Future::poll` if the stream yielded + /// an element. + pub fn get_ref(&self) -> Option<&St> { + self.stream.as_ref() + } + + /// Acquires a mutable reference to the underlying stream that this + /// combinator is pulling from. + /// + /// Note that care must be taken to avoid tampering with the state of the + /// stream which may otherwise confuse this combinator. + /// + /// This method returns an `Option` to account for the fact that `StreamFuture`'s + /// implementation of `Future::poll` consumes the underlying stream during polling + /// in order to return it to the caller of `Future::poll` if the stream yielded + /// an element. + pub fn get_mut(&mut self) -> Option<&mut St> { + self.stream.as_mut() + } + + /// Acquires a pinned mutable reference to the underlying stream that this + /// combinator is pulling from. + /// + /// Note that care must be taken to avoid tampering with the state of the + /// stream which may otherwise confuse this combinator. + /// + /// This method returns an `Option` to account for the fact that `StreamFuture`'s + /// implementation of `Future::poll` consumes the underlying stream during polling + /// in order to return it to the caller of `Future::poll` if the stream yielded + /// an element. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Option<Pin<&mut St>> { + self.get_mut().stream.as_mut().map(Pin::new) + } + + /// Consumes this combinator, returning the underlying stream. + /// + /// Note that this may discard intermediate state of this combinator, so + /// care should be taken to avoid losing resources when this is called. + /// + /// This method returns an `Option` to account for the fact that `StreamFuture`'s + /// implementation of `Future::poll` consumes the underlying stream during polling + /// in order to return it to the caller of `Future::poll` if the stream yielded + /// an element. + pub fn into_inner(self) -> Option<St> { + self.stream + } +} + +impl<St: Stream + Unpin> FusedFuture for StreamFuture<St> { + fn is_terminated(&self) -> bool { + self.stream.is_none() + } +} + +impl<St: Stream + Unpin> Future for StreamFuture<St> { + type Output = (Option<St::Item>, St); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let item = { + let s = self.stream.as_mut().expect("polling StreamFuture twice"); + ready!(s.poll_next_unpin(cx)) + }; + let stream = self.stream.take().unwrap(); + Poll::Ready((item, stream)) + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/map.rs b/third_party/rust/futures-util/src/stream/stream/map.rs new file mode 100644 index 0000000000..88bb6129d4 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/map.rs @@ -0,0 +1,77 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +use crate::fns::FnMut1; + +pin_project! { + /// Stream for the [`map`](super::StreamExt::map) method. + #[must_use = "streams do nothing unless polled"] + pub struct Map<St, F> { + #[pin] + stream: St, + f: F, + } +} + +impl<St, F> fmt::Debug for Map<St, F> +where + St: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Map").field("stream", &self.stream).finish() + } +} + +impl<St, F> Map<St, F> { + pub(crate) fn new(stream: St, f: F) -> Self { + Self { stream, f } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, F> FusedStream for Map<St, F> +where + St: FusedStream, + F: FnMut1<St::Item>, +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St, F> Stream for Map<St, F> +where + St: Stream, + F: FnMut1<St::Item>, +{ + type Item = F::Output; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + let res = ready!(this.stream.as_mut().poll_next(cx)); + Poll::Ready(res.map(|x| this.f.call_mut(x))) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<St, F, Item> Sink<Item> for Map<St, F> +where + St: Stream + Sink<Item>, + F: FnMut1<St::Item>, +{ + type Error = St::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/mod.rs b/third_party/rust/futures-util/src/stream/stream/mod.rs new file mode 100644 index 0000000000..bb5e24907d --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/mod.rs @@ -0,0 +1,1684 @@ +//! Streams +//! +//! This module contains a number of functions for working with `Stream`s, +//! including the `StreamExt` trait which adds methods to `Stream` types. + +use crate::future::{assert_future, Either}; +use crate::stream::assert_stream; +#[cfg(feature = "alloc")] +use alloc::boxed::Box; +#[cfg(feature = "alloc")] +use alloc::vec::Vec; +use core::pin::Pin; +#[cfg(feature = "sink")] +use futures_core::stream::TryStream; +#[cfg(feature = "alloc")] +use futures_core::stream::{BoxStream, LocalBoxStream}; +use futures_core::{ + future::Future, + stream::{FusedStream, Stream}, + task::{Context, Poll}, +}; +#[cfg(feature = "sink")] +use futures_sink::Sink; + +use crate::fns::{inspect_fn, InspectFn}; + +mod chain; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::chain::Chain; + +mod collect; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::collect::Collect; + +mod unzip; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::unzip::Unzip; + +mod concat; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::concat::Concat; + +mod count; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::count::Count; + +mod cycle; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::cycle::Cycle; + +mod enumerate; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::enumerate::Enumerate; + +mod filter; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::filter::Filter; + +mod filter_map; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::filter_map::FilterMap; + +mod flatten; + +delegate_all!( + /// Stream for the [`flatten`](StreamExt::flatten) method. + Flatten<St>( + flatten::Flatten<St, St::Item> + ): Debug + Sink + Stream + FusedStream + AccessInner[St, (.)] + New[|x: St| flatten::Flatten::new(x)] + where St: Stream +); + +mod fold; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::fold::Fold; + +mod any; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::any::Any; + +mod all; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::all::All; + +#[cfg(feature = "sink")] +mod forward; + +#[cfg(feature = "sink")] +delegate_all!( + /// Future for the [`forward`](super::StreamExt::forward) method. + #[cfg_attr(docsrs, doc(cfg(feature = "sink")))] + Forward<St, Si>( + forward::Forward<St, Si, St::Ok> + ): Debug + Future + FusedFuture + New[|x: St, y: Si| forward::Forward::new(x, y)] + where St: TryStream +); + +mod for_each; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::for_each::ForEach; + +mod fuse; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::fuse::Fuse; + +mod into_future; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::into_future::StreamFuture; + +delegate_all!( + /// Stream for the [`inspect`](StreamExt::inspect) method. + Inspect<St, F>( + map::Map<St, InspectFn<F>> + ): Debug + Sink + Stream + FusedStream + AccessInner[St, (.)] + New[|x: St, f: F| map::Map::new(x, inspect_fn(f))] +); + +mod map; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::map::Map; + +delegate_all!( + /// Stream for the [`flat_map`](StreamExt::flat_map) method. + FlatMap<St, U, F>( + flatten::Flatten<Map<St, F>, U> + ): Debug + Sink + Stream + FusedStream + AccessInner[St, (. .)] + New[|x: St, f: F| flatten::Flatten::new(Map::new(x, f))] +); + +mod next; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::next::Next; + +mod select_next_some; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::select_next_some::SelectNextSome; + +mod peek; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::peek::{NextIf, NextIfEq, Peek, PeekMut, Peekable}; + +mod skip; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::skip::Skip; + +mod skip_while; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::skip_while::SkipWhile; + +mod take; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::take::Take; + +mod take_while; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::take_while::TakeWhile; + +mod take_until; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::take_until::TakeUntil; + +mod then; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::then::Then; + +mod zip; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::zip::Zip; + +#[cfg(feature = "alloc")] +mod chunks; +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::chunks::Chunks; + +#[cfg(feature = "alloc")] +mod ready_chunks; +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::ready_chunks::ReadyChunks; + +mod scan; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::scan::Scan; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +mod buffer_unordered; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::buffer_unordered::BufferUnordered; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +mod buffered; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::buffered::Buffered; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +mod flatten_unordered; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] +pub use self::flatten_unordered::FlattenUnordered; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +delegate_all!( + /// Stream for the [`flat_map_unordered`](StreamExt::flat_map_unordered) method. + FlatMapUnordered<St, U, F>( + FlattenUnordered<Map<St, F>> + ): Debug + Sink + Stream + FusedStream + AccessInner[St, (. .)] + New[|x: St, limit: Option<usize>, f: F| FlattenUnordered::new(Map::new(x, f), limit)] + where St: Stream, U: Stream, U: Unpin, F: FnMut(St::Item) -> U +); + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +mod for_each_concurrent; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::for_each_concurrent::ForEachConcurrent; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "sink")] +#[cfg_attr(docsrs, doc(cfg(feature = "sink")))] +#[cfg(feature = "alloc")] +mod split; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "sink")] +#[cfg_attr(docsrs, doc(cfg(feature = "sink")))] +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::split::{ReuniteError, SplitSink, SplitStream}; + +#[cfg(feature = "std")] +mod catch_unwind; +#[cfg(feature = "std")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::catch_unwind::CatchUnwind; + +impl<T: ?Sized> StreamExt for T where T: Stream {} + +/// An extension trait for `Stream`s that provides a variety of convenient +/// combinator functions. +pub trait StreamExt: Stream { + /// Creates a future that resolves to the next item in the stream. + /// + /// Note that because `next` doesn't take ownership over the stream, + /// the [`Stream`] type must be [`Unpin`]. If you want to use `next` with a + /// [`!Unpin`](Unpin) stream, you'll first have to pin the stream. This can + /// be done by boxing the stream using [`Box::pin`] or + /// pinning it to the stack using the `pin_mut!` macro from the `pin_utils` + /// crate. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let mut stream = stream::iter(1..=3); + /// + /// assert_eq!(stream.next().await, Some(1)); + /// assert_eq!(stream.next().await, Some(2)); + /// assert_eq!(stream.next().await, Some(3)); + /// assert_eq!(stream.next().await, None); + /// # }); + /// ``` + fn next(&mut self) -> Next<'_, Self> + where + Self: Unpin, + { + assert_future::<Option<Self::Item>, _>(Next::new(self)) + } + + /// Converts this stream into a future of `(next_item, tail_of_stream)`. + /// If the stream terminates, then the next item is [`None`]. + /// + /// The returned future can be used to compose streams and futures together + /// by placing everything into the "world of futures". + /// + /// Note that because `into_future` moves the stream, the [`Stream`] type + /// must be [`Unpin`]. If you want to use `into_future` with a + /// [`!Unpin`](Unpin) stream, you'll first have to pin the stream. This can + /// be done by boxing the stream using [`Box::pin`] or + /// pinning it to the stack using the `pin_mut!` macro from the `pin_utils` + /// crate. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=3); + /// + /// let (item, stream) = stream.into_future().await; + /// assert_eq!(Some(1), item); + /// + /// let (item, stream) = stream.into_future().await; + /// assert_eq!(Some(2), item); + /// # }); + /// ``` + fn into_future(self) -> StreamFuture<Self> + where + Self: Sized + Unpin, + { + assert_future::<(Option<Self::Item>, Self), _>(StreamFuture::new(self)) + } + + /// Maps this stream's items to a different type, returning a new stream of + /// the resulting type. + /// + /// The provided closure is executed over all elements of this stream as + /// they are made available. It is executed inline with calls to + /// [`poll_next`](Stream::poll_next). + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `map` methods in the + /// standard library. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=3); + /// let stream = stream.map(|x| x + 3); + /// + /// assert_eq!(vec![4, 5, 6], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn map<T, F>(self, f: F) -> Map<Self, F> + where + F: FnMut(Self::Item) -> T, + Self: Sized, + { + assert_stream::<T, _>(Map::new(self, f)) + } + + /// Creates a stream which gives the current iteration count as well as + /// the next value. + /// + /// The stream returned yields pairs `(i, val)`, where `i` is the + /// current index of iteration and `val` is the value returned by the + /// stream. + /// + /// `enumerate()` keeps its count as a [`usize`]. If you want to count by a + /// different sized integer, the [`zip`](StreamExt::zip) function provides similar + /// functionality. + /// + /// # Overflow Behavior + /// + /// The method does no guarding against overflows, so enumerating more than + /// [`prim@usize::max_value()`] elements either produces the wrong result or panics. If + /// debug assertions are enabled, a panic is guaranteed. + /// + /// # Panics + /// + /// The returned stream might panic if the to-be-returned index would + /// overflow a [`usize`]. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(vec!['a', 'b', 'c']); + /// + /// let mut stream = stream.enumerate(); + /// + /// assert_eq!(stream.next().await, Some((0, 'a'))); + /// assert_eq!(stream.next().await, Some((1, 'b'))); + /// assert_eq!(stream.next().await, Some((2, 'c'))); + /// assert_eq!(stream.next().await, None); + /// # }); + /// ``` + fn enumerate(self) -> Enumerate<Self> + where + Self: Sized, + { + assert_stream::<(usize, Self::Item), _>(Enumerate::new(self)) + } + + /// Filters the values produced by this stream according to the provided + /// asynchronous predicate. + /// + /// As values of this stream are made available, the provided predicate `f` + /// will be run against them. If the predicate returns a `Future` which + /// resolves to `true`, then the stream will yield the value, but if the + /// predicate returns a `Future` which resolves to `false`, then the value + /// will be discarded and the next value will be produced. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `filter` methods in the + /// standard library. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// let events = stream.filter(|x| future::ready(x % 2 == 0)); + /// + /// assert_eq!(vec![2, 4, 6, 8, 10], events.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn filter<Fut, F>(self, f: F) -> Filter<Self, Fut, F> + where + F: FnMut(&Self::Item) -> Fut, + Fut: Future<Output = bool>, + Self: Sized, + { + assert_stream::<Self::Item, _>(Filter::new(self, f)) + } + + /// Filters the values produced by this stream while simultaneously mapping + /// them to a different type according to the provided asynchronous closure. + /// + /// As values of this stream are made available, the provided function will + /// be run on them. If the future returned by the predicate `f` resolves to + /// [`Some(item)`](Some) then the stream will yield the value `item`, but if + /// it resolves to [`None`] then the next value will be produced. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `filter_map` methods in + /// the standard library. + /// + /// # Examples + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// let events = stream.filter_map(|x| async move { + /// if x % 2 == 0 { Some(x + 1) } else { None } + /// }); + /// + /// assert_eq!(vec![3, 5, 7, 9, 11], events.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn filter_map<Fut, T, F>(self, f: F) -> FilterMap<Self, Fut, F> + where + F: FnMut(Self::Item) -> Fut, + Fut: Future<Output = Option<T>>, + Self: Sized, + { + assert_stream::<T, _>(FilterMap::new(self, f)) + } + + /// Computes from this stream's items new items of a different type using + /// an asynchronous closure. + /// + /// The provided closure `f` will be called with an `Item` once a value is + /// ready, it returns a future which will then be run to completion + /// to produce the next value on this stream. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=3); + /// let stream = stream.then(|x| async move { x + 3 }); + /// + /// assert_eq!(vec![4, 5, 6], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn then<Fut, F>(self, f: F) -> Then<Self, Fut, F> + where + F: FnMut(Self::Item) -> Fut, + Fut: Future, + Self: Sized, + { + assert_stream::<Fut::Output, _>(Then::new(self, f)) + } + + /// Transforms a stream into a collection, returning a + /// future representing the result of that computation. + /// + /// The returned future will be resolved when the stream terminates. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::StreamExt; + /// use std::thread; + /// + /// let (tx, rx) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// for i in 1..=5 { + /// tx.unbounded_send(i).unwrap(); + /// } + /// }); + /// + /// let output = rx.collect::<Vec<i32>>().await; + /// assert_eq!(output, vec![1, 2, 3, 4, 5]); + /// # }); + /// ``` + fn collect<C: Default + Extend<Self::Item>>(self) -> Collect<Self, C> + where + Self: Sized, + { + assert_future::<C, _>(Collect::new(self)) + } + + /// Converts a stream of pairs into a future, which + /// resolves to pair of containers. + /// + /// `unzip()` produces a future, which resolves to two + /// collections: one from the left elements of the pairs, + /// and one from the right elements. + /// + /// The returned future will be resolved when the stream terminates. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::StreamExt; + /// use std::thread; + /// + /// let (tx, rx) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// tx.unbounded_send((1, 2)).unwrap(); + /// tx.unbounded_send((3, 4)).unwrap(); + /// tx.unbounded_send((5, 6)).unwrap(); + /// }); + /// + /// let (o1, o2): (Vec<_>, Vec<_>) = rx.unzip().await; + /// assert_eq!(o1, vec![1, 3, 5]); + /// assert_eq!(o2, vec![2, 4, 6]); + /// # }); + /// ``` + fn unzip<A, B, FromA, FromB>(self) -> Unzip<Self, FromA, FromB> + where + FromA: Default + Extend<A>, + FromB: Default + Extend<B>, + Self: Sized + Stream<Item = (A, B)>, + { + assert_future::<(FromA, FromB), _>(Unzip::new(self)) + } + + /// Concatenate all items of a stream into a single extendable + /// destination, returning a future representing the end result. + /// + /// This combinator will extend the first item with the contents + /// of all the subsequent results of the stream. If the stream is + /// empty, the default value will be returned. + /// + /// Works with all collections that implement the + /// [`Extend`](std::iter::Extend) trait. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::StreamExt; + /// use std::thread; + /// + /// let (tx, rx) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// for i in (0..3).rev() { + /// let n = i * 3; + /// tx.unbounded_send(vec![n + 1, n + 2, n + 3]).unwrap(); + /// } + /// }); + /// + /// let result = rx.concat().await; + /// + /// assert_eq!(result, vec![7, 8, 9, 4, 5, 6, 1, 2, 3]); + /// # }); + /// ``` + fn concat(self) -> Concat<Self> + where + Self: Sized, + Self::Item: Extend<<<Self as Stream>::Item as IntoIterator>::Item> + IntoIterator + Default, + { + assert_future::<Self::Item, _>(Concat::new(self)) + } + + /// Drives the stream to completion, counting the number of items. + /// + /// # Overflow Behavior + /// + /// The method does no guarding against overflows, so counting elements of a + /// stream with more than [`usize::MAX`] elements either produces the wrong + /// result or panics. If debug assertions are enabled, a panic is guaranteed. + /// + /// # Panics + /// + /// This function might panic if the iterator has more than [`usize::MAX`] + /// elements. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// let count = stream.count().await; + /// + /// assert_eq!(count, 10); + /// # }); + /// ``` + fn count(self) -> Count<Self> + where + Self: Sized, + { + assert_future::<usize, _>(Count::new(self)) + } + + /// Repeats a stream endlessly. + /// + /// The stream never terminates. Note that you likely want to avoid + /// usage of `collect` or such on the returned stream as it will exhaust + /// available memory as it tries to just fill up all RAM. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// let a = [1, 2, 3]; + /// let mut s = stream::iter(a.iter()).cycle(); + /// + /// assert_eq!(s.next().await, Some(&1)); + /// assert_eq!(s.next().await, Some(&2)); + /// assert_eq!(s.next().await, Some(&3)); + /// assert_eq!(s.next().await, Some(&1)); + /// assert_eq!(s.next().await, Some(&2)); + /// assert_eq!(s.next().await, Some(&3)); + /// assert_eq!(s.next().await, Some(&1)); + /// # }); + /// ``` + fn cycle(self) -> Cycle<Self> + where + Self: Sized + Clone, + { + assert_stream::<Self::Item, _>(Cycle::new(self)) + } + + /// Execute an accumulating asynchronous computation over a stream, + /// collecting all the values into one final result. + /// + /// This combinator will accumulate all values returned by this stream + /// according to the closure provided. The initial state is also provided to + /// this method and then is returned again by each execution of the closure. + /// Once the entire stream has been exhausted the returned future will + /// resolve to this value. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let number_stream = stream::iter(0..6); + /// let sum = number_stream.fold(0, |acc, x| async move { acc + x }); + /// assert_eq!(sum.await, 15); + /// # }); + /// ``` + fn fold<T, Fut, F>(self, init: T, f: F) -> Fold<Self, Fut, T, F> + where + F: FnMut(T, Self::Item) -> Fut, + Fut: Future<Output = T>, + Self: Sized, + { + assert_future::<T, _>(Fold::new(self, f, init)) + } + + /// Execute predicate over asynchronous stream, and return `true` if any element in stream satisfied a predicate. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let number_stream = stream::iter(0..10); + /// let contain_three = number_stream.any(|i| async move { i == 3 }); + /// assert_eq!(contain_three.await, true); + /// # }); + /// ``` + fn any<Fut, F>(self, f: F) -> Any<Self, Fut, F> + where + F: FnMut(Self::Item) -> Fut, + Fut: Future<Output = bool>, + Self: Sized, + { + assert_future::<bool, _>(Any::new(self, f)) + } + + /// Execute predicate over asynchronous stream, and return `true` if all element in stream satisfied a predicate. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let number_stream = stream::iter(0..10); + /// let less_then_twenty = number_stream.all(|i| async move { i < 20 }); + /// assert_eq!(less_then_twenty.await, true); + /// # }); + /// ``` + fn all<Fut, F>(self, f: F) -> All<Self, Fut, F> + where + F: FnMut(Self::Item) -> Fut, + Fut: Future<Output = bool>, + Self: Sized, + { + assert_future::<bool, _>(All::new(self, f)) + } + + /// Flattens a stream of streams into just one continuous stream. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::StreamExt; + /// use std::thread; + /// + /// let (tx1, rx1) = mpsc::unbounded(); + /// let (tx2, rx2) = mpsc::unbounded(); + /// let (tx3, rx3) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// tx1.unbounded_send(1).unwrap(); + /// tx1.unbounded_send(2).unwrap(); + /// }); + /// thread::spawn(move || { + /// tx2.unbounded_send(3).unwrap(); + /// tx2.unbounded_send(4).unwrap(); + /// }); + /// thread::spawn(move || { + /// tx3.unbounded_send(rx1).unwrap(); + /// tx3.unbounded_send(rx2).unwrap(); + /// }); + /// + /// let output = rx3.flatten().collect::<Vec<i32>>().await; + /// assert_eq!(output, vec![1, 2, 3, 4]); + /// # }); + /// ``` + fn flatten(self) -> Flatten<Self> + where + Self::Item: Stream, + Self: Sized, + { + assert_stream::<<Self::Item as Stream>::Item, _>(Flatten::new(self)) + } + + /// Flattens a stream of streams into just one continuous stream. Polls + /// inner streams concurrently. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::StreamExt; + /// use std::thread; + /// + /// let (tx1, rx1) = mpsc::unbounded(); + /// let (tx2, rx2) = mpsc::unbounded(); + /// let (tx3, rx3) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// tx1.unbounded_send(1).unwrap(); + /// tx1.unbounded_send(2).unwrap(); + /// }); + /// thread::spawn(move || { + /// tx2.unbounded_send(3).unwrap(); + /// tx2.unbounded_send(4).unwrap(); + /// }); + /// thread::spawn(move || { + /// tx3.unbounded_send(rx1).unwrap(); + /// tx3.unbounded_send(rx2).unwrap(); + /// }); + /// + /// let mut output = rx3.flatten_unordered(None).collect::<Vec<i32>>().await; + /// output.sort(); + /// + /// assert_eq!(output, vec![1, 2, 3, 4]); + /// # }); + /// ``` + #[cfg(not(futures_no_atomic_cas))] + #[cfg(feature = "alloc")] + fn flatten_unordered(self, limit: impl Into<Option<usize>>) -> FlattenUnordered<Self> + where + Self::Item: Stream + Unpin, + Self: Sized, + { + FlattenUnordered::new(self, limit.into()) + } + + /// Maps a stream like [`StreamExt::map`] but flattens nested `Stream`s. + /// + /// [`StreamExt::map`] is very useful, but if it produces a `Stream` instead, + /// you would have to chain combinators like `.map(f).flatten()` while this + /// combinator provides ability to write `.flat_map(f)` instead of chaining. + /// + /// The provided closure which produces inner streams is executed over all elements + /// of stream as last inner stream is terminated and next stream item is available. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `flat_map` methods in the + /// standard library. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=3); + /// let stream = stream.flat_map(|x| stream::iter(vec![x + 3; x])); + /// + /// assert_eq!(vec![4, 5, 5, 6, 6, 6], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn flat_map<U, F>(self, f: F) -> FlatMap<Self, U, F> + where + F: FnMut(Self::Item) -> U, + U: Stream, + Self: Sized, + { + assert_stream::<U::Item, _>(FlatMap::new(self, f)) + } + + /// Maps a stream like [`StreamExt::map`] but flattens nested `Stream`s + /// and polls them concurrently, yielding items in any order, as they made + /// available. + /// + /// [`StreamExt::map`] is very useful, but if it produces `Stream`s + /// instead, and you need to poll all of them concurrently, you would + /// have to use something like `for_each_concurrent` and merge values + /// by hand. This combinator provides ability to collect all values + /// from concurrently polled streams into one stream. + /// + /// The first argument is an optional limit on the number of concurrently + /// polled streams. If this limit is not `None`, no more than `limit` streams + /// will be polled concurrently. The `limit` argument is of type + /// `Into<Option<usize>>`, and so can be provided as either `None`, + /// `Some(10)`, or just `10`. Note: a limit of zero is interpreted as + /// no limit at all, and will have the same result as passing in `None`. + /// + /// The provided closure which produces inner streams is executed over + /// all elements of stream as next stream item is available and limit + /// of concurrently processed streams isn't exceeded. + /// + /// Note that this function consumes the stream passed into it and + /// returns a wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..5); + /// let stream = stream.flat_map_unordered(1, |x| stream::iter(vec![x; x])); + /// let mut values = stream.collect::<Vec<_>>().await; + /// values.sort(); + /// + /// assert_eq!(vec![1usize, 2, 2, 3, 3, 3, 4, 4, 4, 4], values); + /// # }); + /// ``` + #[cfg(not(futures_no_atomic_cas))] + #[cfg(feature = "alloc")] + fn flat_map_unordered<U, F>( + self, + limit: impl Into<Option<usize>>, + f: F, + ) -> FlatMapUnordered<Self, U, F> + where + U: Stream + Unpin, + F: FnMut(Self::Item) -> U, + Self: Sized, + { + FlatMapUnordered::new(self, limit.into(), f) + } + + /// Combinator similar to [`StreamExt::fold`] that holds internal state + /// and produces a new stream. + /// + /// Accepts initial state and closure which will be applied to each element + /// of the stream until provided closure returns `None`. Once `None` is + /// returned, stream will be terminated. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// + /// let stream = stream.scan(0, |state, x| { + /// *state += x; + /// future::ready(if *state < 10 { Some(x) } else { None }) + /// }); + /// + /// assert_eq!(vec![1, 2, 3], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn scan<S, B, Fut, F>(self, initial_state: S, f: F) -> Scan<Self, S, Fut, F> + where + F: FnMut(&mut S, Self::Item) -> Fut, + Fut: Future<Output = Option<B>>, + Self: Sized, + { + assert_stream::<B, _>(Scan::new(self, initial_state, f)) + } + + /// Skip elements on this stream while the provided asynchronous predicate + /// resolves to `true`. + /// + /// This function, like `Iterator::skip_while`, will skip elements on the + /// stream until the predicate `f` resolves to `false`. Once one element + /// returns `false`, all future elements will be returned from the underlying + /// stream. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// + /// let stream = stream.skip_while(|x| future::ready(*x <= 5)); + /// + /// assert_eq!(vec![6, 7, 8, 9, 10], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn skip_while<Fut, F>(self, f: F) -> SkipWhile<Self, Fut, F> + where + F: FnMut(&Self::Item) -> Fut, + Fut: Future<Output = bool>, + Self: Sized, + { + assert_stream::<Self::Item, _>(SkipWhile::new(self, f)) + } + + /// Take elements from this stream while the provided asynchronous predicate + /// resolves to `true`. + /// + /// This function, like `Iterator::take_while`, will take elements from the + /// stream until the predicate `f` resolves to `false`. Once one element + /// returns `false`, it will always return that the stream is done. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// + /// let stream = stream.take_while(|x| future::ready(*x <= 5)); + /// + /// assert_eq!(vec![1, 2, 3, 4, 5], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn take_while<Fut, F>(self, f: F) -> TakeWhile<Self, Fut, F> + where + F: FnMut(&Self::Item) -> Fut, + Fut: Future<Output = bool>, + Self: Sized, + { + assert_stream::<Self::Item, _>(TakeWhile::new(self, f)) + } + + /// Take elements from this stream until the provided future resolves. + /// + /// This function will take elements from the stream until the provided + /// stopping future `fut` resolves. Once the `fut` future becomes ready, + /// this stream combinator will always return that the stream is done. + /// + /// The stopping future may return any type. Once the stream is stopped + /// the result of the stopping future may be accessed with `TakeUntil::take_result()`. + /// The stream may also be resumed with `TakeUntil::take_future()`. + /// See the documentation of [`TakeUntil`] for more information. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// use futures::task::Poll; + /// + /// let stream = stream::iter(1..=10); + /// + /// let mut i = 0; + /// let stop_fut = future::poll_fn(|_cx| { + /// i += 1; + /// if i <= 5 { + /// Poll::Pending + /// } else { + /// Poll::Ready(()) + /// } + /// }); + /// + /// let stream = stream.take_until(stop_fut); + /// + /// assert_eq!(vec![1, 2, 3, 4, 5], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn take_until<Fut>(self, fut: Fut) -> TakeUntil<Self, Fut> + where + Fut: Future, + Self: Sized, + { + assert_stream::<Self::Item, _>(TakeUntil::new(self, fut)) + } + + /// Runs this stream to completion, executing the provided asynchronous + /// closure for each element on the stream. + /// + /// The closure provided will be called for each item this stream produces, + /// yielding a future. That future will then be executed to completion + /// before moving on to the next item. + /// + /// The returned value is a `Future` where the `Output` type is `()`; it is + /// executed entirely for its side effects. + /// + /// To process each item in the stream and produce another stream instead + /// of a single future, use `then` instead. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// + /// let mut x = 0; + /// + /// { + /// let fut = stream::repeat(1).take(3).for_each(|item| { + /// x += item; + /// future::ready(()) + /// }); + /// fut.await; + /// } + /// + /// assert_eq!(x, 3); + /// # }); + /// ``` + fn for_each<Fut, F>(self, f: F) -> ForEach<Self, Fut, F> + where + F: FnMut(Self::Item) -> Fut, + Fut: Future<Output = ()>, + Self: Sized, + { + assert_future::<(), _>(ForEach::new(self, f)) + } + + /// Runs this stream to completion, executing the provided asynchronous + /// closure for each element on the stream concurrently as elements become + /// available. + /// + /// This is similar to [`StreamExt::for_each`], but the futures + /// produced by the closure are run concurrently (but not in parallel-- + /// this combinator does not introduce any threads). + /// + /// The closure provided will be called for each item this stream produces, + /// yielding a future. That future will then be executed to completion + /// concurrently with the other futures produced by the closure. + /// + /// The first argument is an optional limit on the number of concurrent + /// futures. If this limit is not `None`, no more than `limit` futures + /// will be run concurrently. The `limit` argument is of type + /// `Into<Option<usize>>`, and so can be provided as either `None`, + /// `Some(10)`, or just `10`. Note: a limit of zero is interpreted as + /// no limit at all, and will have the same result as passing in `None`. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::oneshot; + /// use futures::stream::{self, StreamExt}; + /// + /// let (tx1, rx1) = oneshot::channel(); + /// let (tx2, rx2) = oneshot::channel(); + /// let (tx3, rx3) = oneshot::channel(); + /// + /// let fut = stream::iter(vec![rx1, rx2, rx3]).for_each_concurrent( + /// /* limit */ 2, + /// |rx| async move { + /// rx.await.unwrap(); + /// } + /// ); + /// tx1.send(()).unwrap(); + /// tx2.send(()).unwrap(); + /// tx3.send(()).unwrap(); + /// fut.await; + /// # }) + /// ``` + #[cfg(not(futures_no_atomic_cas))] + #[cfg(feature = "alloc")] + fn for_each_concurrent<Fut, F>( + self, + limit: impl Into<Option<usize>>, + f: F, + ) -> ForEachConcurrent<Self, Fut, F> + where + F: FnMut(Self::Item) -> Fut, + Fut: Future<Output = ()>, + Self: Sized, + { + assert_future::<(), _>(ForEachConcurrent::new(self, limit.into(), f)) + } + + /// Creates a new stream of at most `n` items of the underlying stream. + /// + /// Once `n` items have been yielded from this stream then it will always + /// return that the stream is done. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10).take(3); + /// + /// assert_eq!(vec![1, 2, 3], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn take(self, n: usize) -> Take<Self> + where + Self: Sized, + { + assert_stream::<Self::Item, _>(Take::new(self, n)) + } + + /// Creates a new stream which skips `n` items of the underlying stream. + /// + /// Once `n` items have been skipped from this stream then it will always + /// return the remaining items on this stream. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10).skip(5); + /// + /// assert_eq!(vec![6, 7, 8, 9, 10], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn skip(self, n: usize) -> Skip<Self> + where + Self: Sized, + { + assert_stream::<Self::Item, _>(Skip::new(self, n)) + } + + /// Fuse a stream such that [`poll_next`](Stream::poll_next) will never + /// again be called once it has finished. This method can be used to turn + /// any `Stream` into a `FusedStream`. + /// + /// Normally, once a stream has returned [`None`] from + /// [`poll_next`](Stream::poll_next) any further calls could exhibit bad + /// behavior such as block forever, panic, never return, etc. If it is known + /// that [`poll_next`](Stream::poll_next) may be called after stream + /// has already finished, then this method can be used to ensure that it has + /// defined semantics. + /// + /// The [`poll_next`](Stream::poll_next) method of a `fuse`d stream + /// is guaranteed to return [`None`] after the underlying stream has + /// finished. + /// + /// # Examples + /// + /// ``` + /// use futures::executor::block_on_stream; + /// use futures::stream::{self, StreamExt}; + /// use futures::task::Poll; + /// + /// let mut x = 0; + /// let stream = stream::poll_fn(|_| { + /// x += 1; + /// match x { + /// 0..=2 => Poll::Ready(Some(x)), + /// 3 => Poll::Ready(None), + /// _ => panic!("should not happen") + /// } + /// }).fuse(); + /// + /// let mut iter = block_on_stream(stream); + /// assert_eq!(Some(1), iter.next()); + /// assert_eq!(Some(2), iter.next()); + /// assert_eq!(None, iter.next()); + /// assert_eq!(None, iter.next()); + /// // ... + /// ``` + fn fuse(self) -> Fuse<Self> + where + Self: Sized, + { + assert_stream::<Self::Item, _>(Fuse::new(self)) + } + + /// Borrows a stream, rather than consuming it. + /// + /// This is useful to allow applying stream adaptors while still retaining + /// ownership of the original stream. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let mut stream = stream::iter(1..5); + /// + /// let sum = stream.by_ref() + /// .take(2) + /// .fold(0, |a, b| async move { a + b }) + /// .await; + /// assert_eq!(sum, 3); + /// + /// // You can use the stream again + /// let sum = stream.take(2) + /// .fold(0, |a, b| async move { a + b }) + /// .await; + /// assert_eq!(sum, 7); + /// # }); + /// ``` + fn by_ref(&mut self) -> &mut Self { + self + } + + /// Catches unwinding panics while polling the stream. + /// + /// Caught panic (if any) will be the last element of the resulting stream. + /// + /// In general, panics within a stream can propagate all the way out to the + /// task level. This combinator makes it possible to halt unwinding within + /// the stream itself. It's most commonly used within task executors. This + /// method should not be used for error handling. + /// + /// Note that this method requires the `UnwindSafe` bound from the standard + /// library. This isn't always applied automatically, and the standard + /// library provides an `AssertUnwindSafe` wrapper type to apply it + /// after-the fact. To assist using this method, the [`Stream`] trait is + /// also implemented for `AssertUnwindSafe<St>` where `St` implements + /// [`Stream`]. + /// + /// This method is only available when the `std` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(vec![Some(10), None, Some(11)]); + /// // Panic on second element + /// let stream_panicking = stream.map(|o| o.unwrap()); + /// // Collect all the results + /// let stream = stream_panicking.catch_unwind(); + /// + /// let results: Vec<Result<i32, _>> = stream.collect().await; + /// match results[0] { + /// Ok(10) => {} + /// _ => panic!("unexpected result!"), + /// } + /// assert!(results[1].is_err()); + /// assert_eq!(results.len(), 2); + /// # }); + /// ``` + #[cfg(feature = "std")] + fn catch_unwind(self) -> CatchUnwind<Self> + where + Self: Sized + std::panic::UnwindSafe, + { + assert_stream(CatchUnwind::new(self)) + } + + /// Wrap the stream in a Box, pinning it. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg(feature = "alloc")] + fn boxed<'a>(self) -> BoxStream<'a, Self::Item> + where + Self: Sized + Send + 'a, + { + assert_stream::<Self::Item, _>(Box::pin(self)) + } + + /// Wrap the stream in a Box, pinning it. + /// + /// Similar to `boxed`, but without the `Send` requirement. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg(feature = "alloc")] + fn boxed_local<'a>(self) -> LocalBoxStream<'a, Self::Item> + where + Self: Sized + 'a, + { + assert_stream::<Self::Item, _>(Box::pin(self)) + } + + /// An adaptor for creating a buffered list of pending futures. + /// + /// If this stream's item can be converted into a future, then this adaptor + /// will buffer up to at most `n` futures and then return the outputs in the + /// same order as the underlying stream. No more than `n` futures will be + /// buffered at any point in time, and less than `n` may also be buffered + /// depending on the state of each future. + /// + /// The returned stream will be a stream of each future's output. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg(not(futures_no_atomic_cas))] + #[cfg(feature = "alloc")] + fn buffered(self, n: usize) -> Buffered<Self> + where + Self::Item: Future, + Self: Sized, + { + assert_stream::<<Self::Item as Future>::Output, _>(Buffered::new(self, n)) + } + + /// An adaptor for creating a buffered list of pending futures (unordered). + /// + /// If this stream's item can be converted into a future, then this adaptor + /// will buffer up to `n` futures and then return the outputs in the order + /// in which they complete. No more than `n` futures will be buffered at + /// any point in time, and less than `n` may also be buffered depending on + /// the state of each future. + /// + /// The returned stream will be a stream of each future's output. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::oneshot; + /// use futures::stream::{self, StreamExt}; + /// + /// let (send_one, recv_one) = oneshot::channel(); + /// let (send_two, recv_two) = oneshot::channel(); + /// + /// let stream_of_futures = stream::iter(vec![recv_one, recv_two]); + /// let mut buffered = stream_of_futures.buffer_unordered(10); + /// + /// send_two.send(2i32)?; + /// assert_eq!(buffered.next().await, Some(Ok(2i32))); + /// + /// send_one.send(1i32)?; + /// assert_eq!(buffered.next().await, Some(Ok(1i32))); + /// + /// assert_eq!(buffered.next().await, None); + /// # Ok::<(), i32>(()) }).unwrap(); + /// ``` + #[cfg(not(futures_no_atomic_cas))] + #[cfg(feature = "alloc")] + fn buffer_unordered(self, n: usize) -> BufferUnordered<Self> + where + Self::Item: Future, + Self: Sized, + { + assert_stream::<<Self::Item as Future>::Output, _>(BufferUnordered::new(self, n)) + } + + /// An adapter for zipping two streams together. + /// + /// The zipped stream waits for both streams to produce an item, and then + /// returns that pair. If either stream ends then the zipped stream will + /// also end. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream1 = stream::iter(1..=3); + /// let stream2 = stream::iter(5..=10); + /// + /// let vec = stream1.zip(stream2) + /// .collect::<Vec<_>>() + /// .await; + /// assert_eq!(vec![(1, 5), (2, 6), (3, 7)], vec); + /// # }); + /// ``` + /// + fn zip<St>(self, other: St) -> Zip<Self, St> + where + St: Stream, + Self: Sized, + { + assert_stream::<(Self::Item, St::Item), _>(Zip::new(self, other)) + } + + /// Adapter for chaining two streams. + /// + /// The resulting stream emits elements from the first stream, and when + /// first stream reaches the end, emits the elements from the second stream. + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream1 = stream::iter(vec![Ok(10), Err(false)]); + /// let stream2 = stream::iter(vec![Err(true), Ok(20)]); + /// + /// let stream = stream1.chain(stream2); + /// + /// let result: Vec<_> = stream.collect().await; + /// assert_eq!(result, vec![ + /// Ok(10), + /// Err(false), + /// Err(true), + /// Ok(20), + /// ]); + /// # }); + /// ``` + fn chain<St>(self, other: St) -> Chain<Self, St> + where + St: Stream<Item = Self::Item>, + Self: Sized, + { + assert_stream::<Self::Item, _>(Chain::new(self, other)) + } + + /// Creates a new stream which exposes a `peek` method. + /// + /// Calling `peek` returns a reference to the next item in the stream. + fn peekable(self) -> Peekable<Self> + where + Self: Sized, + { + assert_stream::<Self::Item, _>(Peekable::new(self)) + } + + /// An adaptor for chunking up items of the stream inside a vector. + /// + /// This combinator will attempt to pull items from this stream and buffer + /// them into a local vector. At most `capacity` items will get buffered + /// before they're yielded from the returned stream. + /// + /// Note that the vectors returned from this iterator may not always have + /// `capacity` elements. If the underlying stream ended and only a partial + /// vector was created, it'll be returned. Additionally if an error happens + /// from the underlying stream then the currently buffered items will be + /// yielded. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Panics + /// + /// This method will panic if `capacity` is zero. + #[cfg(feature = "alloc")] + fn chunks(self, capacity: usize) -> Chunks<Self> + where + Self: Sized, + { + assert_stream::<Vec<Self::Item>, _>(Chunks::new(self, capacity)) + } + + /// An adaptor for chunking up ready items of the stream inside a vector. + /// + /// This combinator will attempt to pull ready items from this stream and + /// buffer them into a local vector. At most `capacity` items will get + /// buffered before they're yielded from the returned stream. If underlying + /// stream returns `Poll::Pending`, and collected chunk is not empty, it will + /// be immediately returned. + /// + /// If the underlying stream ended and only a partial vector was created, + /// it will be returned. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Panics + /// + /// This method will panic if `capacity` is zero. + #[cfg(feature = "alloc")] + fn ready_chunks(self, capacity: usize) -> ReadyChunks<Self> + where + Self: Sized, + { + assert_stream::<Vec<Self::Item>, _>(ReadyChunks::new(self, capacity)) + } + + /// A future that completes after the given stream has been fully processed + /// into the sink and the sink has been flushed and closed. + /// + /// This future will drive the stream to keep producing items until it is + /// exhausted, sending each item to the sink. It will complete once the + /// stream is exhausted, the sink has received and flushed all items, and + /// the sink is closed. Note that neither the original stream nor provided + /// sink will be output by this future. Pass the sink by `Pin<&mut S>` + /// (for example, via `forward(&mut sink)` inside an `async` fn/block) in + /// order to preserve access to the `Sink`. If the stream produces an error, + /// that error will be returned by this future without flushing/closing the sink. + #[cfg(feature = "sink")] + #[cfg_attr(docsrs, doc(cfg(feature = "sink")))] + fn forward<S>(self, sink: S) -> Forward<Self, S> + where + S: Sink<Self::Ok, Error = Self::Error>, + Self: TryStream + Sized, + // Self: TryStream + Sized + Stream<Item = Result<<Self as TryStream>::Ok, <Self as TryStream>::Error>>, + { + // TODO: type mismatch resolving `<Self as futures_core::Stream>::Item == std::result::Result<<Self as futures_core::TryStream>::Ok, <Self as futures_core::TryStream>::Error>` + // assert_future::<Result<(), Self::Error>, _>(Forward::new(self, sink)) + Forward::new(self, sink) + } + + /// Splits this `Stream + Sink` object into separate `Sink` and `Stream` + /// objects. + /// + /// This can be useful when you want to split ownership between tasks, or + /// allow direct interaction between the two objects (e.g. via + /// `Sink::send_all`). + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg(feature = "sink")] + #[cfg_attr(docsrs, doc(cfg(feature = "sink")))] + #[cfg(not(futures_no_atomic_cas))] + #[cfg(feature = "alloc")] + fn split<Item>(self) -> (SplitSink<Self, Item>, SplitStream<Self>) + where + Self: Sink<Item> + Sized, + { + let (sink, stream) = split::split(self); + ( + crate::sink::assert_sink::<Item, Self::Error, _>(sink), + assert_stream::<Self::Item, _>(stream), + ) + } + + /// Do something with each item of this stream, afterwards passing it on. + /// + /// This is similar to the `Iterator::inspect` method in the standard + /// library where it allows easily inspecting each value as it passes + /// through the stream, for example to debug what's going on. + fn inspect<F>(self, f: F) -> Inspect<Self, F> + where + F: FnMut(&Self::Item), + Self: Sized, + { + assert_stream::<Self::Item, _>(Inspect::new(self, f)) + } + + /// Wrap this stream in an `Either` stream, making it the left-hand variant + /// of that `Either`. + /// + /// This can be used in combination with the `right_stream` method to write `if` + /// statements that evaluate to different streams in different branches. + fn left_stream<B>(self) -> Either<Self, B> + where + B: Stream<Item = Self::Item>, + Self: Sized, + { + assert_stream::<Self::Item, _>(Either::Left(self)) + } + + /// Wrap this stream in an `Either` stream, making it the right-hand variant + /// of that `Either`. + /// + /// This can be used in combination with the `left_stream` method to write `if` + /// statements that evaluate to different streams in different branches. + fn right_stream<B>(self) -> Either<B, Self> + where + B: Stream<Item = Self::Item>, + Self: Sized, + { + assert_stream::<Self::Item, _>(Either::Right(self)) + } + + /// A convenience method for calling [`Stream::poll_next`] on [`Unpin`] + /// stream types. + fn poll_next_unpin(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> + where + Self: Unpin, + { + Pin::new(self).poll_next(cx) + } + + /// Returns a [`Future`] that resolves when the next item in this stream is + /// ready. + /// + /// This is similar to the [`next`][StreamExt::next] method, but it won't + /// resolve to [`None`] if used on an empty [`Stream`]. Instead, the + /// returned future type will return `true` from + /// [`FusedFuture::is_terminated`][] when the [`Stream`] is empty, allowing + /// [`select_next_some`][StreamExt::select_next_some] to be easily used with + /// the [`select!`] macro. + /// + /// If the future is polled after this [`Stream`] is empty it will panic. + /// Using the future with a [`FusedFuture`][]-aware primitive like the + /// [`select!`] macro will prevent this. + /// + /// [`FusedFuture`]: futures_core::future::FusedFuture + /// [`FusedFuture::is_terminated`]: futures_core::future::FusedFuture::is_terminated + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::{future, select}; + /// use futures::stream::{StreamExt, FuturesUnordered}; + /// + /// let mut fut = future::ready(1); + /// let mut async_tasks = FuturesUnordered::new(); + /// let mut total = 0; + /// loop { + /// select! { + /// num = fut => { + /// // First, the `ready` future completes. + /// total += num; + /// // Then we spawn a new task onto `async_tasks`, + /// async_tasks.push(async { 5 }); + /// }, + /// // On the next iteration of the loop, the task we spawned + /// // completes. + /// num = async_tasks.select_next_some() => { + /// total += num; + /// } + /// // Finally, both the `ready` future and `async_tasks` have + /// // finished, so we enter the `complete` branch. + /// complete => break, + /// } + /// } + /// assert_eq!(total, 6); + /// # }); + /// ``` + /// + /// [`select!`]: crate::select + fn select_next_some(&mut self) -> SelectNextSome<'_, Self> + where + Self: Unpin + FusedStream, + { + assert_future::<Self::Item, _>(SelectNextSome::new(self)) + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/next.rs b/third_party/rust/futures-util/src/stream/stream/next.rs new file mode 100644 index 0000000000..8d8347aa03 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/next.rs @@ -0,0 +1,34 @@ +use crate::stream::StreamExt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; + +/// Future for the [`next`](super::StreamExt::next) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Next<'a, St: ?Sized> { + stream: &'a mut St, +} + +impl<St: ?Sized + Unpin> Unpin for Next<'_, St> {} + +impl<'a, St: ?Sized + Stream + Unpin> Next<'a, St> { + pub(super) fn new(stream: &'a mut St) -> Self { + Self { stream } + } +} + +impl<St: ?Sized + FusedStream + Unpin> FusedFuture for Next<'_, St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St: ?Sized + Stream + Unpin> Future for Next<'_, St> { + type Output = Option<St::Item>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + self.stream.poll_next_unpin(cx) + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/peek.rs b/third_party/rust/futures-util/src/stream/stream/peek.rs new file mode 100644 index 0000000000..ea3d6243f3 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/peek.rs @@ -0,0 +1,433 @@ +use crate::fns::FnOnce1; +use crate::stream::{Fuse, StreamExt}; +use core::fmt; +use core::marker::PhantomData; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// A `Stream` that implements a `peek` method. + /// + /// The `peek` method can be used to retrieve a reference + /// to the next `Stream::Item` if available. A subsequent + /// call to `poll` will return the owned item. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Peekable<St: Stream> { + #[pin] + stream: Fuse<St>, + peeked: Option<St::Item>, + } +} + +impl<St: Stream> Peekable<St> { + pub(super) fn new(stream: St) -> Self { + Self { stream: stream.fuse(), peeked: None } + } + + delegate_access_inner!(stream, St, (.)); + + /// Produces a future which retrieves a reference to the next item + /// in the stream, or `None` if the underlying stream terminates. + pub fn peek(self: Pin<&mut Self>) -> Peek<'_, St> { + Peek { inner: Some(self) } + } + + /// Peek retrieves a reference to the next item in the stream. + /// + /// This method polls the underlying stream and return either a reference + /// to the next item if the stream is ready or passes through any errors. + pub fn poll_peek(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<&St::Item>> { + let mut this = self.project(); + + Poll::Ready(loop { + if this.peeked.is_some() { + break this.peeked.as_ref(); + } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) { + *this.peeked = Some(item); + } else { + break None; + } + }) + } + + /// Produces a future which retrieves a mutable reference to the next item + /// in the stream, or `None` if the underlying stream terminates. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// use futures::pin_mut; + /// + /// let stream = stream::iter(vec![1, 2, 3]).peekable(); + /// pin_mut!(stream); + /// + /// assert_eq!(stream.as_mut().peek_mut().await, Some(&mut 1)); + /// assert_eq!(stream.as_mut().next().await, Some(1)); + /// + /// // Peek into the stream and modify the value which will be returned next + /// if let Some(p) = stream.as_mut().peek_mut().await { + /// if *p == 2 { + /// *p = 5; + /// } + /// } + /// + /// assert_eq!(stream.collect::<Vec<_>>().await, vec![5, 3]); + /// # }); + /// ``` + pub fn peek_mut(self: Pin<&mut Self>) -> PeekMut<'_, St> { + PeekMut { inner: Some(self) } + } + + /// Peek retrieves a mutable reference to the next item in the stream. + pub fn poll_peek_mut( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<&mut St::Item>> { + let mut this = self.project(); + + Poll::Ready(loop { + if this.peeked.is_some() { + break this.peeked.as_mut(); + } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) { + *this.peeked = Some(item); + } else { + break None; + } + }) + } + + /// Creates a future which will consume and return the next value of this + /// stream if a condition is true. + /// + /// If `func` returns `true` for the next value of this stream, consume and + /// return it. Otherwise, return `None`. + /// + /// # Examples + /// + /// Consume a number if it's equal to 0. + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// use futures::pin_mut; + /// + /// let stream = stream::iter(0..5).peekable(); + /// pin_mut!(stream); + /// // The first item of the stream is 0; consume it. + /// assert_eq!(stream.as_mut().next_if(|&x| x == 0).await, Some(0)); + /// // The next item returned is now 1, so `consume` will return `false`. + /// assert_eq!(stream.as_mut().next_if(|&x| x == 0).await, None); + /// // `next_if` saves the value of the next item if it was not equal to `expected`. + /// assert_eq!(stream.next().await, Some(1)); + /// # }); + /// ``` + /// + /// Consume any number less than 10. + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// use futures::pin_mut; + /// + /// let stream = stream::iter(1..20).peekable(); + /// pin_mut!(stream); + /// // Consume all numbers less than 10 + /// while stream.as_mut().next_if(|&x| x < 10).await.is_some() {} + /// // The next value returned will be 10 + /// assert_eq!(stream.next().await, Some(10)); + /// # }); + /// ``` + pub fn next_if<F>(self: Pin<&mut Self>, func: F) -> NextIf<'_, St, F> + where + F: FnOnce(&St::Item) -> bool, + { + NextIf { inner: Some((self, func)) } + } + + /// Creates a future which will consume and return the next item if it is + /// equal to `expected`. + /// + /// # Example + /// + /// Consume a number if it's equal to 0. + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// use futures::pin_mut; + /// + /// let stream = stream::iter(0..5).peekable(); + /// pin_mut!(stream); + /// // The first item of the stream is 0; consume it. + /// assert_eq!(stream.as_mut().next_if_eq(&0).await, Some(0)); + /// // The next item returned is now 1, so `consume` will return `false`. + /// assert_eq!(stream.as_mut().next_if_eq(&0).await, None); + /// // `next_if_eq` saves the value of the next item if it was not equal to `expected`. + /// assert_eq!(stream.next().await, Some(1)); + /// # }); + /// ``` + pub fn next_if_eq<'a, T>(self: Pin<&'a mut Self>, expected: &'a T) -> NextIfEq<'a, St, T> + where + T: ?Sized, + St::Item: PartialEq<T>, + { + NextIfEq { + inner: NextIf { inner: Some((self, NextIfEqFn { expected, _next: PhantomData })) }, + } + } +} + +impl<St: Stream> FusedStream for Peekable<St> { + fn is_terminated(&self) -> bool { + self.peeked.is_none() && self.stream.is_terminated() + } +} + +impl<S: Stream> Stream for Peekable<S> { + type Item = S::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let this = self.project(); + if let Some(item) = this.peeked.take() { + return Poll::Ready(Some(item)); + } + this.stream.poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let peek_len = usize::from(self.peeked.is_some()); + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(peek_len); + let upper = match upper { + Some(x) => x.checked_add(peek_len), + None => None, + }; + (lower, upper) + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Peekable<S> +where + S: Sink<Item> + Stream, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} + +pin_project! { + /// Future for the [`Peekable::peek`](self::Peekable::peek) method. + #[must_use = "futures do nothing unless polled"] + pub struct Peek<'a, St: Stream> { + inner: Option<Pin<&'a mut Peekable<St>>>, + } +} + +impl<St> fmt::Debug for Peek<'_, St> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Peek").field("inner", &self.inner).finish() + } +} + +impl<St: Stream> FusedFuture for Peek<'_, St> { + fn is_terminated(&self) -> bool { + self.inner.is_none() + } +} + +impl<'a, St> Future for Peek<'a, St> +where + St: Stream, +{ + type Output = Option<&'a St::Item>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let inner = self.project().inner; + if let Some(peekable) = inner { + ready!(peekable.as_mut().poll_peek(cx)); + + inner.take().unwrap().poll_peek(cx) + } else { + panic!("Peek polled after completion") + } + } +} + +pin_project! { + /// Future for the [`Peekable::peek_mut`](self::Peekable::peek_mut) method. + #[must_use = "futures do nothing unless polled"] + pub struct PeekMut<'a, St: Stream> { + inner: Option<Pin<&'a mut Peekable<St>>>, + } +} + +impl<St> fmt::Debug for PeekMut<'_, St> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PeekMut").field("inner", &self.inner).finish() + } +} + +impl<St: Stream> FusedFuture for PeekMut<'_, St> { + fn is_terminated(&self) -> bool { + self.inner.is_none() + } +} + +impl<'a, St> Future for PeekMut<'a, St> +where + St: Stream, +{ + type Output = Option<&'a mut St::Item>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let inner = self.project().inner; + if let Some(peekable) = inner { + ready!(peekable.as_mut().poll_peek_mut(cx)); + + inner.take().unwrap().poll_peek_mut(cx) + } else { + panic!("PeekMut polled after completion") + } + } +} + +pin_project! { + /// Future for the [`Peekable::next_if`](self::Peekable::next_if) method. + #[must_use = "futures do nothing unless polled"] + pub struct NextIf<'a, St: Stream, F> { + inner: Option<(Pin<&'a mut Peekable<St>>, F)>, + } +} + +impl<St, F> fmt::Debug for NextIf<'_, St, F> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("NextIf").field("inner", &self.inner.as_ref().map(|(s, _f)| s)).finish() + } +} + +#[allow(single_use_lifetimes)] // https://github.com/rust-lang/rust/issues/55058 +impl<St, F> FusedFuture for NextIf<'_, St, F> +where + St: Stream, + F: for<'a> FnOnce1<&'a St::Item, Output = bool>, +{ + fn is_terminated(&self) -> bool { + self.inner.is_none() + } +} + +#[allow(single_use_lifetimes)] // https://github.com/rust-lang/rust/issues/55058 +impl<St, F> Future for NextIf<'_, St, F> +where + St: Stream, + F: for<'a> FnOnce1<&'a St::Item, Output = bool>, +{ + type Output = Option<St::Item>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let inner = self.project().inner; + if let Some((peekable, _)) = inner { + let res = ready!(peekable.as_mut().poll_next(cx)); + + let (peekable, func) = inner.take().unwrap(); + match res { + Some(ref matched) if func.call_once(matched) => Poll::Ready(res), + other => { + let peekable = peekable.project(); + // Since we called `self.next()`, we consumed `self.peeked`. + assert!(peekable.peeked.is_none()); + *peekable.peeked = other; + Poll::Ready(None) + } + } + } else { + panic!("NextIf polled after completion") + } + } +} + +pin_project! { + /// Future for the [`Peekable::next_if_eq`](self::Peekable::next_if_eq) method. + #[must_use = "futures do nothing unless polled"] + pub struct NextIfEq<'a, St: Stream, T: ?Sized> { + #[pin] + inner: NextIf<'a, St, NextIfEqFn<'a, T, St::Item>>, + } +} + +impl<St, T> fmt::Debug for NextIfEq<'_, St, T> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, + T: ?Sized, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("NextIfEq") + .field("inner", &self.inner.inner.as_ref().map(|(s, _f)| s)) + .finish() + } +} + +impl<St, T> FusedFuture for NextIfEq<'_, St, T> +where + St: Stream, + T: ?Sized, + St::Item: PartialEq<T>, +{ + fn is_terminated(&self) -> bool { + self.inner.is_terminated() + } +} + +impl<St, T> Future for NextIfEq<'_, St, T> +where + St: Stream, + T: ?Sized, + St::Item: PartialEq<T>, +{ + type Output = Option<St::Item>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + self.project().inner.poll(cx) + } +} + +struct NextIfEqFn<'a, T: ?Sized, Item> { + expected: &'a T, + _next: PhantomData<Item>, +} + +impl<T, Item> FnOnce1<&Item> for NextIfEqFn<'_, T, Item> +where + T: ?Sized, + Item: PartialEq<T>, +{ + type Output = bool; + + fn call_once(self, next: &Item) -> Self::Output { + next == self.expected + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/ready_chunks.rs b/third_party/rust/futures-util/src/stream/stream/ready_chunks.rs new file mode 100644 index 0000000000..192054c4ae --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/ready_chunks.rs @@ -0,0 +1,93 @@ +use crate::stream::{Fuse, StreamExt}; +use alloc::vec::Vec; +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`ready_chunks`](super::StreamExt::ready_chunks) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct ReadyChunks<St: Stream> { + #[pin] + stream: Fuse<St>, + cap: usize, // https://github.com/rust-lang/futures-rs/issues/1475 + } +} + +impl<St: Stream> ReadyChunks<St> { + pub(super) fn new(stream: St, capacity: usize) -> Self { + assert!(capacity > 0); + + Self { stream: stream.fuse(), cap: capacity } + } + + delegate_access_inner!(stream, St, (.)); +} + +impl<St: Stream> Stream for ReadyChunks<St> { + type Item = Vec<St::Item>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + let mut items: Vec<St::Item> = Vec::new(); + + loop { + match this.stream.as_mut().poll_next(cx) { + // Flush all collected data if underlying stream doesn't contain + // more ready values + Poll::Pending => { + return if items.is_empty() { Poll::Pending } else { Poll::Ready(Some(items)) } + } + + // Push the ready item into the buffer and check whether it is full. + // If so, replace our buffer with a new and empty one and return + // the full one. + Poll::Ready(Some(item)) => { + if items.is_empty() { + items.reserve(*this.cap); + } + items.push(item); + if items.len() >= *this.cap { + return Poll::Ready(Some(items)); + } + } + + // Since the underlying stream ran out of values, return what we + // have buffered, if we have anything. + Poll::Ready(None) => { + let last = if items.is_empty() { None } else { Some(items) }; + + return Poll::Ready(last); + } + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let (lower, upper) = self.stream.size_hint(); + let lower = lower / self.cap; + (lower, upper) + } +} + +impl<St: Stream> FusedStream for ReadyChunks<St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for ReadyChunks<S> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/scan.rs b/third_party/rust/futures-util/src/stream/stream/scan.rs new file mode 100644 index 0000000000..f5cfde9c36 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/scan.rs @@ -0,0 +1,128 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +struct StateFn<S, F> { + state: S, + f: F, +} + +pin_project! { + /// Stream for the [`scan`](super::StreamExt::scan) method. + #[must_use = "streams do nothing unless polled"] + pub struct Scan<St: Stream, S, Fut, F> { + #[pin] + stream: St, + state_f: Option<StateFn<S, F>>, + #[pin] + future: Option<Fut>, + } +} + +impl<St, S, Fut, F> fmt::Debug for Scan<St, S, Fut, F> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, + S: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Scan") + .field("stream", &self.stream) + .field("state", &self.state_f.as_ref().map(|s| &s.state)) + .field("future", &self.future) + .field("done_taking", &self.is_done_taking()) + .finish() + } +} + +impl<St: Stream, S, Fut, F> Scan<St, S, Fut, F> { + /// Checks if internal state is `None`. + fn is_done_taking(&self) -> bool { + self.state_f.is_none() + } +} + +impl<B, St, S, Fut, F> Scan<St, S, Fut, F> +where + St: Stream, + F: FnMut(&mut S, St::Item) -> Fut, + Fut: Future<Output = Option<B>>, +{ + pub(super) fn new(stream: St, initial_state: S, f: F) -> Self { + Self { stream, state_f: Some(StateFn { state: initial_state, f }), future: None } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<B, St, S, Fut, F> Stream for Scan<St, S, Fut, F> +where + St: Stream, + F: FnMut(&mut S, St::Item) -> Fut, + Fut: Future<Output = Option<B>>, +{ + type Item = B; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<B>> { + if self.is_done_taking() { + return Poll::Ready(None); + } + + let mut this = self.project(); + + Poll::Ready(loop { + if let Some(fut) = this.future.as_mut().as_pin_mut() { + let item = ready!(fut.poll(cx)); + this.future.set(None); + + if item.is_none() { + *this.state_f = None; + } + + break item; + } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) { + let state_f = this.state_f.as_mut().unwrap(); + this.future.set(Some((state_f.f)(&mut state_f.state, item))) + } else { + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.is_done_taking() { + (0, Some(0)) + } else { + self.stream.size_hint() // can't know a lower bound, due to the predicate + } + } +} + +impl<B, St, S, Fut, F> FusedStream for Scan<St, S, Fut, F> +where + St: FusedStream, + F: FnMut(&mut S, St::Item) -> Fut, + Fut: Future<Output = Option<B>>, +{ + fn is_terminated(&self) -> bool { + self.is_done_taking() || self.future.is_none() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<St, S, Fut, F, Item> Sink<Item> for Scan<St, S, Fut, F> +where + St: Stream + Sink<Item>, +{ + type Error = St::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/select_next_some.rs b/third_party/rust/futures-util/src/stream/stream/select_next_some.rs new file mode 100644 index 0000000000..3115e14d9a --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/select_next_some.rs @@ -0,0 +1,42 @@ +use crate::stream::StreamExt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::FusedStream; +use futures_core::task::{Context, Poll}; + +/// Future for the [`select_next_some`](super::StreamExt::select_next_some) +/// method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct SelectNextSome<'a, St: ?Sized> { + stream: &'a mut St, +} + +impl<'a, St: ?Sized> SelectNextSome<'a, St> { + pub(super) fn new(stream: &'a mut St) -> Self { + Self { stream } + } +} + +impl<St: ?Sized + FusedStream + Unpin> FusedFuture for SelectNextSome<'_, St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St: ?Sized + FusedStream + Unpin> Future for SelectNextSome<'_, St> { + type Output = St::Item; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + assert!(!self.stream.is_terminated(), "SelectNextSome polled after terminated"); + + if let Some(item) = ready!(self.stream.poll_next_unpin(cx)) { + Poll::Ready(item) + } else { + debug_assert!(self.stream.is_terminated()); + cx.waker().wake_by_ref(); + Poll::Pending + } + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/skip.rs b/third_party/rust/futures-util/src/stream/stream/skip.rs new file mode 100644 index 0000000000..f495779521 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/skip.rs @@ -0,0 +1,70 @@ +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`skip`](super::StreamExt::skip) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Skip<St> { + #[pin] + stream: St, + remaining: usize, + } +} + +impl<St: Stream> Skip<St> { + pub(super) fn new(stream: St, n: usize) -> Self { + Self { stream, remaining: n } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St: FusedStream> FusedStream for Skip<St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St: Stream> Stream for Skip<St> { + type Item = St::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St::Item>> { + let mut this = self.project(); + + while *this.remaining > 0 { + if ready!(this.stream.as_mut().poll_next(cx)).is_some() { + *this.remaining -= 1; + } else { + return Poll::Ready(None); + } + } + + this.stream.poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let (lower, upper) = self.stream.size_hint(); + + let lower = lower.saturating_sub(self.remaining); + let upper = upper.map(|x| x.saturating_sub(self.remaining)); + + (lower, upper) + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Skip<S> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/skip_while.rs b/third_party/rust/futures-util/src/stream/stream/skip_while.rs new file mode 100644 index 0000000000..dabd5eefae --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/skip_while.rs @@ -0,0 +1,124 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`skip_while`](super::StreamExt::skip_while) method. + #[must_use = "streams do nothing unless polled"] + pub struct SkipWhile<St, Fut, F> where St: Stream { + #[pin] + stream: St, + f: F, + #[pin] + pending_fut: Option<Fut>, + pending_item: Option<St::Item>, + done_skipping: bool, + } +} + +impl<St, Fut, F> fmt::Debug for SkipWhile<St, Fut, F> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SkipWhile") + .field("stream", &self.stream) + .field("pending_fut", &self.pending_fut) + .field("pending_item", &self.pending_item) + .field("done_skipping", &self.done_skipping) + .finish() + } +} + +impl<St, Fut, F> SkipWhile<St, Fut, F> +where + St: Stream, + F: FnMut(&St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, pending_fut: None, pending_item: None, done_skipping: false } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F> FusedStream for SkipWhile<St, Fut, F> +where + St: FusedStream, + F: FnMut(&St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + fn is_terminated(&self) -> bool { + self.pending_item.is_none() && self.stream.is_terminated() + } +} + +impl<St, Fut, F> Stream for SkipWhile<St, Fut, F> +where + St: Stream, + F: FnMut(&St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + type Item = St::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St::Item>> { + let mut this = self.project(); + + if *this.done_skipping { + return this.stream.poll_next(cx); + } + + Poll::Ready(loop { + if let Some(fut) = this.pending_fut.as_mut().as_pin_mut() { + let skipped = ready!(fut.poll(cx)); + let item = this.pending_item.take(); + this.pending_fut.set(None); + if !skipped { + *this.done_skipping = true; + break item; + } + } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) { + this.pending_fut.set(Some((this.f)(&item))); + *this.pending_item = Some(item); + } else { + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.done_skipping { + self.stream.size_hint() + } else { + let pending_len = usize::from(self.pending_item.is_some()); + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for SkipWhile<S, Fut, F> +where + S: Stream + Sink<Item>, + F: FnMut(&S::Item) -> Fut, + Fut: Future<Output = bool>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/split.rs b/third_party/rust/futures-util/src/stream/stream/split.rs new file mode 100644 index 0000000000..e2034e0c27 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/split.rs @@ -0,0 +1,144 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; + +use crate::lock::BiLock; + +/// A `Stream` part of the split pair +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +#[cfg_attr(docsrs, doc(cfg(feature = "sink")))] +pub struct SplitStream<S>(BiLock<S>); + +impl<S> Unpin for SplitStream<S> {} + +impl<S: Unpin> SplitStream<S> { + /// Attempts to put the two "halves" of a split `Stream + Sink` back + /// together. Succeeds only if the `SplitStream<S>` and `SplitSink<S>` are + /// a matching pair originating from the same call to `StreamExt::split`. + pub fn reunite<Item>(self, other: SplitSink<S, Item>) -> Result<S, ReuniteError<S, Item>> + where + S: Sink<Item>, + { + other.reunite(self) + } +} + +impl<S: Stream> Stream for SplitStream<S> { + type Item = S::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<S::Item>> { + ready!(self.0.poll_lock(cx)).as_pin_mut().poll_next(cx) + } +} + +#[allow(non_snake_case)] +fn SplitSink<S: Sink<Item>, Item>(lock: BiLock<S>) -> SplitSink<S, Item> { + SplitSink { lock, slot: None } +} + +/// A `Sink` part of the split pair +#[derive(Debug)] +#[must_use = "sinks do nothing unless polled"] +#[cfg_attr(docsrs, doc(cfg(feature = "sink")))] +pub struct SplitSink<S, Item> { + lock: BiLock<S>, + slot: Option<Item>, +} + +impl<S, Item> Unpin for SplitSink<S, Item> {} + +impl<S: Sink<Item> + Unpin, Item> SplitSink<S, Item> { + /// Attempts to put the two "halves" of a split `Stream + Sink` back + /// together. Succeeds only if the `SplitStream<S>` and `SplitSink<S>` are + /// a matching pair originating from the same call to `StreamExt::split`. + pub fn reunite(self, other: SplitStream<S>) -> Result<S, ReuniteError<S, Item>> { + self.lock.reunite(other.0).map_err(|err| ReuniteError(SplitSink(err.0), SplitStream(err.1))) + } +} + +impl<S: Sink<Item>, Item> SplitSink<S, Item> { + fn poll_flush_slot( + mut inner: Pin<&mut S>, + slot: &mut Option<Item>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), S::Error>> { + if slot.is_some() { + ready!(inner.as_mut().poll_ready(cx))?; + Poll::Ready(inner.start_send(slot.take().unwrap())) + } else { + Poll::Ready(Ok(())) + } + } + + fn poll_lock_and_flush_slot( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), S::Error>> { + let this = &mut *self; + let mut inner = ready!(this.lock.poll_lock(cx)); + Self::poll_flush_slot(inner.as_pin_mut(), &mut this.slot, cx) + } +} + +impl<S: Sink<Item>, Item> Sink<Item> for SplitSink<S, Item> { + type Error = S::Error; + + fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> { + loop { + if self.slot.is_none() { + return Poll::Ready(Ok(())); + } + ready!(self.as_mut().poll_lock_and_flush_slot(cx))?; + } + } + + fn start_send(mut self: Pin<&mut Self>, item: Item) -> Result<(), S::Error> { + self.slot = Some(item); + Ok(()) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> { + let this = &mut *self; + let mut inner = ready!(this.lock.poll_lock(cx)); + ready!(Self::poll_flush_slot(inner.as_pin_mut(), &mut this.slot, cx))?; + inner.as_pin_mut().poll_flush(cx) + } + + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> { + let this = &mut *self; + let mut inner = ready!(this.lock.poll_lock(cx)); + ready!(Self::poll_flush_slot(inner.as_pin_mut(), &mut this.slot, cx))?; + inner.as_pin_mut().poll_close(cx) + } +} + +pub(super) fn split<S: Stream + Sink<Item>, Item>(s: S) -> (SplitSink<S, Item>, SplitStream<S>) { + let (a, b) = BiLock::new(s); + let read = SplitStream(a); + let write = SplitSink(b); + (write, read) +} + +/// Error indicating a `SplitSink<S>` and `SplitStream<S>` were not two halves +/// of a `Stream + Split`, and thus could not be `reunite`d. +#[cfg_attr(docsrs, doc(cfg(feature = "sink")))] +pub struct ReuniteError<T, Item>(pub SplitSink<T, Item>, pub SplitStream<T>); + +impl<T, Item> fmt::Debug for ReuniteError<T, Item> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("ReuniteError").field(&"...").finish() + } +} + +impl<T, Item> fmt::Display for ReuniteError<T, Item> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "tried to reunite a SplitStream and SplitSink that don't form a pair") + } +} + +#[cfg(feature = "std")] +impl<T: core::any::Any, Item> std::error::Error for ReuniteError<T, Item> {} diff --git a/third_party/rust/futures-util/src/stream/stream/take.rs b/third_party/rust/futures-util/src/stream/stream/take.rs new file mode 100644 index 0000000000..29d6c39ee3 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/take.rs @@ -0,0 +1,86 @@ +use core::cmp; +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`take`](super::StreamExt::take) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Take<St> { + #[pin] + stream: St, + remaining: usize, + } +} + +impl<St: Stream> Take<St> { + pub(super) fn new(stream: St, n: usize) -> Self { + Self { stream, remaining: n } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St> Stream for Take<St> +where + St: Stream, +{ + type Item = St::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St::Item>> { + if self.remaining == 0 { + Poll::Ready(None) + } else { + let this = self.project(); + let next = ready!(this.stream.poll_next(cx)); + if next.is_some() { + *this.remaining -= 1; + } else { + *this.remaining = 0; + } + Poll::Ready(next) + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.remaining == 0 { + return (0, Some(0)); + } + + let (lower, upper) = self.stream.size_hint(); + + let lower = cmp::min(lower, self.remaining); + + let upper = match upper { + Some(x) if x < self.remaining => Some(x), + _ => Some(self.remaining), + }; + + (lower, upper) + } +} + +impl<St> FusedStream for Take<St> +where + St: FusedStream, +{ + fn is_terminated(&self) -> bool { + self.remaining == 0 || self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Take<S> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/take_until.rs b/third_party/rust/futures-util/src/stream/stream/take_until.rs new file mode 100644 index 0000000000..d14f9ce100 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/take_until.rs @@ -0,0 +1,170 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +// FIXME: docs, tests + +pin_project! { + /// Stream for the [`take_until`](super::StreamExt::take_until) method. + #[must_use = "streams do nothing unless polled"] + pub struct TakeUntil<St: Stream, Fut: Future> { + #[pin] + stream: St, + // Contains the inner Future on start and None once the inner Future is resolved + // or taken out by the user. + #[pin] + fut: Option<Fut>, + // Contains fut's return value once fut is resolved + fut_result: Option<Fut::Output>, + // Whether the future was taken out by the user. + free: bool, + } +} + +impl<St, Fut> fmt::Debug for TakeUntil<St, Fut> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, + Fut: Future + fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TakeUntil").field("stream", &self.stream).field("fut", &self.fut).finish() + } +} + +impl<St, Fut> TakeUntil<St, Fut> +where + St: Stream, + Fut: Future, +{ + pub(super) fn new(stream: St, fut: Fut) -> Self { + Self { stream, fut: Some(fut), fut_result: None, free: false } + } + + delegate_access_inner!(stream, St, ()); + + /// Extract the stopping future out of the combinator. + /// The future is returned only if it isn't resolved yet, ie. if the stream isn't stopped yet. + /// Taking out the future means the combinator will be yielding + /// elements from the wrapped stream without ever stopping it. + pub fn take_future(&mut self) -> Option<Fut> { + if self.fut.is_some() { + self.free = true; + } + + self.fut.take() + } + + /// Once the stopping future is resolved, this method can be used + /// to extract the value returned by the stopping future. + /// + /// This may be used to retrieve arbitrary data from the stopping + /// future, for example a reason why the stream was stopped. + /// + /// This method will return `None` if the future isn't resolved yet, + /// or if the result was already taken out. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// use futures::task::Poll; + /// + /// let stream = stream::iter(1..=10); + /// + /// let mut i = 0; + /// let stop_fut = future::poll_fn(|_cx| { + /// i += 1; + /// if i <= 5 { + /// Poll::Pending + /// } else { + /// Poll::Ready("reason") + /// } + /// }); + /// + /// let mut stream = stream.take_until(stop_fut); + /// let _ = stream.by_ref().collect::<Vec<_>>().await; + /// + /// let result = stream.take_result().unwrap(); + /// assert_eq!(result, "reason"); + /// # }); + /// ``` + pub fn take_result(&mut self) -> Option<Fut::Output> { + self.fut_result.take() + } + + /// Whether the stream was stopped yet by the stopping future + /// being resolved. + pub fn is_stopped(&self) -> bool { + !self.free && self.fut.is_none() + } +} + +impl<St, Fut> Stream for TakeUntil<St, Fut> +where + St: Stream, + Fut: Future, +{ + type Item = St::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St::Item>> { + let mut this = self.project(); + + if let Some(f) = this.fut.as_mut().as_pin_mut() { + if let Poll::Ready(result) = f.poll(cx) { + this.fut.set(None); + *this.fut_result = Some(result); + } + } + + if !*this.free && this.fut.is_none() { + // Future resolved, inner stream stopped + Poll::Ready(None) + } else { + // Future either not resolved yet or taken out by the user + let item = ready!(this.stream.poll_next(cx)); + if item.is_none() { + this.fut.set(None); + } + Poll::Ready(item) + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.is_stopped() { + return (0, Some(0)); + } + + self.stream.size_hint() + } +} + +impl<St, Fut> FusedStream for TakeUntil<St, Fut> +where + St: Stream, + Fut: Future, +{ + fn is_terminated(&self) -> bool { + self.is_stopped() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, Item> Sink<Item> for TakeUntil<S, Fut> +where + S: Stream + Sink<Item>, + Fut: Future, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/take_while.rs b/third_party/rust/futures-util/src/stream/stream/take_while.rs new file mode 100644 index 0000000000..9256943010 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/take_while.rs @@ -0,0 +1,124 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`take_while`](super::StreamExt::take_while) method. + #[must_use = "streams do nothing unless polled"] + pub struct TakeWhile<St: Stream, Fut, F> { + #[pin] + stream: St, + f: F, + #[pin] + pending_fut: Option<Fut>, + pending_item: Option<St::Item>, + done_taking: bool, + } +} + +impl<St, Fut, F> fmt::Debug for TakeWhile<St, Fut, F> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TakeWhile") + .field("stream", &self.stream) + .field("pending_fut", &self.pending_fut) + .field("pending_item", &self.pending_item) + .field("done_taking", &self.done_taking) + .finish() + } +} + +impl<St, Fut, F> TakeWhile<St, Fut, F> +where + St: Stream, + F: FnMut(&St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, pending_fut: None, pending_item: None, done_taking: false } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F> Stream for TakeWhile<St, Fut, F> +where + St: Stream, + F: FnMut(&St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + type Item = St::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St::Item>> { + if self.done_taking { + return Poll::Ready(None); + } + + let mut this = self.project(); + + Poll::Ready(loop { + if let Some(fut) = this.pending_fut.as_mut().as_pin_mut() { + let take = ready!(fut.poll(cx)); + let item = this.pending_item.take(); + this.pending_fut.set(None); + if take { + break item; + } else { + *this.done_taking = true; + break None; + } + } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) { + this.pending_fut.set(Some((this.f)(&item))); + *this.pending_item = Some(item); + } else { + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.done_taking { + return (0, Some(0)); + } + + let pending_len = usize::from(self.pending_item.is_some()); + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +impl<St, Fut, F> FusedStream for TakeWhile<St, Fut, F> +where + St: FusedStream, + F: FnMut(&St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + fn is_terminated(&self) -> bool { + self.done_taking || self.pending_item.is_none() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for TakeWhile<S, Fut, F> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/then.rs b/third_party/rust/futures-util/src/stream/stream/then.rs new file mode 100644 index 0000000000..9192c0b0cf --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/then.rs @@ -0,0 +1,101 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`then`](super::StreamExt::then) method. + #[must_use = "streams do nothing unless polled"] + pub struct Then<St, Fut, F> { + #[pin] + stream: St, + #[pin] + future: Option<Fut>, + f: F, + } +} + +impl<St, Fut, F> fmt::Debug for Then<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Then").field("stream", &self.stream).field("future", &self.future).finish() + } +} + +impl<St, Fut, F> Then<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, future: None, f } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F> FusedStream for Then<St, Fut, F> +where + St: FusedStream, + F: FnMut(St::Item) -> Fut, + Fut: Future, +{ + fn is_terminated(&self) -> bool { + self.future.is_none() && self.stream.is_terminated() + } +} + +impl<St, Fut, F> Stream for Then<St, Fut, F> +where + St: Stream, + F: FnMut(St::Item) -> Fut, + Fut: Future, +{ + type Item = Fut::Output; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + Poll::Ready(loop { + if let Some(fut) = this.future.as_mut().as_pin_mut() { + let item = ready!(fut.poll(cx)); + this.future.set(None); + break Some(item); + } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) { + this.future.set(Some((this.f)(item))); + } else { + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let future_len = usize::from(self.future.is_some()); + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(future_len); + let upper = match upper { + Some(x) => x.checked_add(future_len), + None => None, + }; + (lower, upper) + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for Then<S, Fut, F> +where + S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/stream/unzip.rs b/third_party/rust/futures-util/src/stream/stream/unzip.rs new file mode 100644 index 0000000000..a88cf03266 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/unzip.rs @@ -0,0 +1,63 @@ +use core::mem; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`unzip`](super::StreamExt::unzip) method. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Unzip<St, FromA, FromB> { + #[pin] + stream: St, + left: FromA, + right: FromB, + } +} + +impl<St: Stream, FromA: Default, FromB: Default> Unzip<St, FromA, FromB> { + fn finish(self: Pin<&mut Self>) -> (FromA, FromB) { + let this = self.project(); + (mem::take(this.left), mem::take(this.right)) + } + + pub(super) fn new(stream: St) -> Self { + Self { stream, left: Default::default(), right: Default::default() } + } +} + +impl<St, A, B, FromA, FromB> FusedFuture for Unzip<St, FromA, FromB> +where + St: FusedStream<Item = (A, B)>, + FromA: Default + Extend<A>, + FromB: Default + Extend<B>, +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St, A, B, FromA, FromB> Future for Unzip<St, FromA, FromB> +where + St: Stream<Item = (A, B)>, + FromA: Default + Extend<A>, + FromB: Default + Extend<B>, +{ + type Output = (FromA, FromB); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<(FromA, FromB)> { + let mut this = self.as_mut().project(); + loop { + match ready!(this.stream.as_mut().poll_next(cx)) { + Some(e) => { + this.left.extend(Some(e.0)); + this.right.extend(Some(e.1)); + } + None => return Poll::Ready(self.finish()), + } + } + } +} diff --git a/third_party/rust/futures-util/src/stream/stream/zip.rs b/third_party/rust/futures-util/src/stream/stream/zip.rs new file mode 100644 index 0000000000..25a47e96be --- /dev/null +++ b/third_party/rust/futures-util/src/stream/stream/zip.rs @@ -0,0 +1,128 @@ +use crate::stream::{Fuse, StreamExt}; +use core::cmp; +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`zip`](super::StreamExt::zip) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Zip<St1: Stream, St2: Stream> { + #[pin] + stream1: Fuse<St1>, + #[pin] + stream2: Fuse<St2>, + queued1: Option<St1::Item>, + queued2: Option<St2::Item>, + } +} + +impl<St1: Stream, St2: Stream> Zip<St1, St2> { + pub(super) fn new(stream1: St1, stream2: St2) -> Self { + Self { stream1: stream1.fuse(), stream2: stream2.fuse(), queued1: None, queued2: None } + } + + /// Acquires a reference to the underlying streams that this combinator is + /// pulling from. + pub fn get_ref(&self) -> (&St1, &St2) { + (self.stream1.get_ref(), self.stream2.get_ref()) + } + + /// Acquires a mutable reference to the underlying streams that this + /// combinator is pulling from. + /// + /// Note that care must be taken to avoid tampering with the state of the + /// stream which may otherwise confuse this combinator. + pub fn get_mut(&mut self) -> (&mut St1, &mut St2) { + (self.stream1.get_mut(), self.stream2.get_mut()) + } + + /// Acquires a pinned mutable reference to the underlying streams that this + /// combinator is pulling from. + /// + /// Note that care must be taken to avoid tampering with the state of the + /// stream which may otherwise confuse this combinator. + pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut St1>, Pin<&mut St2>) { + let this = self.project(); + (this.stream1.get_pin_mut(), this.stream2.get_pin_mut()) + } + + /// Consumes this combinator, returning the underlying streams. + /// + /// Note that this may discard intermediate state of this combinator, so + /// care should be taken to avoid losing resources when this is called. + pub fn into_inner(self) -> (St1, St2) { + (self.stream1.into_inner(), self.stream2.into_inner()) + } +} + +impl<St1, St2> FusedStream for Zip<St1, St2> +where + St1: Stream, + St2: Stream, +{ + fn is_terminated(&self) -> bool { + self.stream1.is_terminated() && self.stream2.is_terminated() + } +} + +impl<St1, St2> Stream for Zip<St1, St2> +where + St1: Stream, + St2: Stream, +{ + type Item = (St1::Item, St2::Item); + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + if this.queued1.is_none() { + match this.stream1.as_mut().poll_next(cx) { + Poll::Ready(Some(item1)) => *this.queued1 = Some(item1), + Poll::Ready(None) | Poll::Pending => {} + } + } + if this.queued2.is_none() { + match this.stream2.as_mut().poll_next(cx) { + Poll::Ready(Some(item2)) => *this.queued2 = Some(item2), + Poll::Ready(None) | Poll::Pending => {} + } + } + + if this.queued1.is_some() && this.queued2.is_some() { + let pair = (this.queued1.take().unwrap(), this.queued2.take().unwrap()); + Poll::Ready(Some(pair)) + } else if this.stream1.is_done() || this.stream2.is_done() { + Poll::Ready(None) + } else { + Poll::Pending + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let queued1_len = usize::from(self.queued1.is_some()); + let queued2_len = usize::from(self.queued2.is_some()); + let (stream1_lower, stream1_upper) = self.stream1.size_hint(); + let (stream2_lower, stream2_upper) = self.stream2.size_hint(); + + let stream1_lower = stream1_lower.saturating_add(queued1_len); + let stream2_lower = stream2_lower.saturating_add(queued2_len); + + let lower = cmp::min(stream1_lower, stream2_lower); + + let upper = match (stream1_upper, stream2_upper) { + (Some(x), Some(y)) => { + let x = x.saturating_add(queued1_len); + let y = y.saturating_add(queued2_len); + Some(cmp::min(x, y)) + } + (Some(x), None) => x.checked_add(queued1_len), + (None, Some(y)) => y.checked_add(queued2_len), + (None, None) => None, + }; + + (lower, upper) + } +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/and_then.rs b/third_party/rust/futures-util/src/stream/try_stream/and_then.rs new file mode 100644 index 0000000000..2f8b6f2589 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/and_then.rs @@ -0,0 +1,105 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`and_then`](super::TryStreamExt::and_then) method. + #[must_use = "streams do nothing unless polled"] + pub struct AndThen<St, Fut, F> { + #[pin] + stream: St, + #[pin] + future: Option<Fut>, + f: F, + } +} + +impl<St, Fut, F> fmt::Debug for AndThen<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AndThen") + .field("stream", &self.stream) + .field("future", &self.future) + .finish() + } +} + +impl<St, Fut, F> AndThen<St, Fut, F> +where + St: TryStream, + F: FnMut(St::Ok) -> Fut, + Fut: TryFuture<Error = St::Error>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, future: None, f } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F> Stream for AndThen<St, Fut, F> +where + St: TryStream, + F: FnMut(St::Ok) -> Fut, + Fut: TryFuture<Error = St::Error>, +{ + type Item = Result<Fut::Ok, St::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + Poll::Ready(loop { + if let Some(fut) = this.future.as_mut().as_pin_mut() { + let item = ready!(fut.try_poll(cx)); + this.future.set(None); + break Some(item); + } else if let Some(item) = ready!(this.stream.as_mut().try_poll_next(cx)?) { + this.future.set(Some((this.f)(item))); + } else { + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let future_len = usize::from(self.future.is_some()); + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(future_len); + let upper = match upper { + Some(x) => x.checked_add(future_len), + None => None, + }; + (lower, upper) + } +} + +impl<St, Fut, F> FusedStream for AndThen<St, Fut, F> +where + St: TryStream + FusedStream, + F: FnMut(St::Ok) -> Fut, + Fut: TryFuture<Error = St::Error>, +{ + fn is_terminated(&self) -> bool { + self.future.is_none() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for AndThen<S, Fut, F> +where + S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/into_async_read.rs b/third_party/rust/futures-util/src/stream/try_stream/into_async_read.rs new file mode 100644 index 0000000000..ffbfc7eae9 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/into_async_read.rs @@ -0,0 +1,166 @@ +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::TryStream; +use futures_core::task::{Context, Poll}; +use futures_io::{AsyncBufRead, AsyncRead, AsyncWrite}; +use pin_project_lite::pin_project; +use std::cmp; +use std::io::{Error, Result}; + +pin_project! { + /// Reader for the [`into_async_read`](super::TryStreamExt::into_async_read) method. + #[derive(Debug)] + #[must_use = "readers do nothing unless polled"] + #[cfg_attr(docsrs, doc(cfg(feature = "io")))] + pub struct IntoAsyncRead<St> + where + St: TryStream<Error = Error>, + St::Ok: AsRef<[u8]>, + { + #[pin] + stream: St, + state: ReadState<St::Ok>, + } +} + +#[derive(Debug)] +enum ReadState<T: AsRef<[u8]>> { + Ready { chunk: T, chunk_start: usize }, + PendingChunk, + Eof, +} + +impl<St> IntoAsyncRead<St> +where + St: TryStream<Error = Error>, + St::Ok: AsRef<[u8]>, +{ + pub(super) fn new(stream: St) -> Self { + Self { stream, state: ReadState::PendingChunk } + } +} + +impl<St> AsyncRead for IntoAsyncRead<St> +where + St: TryStream<Error = Error>, + St::Ok: AsRef<[u8]>, +{ + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll<Result<usize>> { + let mut this = self.project(); + + loop { + match this.state { + ReadState::Ready { chunk, chunk_start } => { + let chunk = chunk.as_ref(); + let len = cmp::min(buf.len(), chunk.len() - *chunk_start); + + buf[..len].copy_from_slice(&chunk[*chunk_start..*chunk_start + len]); + *chunk_start += len; + + if chunk.len() == *chunk_start { + *this.state = ReadState::PendingChunk; + } + + return Poll::Ready(Ok(len)); + } + ReadState::PendingChunk => match ready!(this.stream.as_mut().try_poll_next(cx)) { + Some(Ok(chunk)) => { + if !chunk.as_ref().is_empty() { + *this.state = ReadState::Ready { chunk, chunk_start: 0 }; + } + } + Some(Err(err)) => { + *this.state = ReadState::Eof; + return Poll::Ready(Err(err)); + } + None => { + *this.state = ReadState::Eof; + return Poll::Ready(Ok(0)); + } + }, + ReadState::Eof => { + return Poll::Ready(Ok(0)); + } + } + } + } +} + +impl<St> AsyncWrite for IntoAsyncRead<St> +where + St: TryStream<Error = Error> + AsyncWrite, + St::Ok: AsRef<[u8]>, +{ + fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize>> { + let this = self.project(); + this.stream.poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> { + let this = self.project(); + this.stream.poll_flush(cx) + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> { + let this = self.project(); + this.stream.poll_close(cx) + } +} + +impl<St> AsyncBufRead for IntoAsyncRead<St> +where + St: TryStream<Error = Error>, + St::Ok: AsRef<[u8]>, +{ + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<&[u8]>> { + let mut this = self.project(); + + while let ReadState::PendingChunk = this.state { + match ready!(this.stream.as_mut().try_poll_next(cx)) { + Some(Ok(chunk)) => { + if !chunk.as_ref().is_empty() { + *this.state = ReadState::Ready { chunk, chunk_start: 0 }; + } + } + Some(Err(err)) => { + *this.state = ReadState::Eof; + return Poll::Ready(Err(err)); + } + None => { + *this.state = ReadState::Eof; + return Poll::Ready(Ok(&[])); + } + } + } + + if let &mut ReadState::Ready { ref chunk, chunk_start } = this.state { + let chunk = chunk.as_ref(); + return Poll::Ready(Ok(&chunk[chunk_start..])); + } + + // To get to this point we must be in ReadState::Eof + Poll::Ready(Ok(&[])) + } + + fn consume(self: Pin<&mut Self>, amount: usize) { + let this = self.project(); + + // https://github.com/rust-lang/futures-rs/pull/1556#discussion_r281644295 + if amount == 0 { + return; + } + if let ReadState::Ready { chunk, chunk_start } = this.state { + *chunk_start += amount; + debug_assert!(*chunk_start <= chunk.as_ref().len()); + if *chunk_start >= chunk.as_ref().len() { + *this.state = ReadState::PendingChunk; + } + } else { + debug_assert!(false, "Attempted to consume from IntoAsyncRead without chunk"); + } + } +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/into_stream.rs b/third_party/rust/futures-util/src/stream/try_stream/into_stream.rs new file mode 100644 index 0000000000..2126258af7 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/into_stream.rs @@ -0,0 +1,52 @@ +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`into_stream`](super::TryStreamExt::into_stream) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct IntoStream<St> { + #[pin] + stream: St, + } +} + +impl<St> IntoStream<St> { + #[inline] + pub(super) fn new(stream: St) -> Self { + Self { stream } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St: TryStream + FusedStream> FusedStream for IntoStream<St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St: TryStream> Stream for IntoStream<St> { + type Item = Result<St::Ok, St::Error>; + + #[inline] + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + self.project().stream.try_poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S: Sink<Item>, Item> Sink<Item> for IntoStream<S> { + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/mod.rs b/third_party/rust/futures-util/src/stream/try_stream/mod.rs new file mode 100644 index 0000000000..bc4c6e4f6a --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/mod.rs @@ -0,0 +1,1064 @@ +//! Streams +//! +//! This module contains a number of functions for working with `Streams`s +//! that return `Result`s, allowing for short-circuiting computations. + +#[cfg(feature = "compat")] +use crate::compat::Compat; +use crate::fns::{ + inspect_err_fn, inspect_ok_fn, into_fn, map_err_fn, map_ok_fn, InspectErrFn, InspectOkFn, + IntoFn, MapErrFn, MapOkFn, +}; +use crate::future::assert_future; +use crate::stream::assert_stream; +use crate::stream::{Inspect, Map}; +#[cfg(feature = "alloc")] +use alloc::vec::Vec; +use core::pin::Pin; +use futures_core::{ + future::{Future, TryFuture}, + stream::TryStream, + task::{Context, Poll}, +}; + +mod and_then; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::and_then::AndThen; + +delegate_all!( + /// Stream for the [`err_into`](super::TryStreamExt::err_into) method. + ErrInto<St, E>( + MapErr<St, IntoFn<E>> + ): Debug + Sink + Stream + FusedStream + AccessInner[St, (.)] + New[|x: St| MapErr::new(x, into_fn())] +); + +delegate_all!( + /// Stream for the [`inspect_ok`](super::TryStreamExt::inspect_ok) method. + InspectOk<St, F>( + Inspect<IntoStream<St>, InspectOkFn<F>> + ): Debug + Sink + Stream + FusedStream + AccessInner[St, (. .)] + New[|x: St, f: F| Inspect::new(IntoStream::new(x), inspect_ok_fn(f))] +); + +delegate_all!( + /// Stream for the [`inspect_err`](super::TryStreamExt::inspect_err) method. + InspectErr<St, F>( + Inspect<IntoStream<St>, InspectErrFn<F>> + ): Debug + Sink + Stream + FusedStream + AccessInner[St, (. .)] + New[|x: St, f: F| Inspect::new(IntoStream::new(x), inspect_err_fn(f))] +); + +mod into_stream; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::into_stream::IntoStream; + +delegate_all!( + /// Stream for the [`map_ok`](super::TryStreamExt::map_ok) method. + MapOk<St, F>( + Map<IntoStream<St>, MapOkFn<F>> + ): Debug + Sink + Stream + FusedStream + AccessInner[St, (. .)] + New[|x: St, f: F| Map::new(IntoStream::new(x), map_ok_fn(f))] +); + +delegate_all!( + /// Stream for the [`map_err`](super::TryStreamExt::map_err) method. + MapErr<St, F>( + Map<IntoStream<St>, MapErrFn<F>> + ): Debug + Sink + Stream + FusedStream + AccessInner[St, (. .)] + New[|x: St, f: F| Map::new(IntoStream::new(x), map_err_fn(f))] +); + +mod or_else; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::or_else::OrElse; + +mod try_next; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_next::TryNext; + +mod try_for_each; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_for_each::TryForEach; + +mod try_filter; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_filter::TryFilter; + +mod try_filter_map; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_filter_map::TryFilterMap; + +mod try_flatten; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_flatten::TryFlatten; + +mod try_collect; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_collect::TryCollect; + +mod try_concat; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_concat::TryConcat; + +#[cfg(feature = "alloc")] +mod try_chunks; +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_chunks::{TryChunks, TryChunksError}; + +mod try_fold; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_fold::TryFold; + +mod try_unfold; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_unfold::{try_unfold, TryUnfold}; + +mod try_skip_while; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_skip_while::TrySkipWhile; + +mod try_take_while; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_take_while::TryTakeWhile; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +mod try_buffer_unordered; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_buffer_unordered::TryBufferUnordered; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +mod try_buffered; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_buffered::TryBuffered; + +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +mod try_for_each_concurrent; +#[cfg(not(futures_no_atomic_cas))] +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_for_each_concurrent::TryForEachConcurrent; + +#[cfg(feature = "io")] +#[cfg(feature = "std")] +mod into_async_read; +#[cfg(feature = "io")] +#[cfg_attr(docsrs, doc(cfg(feature = "io")))] +#[cfg(feature = "std")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::into_async_read::IntoAsyncRead; + +impl<S: ?Sized + TryStream> TryStreamExt for S {} + +/// Adapters specific to `Result`-returning streams +pub trait TryStreamExt: TryStream { + /// Wraps the current stream in a new stream which converts the error type + /// into the one provided. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// + /// let mut stream = + /// stream::iter(vec![Ok(()), Err(5i32)]) + /// .err_into::<i64>(); + /// + /// assert_eq!(stream.try_next().await, Ok(Some(()))); + /// assert_eq!(stream.try_next().await, Err(5i64)); + /// # }) + /// ``` + fn err_into<E>(self) -> ErrInto<Self, E> + where + Self: Sized, + Self::Error: Into<E>, + { + assert_stream::<Result<Self::Ok, E>, _>(ErrInto::new(self)) + } + + /// Wraps the current stream in a new stream which maps the success value + /// using the provided closure. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// + /// let mut stream = + /// stream::iter(vec![Ok(5), Err(0)]) + /// .map_ok(|x| x + 2); + /// + /// assert_eq!(stream.try_next().await, Ok(Some(7))); + /// assert_eq!(stream.try_next().await, Err(0)); + /// # }) + /// ``` + fn map_ok<T, F>(self, f: F) -> MapOk<Self, F> + where + Self: Sized, + F: FnMut(Self::Ok) -> T, + { + assert_stream::<Result<T, Self::Error>, _>(MapOk::new(self, f)) + } + + /// Wraps the current stream in a new stream which maps the error value + /// using the provided closure. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// + /// let mut stream = + /// stream::iter(vec![Ok(5), Err(0)]) + /// .map_err(|x| x + 2); + /// + /// assert_eq!(stream.try_next().await, Ok(Some(5))); + /// assert_eq!(stream.try_next().await, Err(2)); + /// # }) + /// ``` + fn map_err<E, F>(self, f: F) -> MapErr<Self, F> + where + Self: Sized, + F: FnMut(Self::Error) -> E, + { + assert_stream::<Result<Self::Ok, E>, _>(MapErr::new(self, f)) + } + + /// Chain on a computation for when a value is ready, passing the successful + /// results to the provided closure `f`. + /// + /// This function can be used to run a unit of work when the next successful + /// value on a stream is ready. The closure provided will be yielded a value + /// when ready, and the returned future will then be run to completion to + /// produce the next value on this stream. + /// + /// Any errors produced by this stream will not be passed to the closure, + /// and will be passed through. + /// + /// The returned value of the closure must implement the `TryFuture` trait + /// and can represent some more work to be done before the composed stream + /// is finished. + /// + /// Note that this function consumes the receiving stream and returns a + /// wrapped version of it. + /// + /// To process the entire stream and return a single future representing + /// success or error, use `try_for_each` instead. + /// + /// # Examples + /// + /// ``` + /// use futures::channel::mpsc; + /// use futures::future; + /// use futures::stream::TryStreamExt; + /// + /// let (_tx, rx) = mpsc::channel::<Result<i32, ()>>(1); + /// + /// let rx = rx.and_then(|result| { + /// future::ok(if result % 2 == 0 { + /// Some(result) + /// } else { + /// None + /// }) + /// }); + /// ``` + fn and_then<Fut, F>(self, f: F) -> AndThen<Self, Fut, F> + where + F: FnMut(Self::Ok) -> Fut, + Fut: TryFuture<Error = Self::Error>, + Self: Sized, + { + assert_stream::<Result<Fut::Ok, Fut::Error>, _>(AndThen::new(self, f)) + } + + /// Chain on a computation for when an error happens, passing the + /// erroneous result to the provided closure `f`. + /// + /// This function can be used to run a unit of work and attempt to recover from + /// an error if one happens. The closure provided will be yielded an error + /// when one appears, and the returned future will then be run to completion + /// to produce the next value on this stream. + /// + /// Any successful values produced by this stream will not be passed to the + /// closure, and will be passed through. + /// + /// The returned value of the closure must implement the [`TryFuture`](futures_core::future::TryFuture) trait + /// and can represent some more work to be done before the composed stream + /// is finished. + /// + /// Note that this function consumes the receiving stream and returns a + /// wrapped version of it. + fn or_else<Fut, F>(self, f: F) -> OrElse<Self, Fut, F> + where + F: FnMut(Self::Error) -> Fut, + Fut: TryFuture<Ok = Self::Ok>, + Self: Sized, + { + assert_stream::<Result<Self::Ok, Fut::Error>, _>(OrElse::new(self, f)) + } + + /// Do something with the success value of this stream, afterwards passing + /// it on. + /// + /// This is similar to the `StreamExt::inspect` method where it allows + /// easily inspecting the success value as it passes through the stream, for + /// example to debug what's going on. + fn inspect_ok<F>(self, f: F) -> InspectOk<Self, F> + where + F: FnMut(&Self::Ok), + Self: Sized, + { + assert_stream::<Result<Self::Ok, Self::Error>, _>(InspectOk::new(self, f)) + } + + /// Do something with the error value of this stream, afterwards passing it on. + /// + /// This is similar to the `StreamExt::inspect` method where it allows + /// easily inspecting the error value as it passes through the stream, for + /// example to debug what's going on. + fn inspect_err<F>(self, f: F) -> InspectErr<Self, F> + where + F: FnMut(&Self::Error), + Self: Sized, + { + assert_stream::<Result<Self::Ok, Self::Error>, _>(InspectErr::new(self, f)) + } + + /// Wraps a [`TryStream`] into a type that implements + /// [`Stream`](futures_core::stream::Stream) + /// + /// [`TryStream`]s currently do not implement the + /// [`Stream`](futures_core::stream::Stream) trait because of limitations + /// of the compiler. + /// + /// # Examples + /// + /// ``` + /// use futures::stream::{Stream, TryStream, TryStreamExt}; + /// + /// # type T = i32; + /// # type E = (); + /// fn make_try_stream() -> impl TryStream<Ok = T, Error = E> { // ... } + /// # futures::stream::empty() + /// # } + /// fn take_stream(stream: impl Stream<Item = Result<T, E>>) { /* ... */ } + /// + /// take_stream(make_try_stream().into_stream()); + /// ``` + fn into_stream(self) -> IntoStream<Self> + where + Self: Sized, + { + assert_stream::<Result<Self::Ok, Self::Error>, _>(IntoStream::new(self)) + } + + /// Creates a future that attempts to resolve the next item in the stream. + /// If an error is encountered before the next item, the error is returned + /// instead. + /// + /// This is similar to the `Stream::next` combinator, but returns a + /// `Result<Option<T>, E>` rather than an `Option<Result<T, E>>`, making + /// for easy use with the `?` operator. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// + /// let mut stream = stream::iter(vec![Ok(()), Err(())]); + /// + /// assert_eq!(stream.try_next().await, Ok(Some(()))); + /// assert_eq!(stream.try_next().await, Err(())); + /// # }) + /// ``` + fn try_next(&mut self) -> TryNext<'_, Self> + where + Self: Unpin, + { + assert_future::<Result<Option<Self::Ok>, Self::Error>, _>(TryNext::new(self)) + } + + /// Attempts to run this stream to completion, executing the provided + /// asynchronous closure for each element on the stream. + /// + /// The provided closure will be called for each item this stream produces, + /// yielding a future. That future will then be executed to completion + /// before moving on to the next item. + /// + /// The returned value is a [`Future`](futures_core::future::Future) where the + /// [`Output`](futures_core::future::Future::Output) type is + /// `Result<(), Self::Error>`. If any of the intermediate + /// futures or the stream returns an error, this future will return + /// immediately with an error. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, TryStreamExt}; + /// + /// let mut x = 0i32; + /// + /// { + /// let fut = stream::repeat(Ok(1)).try_for_each(|item| { + /// x += item; + /// future::ready(if x == 3 { Err(()) } else { Ok(()) }) + /// }); + /// assert_eq!(fut.await, Err(())); + /// } + /// + /// assert_eq!(x, 3); + /// # }) + /// ``` + fn try_for_each<Fut, F>(self, f: F) -> TryForEach<Self, Fut, F> + where + F: FnMut(Self::Ok) -> Fut, + Fut: TryFuture<Ok = (), Error = Self::Error>, + Self: Sized, + { + assert_future::<Result<(), Self::Error>, _>(TryForEach::new(self, f)) + } + + /// Skip elements on this stream while the provided asynchronous predicate + /// resolves to `true`. + /// + /// This function is similar to + /// [`StreamExt::skip_while`](crate::stream::StreamExt::skip_while) but exits + /// early if an error occurs. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, TryStreamExt}; + /// + /// let stream = stream::iter(vec![Ok::<i32, i32>(1), Ok(3), Ok(2)]); + /// let stream = stream.try_skip_while(|x| future::ready(Ok(*x < 3))); + /// + /// let output: Result<Vec<i32>, i32> = stream.try_collect().await; + /// assert_eq!(output, Ok(vec![3, 2])); + /// # }) + /// ``` + fn try_skip_while<Fut, F>(self, f: F) -> TrySkipWhile<Self, Fut, F> + where + F: FnMut(&Self::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = Self::Error>, + Self: Sized, + { + assert_stream::<Result<Self::Ok, Self::Error>, _>(TrySkipWhile::new(self, f)) + } + + /// Take elements on this stream while the provided asynchronous predicate + /// resolves to `true`. + /// + /// This function is similar to + /// [`StreamExt::take_while`](crate::stream::StreamExt::take_while) but exits + /// early if an error occurs. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, TryStreamExt}; + /// + /// let stream = stream::iter(vec![Ok::<i32, i32>(1), Ok(2), Ok(3), Ok(2)]); + /// let stream = stream.try_take_while(|x| future::ready(Ok(*x < 3))); + /// + /// let output: Result<Vec<i32>, i32> = stream.try_collect().await; + /// assert_eq!(output, Ok(vec![1, 2])); + /// # }) + /// ``` + fn try_take_while<Fut, F>(self, f: F) -> TryTakeWhile<Self, Fut, F> + where + F: FnMut(&Self::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = Self::Error>, + Self: Sized, + { + assert_stream::<Result<Self::Ok, Self::Error>, _>(TryTakeWhile::new(self, f)) + } + + /// Attempts to run this stream to completion, executing the provided asynchronous + /// closure for each element on the stream concurrently as elements become + /// available, exiting as soon as an error occurs. + /// + /// This is similar to + /// [`StreamExt::for_each_concurrent`](crate::stream::StreamExt::for_each_concurrent), + /// but will resolve to an error immediately if the underlying stream or the provided + /// closure return an error. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::oneshot; + /// use futures::stream::{self, StreamExt, TryStreamExt}; + /// + /// let (tx1, rx1) = oneshot::channel(); + /// let (tx2, rx2) = oneshot::channel(); + /// let (_tx3, rx3) = oneshot::channel(); + /// + /// let stream = stream::iter(vec![rx1, rx2, rx3]); + /// let fut = stream.map(Ok).try_for_each_concurrent( + /// /* limit */ 2, + /// |rx| async move { + /// let res: Result<(), oneshot::Canceled> = rx.await; + /// res + /// } + /// ); + /// + /// tx1.send(()).unwrap(); + /// // Drop the second sender so that `rx2` resolves to `Canceled`. + /// drop(tx2); + /// + /// // The final result is an error because the second future + /// // resulted in an error. + /// assert_eq!(Err(oneshot::Canceled), fut.await); + /// # }) + /// ``` + #[cfg(not(futures_no_atomic_cas))] + #[cfg(feature = "alloc")] + fn try_for_each_concurrent<Fut, F>( + self, + limit: impl Into<Option<usize>>, + f: F, + ) -> TryForEachConcurrent<Self, Fut, F> + where + F: FnMut(Self::Ok) -> Fut, + Fut: Future<Output = Result<(), Self::Error>>, + Self: Sized, + { + assert_future::<Result<(), Self::Error>, _>(TryForEachConcurrent::new( + self, + limit.into(), + f, + )) + } + + /// Attempt to transform a stream into a collection, + /// returning a future representing the result of that computation. + /// + /// This combinator will collect all successful results of this stream and + /// collect them into the specified collection type. If an error happens then all + /// collected elements will be dropped and the error will be returned. + /// + /// The returned future will be resolved when the stream terminates. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::TryStreamExt; + /// use std::thread; + /// + /// let (tx, rx) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// for i in 1..=5 { + /// tx.unbounded_send(Ok(i)).unwrap(); + /// } + /// tx.unbounded_send(Err(6)).unwrap(); + /// }); + /// + /// let output: Result<Vec<i32>, i32> = rx.try_collect().await; + /// assert_eq!(output, Err(6)); + /// # }) + /// ``` + fn try_collect<C: Default + Extend<Self::Ok>>(self) -> TryCollect<Self, C> + where + Self: Sized, + { + assert_future::<Result<C, Self::Error>, _>(TryCollect::new(self)) + } + + /// An adaptor for chunking up successful items of the stream inside a vector. + /// + /// This combinator will attempt to pull successful items from this stream and buffer + /// them into a local vector. At most `capacity` items will get buffered + /// before they're yielded from the returned stream. + /// + /// Note that the vectors returned from this iterator may not always have + /// `capacity` elements. If the underlying stream ended and only a partial + /// vector was created, it'll be returned. Additionally if an error happens + /// from the underlying stream then the currently buffered items will be + /// yielded. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// This function is similar to + /// [`StreamExt::chunks`](crate::stream::StreamExt::chunks) but exits + /// early if an error occurs. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryChunksError, TryStreamExt}; + /// + /// let stream = stream::iter(vec![Ok::<i32, i32>(1), Ok(2), Ok(3), Err(4), Ok(5), Ok(6)]); + /// let mut stream = stream.try_chunks(2); + /// + /// assert_eq!(stream.try_next().await, Ok(Some(vec![1, 2]))); + /// assert_eq!(stream.try_next().await, Err(TryChunksError(vec![3], 4))); + /// assert_eq!(stream.try_next().await, Ok(Some(vec![5, 6]))); + /// # }) + /// ``` + /// + /// # Panics + /// + /// This method will panic if `capacity` is zero. + #[cfg(feature = "alloc")] + fn try_chunks(self, capacity: usize) -> TryChunks<Self> + where + Self: Sized, + { + assert_stream::<Result<Vec<Self::Ok>, TryChunksError<Self::Ok, Self::Error>>, _>( + TryChunks::new(self, capacity), + ) + } + + /// Attempt to filter the values produced by this stream according to the + /// provided asynchronous closure. + /// + /// As values of this stream are made available, the provided predicate `f` + /// will be run on them. If the predicate returns a `Future` which resolves + /// to `true`, then the stream will yield the value, but if the predicate + /// return a `Future` which resolves to `false`, then the value will be + /// discarded and the next value will be produced. + /// + /// All errors are passed through without filtering in this combinator. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `filter` methods in + /// the standard library. + /// + /// # Examples + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt, TryStreamExt}; + /// + /// let stream = stream::iter(vec![Ok(1i32), Ok(2i32), Ok(3i32), Err("error")]); + /// let mut evens = stream.try_filter(|x| { + /// future::ready(x % 2 == 0) + /// }); + /// + /// assert_eq!(evens.next().await, Some(Ok(2))); + /// assert_eq!(evens.next().await, Some(Err("error"))); + /// # }) + /// ``` + fn try_filter<Fut, F>(self, f: F) -> TryFilter<Self, Fut, F> + where + Fut: Future<Output = bool>, + F: FnMut(&Self::Ok) -> Fut, + Self: Sized, + { + assert_stream::<Result<Self::Ok, Self::Error>, _>(TryFilter::new(self, f)) + } + + /// Attempt to filter the values produced by this stream while + /// simultaneously mapping them to a different type according to the + /// provided asynchronous closure. + /// + /// As values of this stream are made available, the provided function will + /// be run on them. If the future returned by the predicate `f` resolves to + /// [`Some(item)`](Some) then the stream will yield the value `item`, but if + /// it resolves to [`None`] then the next value will be produced. + /// + /// All errors are passed through without filtering in this combinator. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `filter_map` methods in + /// the standard library. + /// + /// # Examples + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt, TryStreamExt}; + /// use futures::pin_mut; + /// + /// let stream = stream::iter(vec![Ok(1i32), Ok(6i32), Err("error")]); + /// let halves = stream.try_filter_map(|x| async move { + /// let ret = if x % 2 == 0 { Some(x / 2) } else { None }; + /// Ok(ret) + /// }); + /// + /// pin_mut!(halves); + /// assert_eq!(halves.next().await, Some(Ok(3))); + /// assert_eq!(halves.next().await, Some(Err("error"))); + /// # }) + /// ``` + fn try_filter_map<Fut, F, T>(self, f: F) -> TryFilterMap<Self, Fut, F> + where + Fut: TryFuture<Ok = Option<T>, Error = Self::Error>, + F: FnMut(Self::Ok) -> Fut, + Self: Sized, + { + assert_stream::<Result<T, Self::Error>, _>(TryFilterMap::new(self, f)) + } + + /// Flattens a stream of streams into just one continuous stream. + /// + /// If this stream's elements are themselves streams then this combinator + /// will flatten out the entire stream to one long chain of elements. Any + /// errors are passed through without looking at them, but otherwise each + /// individual stream will get exhausted before moving on to the next. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::{StreamExt, TryStreamExt}; + /// use std::thread; + /// + /// let (tx1, rx1) = mpsc::unbounded(); + /// let (tx2, rx2) = mpsc::unbounded(); + /// let (tx3, rx3) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// tx1.unbounded_send(Ok(1)).unwrap(); + /// }); + /// thread::spawn(move || { + /// tx2.unbounded_send(Ok(2)).unwrap(); + /// tx2.unbounded_send(Err(3)).unwrap(); + /// tx2.unbounded_send(Ok(4)).unwrap(); + /// }); + /// thread::spawn(move || { + /// tx3.unbounded_send(Ok(rx1)).unwrap(); + /// tx3.unbounded_send(Ok(rx2)).unwrap(); + /// tx3.unbounded_send(Err(5)).unwrap(); + /// }); + /// + /// let mut stream = rx3.try_flatten(); + /// assert_eq!(stream.next().await, Some(Ok(1))); + /// assert_eq!(stream.next().await, Some(Ok(2))); + /// assert_eq!(stream.next().await, Some(Err(3))); + /// assert_eq!(stream.next().await, Some(Ok(4))); + /// assert_eq!(stream.next().await, Some(Err(5))); + /// assert_eq!(stream.next().await, None); + /// # }); + /// ``` + fn try_flatten(self) -> TryFlatten<Self> + where + Self::Ok: TryStream, + <Self::Ok as TryStream>::Error: From<Self::Error>, + Self: Sized, + { + assert_stream::<Result<<Self::Ok as TryStream>::Ok, <Self::Ok as TryStream>::Error>, _>( + TryFlatten::new(self), + ) + } + + /// Attempt to execute an accumulating asynchronous computation over a + /// stream, collecting all the values into one final result. + /// + /// This combinator will accumulate all values returned by this stream + /// according to the closure provided. The initial state is also provided to + /// this method and then is returned again by each execution of the closure. + /// Once the entire stream has been exhausted the returned future will + /// resolve to this value. + /// + /// This method is similar to [`fold`](crate::stream::StreamExt::fold), but will + /// exit early if an error is encountered in either the stream or the + /// provided closure. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// + /// let number_stream = stream::iter(vec![Ok::<i32, i32>(1), Ok(2)]); + /// let sum = number_stream.try_fold(0, |acc, x| async move { Ok(acc + x) }); + /// assert_eq!(sum.await, Ok(3)); + /// + /// let number_stream_with_err = stream::iter(vec![Ok::<i32, i32>(1), Err(2), Ok(1)]); + /// let sum = number_stream_with_err.try_fold(0, |acc, x| async move { Ok(acc + x) }); + /// assert_eq!(sum.await, Err(2)); + /// # }) + /// ``` + fn try_fold<T, Fut, F>(self, init: T, f: F) -> TryFold<Self, Fut, T, F> + where + F: FnMut(T, Self::Ok) -> Fut, + Fut: TryFuture<Ok = T, Error = Self::Error>, + Self: Sized, + { + assert_future::<Result<T, Self::Error>, _>(TryFold::new(self, f, init)) + } + + /// Attempt to concatenate all items of a stream into a single + /// extendable destination, returning a future representing the end result. + /// + /// This combinator will extend the first item with the contents of all + /// the subsequent successful results of the stream. If the stream is empty, + /// the default value will be returned. + /// + /// Works with all collections that implement the [`Extend`](std::iter::Extend) trait. + /// + /// This method is similar to [`concat`](crate::stream::StreamExt::concat), but will + /// exit early if an error is encountered in the stream. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::TryStreamExt; + /// use std::thread; + /// + /// let (tx, rx) = mpsc::unbounded::<Result<Vec<i32>, ()>>(); + /// + /// thread::spawn(move || { + /// for i in (0..3).rev() { + /// let n = i * 3; + /// tx.unbounded_send(Ok(vec![n + 1, n + 2, n + 3])).unwrap(); + /// } + /// }); + /// + /// let result = rx.try_concat().await; + /// + /// assert_eq!(result, Ok(vec![7, 8, 9, 4, 5, 6, 1, 2, 3])); + /// # }); + /// ``` + fn try_concat(self) -> TryConcat<Self> + where + Self: Sized, + Self::Ok: Extend<<<Self as TryStream>::Ok as IntoIterator>::Item> + IntoIterator + Default, + { + assert_future::<Result<Self::Ok, Self::Error>, _>(TryConcat::new(self)) + } + + /// Attempt to execute several futures from a stream concurrently (unordered). + /// + /// This stream's `Ok` type must be a [`TryFuture`](futures_core::future::TryFuture) with an `Error` type + /// that matches the stream's `Error` type. + /// + /// This adaptor will buffer up to `n` futures and then return their + /// outputs in the order in which they complete. If the underlying stream + /// returns an error, it will be immediately propagated. + /// + /// The returned stream will be a stream of results, each containing either + /// an error or a future's output. An error can be produced either by the + /// underlying stream itself or by one of the futures it yielded. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// Results are returned in the order of completion: + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::oneshot; + /// use futures::stream::{self, StreamExt, TryStreamExt}; + /// + /// let (send_one, recv_one) = oneshot::channel(); + /// let (send_two, recv_two) = oneshot::channel(); + /// + /// let stream_of_futures = stream::iter(vec![Ok(recv_one), Ok(recv_two)]); + /// + /// let mut buffered = stream_of_futures.try_buffer_unordered(10); + /// + /// send_two.send(2i32)?; + /// assert_eq!(buffered.next().await, Some(Ok(2i32))); + /// + /// send_one.send(1i32)?; + /// assert_eq!(buffered.next().await, Some(Ok(1i32))); + /// + /// assert_eq!(buffered.next().await, None); + /// # Ok::<(), i32>(()) }).unwrap(); + /// ``` + /// + /// Errors from the underlying stream itself are propagated: + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::{StreamExt, TryStreamExt}; + /// + /// let (sink, stream_of_futures) = mpsc::unbounded(); + /// let mut buffered = stream_of_futures.try_buffer_unordered(10); + /// + /// sink.unbounded_send(Ok(async { Ok(7i32) }))?; + /// assert_eq!(buffered.next().await, Some(Ok(7i32))); + /// + /// sink.unbounded_send(Err("error in the stream"))?; + /// assert_eq!(buffered.next().await, Some(Err("error in the stream"))); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + #[cfg(not(futures_no_atomic_cas))] + #[cfg(feature = "alloc")] + fn try_buffer_unordered(self, n: usize) -> TryBufferUnordered<Self> + where + Self::Ok: TryFuture<Error = Self::Error>, + Self: Sized, + { + assert_stream::<Result<<Self::Ok as TryFuture>::Ok, Self::Error>, _>( + TryBufferUnordered::new(self, n), + ) + } + + /// Attempt to execute several futures from a stream concurrently. + /// + /// This stream's `Ok` type must be a [`TryFuture`](futures_core::future::TryFuture) with an `Error` type + /// that matches the stream's `Error` type. + /// + /// This adaptor will buffer up to `n` futures and then return their + /// outputs in the same order as the underlying stream. If the underlying stream returns an error, it will + /// be immediately propagated. + /// + /// The returned stream will be a stream of results, each containing either + /// an error or a future's output. An error can be produced either by the + /// underlying stream itself or by one of the futures it yielded. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// Results are returned in the order of addition: + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::oneshot; + /// use futures::future::lazy; + /// use futures::stream::{self, StreamExt, TryStreamExt}; + /// + /// let (send_one, recv_one) = oneshot::channel(); + /// let (send_two, recv_two) = oneshot::channel(); + /// + /// let mut buffered = lazy(move |cx| { + /// let stream_of_futures = stream::iter(vec![Ok(recv_one), Ok(recv_two)]); + /// + /// let mut buffered = stream_of_futures.try_buffered(10); + /// + /// assert!(buffered.try_poll_next_unpin(cx).is_pending()); + /// + /// send_two.send(2i32)?; + /// assert!(buffered.try_poll_next_unpin(cx).is_pending()); + /// Ok::<_, i32>(buffered) + /// }).await?; + /// + /// send_one.send(1i32)?; + /// assert_eq!(buffered.next().await, Some(Ok(1i32))); + /// assert_eq!(buffered.next().await, Some(Ok(2i32))); + /// + /// assert_eq!(buffered.next().await, None); + /// # Ok::<(), i32>(()) }).unwrap(); + /// ``` + /// + /// Errors from the underlying stream itself are propagated: + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::{StreamExt, TryStreamExt}; + /// + /// let (sink, stream_of_futures) = mpsc::unbounded(); + /// let mut buffered = stream_of_futures.try_buffered(10); + /// + /// sink.unbounded_send(Ok(async { Ok(7i32) }))?; + /// assert_eq!(buffered.next().await, Some(Ok(7i32))); + /// + /// sink.unbounded_send(Err("error in the stream"))?; + /// assert_eq!(buffered.next().await, Some(Err("error in the stream"))); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + #[cfg(not(futures_no_atomic_cas))] + #[cfg(feature = "alloc")] + fn try_buffered(self, n: usize) -> TryBuffered<Self> + where + Self::Ok: TryFuture<Error = Self::Error>, + Self: Sized, + { + assert_stream::<Result<<Self::Ok as TryFuture>::Ok, Self::Error>, _>(TryBuffered::new( + self, n, + )) + } + + // TODO: false positive warning from rustdoc. Verify once #43466 settles + // + /// A convenience method for calling [`TryStream::try_poll_next`] on [`Unpin`] + /// stream types. + fn try_poll_next_unpin( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<Self::Ok, Self::Error>>> + where + Self: Unpin, + { + Pin::new(self).try_poll_next(cx) + } + + /// Wraps a [`TryStream`] into a stream compatible with libraries using + /// futures 0.1 `Stream`. Requires the `compat` feature to be enabled. + /// ``` + /// # if cfg!(miri) { return; } // Miri does not support epoll + /// use futures::future::{FutureExt, TryFutureExt}; + /// # let (tx, rx) = futures::channel::oneshot::channel(); + /// + /// let future03 = async { + /// println!("Running on the pool"); + /// tx.send(42).unwrap(); + /// }; + /// + /// let future01 = future03 + /// .unit_error() // Make it a TryFuture + /// .boxed() // Make it Unpin + /// .compat(); + /// + /// tokio::run(future01); + /// # assert_eq!(42, futures::executor::block_on(rx).unwrap()); + /// ``` + #[cfg(feature = "compat")] + #[cfg_attr(docsrs, doc(cfg(feature = "compat")))] + fn compat(self) -> Compat<Self> + where + Self: Sized + Unpin, + { + Compat::new(self) + } + + /// Adapter that converts this stream into an [`AsyncBufRead`](crate::io::AsyncBufRead). + /// + /// This method is only available when the `std` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// use futures::io::AsyncReadExt; + /// + /// let stream = stream::iter([Ok(vec![1, 2, 3]), Ok(vec![4, 5])]); + /// let mut reader = stream.into_async_read(); + /// + /// let mut buf = Vec::new(); + /// reader.read_to_end(&mut buf).await.unwrap(); + /// assert_eq!(buf, [1, 2, 3, 4, 5]); + /// # }) + /// ``` + #[cfg(feature = "io")] + #[cfg_attr(docsrs, doc(cfg(feature = "io")))] + #[cfg(feature = "std")] + fn into_async_read(self) -> IntoAsyncRead<Self> + where + Self: Sized + TryStreamExt<Error = std::io::Error>, + Self::Ok: AsRef<[u8]>, + { + crate::io::assert_read(IntoAsyncRead::new(self)) + } +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/or_else.rs b/third_party/rust/futures-util/src/stream/try_stream/or_else.rs new file mode 100644 index 0000000000..53aceb8e64 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/or_else.rs @@ -0,0 +1,109 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`or_else`](super::TryStreamExt::or_else) method. + #[must_use = "streams do nothing unless polled"] + pub struct OrElse<St, Fut, F> { + #[pin] + stream: St, + #[pin] + future: Option<Fut>, + f: F, + } +} + +impl<St, Fut, F> fmt::Debug for OrElse<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("OrElse") + .field("stream", &self.stream) + .field("future", &self.future) + .finish() + } +} + +impl<St, Fut, F> OrElse<St, Fut, F> +where + St: TryStream, + F: FnMut(St::Error) -> Fut, + Fut: TryFuture<Ok = St::Ok>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, future: None, f } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F> Stream for OrElse<St, Fut, F> +where + St: TryStream, + F: FnMut(St::Error) -> Fut, + Fut: TryFuture<Ok = St::Ok>, +{ + type Item = Result<St::Ok, Fut::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + Poll::Ready(loop { + if let Some(fut) = this.future.as_mut().as_pin_mut() { + let item = ready!(fut.try_poll(cx)); + this.future.set(None); + break Some(item); + } else { + match ready!(this.stream.as_mut().try_poll_next(cx)) { + Some(Ok(item)) => break Some(Ok(item)), + Some(Err(e)) => { + this.future.set(Some((this.f)(e))); + } + None => break None, + } + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let future_len = usize::from(self.future.is_some()); + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(future_len); + let upper = match upper { + Some(x) => x.checked_add(future_len), + None => None, + }; + (lower, upper) + } +} + +impl<St, Fut, F> FusedStream for OrElse<St, Fut, F> +where + St: TryStream + FusedStream, + F: FnMut(St::Error) -> Fut, + Fut: TryFuture<Ok = St::Ok>, +{ + fn is_terminated(&self) -> bool { + self.future.is_none() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for OrElse<S, Fut, F> +where + S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_buffer_unordered.rs b/third_party/rust/futures-util/src/stream/try_stream/try_buffer_unordered.rs new file mode 100644 index 0000000000..9a899d4ea6 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_buffer_unordered.rs @@ -0,0 +1,86 @@ +use crate::future::{IntoFuture, TryFutureExt}; +use crate::stream::{Fuse, FuturesUnordered, IntoStream, StreamExt}; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::stream::{Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the + /// [`try_buffer_unordered`](super::TryStreamExt::try_buffer_unordered) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct TryBufferUnordered<St> + where St: TryStream + { + #[pin] + stream: Fuse<IntoStream<St>>, + in_progress_queue: FuturesUnordered<IntoFuture<St::Ok>>, + max: usize, + } +} + +impl<St> TryBufferUnordered<St> +where + St: TryStream, + St::Ok: TryFuture, +{ + pub(super) fn new(stream: St, n: usize) -> Self { + Self { + stream: IntoStream::new(stream).fuse(), + in_progress_queue: FuturesUnordered::new(), + max: n, + } + } + + delegate_access_inner!(stream, St, (. .)); +} + +impl<St> Stream for TryBufferUnordered<St> +where + St: TryStream, + St::Ok: TryFuture<Error = St::Error>, +{ + type Item = Result<<St::Ok as TryFuture>::Ok, St::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + // First up, try to spawn off as many futures as possible by filling up + // our queue of futures. Propagate errors from the stream immediately. + while this.in_progress_queue.len() < *this.max { + match this.stream.as_mut().poll_next(cx)? { + Poll::Ready(Some(fut)) => this.in_progress_queue.push(fut.into_future()), + Poll::Ready(None) | Poll::Pending => break, + } + } + + // Attempt to pull the next value from the in_progress_queue + match this.in_progress_queue.poll_next_unpin(cx) { + x @ Poll::Pending | x @ Poll::Ready(Some(_)) => return x, + Poll::Ready(None) => {} + } + + // If more values are still coming from the stream, we're not done yet + if this.stream.is_done() { + Poll::Ready(None) + } else { + Poll::Pending + } + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item, E> Sink<Item> for TryBufferUnordered<S> +where + S: TryStream + Sink<Item, Error = E>, + S::Ok: TryFuture<Error = E>, +{ + type Error = E; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_buffered.rs b/third_party/rust/futures-util/src/stream/try_stream/try_buffered.rs new file mode 100644 index 0000000000..9f48e5c0a7 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_buffered.rs @@ -0,0 +1,87 @@ +use crate::future::{IntoFuture, TryFutureExt}; +use crate::stream::{Fuse, FuturesOrdered, IntoStream, StreamExt}; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::stream::{Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`try_buffered`](super::TryStreamExt::try_buffered) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct TryBuffered<St> + where + St: TryStream, + St::Ok: TryFuture, + { + #[pin] + stream: Fuse<IntoStream<St>>, + in_progress_queue: FuturesOrdered<IntoFuture<St::Ok>>, + max: usize, + } +} + +impl<St> TryBuffered<St> +where + St: TryStream, + St::Ok: TryFuture, +{ + pub(super) fn new(stream: St, n: usize) -> Self { + Self { + stream: IntoStream::new(stream).fuse(), + in_progress_queue: FuturesOrdered::new(), + max: n, + } + } + + delegate_access_inner!(stream, St, (. .)); +} + +impl<St> Stream for TryBuffered<St> +where + St: TryStream, + St::Ok: TryFuture<Error = St::Error>, +{ + type Item = Result<<St::Ok as TryFuture>::Ok, St::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + // First up, try to spawn off as many futures as possible by filling up + // our queue of futures. Propagate errors from the stream immediately. + while this.in_progress_queue.len() < *this.max { + match this.stream.as_mut().poll_next(cx)? { + Poll::Ready(Some(fut)) => this.in_progress_queue.push_back(fut.into_future()), + Poll::Ready(None) | Poll::Pending => break, + } + } + + // Attempt to pull the next value from the in_progress_queue + match this.in_progress_queue.poll_next_unpin(cx) { + x @ Poll::Pending | x @ Poll::Ready(Some(_)) => return x, + Poll::Ready(None) => {} + } + + // If more values are still coming from the stream, we're not done yet + if this.stream.is_done() { + Poll::Ready(None) + } else { + Poll::Pending + } + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item, E> Sink<Item> for TryBuffered<S> +where + S: TryStream + Sink<Item, Error = E>, + S::Ok: TryFuture<Error = E>, +{ + type Error = E; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_chunks.rs b/third_party/rust/futures-util/src/stream/try_stream/try_chunks.rs new file mode 100644 index 0000000000..3bb253a714 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_chunks.rs @@ -0,0 +1,131 @@ +use crate::stream::{Fuse, IntoStream, StreamExt}; + +use alloc::vec::Vec; +use core::pin::Pin; +use core::{fmt, mem}; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`try_chunks`](super::TryStreamExt::try_chunks) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct TryChunks<St: TryStream> { + #[pin] + stream: Fuse<IntoStream<St>>, + items: Vec<St::Ok>, + cap: usize, // https://github.com/rust-lang/futures-rs/issues/1475 + } +} + +impl<St: TryStream> TryChunks<St> { + pub(super) fn new(stream: St, capacity: usize) -> Self { + assert!(capacity > 0); + + Self { + stream: IntoStream::new(stream).fuse(), + items: Vec::with_capacity(capacity), + cap: capacity, + } + } + + fn take(self: Pin<&mut Self>) -> Vec<St::Ok> { + let cap = self.cap; + mem::replace(self.project().items, Vec::with_capacity(cap)) + } + + delegate_access_inner!(stream, St, (. .)); +} + +impl<St: TryStream> Stream for TryChunks<St> { + #[allow(clippy::type_complexity)] + type Item = Result<Vec<St::Ok>, TryChunksError<St::Ok, St::Error>>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.as_mut().project(); + loop { + match ready!(this.stream.as_mut().try_poll_next(cx)) { + // Push the item into the buffer and check whether it is full. + // If so, replace our buffer with a new and empty one and return + // the full one. + Some(item) => match item { + Ok(item) => { + this.items.push(item); + if this.items.len() >= *this.cap { + return Poll::Ready(Some(Ok(self.take()))); + } + } + Err(e) => { + return Poll::Ready(Some(Err(TryChunksError(self.take(), e)))); + } + }, + + // Since the underlying stream ran out of values, return what we + // have buffered, if we have anything. + None => { + let last = if this.items.is_empty() { + None + } else { + let full_buf = mem::take(this.items); + Some(full_buf) + }; + + return Poll::Ready(last.map(Ok)); + } + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let chunk_len = usize::from(!self.items.is_empty()); + let (lower, upper) = self.stream.size_hint(); + let lower = (lower / self.cap).saturating_add(chunk_len); + let upper = match upper { + Some(x) => x.checked_add(chunk_len), + None => None, + }; + (lower, upper) + } +} + +impl<St: TryStream + FusedStream> FusedStream for TryChunks<St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() && self.items.is_empty() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for TryChunks<S> +where + S: TryStream + Sink<Item>, +{ + type Error = <S as Sink<Item>>::Error; + + delegate_sink!(stream, Item); +} + +/// Error indicating, that while chunk was collected inner stream produced an error. +/// +/// Contains all items that were collected before an error occurred, and the stream error itself. +#[derive(PartialEq, Eq)] +pub struct TryChunksError<T, E>(pub Vec<T>, pub E); + +impl<T, E: fmt::Debug> fmt::Debug for TryChunksError<T, E> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.1.fmt(f) + } +} + +impl<T, E: fmt::Display> fmt::Display for TryChunksError<T, E> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.1.fmt(f) + } +} + +#[cfg(feature = "std")] +impl<T, E: fmt::Debug + fmt::Display> std::error::Error for TryChunksError<T, E> {} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_collect.rs b/third_party/rust/futures-util/src/stream/try_stream/try_collect.rs new file mode 100644 index 0000000000..3e5963f033 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_collect.rs @@ -0,0 +1,52 @@ +use core::mem; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::ready; +use futures_core::stream::{FusedStream, TryStream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`try_collect`](super::TryStreamExt::try_collect) method. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct TryCollect<St, C> { + #[pin] + stream: St, + items: C, + } +} + +impl<St: TryStream, C: Default> TryCollect<St, C> { + pub(super) fn new(s: St) -> Self { + Self { stream: s, items: Default::default() } + } +} + +impl<St, C> FusedFuture for TryCollect<St, C> +where + St: TryStream + FusedStream, + C: Default + Extend<St::Ok>, +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St, C> Future for TryCollect<St, C> +where + St: TryStream, + C: Default + Extend<St::Ok>, +{ + type Output = Result<C, St::Error>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut this = self.project(); + Poll::Ready(Ok(loop { + match ready!(this.stream.as_mut().try_poll_next(cx)?) { + Some(x) => this.items.extend(Some(x)), + None => break mem::take(this.items), + } + })) + } +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_concat.rs b/third_party/rust/futures-util/src/stream/try_stream/try_concat.rs new file mode 100644 index 0000000000..58fb6a5413 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_concat.rs @@ -0,0 +1,51 @@ +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::TryStream; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`try_concat`](super::TryStreamExt::try_concat) method. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct TryConcat<St: TryStream> { + #[pin] + stream: St, + accum: Option<St::Ok>, + } +} + +impl<St> TryConcat<St> +where + St: TryStream, + St::Ok: Extend<<St::Ok as IntoIterator>::Item> + IntoIterator + Default, +{ + pub(super) fn new(stream: St) -> Self { + Self { stream, accum: None } + } +} + +impl<St> Future for TryConcat<St> +where + St: TryStream, + St::Ok: Extend<<St::Ok as IntoIterator>::Item> + IntoIterator + Default, +{ + type Output = Result<St::Ok, St::Error>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut this = self.project(); + + Poll::Ready(Ok(loop { + if let Some(x) = ready!(this.stream.as_mut().try_poll_next(cx)?) { + if let Some(a) = this.accum { + a.extend(x) + } else { + *this.accum = Some(x) + } + } else { + break this.accum.take().unwrap_or_default(); + } + })) + } +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_filter.rs b/third_party/rust/futures-util/src/stream/try_stream/try_filter.rs new file mode 100644 index 0000000000..11d58243fd --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_filter.rs @@ -0,0 +1,112 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`try_filter`](super::TryStreamExt::try_filter) + /// method. + #[must_use = "streams do nothing unless polled"] + pub struct TryFilter<St, Fut, F> + where St: TryStream + { + #[pin] + stream: St, + f: F, + #[pin] + pending_fut: Option<Fut>, + pending_item: Option<St::Ok>, + } +} + +impl<St, Fut, F> fmt::Debug for TryFilter<St, Fut, F> +where + St: TryStream + fmt::Debug, + St::Ok: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TryFilter") + .field("stream", &self.stream) + .field("pending_fut", &self.pending_fut) + .field("pending_item", &self.pending_item) + .finish() + } +} + +impl<St, Fut, F> TryFilter<St, Fut, F> +where + St: TryStream, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, pending_fut: None, pending_item: None } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F> FusedStream for TryFilter<St, Fut, F> +where + St: TryStream + FusedStream, + F: FnMut(&St::Ok) -> Fut, + Fut: Future<Output = bool>, +{ + fn is_terminated(&self) -> bool { + self.pending_fut.is_none() && self.stream.is_terminated() + } +} + +impl<St, Fut, F> Stream for TryFilter<St, Fut, F> +where + St: TryStream, + Fut: Future<Output = bool>, + F: FnMut(&St::Ok) -> Fut, +{ + type Item = Result<St::Ok, St::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + Poll::Ready(loop { + if let Some(fut) = this.pending_fut.as_mut().as_pin_mut() { + let res = ready!(fut.poll(cx)); + this.pending_fut.set(None); + if res { + break this.pending_item.take().map(Ok); + } + *this.pending_item = None; + } else if let Some(item) = ready!(this.stream.as_mut().try_poll_next(cx)?) { + this.pending_fut.set(Some((this.f)(&item))); + *this.pending_item = Some(item); + } else { + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = usize::from(self.pending_fut.is_some()); + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item, E> Sink<Item> for TryFilter<S, Fut, F> +where + S: TryStream + Sink<Item, Error = E>, +{ + type Error = E; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_filter_map.rs b/third_party/rust/futures-util/src/stream/try_stream/try_filter_map.rs new file mode 100644 index 0000000000..ed1201732b --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_filter_map.rs @@ -0,0 +1,106 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`try_filter_map`](super::TryStreamExt::try_filter_map) + /// method. + #[must_use = "streams do nothing unless polled"] + pub struct TryFilterMap<St, Fut, F> { + #[pin] + stream: St, + f: F, + #[pin] + pending: Option<Fut>, + } +} + +impl<St, Fut, F> fmt::Debug for TryFilterMap<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TryFilterMap") + .field("stream", &self.stream) + .field("pending", &self.pending) + .finish() + } +} + +impl<St, Fut, F> TryFilterMap<St, Fut, F> { + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, pending: None } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F, T> FusedStream for TryFilterMap<St, Fut, F> +where + St: TryStream + FusedStream, + Fut: TryFuture<Ok = Option<T>, Error = St::Error>, + F: FnMut(St::Ok) -> Fut, +{ + fn is_terminated(&self) -> bool { + self.pending.is_none() && self.stream.is_terminated() + } +} + +impl<St, Fut, F, T> Stream for TryFilterMap<St, Fut, F> +where + St: TryStream, + Fut: TryFuture<Ok = Option<T>, Error = St::Error>, + F: FnMut(St::Ok) -> Fut, +{ + type Item = Result<T, St::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + Poll::Ready(loop { + if let Some(p) = this.pending.as_mut().as_pin_mut() { + // We have an item in progress, poll that until it's done + let res = ready!(p.try_poll(cx)); + this.pending.set(None); + let item = res?; + if item.is_some() { + break item.map(Ok); + } + } else if let Some(item) = ready!(this.stream.as_mut().try_poll_next(cx)?) { + // No item in progress, but the stream is still going + this.pending.set(Some((this.f)(item))); + } else { + // The stream is done + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = usize::from(self.pending.is_some()); + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for TryFilterMap<S, Fut, F> +where + S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_flatten.rs b/third_party/rust/futures-util/src/stream/try_stream/try_flatten.rs new file mode 100644 index 0000000000..4fc04a07bb --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_flatten.rs @@ -0,0 +1,84 @@ +use core::pin::Pin; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`try_flatten`](super::TryStreamExt::try_flatten) method. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct TryFlatten<St> + where + St: TryStream, + { + #[pin] + stream: St, + #[pin] + next: Option<St::Ok>, + } +} + +impl<St> TryFlatten<St> +where + St: TryStream, + St::Ok: TryStream, + <St::Ok as TryStream>::Error: From<St::Error>, +{ + pub(super) fn new(stream: St) -> Self { + Self { stream, next: None } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St> FusedStream for TryFlatten<St> +where + St: TryStream + FusedStream, + St::Ok: TryStream, + <St::Ok as TryStream>::Error: From<St::Error>, +{ + fn is_terminated(&self) -> bool { + self.next.is_none() && self.stream.is_terminated() + } +} + +impl<St> Stream for TryFlatten<St> +where + St: TryStream, + St::Ok: TryStream, + <St::Ok as TryStream>::Error: From<St::Error>, +{ + type Item = Result<<St::Ok as TryStream>::Ok, <St::Ok as TryStream>::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + Poll::Ready(loop { + if let Some(s) = this.next.as_mut().as_pin_mut() { + if let Some(item) = ready!(s.try_poll_next(cx)?) { + break Some(Ok(item)); + } else { + this.next.set(None); + } + } else if let Some(s) = ready!(this.stream.as_mut().try_poll_next(cx)?) { + this.next.set(Some(s)); + } else { + break None; + } + }) + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for TryFlatten<S> +where + S: TryStream + Sink<Item>, +{ + type Error = <S as Sink<Item>>::Error; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_fold.rs b/third_party/rust/futures-util/src/stream/try_stream/try_fold.rs new file mode 100644 index 0000000000..d344d96e7d --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_fold.rs @@ -0,0 +1,93 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::ready; +use futures_core::stream::TryStream; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`try_fold`](super::TryStreamExt::try_fold) method. + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct TryFold<St, Fut, T, F> { + #[pin] + stream: St, + f: F, + accum: Option<T>, + #[pin] + future: Option<Fut>, + } +} + +impl<St, Fut, T, F> fmt::Debug for TryFold<St, Fut, T, F> +where + St: fmt::Debug, + Fut: fmt::Debug, + T: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TryFold") + .field("stream", &self.stream) + .field("accum", &self.accum) + .field("future", &self.future) + .finish() + } +} + +impl<St, Fut, T, F> TryFold<St, Fut, T, F> +where + St: TryStream, + F: FnMut(T, St::Ok) -> Fut, + Fut: TryFuture<Ok = T, Error = St::Error>, +{ + pub(super) fn new(stream: St, f: F, t: T) -> Self { + Self { stream, f, accum: Some(t), future: None } + } +} + +impl<St, Fut, T, F> FusedFuture for TryFold<St, Fut, T, F> +where + St: TryStream, + F: FnMut(T, St::Ok) -> Fut, + Fut: TryFuture<Ok = T, Error = St::Error>, +{ + fn is_terminated(&self) -> bool { + self.accum.is_none() && self.future.is_none() + } +} + +impl<St, Fut, T, F> Future for TryFold<St, Fut, T, F> +where + St: TryStream, + F: FnMut(T, St::Ok) -> Fut, + Fut: TryFuture<Ok = T, Error = St::Error>, +{ + type Output = Result<T, St::Error>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut this = self.project(); + + Poll::Ready(loop { + if let Some(fut) = this.future.as_mut().as_pin_mut() { + // we're currently processing a future to produce a new accum value + let res = ready!(fut.try_poll(cx)); + this.future.set(None); + match res { + Ok(a) => *this.accum = Some(a), + Err(e) => break Err(e), + } + } else if this.accum.is_some() { + // we're waiting on a new item from the stream + let res = ready!(this.stream.as_mut().try_poll_next(cx)); + let a = this.accum.take().unwrap(); + match res { + Some(Ok(item)) => this.future.set(Some((this.f)(a, item))), + Some(Err(e)) => break Err(e), + None => break Ok(a), + } + } else { + panic!("Fold polled after completion") + } + }) + } +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_for_each.rs b/third_party/rust/futures-util/src/stream/try_stream/try_for_each.rs new file mode 100644 index 0000000000..6a081d84e7 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_for_each.rs @@ -0,0 +1,68 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{Future, TryFuture}; +use futures_core::ready; +use futures_core::stream::TryStream; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the [`try_for_each`](super::TryStreamExt::try_for_each) method. + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct TryForEach<St, Fut, F> { + #[pin] + stream: St, + f: F, + #[pin] + future: Option<Fut>, + } +} + +impl<St, Fut, F> fmt::Debug for TryForEach<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TryForEach") + .field("stream", &self.stream) + .field("future", &self.future) + .finish() + } +} + +impl<St, Fut, F> TryForEach<St, Fut, F> +where + St: TryStream, + F: FnMut(St::Ok) -> Fut, + Fut: TryFuture<Ok = (), Error = St::Error>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, future: None } + } +} + +impl<St, Fut, F> Future for TryForEach<St, Fut, F> +where + St: TryStream, + F: FnMut(St::Ok) -> Fut, + Fut: TryFuture<Ok = (), Error = St::Error>, +{ + type Output = Result<(), St::Error>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut this = self.project(); + loop { + if let Some(fut) = this.future.as_mut().as_pin_mut() { + ready!(fut.try_poll(cx))?; + this.future.set(None); + } else { + match ready!(this.stream.as_mut().try_poll_next(cx)?) { + Some(e) => this.future.set(Some((this.f)(e))), + None => break, + } + } + } + Poll::Ready(Ok(())) + } +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_for_each_concurrent.rs b/third_party/rust/futures-util/src/stream/try_stream/try_for_each_concurrent.rs new file mode 100644 index 0000000000..62734c746b --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_for_each_concurrent.rs @@ -0,0 +1,133 @@ +use crate::stream::{FuturesUnordered, StreamExt}; +use core::fmt; +use core::mem; +use core::num::NonZeroUsize; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::TryStream; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Future for the + /// [`try_for_each_concurrent`](super::TryStreamExt::try_for_each_concurrent) + /// method. + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct TryForEachConcurrent<St, Fut, F> { + #[pin] + stream: Option<St>, + f: F, + futures: FuturesUnordered<Fut>, + limit: Option<NonZeroUsize>, + } +} + +impl<St, Fut, F> fmt::Debug for TryForEachConcurrent<St, Fut, F> +where + St: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TryForEachConcurrent") + .field("stream", &self.stream) + .field("futures", &self.futures) + .field("limit", &self.limit) + .finish() + } +} + +impl<St, Fut, F> FusedFuture for TryForEachConcurrent<St, Fut, F> +where + St: TryStream, + F: FnMut(St::Ok) -> Fut, + Fut: Future<Output = Result<(), St::Error>>, +{ + fn is_terminated(&self) -> bool { + self.stream.is_none() && self.futures.is_empty() + } +} + +impl<St, Fut, F> TryForEachConcurrent<St, Fut, F> +where + St: TryStream, + F: FnMut(St::Ok) -> Fut, + Fut: Future<Output = Result<(), St::Error>>, +{ + pub(super) fn new(stream: St, limit: Option<usize>, f: F) -> Self { + Self { + stream: Some(stream), + // Note: `limit` = 0 gets ignored. + limit: limit.and_then(NonZeroUsize::new), + f, + futures: FuturesUnordered::new(), + } + } +} + +impl<St, Fut, F> Future for TryForEachConcurrent<St, Fut, F> +where + St: TryStream, + F: FnMut(St::Ok) -> Fut, + Fut: Future<Output = Result<(), St::Error>>, +{ + type Output = Result<(), St::Error>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut this = self.project(); + loop { + let mut made_progress_this_iter = false; + + // Check if we've already created a number of futures greater than `limit` + if this.limit.map(|limit| limit.get() > this.futures.len()).unwrap_or(true) { + let poll_res = match this.stream.as_mut().as_pin_mut() { + Some(stream) => stream.try_poll_next(cx), + None => Poll::Ready(None), + }; + + let elem = match poll_res { + Poll::Ready(Some(Ok(elem))) => { + made_progress_this_iter = true; + Some(elem) + } + Poll::Ready(None) => { + this.stream.set(None); + None + } + Poll::Pending => None, + Poll::Ready(Some(Err(e))) => { + // Empty the stream and futures so that we know + // the future has completed. + this.stream.set(None); + drop(mem::replace(this.futures, FuturesUnordered::new())); + return Poll::Ready(Err(e)); + } + }; + + if let Some(elem) = elem { + this.futures.push((this.f)(elem)); + } + } + + match this.futures.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(()))) => made_progress_this_iter = true, + Poll::Ready(None) => { + if this.stream.is_none() { + return Poll::Ready(Ok(())); + } + } + Poll::Pending => {} + Poll::Ready(Some(Err(e))) => { + // Empty the stream and futures so that we know + // the future has completed. + this.stream.set(None); + drop(mem::replace(this.futures, FuturesUnordered::new())); + return Poll::Ready(Err(e)); + } + } + + if !made_progress_this_iter { + return Poll::Pending; + } + } + } +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_next.rs b/third_party/rust/futures-util/src/stream/try_stream/try_next.rs new file mode 100644 index 0000000000..13fcf80cae --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_next.rs @@ -0,0 +1,34 @@ +use crate::stream::TryStreamExt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::{FusedStream, TryStream}; +use futures_core::task::{Context, Poll}; + +/// Future for the [`try_next`](super::TryStreamExt::try_next) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct TryNext<'a, St: ?Sized> { + stream: &'a mut St, +} + +impl<St: ?Sized + Unpin> Unpin for TryNext<'_, St> {} + +impl<'a, St: ?Sized + TryStream + Unpin> TryNext<'a, St> { + pub(super) fn new(stream: &'a mut St) -> Self { + Self { stream } + } +} + +impl<St: ?Sized + TryStream + Unpin + FusedStream> FusedFuture for TryNext<'_, St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St: ?Sized + TryStream + Unpin> Future for TryNext<'_, St> { + type Output = Result<Option<St::Ok>, St::Error>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + self.stream.try_poll_next_unpin(cx)?.map(Ok) + } +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_skip_while.rs b/third_party/rust/futures-util/src/stream/try_stream/try_skip_while.rs new file mode 100644 index 0000000000..52aa2d478b --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_skip_while.rs @@ -0,0 +1,120 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`try_skip_while`](super::TryStreamExt::try_skip_while) + /// method. + #[must_use = "streams do nothing unless polled"] + pub struct TrySkipWhile<St, Fut, F> where St: TryStream { + #[pin] + stream: St, + f: F, + #[pin] + pending_fut: Option<Fut>, + pending_item: Option<St::Ok>, + done_skipping: bool, + } +} + +impl<St, Fut, F> fmt::Debug for TrySkipWhile<St, Fut, F> +where + St: TryStream + fmt::Debug, + St::Ok: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TrySkipWhile") + .field("stream", &self.stream) + .field("pending_fut", &self.pending_fut) + .field("pending_item", &self.pending_item) + .field("done_skipping", &self.done_skipping) + .finish() + } +} + +impl<St, Fut, F> TrySkipWhile<St, Fut, F> +where + St: TryStream, + F: FnMut(&St::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = St::Error>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, pending_fut: None, pending_item: None, done_skipping: false } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F> Stream for TrySkipWhile<St, Fut, F> +where + St: TryStream, + F: FnMut(&St::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = St::Error>, +{ + type Item = Result<St::Ok, St::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + if *this.done_skipping { + return this.stream.try_poll_next(cx); + } + + Poll::Ready(loop { + if let Some(fut) = this.pending_fut.as_mut().as_pin_mut() { + let res = ready!(fut.try_poll(cx)); + this.pending_fut.set(None); + let skipped = res?; + let item = this.pending_item.take(); + if !skipped { + *this.done_skipping = true; + break item.map(Ok); + } + } else if let Some(item) = ready!(this.stream.as_mut().try_poll_next(cx)?) { + this.pending_fut.set(Some((this.f)(&item))); + *this.pending_item = Some(item); + } else { + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = usize::from(self.pending_item.is_some()); + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +impl<St, Fut, F> FusedStream for TrySkipWhile<St, Fut, F> +where + St: TryStream + FusedStream, + F: FnMut(&St::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = St::Error>, +{ + fn is_terminated(&self) -> bool { + self.pending_item.is_none() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item, E> Sink<Item> for TrySkipWhile<S, Fut, F> +where + S: TryStream + Sink<Item, Error = E>, +{ + type Error = E; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_take_while.rs b/third_party/rust/futures-util/src/stream/try_stream/try_take_while.rs new file mode 100644 index 0000000000..4b5ff1ad38 --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_take_while.rs @@ -0,0 +1,129 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`try_take_while`](super::TryStreamExt::try_take_while) + /// method. + #[must_use = "streams do nothing unless polled"] + pub struct TryTakeWhile<St, Fut, F> + where + St: TryStream, + { + #[pin] + stream: St, + f: F, + #[pin] + pending_fut: Option<Fut>, + pending_item: Option<St::Ok>, + done_taking: bool, + } +} + +impl<St, Fut, F> fmt::Debug for TryTakeWhile<St, Fut, F> +where + St: TryStream + fmt::Debug, + St::Ok: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TryTakeWhile") + .field("stream", &self.stream) + .field("pending_fut", &self.pending_fut) + .field("pending_item", &self.pending_item) + .field("done_taking", &self.done_taking) + .finish() + } +} + +impl<St, Fut, F> TryTakeWhile<St, Fut, F> +where + St: TryStream, + F: FnMut(&St::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = St::Error>, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f, pending_fut: None, pending_item: None, done_taking: false } + } + + delegate_access_inner!(stream, St, ()); +} + +impl<St, Fut, F> Stream for TryTakeWhile<St, Fut, F> +where + St: TryStream, + F: FnMut(&St::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = St::Error>, +{ + type Item = Result<St::Ok, St::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + if *this.done_taking { + return Poll::Ready(None); + } + + Poll::Ready(loop { + if let Some(fut) = this.pending_fut.as_mut().as_pin_mut() { + let res = ready!(fut.try_poll(cx)); + this.pending_fut.set(None); + let take = res?; + let item = this.pending_item.take(); + if take { + break item.map(Ok); + } else { + *this.done_taking = true; + break None; + } + } else if let Some(item) = ready!(this.stream.as_mut().try_poll_next(cx)?) { + this.pending_fut.set(Some((this.f)(&item))); + *this.pending_item = Some(item); + } else { + break None; + } + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.done_taking { + return (0, Some(0)); + } + + let pending_len = usize::from(self.pending_item.is_some()); + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +impl<St, Fut, F> FusedStream for TryTakeWhile<St, Fut, F> +where + St: TryStream + FusedStream, + F: FnMut(&St::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = St::Error>, +{ + fn is_terminated(&self) -> bool { + self.done_taking || self.pending_item.is_none() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item, E> Sink<Item> for TryTakeWhile<S, Fut, F> +where + S: TryStream + Sink<Item, Error = E>, +{ + type Error = E; + + delegate_sink!(stream, Item); +} diff --git a/third_party/rust/futures-util/src/stream/try_stream/try_unfold.rs b/third_party/rust/futures-util/src/stream/try_stream/try_unfold.rs new file mode 100644 index 0000000000..fd9cdf1d8c --- /dev/null +++ b/third_party/rust/futures-util/src/stream/try_stream/try_unfold.rs @@ -0,0 +1,122 @@ +use super::assert_stream; +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::ready; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +/// Creates a `TryStream` from a seed and a closure returning a `TryFuture`. +/// +/// This function is the dual for the `TryStream::try_fold()` adapter: while +/// `TryStream::try_fold()` reduces a `TryStream` to one single value, +/// `try_unfold()` creates a `TryStream` from a seed value. +/// +/// `try_unfold()` will call the provided closure with the provided seed, then +/// wait for the returned `TryFuture` to complete with `(a, b)`. It will then +/// yield the value `a`, and use `b` as the next internal state. +/// +/// If the closure returns `None` instead of `Some(TryFuture)`, then the +/// `try_unfold()` will stop producing items and return `Poll::Ready(None)` in +/// future calls to `poll()`. +/// +/// In case of error generated by the returned `TryFuture`, the error will be +/// returned by the `TryStream`. The `TryStream` will then yield +/// `Poll::Ready(None)` in future calls to `poll()`. +/// +/// This function can typically be used when wanting to go from the "world of +/// futures" to the "world of streams": the provided closure can build a +/// `TryFuture` using other library functions working on futures, and +/// `try_unfold()` will turn it into a `TryStream` by repeating the operation. +/// +/// # Example +/// +/// ``` +/// # #[derive(Debug, PartialEq)] +/// # struct SomeError; +/// # futures::executor::block_on(async { +/// use futures::stream::{self, TryStreamExt}; +/// +/// let stream = stream::try_unfold(0, |state| async move { +/// if state < 0 { +/// return Err(SomeError); +/// } +/// +/// if state <= 2 { +/// let next_state = state + 1; +/// let yielded = state * 2; +/// Ok(Some((yielded, next_state))) +/// } else { +/// Ok(None) +/// } +/// }); +/// +/// let result: Result<Vec<i32>, _> = stream.try_collect().await; +/// assert_eq!(result, Ok(vec![0, 2, 4])); +/// # }); +/// ``` +pub fn try_unfold<T, F, Fut, Item>(init: T, f: F) -> TryUnfold<T, F, Fut> +where + F: FnMut(T) -> Fut, + Fut: TryFuture<Ok = Option<(Item, T)>>, +{ + assert_stream::<Result<Item, Fut::Error>, _>(TryUnfold { f, state: Some(init), fut: None }) +} + +pin_project! { + /// Stream for the [`try_unfold`] function. + #[must_use = "streams do nothing unless polled"] + pub struct TryUnfold<T, F, Fut> { + f: F, + state: Option<T>, + #[pin] + fut: Option<Fut>, + } +} + +impl<T, F, Fut> fmt::Debug for TryUnfold<T, F, Fut> +where + T: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TryUnfold").field("state", &self.state).field("fut", &self.fut).finish() + } +} + +impl<T, F, Fut, Item> Stream for TryUnfold<T, F, Fut> +where + F: FnMut(T) -> Fut, + Fut: TryFuture<Ok = Option<(Item, T)>>, +{ + type Item = Result<Item, Fut::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + if let Some(state) = this.state.take() { + this.fut.set(Some((this.f)(state))); + } + + match this.fut.as_mut().as_pin_mut() { + None => { + // The future previously errored + Poll::Ready(None) + } + Some(future) => { + let step = ready!(future.try_poll(cx)); + this.fut.set(None); + + match step { + Ok(Some((item, next_state))) => { + *this.state = Some(next_state); + Poll::Ready(Some(Ok(item))) + } + Ok(None) => Poll::Ready(None), + Err(e) => Poll::Ready(Some(Err(e))), + } + } + } + } +} diff --git a/third_party/rust/futures-util/src/stream/unfold.rs b/third_party/rust/futures-util/src/stream/unfold.rs new file mode 100644 index 0000000000..7d8ef6babc --- /dev/null +++ b/third_party/rust/futures-util/src/stream/unfold.rs @@ -0,0 +1,119 @@ +use super::assert_stream; +use crate::unfold_state::UnfoldState; +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::ready; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +/// Creates a `Stream` from a seed and a closure returning a `Future`. +/// +/// This function is the dual for the `Stream::fold()` adapter: while +/// `Stream::fold()` reduces a `Stream` to one single value, `unfold()` creates a +/// `Stream` from a seed value. +/// +/// `unfold()` will call the provided closure with the provided seed, then wait +/// for the returned `Future` to complete with `(a, b)`. It will then yield the +/// value `a`, and use `b` as the next internal state. +/// +/// If the closure returns `None` instead of `Some(Future)`, then the `unfold()` +/// will stop producing items and return `Poll::Ready(None)` in future +/// calls to `poll()`. +/// +/// This function can typically be used when wanting to go from the "world of +/// futures" to the "world of streams": the provided closure can build a +/// `Future` using other library functions working on futures, and `unfold()` +/// will turn it into a `Stream` by repeating the operation. +/// +/// # Example +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::stream::{self, StreamExt}; +/// +/// let stream = stream::unfold(0, |state| async move { +/// if state <= 2 { +/// let next_state = state + 1; +/// let yielded = state * 2; +/// Some((yielded, next_state)) +/// } else { +/// None +/// } +/// }); +/// +/// let result = stream.collect::<Vec<i32>>().await; +/// assert_eq!(result, vec![0, 2, 4]); +/// # }); +/// ``` +pub fn unfold<T, F, Fut, Item>(init: T, f: F) -> Unfold<T, F, Fut> +where + F: FnMut(T) -> Fut, + Fut: Future<Output = Option<(Item, T)>>, +{ + assert_stream::<Item, _>(Unfold { f, state: UnfoldState::Value { value: init } }) +} + +pin_project! { + /// Stream for the [`unfold`] function. + #[must_use = "streams do nothing unless polled"] + pub struct Unfold<T, F, Fut> { + f: F, + #[pin] + state: UnfoldState<T, Fut>, + } +} + +impl<T, F, Fut> fmt::Debug for Unfold<T, F, Fut> +where + T: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Unfold").field("state", &self.state).finish() + } +} + +impl<T, F, Fut, Item> FusedStream for Unfold<T, F, Fut> +where + F: FnMut(T) -> Fut, + Fut: Future<Output = Option<(Item, T)>>, +{ + fn is_terminated(&self) -> bool { + if let UnfoldState::Empty = self.state { + true + } else { + false + } + } +} + +impl<T, F, Fut, Item> Stream for Unfold<T, F, Fut> +where + F: FnMut(T) -> Fut, + Fut: Future<Output = Option<(Item, T)>>, +{ + type Item = Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + + if let Some(state) = this.state.as_mut().take_value() { + this.state.set(UnfoldState::Future { future: (this.f)(state) }); + } + + let step = match this.state.as_mut().project_future() { + Some(fut) => ready!(fut.poll(cx)), + None => panic!("Unfold must not be polled after it returned `Poll::Ready(None)`"), + }; + + if let Some((item, next_state)) = step { + this.state.set(UnfoldState::Value { value: next_state }); + Poll::Ready(Some(item)) + } else { + this.state.set(UnfoldState::Empty); + Poll::Ready(None) + } + } +} |