feat: add display UI with splash screen, dashboard, and RSSI sparkline
Implement full SSD1306 OLED display support: - Splash screen on boot (1.5s) with project name, version, board - Dashboard showing radio state, config, counters, RSSI/SNR - RSSI history sparkline (64-sample ring buffer, 2px-wide bars) - DisplayOn/DisplayOff host commands routed via dedicated channel - Pure render functions against generic DrawTarget (no hardware coupling) - Display on/off control bypasses radio task entirely Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
17624 files changed
+4564928
-38
No files matched your search
+19
@@ -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)
|
||||
}
|
||||
+45
@@ -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()
|
||||
}
|
||||
}
|
||||
+260
@@ -0,0 +1,260 @@
|
||||
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
|
||||
// Use i64 for index since isize may overflow in 32-bit targets.
|
||||
index: i64,
|
||||
}
|
||||
}
|
||||
|
||||
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_next`](FuturesOrdered::poll_next) in this state will result
|
||||
/// in [`Poll::Ready(None)`](Poll::Ready) to be returned. Futures are submitted
|
||||
/// to the queue using [`push_back`](FuturesOrdered::push_back) (or
|
||||
/// [`push_front`](FuturesOrdered::push_front)); however, the future will
|
||||
/// **not** be polled at this point. [`FuturesOrdered`] will only poll managed
|
||||
/// futures when [`FuturesOrdered::poll_next`] is called. As such, it
|
||||
/// is important to call [`poll_next`](FuturesOrdered::poll_next) after pushing
|
||||
/// new futures.
|
||||
///
|
||||
/// If [`FuturesOrdered::poll_next`] returns [`Poll::Ready(None)`](Poll::Ready)
|
||||
/// 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_next`] 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_next`]
|
||||
/// 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: i64,
|
||||
next_outgoing_index: i64,
|
||||
}
|
||||
|
||||
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)`](Poll::Ready).
|
||||
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`](Future::poll) on the submitted
|
||||
/// future. The caller must ensure that [`FuturesOrdered::poll_next`] 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`](Future::poll) on the submitted
|
||||
/// future. The caller must ensure that [`FuturesOrdered::poll_next`] 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`](Future::poll) on the submitted
|
||||
/// future. The caller must ensure that [`FuturesOrdered::poll_next`] 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);
|
||||
}
|
||||
|
||||
/// Clear the whole `FuturesOrdered` queue.
|
||||
///
|
||||
/// This function clears the pending futures and the queued outputs
|
||||
/// to make it fully empty.
|
||||
pub fn clear(&mut self) {
|
||||
self.in_progress_queue.clear();
|
||||
self.queued_outputs.clear();
|
||||
self.next_incoming_index = 0;
|
||||
self.next_outgoing_index = 0;
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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> {}
|
||||
@@ -0,0 +1,674 @@
|
||||
//! 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: Send + 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);
|
||||
|
||||
// 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.
|
||||
//
|
||||
// Use ManuallyDrop to transfer the reference count ownership before
|
||||
// dropping the future so unwinding won't release the reference count.
|
||||
let md_slot;
|
||||
let task = if prev {
|
||||
md_slot = mem::ManuallyDrop::new(task);
|
||||
&*md_slot
|
||||
} else {
|
||||
&task
|
||||
};
|
||||
|
||||
// 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;
|
||||
}
|
||||
}
|
||||
|
||||
/// 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>> {
|
||||
unsafe {
|
||||
// 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);
|
||||
// SAFETY: see the comments of Bomb and this block.
|
||||
let waker = unsafe { 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 = Self::new();
|
||||
}
|
||||
}
|
||||
|
||||
impl<Fut> Drop for FuturesUnordered<Fut> {
|
||||
fn drop(&mut self) {
|
||||
// Before the strong reference to the queue is dropped we need all
|
||||
// futures to be dropped. See note at the bottom of this method.
|
||||
//
|
||||
// If there is a panic before this completes, we leak the queue.
|
||||
struct LeakQueueOnDrop<'a, Fut>(&'a mut FuturesUnordered<Fut>);
|
||||
impl<Fut> Drop for LeakQueueOnDrop<'_, Fut> {
|
||||
fn drop(&mut self) {
|
||||
mem::forget(Arc::clone(&self.0.ready_to_run_queue));
|
||||
}
|
||||
}
|
||||
let guard = LeakQueueOnDrop(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.
|
||||
while !guard.0.head_all.get_mut().is_null() {
|
||||
let head = *guard.0.head_all.get_mut();
|
||||
let task = unsafe { guard.0.unlink(head) };
|
||||
guard.0.release_task(task);
|
||||
}
|
||||
mem::forget(guard); // safe to release strong reference to queue
|
||||
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
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> {
|
||||
// FIXME: this takes raw pointer without safety conditions.
|
||||
|
||||
/// 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> {
|
||||
unsafe {
|
||||
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 !core::ptr::eq(self.head.load(Acquire), 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)
|
||||
}
|
||||
}
|
||||
|
||||
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.
|
||||
//
|
||||
// Note that each task has a strong reference count associated with it
|
||||
// which is owned by the ready to run queue. All tasks should have had
|
||||
// their futures dropped already by the `FuturesUnordered` destructor
|
||||
// above, so we're just pulling out tasks and dropping their refcounts.
|
||||
unsafe {
|
||||
loop {
|
||||
match self.dequeue() {
|
||||
Dequeue::Empty => break,
|
||||
Dequeue::Inconsistent => abort("inconsistent in drop"),
|
||||
Dequeue::Data(ptr) => drop(Arc::from_raw(ptr)),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,214 @@
|
||||
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::ArcWake;
|
||||
|
||||
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) unsafe fn waker_ref(this: &Arc<Self>) -> waker_ref::WakerRef<'_> {
|
||||
unsafe { waker_ref::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");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mod waker_ref {
|
||||
use alloc::sync::Arc;
|
||||
use core::marker::PhantomData;
|
||||
use core::mem;
|
||||
use core::mem::ManuallyDrop;
|
||||
use core::ops::Deref;
|
||||
use core::task::{RawWaker, RawWakerVTable, Waker};
|
||||
use futures_task::ArcWake;
|
||||
|
||||
pub(crate) struct WakerRef<'a> {
|
||||
waker: ManuallyDrop<Waker>,
|
||||
_marker: PhantomData<&'a ()>,
|
||||
}
|
||||
|
||||
impl WakerRef<'_> {
|
||||
#[inline]
|
||||
fn new_unowned(waker: ManuallyDrop<Waker>) -> Self {
|
||||
Self { waker, _marker: PhantomData }
|
||||
}
|
||||
}
|
||||
|
||||
impl Deref for WakerRef<'_> {
|
||||
type Target = Waker;
|
||||
|
||||
#[inline]
|
||||
fn deref(&self) -> &Waker {
|
||||
&self.waker
|
||||
}
|
||||
}
|
||||
|
||||
/// Copy of `future_task::waker_ref` without `W: 'static` bound.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// The caller must guarantee that use-after-free will not occur.
|
||||
#[inline]
|
||||
pub(crate) unsafe fn waker_ref<W>(wake: &Arc<W>) -> WakerRef<'_>
|
||||
where
|
||||
W: ArcWake,
|
||||
{
|
||||
// simply copy the pointer instead of using Arc::into_raw,
|
||||
// as we don't actually keep a refcount by using ManuallyDrop.<
|
||||
let ptr = Arc::as_ptr(wake).cast::<()>();
|
||||
|
||||
let waker =
|
||||
ManuallyDrop::new(unsafe { Waker::from_raw(RawWaker::new(ptr, waker_vtable::<W>())) });
|
||||
WakerRef::new_unowned(waker)
|
||||
}
|
||||
|
||||
fn waker_vtable<W: ArcWake>() -> &'static RawWakerVTable {
|
||||
&RawWakerVTable::new(
|
||||
clone_arc_raw::<W>,
|
||||
wake_arc_raw::<W>,
|
||||
wake_by_ref_arc_raw::<W>,
|
||||
drop_arc_raw::<W>,
|
||||
)
|
||||
}
|
||||
|
||||
// FIXME: panics on Arc::clone / refcount changes could wreak havoc on the
|
||||
// code here. We should guard against this by aborting.
|
||||
|
||||
unsafe fn increase_refcount<T: ArcWake>(data: *const ()) {
|
||||
// Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
|
||||
let arc = mem::ManuallyDrop::new(unsafe { Arc::<T>::from_raw(data.cast::<T>()) });
|
||||
// Now increase refcount, but don't drop new refcount either
|
||||
let _arc_clone: mem::ManuallyDrop<_> = arc.clone();
|
||||
}
|
||||
|
||||
unsafe fn clone_arc_raw<T: ArcWake>(data: *const ()) -> RawWaker {
|
||||
unsafe { increase_refcount::<T>(data) }
|
||||
RawWaker::new(data, waker_vtable::<T>())
|
||||
}
|
||||
|
||||
unsafe fn wake_arc_raw<T: ArcWake>(data: *const ()) {
|
||||
let arc: Arc<T> = unsafe { Arc::from_raw(data.cast::<T>()) };
|
||||
ArcWake::wake(arc);
|
||||
}
|
||||
|
||||
unsafe fn wake_by_ref_arc_raw<T: ArcWake>(data: *const ()) {
|
||||
// Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
|
||||
let arc = mem::ManuallyDrop::new(unsafe { Arc::<T>::from_raw(data.cast::<T>()) });
|
||||
ArcWake::wake_by_ref(&arc);
|
||||
}
|
||||
|
||||
unsafe fn drop_arc_raw<T: ArcWake>(data: *const ()) {
|
||||
drop(unsafe { Arc::<T>::from_raw(data.cast::<T>()) })
|
||||
}
|
||||
}
|
||||
+66
@@ -0,0 +1,66 @@
|
||||
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> Iter<I> {
|
||||
/// Acquires a reference to the underlying iterator that this stream is pulling from.
|
||||
pub fn get_ref(&self) -> &I {
|
||||
&self.iter
|
||||
}
|
||||
|
||||
/// Acquires a mutable reference to the underlying iterator that this stream is pulling from.
|
||||
pub fn get_mut(&mut self) -> &mut I {
|
||||
&mut self.iter
|
||||
}
|
||||
|
||||
/// Consumes this stream, returning the underlying iterator.
|
||||
pub fn into_inner(self) -> I {
|
||||
self.iter
|
||||
}
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
}
|
||||
+148
@@ -0,0 +1,148 @@
|
||||
//! 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::{
|
||||
All, Any, Chain, Collect, Concat, Count, 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_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[cfg(feature = "alloc")]
|
||||
pub use self::stream::{
|
||||
BufferUnordered, Buffered, FlatMapUnordered, FlattenUnordered, ForEachConcurrent,
|
||||
};
|
||||
|
||||
#[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[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, TryAll,
|
||||
TryAny, 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_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[cfg(feature = "alloc")]
|
||||
pub use self::try_stream::{
|
||||
TryBufferUnordered, TryBuffered, TryFlattenUnordered, TryForEachConcurrent,
|
||||
};
|
||||
|
||||
#[cfg(feature = "alloc")]
|
||||
pub use self::try_stream::{TryChunks, TryChunksError, TryReadyChunks, TryReadyChunksError};
|
||||
|
||||
// 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_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[cfg(feature = "alloc")]
|
||||
mod futures_ordered;
|
||||
#[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[cfg(feature = "alloc")]
|
||||
pub use self::futures_ordered::FuturesOrdered;
|
||||
|
||||
#[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[cfg(feature = "alloc")]
|
||||
pub mod futures_unordered;
|
||||
#[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[cfg(feature = "alloc")]
|
||||
#[doc(inline)]
|
||||
pub use self::futures_unordered::FuturesUnordered;
|
||||
|
||||
#[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[cfg(feature = "alloc")]
|
||||
pub mod select_all;
|
||||
#[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[cfg(feature = "alloc")]
|
||||
#[doc(inline)]
|
||||
pub use self::select_all::{select_all, SelectAll};
|
||||
|
||||
#[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[cfg(feature = "alloc")]
|
||||
mod abortable;
|
||||
#[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[cfg(feature = "alloc")]
|
||||
pub use crate::abortable::{AbortHandle, AbortRegistration, Abortable, Aborted};
|
||||
#[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
|
||||
#[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
|
||||
}
|
||||
+67
@@ -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()
|
||||
}
|
||||
}
|
||||
+45
@@ -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()
|
||||
}
|
||||
}
|
||||
+57
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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) })
|
||||
}
|
||||
+58
@@ -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, None)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> FusedStream for Repeat<T>
|
||||
where
|
||||
T: Clone,
|
||||
{
|
||||
fn is_terminated(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
+93
@@ -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, 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()`](crate::stream::StreamExt::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()`](crate::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 })
|
||||
}
|
||||
+117
@@ -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)
|
||||
}
|
||||
}
|
||||
+249
@@ -0,0 +1,249 @@
|
||||
//! 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 super::assert_stream;
|
||||
use crate::stream::{futures_unordered, FuturesUnordered, StreamExt, StreamFuture};
|
||||
|
||||
/// 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> {
|
||||
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> {}
|
||||
@@ -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) -> Self {
|
||||
match self {
|
||||
Self::Left => Self::Right,
|
||||
Self::Right => Self::Left,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for PollNext {
|
||||
fn default() -> Self {
|
||||
Self::Left
|
||||
}
|
||||
}
|
||||
|
||||
enum InternalState {
|
||||
Start,
|
||||
LeftFinished,
|
||||
RightFinished,
|
||||
BothFinished,
|
||||
}
|
||||
|
||||
impl InternalState {
|
||||
fn finish(&mut self, ps: PollNext) {
|
||||
match (&self, ps) {
|
||||
(Self::Start, PollNext::Left) => {
|
||||
*self = Self::LeftFinished;
|
||||
}
|
||||
(Self::Start, PollNext::Right) => {
|
||||
*self = Self::RightFinished;
|
||||
}
|
||||
(Self::LeftFinished, PollNext::Right) | (Self::RightFinished, PollNext::Left) => {
|
||||
*self = Self::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,
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_long_first_doc_paragraph)]
|
||||
/// 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 [`stream::select`](crate::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()
|
||||
}
|
||||
}
|
||||
+93
@@ -0,0 +1,93 @@
|
||||
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,
|
||||
done: 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("done", &self.done)
|
||||
.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, done: false, 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.done && 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 value
|
||||
let res = ready!(fut.poll(cx));
|
||||
this.future.set(None);
|
||||
if !res {
|
||||
*this.done = true;
|
||||
break false;
|
||||
} // early exit
|
||||
} else if !*this.done {
|
||||
// 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 => {
|
||||
*this.done = true;
|
||||
break true;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
panic!("All polled after completion")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+93
@@ -0,0 +1,93 @@
|
||||
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,
|
||||
done: 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("done", &self.done)
|
||||
.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, done: 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.done && 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 value
|
||||
let res = ready!(fut.poll(cx));
|
||||
this.future.set(None);
|
||||
if res {
|
||||
*this.done = true;
|
||||
break true;
|
||||
} // early exit
|
||||
} else if !*this.done {
|
||||
// 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 => {
|
||||
*this.done = true;
|
||||
break false;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
panic!("Any polled after completion")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
+118
@@ -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);
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
use futures_core::stream::{FusedStream, Stream};
|
||||
use futures_core::task::{Context, Poll};
|
||||
use pin_project_lite::pin_project;
|
||||
use std::any::Any;
|
||||
use std::boxed::Box;
|
||||
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()
|
||||
}
|
||||
}
|
||||
+76
@@ -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()
|
||||
}
|
||||
}
|
||||
}
|
||||
+103
@@ -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);
|
||||
}
|
||||
@@ -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()),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+62
@@ -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()
|
||||
}
|
||||
}
|
||||
+53
@@ -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,
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+67
@@ -0,0 +1,67 @@
|
||||
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 [`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, 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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
+117
@@ -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);
|
||||
}
|
||||
+111
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -0,0 +1,531 @@
|
||||
use alloc::sync::Arc;
|
||||
use core::{
|
||||
cell::UnsafeCell,
|
||||
convert::identity,
|
||||
fmt,
|
||||
marker::PhantomData,
|
||||
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;
|
||||
|
||||
/// Stream for the [`flatten_unordered`](super::StreamExt::flatten_unordered)
|
||||
/// method.
|
||||
pub type FlattenUnordered<St> = FlattenUnorderedWithFlowController<St, ()>;
|
||||
|
||||
/// There is nothing to poll and stream isn't being polled/waking/woken 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;
|
||||
|
||||
/// Both base stream and inner streams need to be polled.
|
||||
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;
|
||||
|
||||
/// Stream is being woken at the moment.
|
||||
const WAKING: u8 = 0b1000;
|
||||
|
||||
/// The stream was waked and will be polled.
|
||||
const WOKEN: u8 = 0b10000;
|
||||
|
||||
/// Internal polling state of the stream.
|
||||
#[derive(Clone, Debug)]
|
||||
struct SharedPollState {
|
||||
state: Arc<AtomicU8>,
|
||||
}
|
||||
|
||||
impl SharedPollState {
|
||||
/// Constructs new `SharedPollState` with the given state.
|
||||
fn new(value: u8) -> Self {
|
||||
Self { state: Arc::new(AtomicU8::new(value)) }
|
||||
}
|
||||
|
||||
/// Attempts to start polling, returning stored state in case of success.
|
||||
/// Returns `None` if either waker is waking at the moment.
|
||||
fn start_polling(&self) -> Option<(u8, PollStateBomb<'_, impl FnOnce(&Self) -> u8>)> {
|
||||
let value = self
|
||||
.state
|
||||
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |value| {
|
||||
if value & WAKING == NONE {
|
||||
Some(POLLING)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.ok()?;
|
||||
let bomb = PollStateBomb::new(self, Self::reset);
|
||||
|
||||
Some((value, bomb))
|
||||
}
|
||||
|
||||
/// Attempts to start the waking process and performs bitwise or with the given value.
|
||||
///
|
||||
/// If some waker is already in progress or stream is already woken/being polled, waking process won't start, however
|
||||
/// state will be disjuncted with the given value.
|
||||
fn start_waking(
|
||||
&self,
|
||||
to_poll: u8,
|
||||
) -> Option<(u8, PollStateBomb<'_, impl FnOnce(&Self) -> u8>)> {
|
||||
let value = self
|
||||
.state
|
||||
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |value| {
|
||||
let mut next_value = value | to_poll;
|
||||
if value & (WOKEN | POLLING) == NONE {
|
||||
next_value |= WAKING;
|
||||
}
|
||||
|
||||
if next_value != value {
|
||||
Some(next_value)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.ok()?;
|
||||
|
||||
// Only start the waking process if we're not in the polling/waking phase and the stream isn't woken already
|
||||
if value & (WOKEN | POLLING | WAKING) == NONE {
|
||||
let bomb = PollStateBomb::new(self, Self::stop_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` 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)
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Toggles state to non-waking, allowing to start polling.
|
||||
fn stop_waking(&self) -> u8 {
|
||||
let value = self
|
||||
.state
|
||||
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |value| {
|
||||
let next_value = value & !WAKING | WOKEN;
|
||||
|
||||
if next_value != value {
|
||||
Some(next_value)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.unwrap_or_else(identity);
|
||||
|
||||
debug_assert!(value & (WOKEN | POLLING | WAKING) == WAKING);
|
||||
value
|
||||
}
|
||||
|
||||
/// Resets current state allowing to poll the stream and wake up wakers.
|
||||
fn reset(&self) -> u8 {
|
||||
self.state.swap(NEED_TO_POLL_ALL, Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
/// 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();
|
||||
}
|
||||
}
|
||||
|
||||
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 WrappedWaker {
|
||||
inner_waker: UnsafeCell<Option<Waker>>,
|
||||
poll_state: SharedPollState,
|
||||
need_to_poll: u8,
|
||||
}
|
||||
|
||||
unsafe impl Send for WrappedWaker {}
|
||||
unsafe impl Sync for WrappedWaker {}
|
||||
|
||||
impl WrappedWaker {
|
||||
/// 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 by one thread at the time.
|
||||
unsafe fn replace_waker(self_arc: &mut Arc<Self>, cx: &Context<'_>) {
|
||||
unsafe { *self_arc.inner_waker.get() = cx.waker().clone().into() }
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
}
|
||||
|
||||
impl ArcWake for WrappedWaker {
|
||||
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
|
||||
drop(state_bomb);
|
||||
|
||||
// Wake up inner waker
|
||||
inner_waker.wake();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pin_project! {
|
||||
/// Future which polls optional inner 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, Self)>;
|
||||
|
||||
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 = Self::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 with ability to specify flow controller.
|
||||
#[project = FlattenUnorderedWithFlowControllerProj]
|
||||
#[must_use = "streams do nothing unless polled"]
|
||||
pub struct FlattenUnorderedWithFlowController<St, Fc> 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<WrappedWaker>,
|
||||
stream_waker: Arc<WrappedWaker>,
|
||||
flow_controller: PhantomData<Fc>
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fc> fmt::Debug for FlattenUnorderedWithFlowController<St, Fc>
|
||||
where
|
||||
St: Stream + fmt::Debug,
|
||||
St::Item: Stream + fmt::Debug,
|
||||
{
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("FlattenUnorderedWithFlowController")
|
||||
.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)
|
||||
.field("flow_controller", &self.flow_controller)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fc> FlattenUnorderedWithFlowController<St, Fc>
|
||||
where
|
||||
St: Stream,
|
||||
Fc: FlowController<St::Item, <St::Item as Stream>::Item>,
|
||||
St::Item: Stream + Unpin,
|
||||
{
|
||||
pub(crate) fn new(stream: St, limit: Option<usize>) -> Self {
|
||||
let poll_state = SharedPollState::new(NEED_TO_POLL_STREAM);
|
||||
|
||||
Self {
|
||||
inner_streams: FuturesUnordered::new(),
|
||||
stream,
|
||||
is_stream_done: false,
|
||||
limit: limit.and_then(NonZeroUsize::new),
|
||||
inner_streams_waker: Arc::new(WrappedWaker {
|
||||
inner_waker: UnsafeCell::new(None),
|
||||
poll_state: poll_state.clone(),
|
||||
need_to_poll: NEED_TO_POLL_INNER_STREAMS,
|
||||
}),
|
||||
stream_waker: Arc::new(WrappedWaker {
|
||||
inner_waker: UnsafeCell::new(None),
|
||||
poll_state: poll_state.clone(),
|
||||
need_to_poll: NEED_TO_POLL_STREAM,
|
||||
}),
|
||||
poll_state,
|
||||
flow_controller: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
delegate_access_inner!(stream, St, ());
|
||||
}
|
||||
|
||||
/// Returns the next flow step based on the received item.
|
||||
pub trait FlowController<I, O> {
|
||||
/// Handles an item producing `FlowStep` describing the next flow step.
|
||||
fn next_step(item: I) -> FlowStep<I, O>;
|
||||
}
|
||||
|
||||
impl<I, O> FlowController<I, O> for () {
|
||||
fn next_step(item: I) -> FlowStep<I, O> {
|
||||
FlowStep::Continue(item)
|
||||
}
|
||||
}
|
||||
|
||||
/// Describes the next flow step.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum FlowStep<C, R> {
|
||||
/// Just yields an item and continues standard flow.
|
||||
Continue(C),
|
||||
/// Immediately returns an underlying item from the function.
|
||||
Return(R),
|
||||
}
|
||||
|
||||
impl<St, Fc> FlattenUnorderedWithFlowControllerProj<'_, St, Fc>
|
||||
where
|
||||
St: Stream,
|
||||
{
|
||||
/// Checks if current `inner_streams` bucket size is greater than optional limit.
|
||||
fn is_exceeded_limit(&self) -> bool {
|
||||
self.limit.map_or(false, |limit| self.inner_streams.len() >= limit.get())
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fc> FusedStream for FlattenUnorderedWithFlowController<St, Fc>
|
||||
where
|
||||
St: FusedStream,
|
||||
Fc: FlowController<St::Item, <St::Item as Stream>::Item>,
|
||||
St::Item: Stream + Unpin,
|
||||
{
|
||||
fn is_terminated(&self) -> bool {
|
||||
self.stream.is_terminated() && self.inner_streams.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fc> Stream for FlattenUnorderedWithFlowController<St, Fc>
|
||||
where
|
||||
St: Stream,
|
||||
Fc: FlowController<St::Item, <St::Item as Stream>::Item>,
|
||||
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();
|
||||
|
||||
// Attempt to start polling, in case some waker is holding the lock, wait in loop
|
||||
let (mut poll_state_value, state_bomb) = loop {
|
||||
if let Some(value) = this.poll_state.start_polling() {
|
||||
break value;
|
||||
}
|
||||
};
|
||||
|
||||
// Safety: now state is `POLLING`.
|
||||
unsafe {
|
||||
WrappedWaker::replace_waker(this.stream_waker, cx);
|
||||
WrappedWaker::replace_waker(this.inner_streams_waker, cx)
|
||||
};
|
||||
|
||||
if poll_state_value & NEED_TO_POLL_STREAM != NONE {
|
||||
let mut stream_waker = None;
|
||||
|
||||
// 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 {
|
||||
let mut cx = Context::from_waker(
|
||||
stream_waker.get_or_insert_with(|| waker(this.stream_waker.clone())),
|
||||
);
|
||||
|
||||
match this.stream.as_mut().poll_next(&mut cx) {
|
||||
Poll::Ready(Some(item)) => {
|
||||
let next_item_fut = match Fc::next_step(item) {
|
||||
// Propagates an item immediately (the main use-case is for errors)
|
||||
FlowStep::Return(item) => {
|
||||
need_to_poll_next |= NEED_TO_POLL_STREAM
|
||||
| (poll_state_value & NEED_TO_POLL_INNER_STREAMS);
|
||||
poll_state_value &= !NEED_TO_POLL_INNER_STREAMS;
|
||||
|
||||
next_item = Some(item);
|
||||
|
||||
break;
|
||||
}
|
||||
// Yields an item and continues processing (normal case)
|
||||
FlowStep::Continue(inner_stream) => {
|
||||
PollStreamFut::new(inner_stream)
|
||||
}
|
||||
};
|
||||
// Add new stream to the inner streams bucket
|
||||
this.inner_streams.as_mut().push(next_item_fut);
|
||||
// Inner streams must be polled afterward
|
||||
poll_state_value |= NEED_TO_POLL_INNER_STREAMS;
|
||||
}
|
||||
Poll::Ready(None) => {
|
||||
// Mark the base stream as done
|
||||
*this.is_stream_done = true;
|
||||
}
|
||||
Poll::Pending => {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if poll_state_value & NEED_TO_POLL_INNER_STREAMS != NONE {
|
||||
let inner_streams_waker = waker(this.inner_streams_waker.clone());
|
||||
let mut cx = Context::from_waker(&inner_streams_waker);
|
||||
|
||||
match this.inner_streams.as_mut().poll_next(&mut cx) {
|
||||
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();
|
||||
|
||||
// Call the waker at the end of polling if
|
||||
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 the 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, we also 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, Fc> Sink<Item> for FlattenUnorderedWithFlowController<St, Fc>
|
||||
where
|
||||
St: Stream + Sink<Item>,
|
||||
{
|
||||
type Error = St::Error;
|
||||
|
||||
delegate_sink!(stream, Item);
|
||||
}
|
||||
+88
@@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+75
@@ -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);
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
+77
@@ -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);
|
||||
}
|
||||
+1697
File diff suppressed because it is too large.
Load diff
+34
@@ -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)
|
||||
}
|
||||
}
|
||||
+441
@@ -0,0 +1,441 @@
|
||||
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 core::pin::pin;
|
||||
///
|
||||
/// use futures::stream;
|
||||
/// use futures::stream::StreamExt;
|
||||
///
|
||||
/// let stream = stream::iter(vec![1, 2, 3]).peekable();
|
||||
/// let mut stream = pin!(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 core::pin::pin;
|
||||
///
|
||||
/// use futures::stream;
|
||||
/// use futures::stream::StreamExt;
|
||||
///
|
||||
/// let stream = stream::iter(0..5).peekable();
|
||||
/// let mut stream = pin!(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 core::pin::pin;
|
||||
///
|
||||
/// use futures::stream;
|
||||
/// use futures::stream::StreamExt;
|
||||
///
|
||||
/// let stream = stream::iter(1..20).peekable();
|
||||
/// let mut stream = pin!(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 core::pin::pin;
|
||||
///
|
||||
/// use futures::stream;
|
||||
/// use futures::stream::StreamExt;
|
||||
///
|
||||
/// let stream = stream::iter(0..5).peekable();
|
||||
/// let mut stream = pin!(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
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
+128
@@ -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);
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
+70
@@ -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);
|
||||
}
|
||||
+124
@@ -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);
|
||||
}
|
||||
+224
@@ -0,0 +1,224 @@
|
||||
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> SplitStream<S> {
|
||||
/// Returns `true` if the `SplitStream<S>` and `SplitSink<S>` originate from the same call to `StreamExt::split`.
|
||||
pub fn is_pair_of<Item>(&self, other: &SplitSink<S, Item>) -> bool {
|
||||
other.is_pair_of(&self)
|
||||
}
|
||||
}
|
||||
|
||||
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, Item> SplitSink<S, Item> {
|
||||
/// Returns `true` if the `SplitStream<S>` and `SplitSink<S>` originate from the same call to `StreamExt::split`.
|
||||
pub fn is_pair_of(&self, other: &SplitStream<S>) -> bool {
|
||||
self.lock.is_pair_of(&other.0)
|
||||
}
|
||||
}
|
||||
|
||||
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> {}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::stream::StreamExt;
|
||||
use core::marker::PhantomData;
|
||||
|
||||
struct NopStream<Item> {
|
||||
phantom: PhantomData<Item>,
|
||||
}
|
||||
|
||||
impl<Item> Stream for NopStream<Item> {
|
||||
type Item = Item;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
|
||||
impl<Item> Sink<Item> for NopStream<Item> {
|
||||
type Error = ();
|
||||
|
||||
fn poll_ready(
|
||||
self: Pin<&mut Self>,
|
||||
_cx: &mut Context<'_>,
|
||||
) -> Poll<Result<(), Self::Error>> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
fn start_send(self: Pin<&mut Self>, _item: Item) -> Result<(), Self::Error> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
fn poll_flush(
|
||||
self: Pin<&mut Self>,
|
||||
_cx: &mut Context<'_>,
|
||||
) -> Poll<Result<(), Self::Error>> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
fn poll_close(
|
||||
self: Pin<&mut Self>,
|
||||
_cx: &mut Context<'_>,
|
||||
) -> Poll<Result<(), Self::Error>> {
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pairing() {
|
||||
let s1 = NopStream::<()> { phantom: PhantomData };
|
||||
let (sink1, stream1) = s1.split();
|
||||
assert!(sink1.is_pair_of(&stream1));
|
||||
assert!(stream1.is_pair_of(&sink1));
|
||||
|
||||
let s2 = NopStream::<()> { phantom: PhantomData };
|
||||
let (sink2, stream2) = s2.split();
|
||||
assert!(sink2.is_pair_of(&stream2));
|
||||
assert!(stream2.is_pair_of(&sink2));
|
||||
|
||||
assert!(!sink1.is_pair_of(&stream2));
|
||||
assert!(!stream1.is_pair_of(&sink2));
|
||||
assert!(!sink2.is_pair_of(&stream1));
|
||||
assert!(!stream2.is_pair_of(&sink1));
|
||||
}
|
||||
}
|
||||
+86
@@ -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);
|
||||
}
|
||||
+170
@@ -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);
|
||||
}
|
||||
+124
@@ -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);
|
||||
}
|
||||
+101
@@ -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);
|
||||
}
|
||||
+63
@@ -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()),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+128
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
+1254
File diff suppressed because it is too large.
Load diff
@@ -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);
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
use core::fmt;
|
||||
use core::pin::Pin;
|
||||
use futures_core::future::{FusedFuture, 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_all`](super::TryStreamExt::try_all) method.
|
||||
#[must_use = "futures do nothing unless you `.await` or poll them"]
|
||||
pub struct TryAll<St, Fut, F> {
|
||||
#[pin]
|
||||
stream: St,
|
||||
f: F,
|
||||
done: bool,
|
||||
#[pin]
|
||||
future: Option<Fut>,
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fut, F> fmt::Debug for TryAll<St, Fut, F>
|
||||
where
|
||||
St: fmt::Debug,
|
||||
Fut: fmt::Debug,
|
||||
{
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("TryAll")
|
||||
.field("stream", &self.stream)
|
||||
.field("done", &self.done)
|
||||
.field("future", &self.future)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fut, F> TryAll<St, Fut, F>
|
||||
where
|
||||
St: TryStream,
|
||||
F: FnMut(St::Ok) -> Fut,
|
||||
Fut: Future<Output = bool>,
|
||||
{
|
||||
pub(super) fn new(stream: St, f: F) -> Self {
|
||||
Self { stream, f, done: false, future: None }
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fut, F> FusedFuture for TryAll<St, Fut, F>
|
||||
where
|
||||
St: TryStream,
|
||||
F: FnMut(St::Ok) -> Fut,
|
||||
Fut: Future<Output = bool>,
|
||||
{
|
||||
fn is_terminated(&self) -> bool {
|
||||
self.done && self.future.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fut, F> Future for TryAll<St, Fut, F>
|
||||
where
|
||||
St: TryStream,
|
||||
F: FnMut(St::Ok) -> Fut,
|
||||
Fut: Future<Output = bool>,
|
||||
{
|
||||
type Output = Result<bool, St::Error>;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<bool, St::Error>> {
|
||||
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 value
|
||||
let acc = ready!(fut.poll(cx));
|
||||
this.future.set(None);
|
||||
if !acc {
|
||||
*this.done = true;
|
||||
break Ok(false);
|
||||
} // early exit
|
||||
} else if !*this.done {
|
||||
// we're waiting on a new item from the stream
|
||||
match ready!(this.stream.as_mut().try_poll_next(cx)) {
|
||||
Some(Ok(item)) => {
|
||||
this.future.set(Some((this.f)(item)));
|
||||
}
|
||||
Some(Err(err)) => {
|
||||
*this.done = true;
|
||||
break Err(err);
|
||||
}
|
||||
None => {
|
||||
*this.done = true;
|
||||
break Ok(true);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
panic!("TryAll polled after completion")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
use core::fmt;
|
||||
use core::pin::Pin;
|
||||
use futures_core::future::{FusedFuture, 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_any`](super::TryStreamExt::try_any) method.
|
||||
#[must_use = "futures do nothing unless you `.await` or poll them"]
|
||||
pub struct TryAny<St, Fut, F> {
|
||||
#[pin]
|
||||
stream: St,
|
||||
f: F,
|
||||
done: bool,
|
||||
#[pin]
|
||||
future: Option<Fut>,
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fut, F> fmt::Debug for TryAny<St, Fut, F>
|
||||
where
|
||||
St: fmt::Debug,
|
||||
Fut: fmt::Debug,
|
||||
{
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("TryAny")
|
||||
.field("stream", &self.stream)
|
||||
.field("done", &self.done)
|
||||
.field("future", &self.future)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fut, F> TryAny<St, Fut, F>
|
||||
where
|
||||
St: TryStream,
|
||||
F: FnMut(St::Ok) -> Fut,
|
||||
Fut: Future<Output = bool>,
|
||||
{
|
||||
pub(super) fn new(stream: St, f: F) -> Self {
|
||||
Self { stream, f, done: false, future: None }
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fut, F> FusedFuture for TryAny<St, Fut, F>
|
||||
where
|
||||
St: TryStream,
|
||||
F: FnMut(St::Ok) -> Fut,
|
||||
Fut: Future<Output = bool>,
|
||||
{
|
||||
fn is_terminated(&self) -> bool {
|
||||
self.done && self.future.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
impl<St, Fut, F> Future for TryAny<St, Fut, F>
|
||||
where
|
||||
St: TryStream,
|
||||
F: FnMut(St::Ok) -> Fut,
|
||||
Fut: Future<Output = bool>,
|
||||
{
|
||||
type Output = Result<bool, St::Error>;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<bool, St::Error>> {
|
||||
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 value
|
||||
let acc = ready!(fut.poll(cx));
|
||||
this.future.set(None);
|
||||
if acc {
|
||||
*this.done = true;
|
||||
break Ok(true);
|
||||
} // early exit
|
||||
} else if !*this.done {
|
||||
// we're waiting on a new item from the stream
|
||||
match ready!(this.stream.as_mut().try_poll_next(cx)) {
|
||||
Some(Ok(item)) => {
|
||||
this.future.set(Some((this.f)(item)));
|
||||
}
|
||||
Some(Err(err)) => {
|
||||
*this.done = true;
|
||||
break Err(err);
|
||||
}
|
||||
None => {
|
||||
*this.done = true;
|
||||
break Ok(false);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
panic!("TryAny polled after completion")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
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, (. .));
|
||||
}
|
||||
|
||||
type TryChunksStreamError<St> = TryChunksError<<St as TryStream>::Ok, <St as TryStream>::Error>;
|
||||
|
||||
impl<St: TryStream> Stream for TryChunks<St> {
|
||||
type Item = Result<Vec<St::Ok>, TryChunksStreamError<St>>;
|
||||
|
||||
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> {}
|
||||
@@ -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),
|
||||
}
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
use core::marker::PhantomData;
|
||||
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;
|
||||
|
||||
use crate::future::Either;
|
||||
use crate::stream::stream::flatten_unordered::{
|
||||
FlattenUnorderedWithFlowController, FlowController, FlowStep,
|
||||
};
|
||||
use crate::stream::IntoStream;
|
||||
use crate::TryStreamExt;
|
||||
|
||||
delegate_all!(
|
||||
/// Stream for the [`try_flatten_unordered`](super::TryStreamExt::try_flatten_unordered) method.
|
||||
TryFlattenUnordered<St>(
|
||||
FlattenUnorderedWithFlowController<NestedTryStreamIntoEitherTryStream<St>, PropagateBaseStreamError<St>>
|
||||
): Debug + Sink + Stream + FusedStream + AccessInner[St, (. .)]
|
||||
+ New[
|
||||
|stream: St, limit: impl Into<Option<usize>>|
|
||||
FlattenUnorderedWithFlowController::new(
|
||||
NestedTryStreamIntoEitherTryStream::new(stream),
|
||||
limit.into()
|
||||
)
|
||||
]
|
||||
where
|
||||
St: TryStream,
|
||||
St::Ok: TryStream,
|
||||
St::Ok: Unpin,
|
||||
<St::Ok as TryStream>::Error: From<St::Error>
|
||||
);
|
||||
|
||||
pin_project! {
|
||||
/// Emits either successful streams or single-item streams containing the underlying errors.
|
||||
/// This's a wrapper for `FlattenUnordered` to reuse its logic over `TryStream`.
|
||||
#[derive(Debug)]
|
||||
#[must_use = "streams do nothing unless polled"]
|
||||
pub struct NestedTryStreamIntoEitherTryStream<St>
|
||||
where
|
||||
St: TryStream,
|
||||
St::Ok: TryStream,
|
||||
St::Ok: Unpin,
|
||||
<St::Ok as TryStream>::Error: From<St::Error>
|
||||
{
|
||||
#[pin]
|
||||
stream: St
|
||||
}
|
||||
}
|
||||
|
||||
impl<St> NestedTryStreamIntoEitherTryStream<St>
|
||||
where
|
||||
St: TryStream,
|
||||
St::Ok: TryStream + Unpin,
|
||||
<St::Ok as TryStream>::Error: From<St::Error>,
|
||||
{
|
||||
fn new(stream: St) -> Self {
|
||||
Self { stream }
|
||||
}
|
||||
|
||||
delegate_access_inner!(stream, St, ());
|
||||
}
|
||||
|
||||
/// Emits a single item immediately, then stream will be terminated.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Single<T>(Option<T>);
|
||||
|
||||
impl<T> Single<T> {
|
||||
/// Constructs new `Single` with the given value.
|
||||
fn new(val: T) -> Self {
|
||||
Self(Some(val))
|
||||
}
|
||||
|
||||
/// Attempts to take inner item immediately. Will always succeed if the stream isn't terminated.
|
||||
fn next_immediate(&mut self) -> Option<T> {
|
||||
self.0.take()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Unpin for Single<T> {}
|
||||
|
||||
impl<T> Stream for Single<T> {
|
||||
type Item = T;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
Poll::Ready(self.0.take())
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> (usize, Option<usize>) {
|
||||
self.0.as_ref().map_or((0, Some(0)), |_| (1, Some(1)))
|
||||
}
|
||||
}
|
||||
|
||||
/// Immediately propagates errors occurred in the base stream.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct PropagateBaseStreamError<St>(PhantomData<St>);
|
||||
|
||||
type BaseStreamItem<St> = <NestedTryStreamIntoEitherTryStream<St> as Stream>::Item;
|
||||
type InnerStreamItem<St> = <BaseStreamItem<St> as Stream>::Item;
|
||||
|
||||
impl<St> FlowController<BaseStreamItem<St>, InnerStreamItem<St>> for PropagateBaseStreamError<St>
|
||||
where
|
||||
St: TryStream,
|
||||
St::Ok: TryStream + Unpin,
|
||||
<St::Ok as TryStream>::Error: From<St::Error>,
|
||||
{
|
||||
fn next_step(item: BaseStreamItem<St>) -> FlowStep<BaseStreamItem<St>, InnerStreamItem<St>> {
|
||||
match item {
|
||||
// A new successful inner stream received
|
||||
st @ Either::Left(_) => FlowStep::Continue(st),
|
||||
// An error encountered
|
||||
Either::Right(mut err) => FlowStep::Return(err.next_immediate().unwrap()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type SingleStreamResult<St> = Single<Result<<St as TryStream>::Ok, <St as TryStream>::Error>>;
|
||||
|
||||
impl<St> Stream for NestedTryStreamIntoEitherTryStream<St>
|
||||
where
|
||||
St: TryStream,
|
||||
St::Ok: TryStream + Unpin,
|
||||
<St::Ok as TryStream>::Error: From<St::Error>,
|
||||
{
|
||||
// Item is either an inner stream or a stream containing a single error.
|
||||
// This will allow using `Either`'s `Stream` implementation as both branches are actually streams of `Result`'s.
|
||||
type Item = Either<IntoStream<St::Ok>, SingleStreamResult<St::Ok>>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
let item = ready!(self.project().stream.try_poll_next(cx));
|
||||
|
||||
let out = match item {
|
||||
Some(res) => match res {
|
||||
// Emit successful inner stream as is
|
||||
Ok(stream) => Either::Left(stream.into_stream()),
|
||||
// Wrap an error into a stream containing a single item
|
||||
err @ Err(_) => {
|
||||
let res = err.map(|_: St::Ok| unreachable!()).map_err(Into::into);
|
||||
|
||||
Either::Right(Single::new(res))
|
||||
}
|
||||
},
|
||||
None => return Poll::Ready(None),
|
||||
};
|
||||
|
||||
Poll::Ready(Some(out))
|
||||
}
|
||||
}
|
||||
|
||||
impl<St> FusedStream for NestedTryStreamIntoEitherTryStream<St>
|
||||
where
|
||||
St: TryStream + FusedStream,
|
||||
St::Ok: TryStream + Unpin,
|
||||
<St::Ok as TryStream>::Error: From<St::Error>,
|
||||
{
|
||||
fn is_terminated(&self) -> bool {
|
||||
self.stream.is_terminated()
|
||||
}
|
||||
}
|
||||
|
||||
// Forwarding impl of Sink from the underlying stream
|
||||
#[cfg(feature = "sink")]
|
||||
impl<St, Item> Sink<Item> for NestedTryStreamIntoEitherTryStream<St>
|
||||
where
|
||||
St: TryStream + Sink<Item>,
|
||||
St::Ok: TryStream + Unpin,
|
||||
<St::Ok as TryStream>::Error: From<<St as TryStream>::Error>,
|
||||
{
|
||||
type Error = <St as Sink<Item>>::Error;
|
||||
|
||||
delegate_sink!(stream, Item);
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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(()))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
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::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(core::mem::take(this.futures));
|
||||
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(core::mem::take(this.futures));
|
||||
return Poll::Ready(Err(e));
|
||||
}
|
||||
}
|
||||
|
||||
if !made_progress_this_iter {
|
||||
return Poll::Pending;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
use crate::stream::{Fuse, IntoStream, StreamExt};
|
||||
|
||||
use alloc::vec::Vec;
|
||||
use core::fmt;
|
||||
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 [`try_ready_chunks`](super::TryStreamExt::try_ready_chunks) method.
|
||||
#[derive(Debug)]
|
||||
#[must_use = "streams do nothing unless polled"]
|
||||
pub struct TryReadyChunks<St: TryStream> {
|
||||
#[pin]
|
||||
stream: Fuse<IntoStream<St>>,
|
||||
cap: usize, // https://github.com/rust-lang/futures-rs/issues/1475
|
||||
}
|
||||
}
|
||||
|
||||
impl<St: TryStream> TryReadyChunks<St> {
|
||||
pub(super) fn new(stream: St, capacity: usize) -> Self {
|
||||
assert!(capacity > 0);
|
||||
|
||||
Self { stream: IntoStream::new(stream).fuse(), cap: capacity }
|
||||
}
|
||||
|
||||
delegate_access_inner!(stream, St, (. .));
|
||||
}
|
||||
|
||||
type TryReadyChunksStreamError<St> =
|
||||
TryReadyChunksError<<St as TryStream>::Ok, <St as TryStream>::Error>;
|
||||
|
||||
impl<St: TryStream> Stream for TryReadyChunks<St> {
|
||||
type Item = Result<Vec<St::Ok>, TryReadyChunksStreamError<St>>;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
let mut this = self.as_mut().project();
|
||||
|
||||
let mut items: Vec<St::Ok> = Vec::new();
|
||||
|
||||
loop {
|
||||
match this.stream.as_mut().poll_next(cx) {
|
||||
// Flush all the collected data if the underlying stream doesn't
|
||||
// contain more ready values
|
||||
Poll::Pending => {
|
||||
return if items.is_empty() {
|
||||
Poll::Pending
|
||||
} else {
|
||||
Poll::Ready(Some(Ok(items)))
|
||||
}
|
||||
}
|
||||
|
||||
// Push the ready item into the buffer and check whether it is full.
|
||||
// If so, return the buffer.
|
||||
Poll::Ready(Some(Ok(item))) => {
|
||||
if items.is_empty() {
|
||||
items.reserve_exact(*this.cap);
|
||||
}
|
||||
items.push(item);
|
||||
if items.len() >= *this.cap {
|
||||
return Poll::Ready(Some(Ok(items)));
|
||||
}
|
||||
}
|
||||
|
||||
// Return the already collected items and the error.
|
||||
Poll::Ready(Some(Err(e))) => {
|
||||
return Poll::Ready(Some(Err(TryReadyChunksError(items, e))));
|
||||
}
|
||||
|
||||
// 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(Ok(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: TryStream + FusedStream> FusedStream for TryReadyChunks<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 TryReadyChunks<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 TryReadyChunksError<T, E>(pub Vec<T>, pub E);
|
||||
|
||||
impl<T, E: fmt::Debug> fmt::Debug for TryReadyChunksError<T, E> {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
self.1.fmt(f)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, E: fmt::Display> fmt::Display for TryReadyChunksError<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 TryReadyChunksError<T, E> {}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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))),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+119
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user