diff --git a/Cargo.toml b/Cargo.toml index f215ef7..27f88da 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,13 +10,15 @@ exclude = [".gitignore"] [features] default = ["std", "alloc"] -std = ["futures/std", "futures-timeout"] +std = ["futures/std"] alloc = ["futures/alloc"] +timeout = ["std", "dep:futures-timeout"] [dependencies] futures = { version = "0.3.31", default-features = false, features = ["async-await"] } futures-timeout = { version = "0.1.3", optional = true } +pin-project = "1.1.10" [dev-dependencies] futures = { version = "0.3.31", default-features = false, features = ["std", "alloc", "async-await", "thread-pool", "executor"] } \ No newline at end of file diff --git a/src/common.rs b/src/common.rs index aae4dc9..e24e1e7 100644 --- a/src/common.rs +++ b/src/common.rs @@ -3,11 +3,13 @@ use core::pin::Pin; use core::task::{Context, Poll}; use futures::Stream; -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] use futures_timeout::{Timeout, TimeoutExt}; +#[pin_project::pin_project] pub struct InnerMap { key: K, + #[pin] inner: Option, wake_on_success: bool, } @@ -21,19 +23,30 @@ impl InnerMap { } } - pub fn set_wake_on_success(&mut self, wake_on_success: bool) -> bool { - let prev = self.wake_on_success; - self.wake_on_success = wake_on_success; - wake_on_success != prev + pub fn set_wake_on_success_pin(self: Pin<&mut Self>, wake_on_success: bool) -> bool { + let this = self.project(); + let previous = *this.wake_on_success; + + *this.wake_on_success = wake_on_success; + previous != wake_on_success } pub fn key(&self) -> &K { &self.key } - pub fn key_value(&self) -> Option<(&K, &S)> { - let Self { key, inner, .. } = self; - inner.as_ref().map(|st| (key, st)) + pub fn key_pin(self: Pin<&Self>) -> &K { + self.project_ref().key + } + + pub fn inner_pin_ref(self: Pin<&Self>) -> Option> { + self.project_ref().inner.as_pin_ref() + } + + pub fn key_value_pin_ref(self: Pin<&Self>) -> Option<(&K, Pin<&S>)> { + let this = self.project_ref(); + + this.inner.as_pin_ref().map(|inner| (this.key, inner)) } pub fn key_value_mut(&mut self) -> Option<(&K, &mut S)> { @@ -41,10 +54,6 @@ impl InnerMap { inner.as_mut().map(|s| (key, s)) } - pub fn inner(&self) -> Option<&S> { - self.inner.as_ref() - } - pub fn inner_mut(&mut self) -> Option<&mut S> { self.inner.as_mut() } @@ -53,49 +62,54 @@ impl InnerMap { self.inner.take() } - pub fn key_value_pin(&mut self) -> Option<(&K, Pin<&mut S>)> + pub fn take_inner_pin(self: Pin<&mut Self>) -> Option where S: Unpin, { - let Self { ref key, inner, .. } = self; - inner.as_mut().map(|s| (key, Pin::new(s))) + self.project().inner.get_mut().take() } - pub fn inner_pin(&mut self) -> Option> - where - S: Unpin, - { - self.inner_mut().map(Pin::new) + pub fn key_value_pin(self: Pin<&mut Self>) -> Option<(&K, Pin<&mut S>)> { + let this = self.project(); + this.inner.as_pin_mut().map(|inner| (&*this.key, inner)) + } + + pub fn inner_pin(self: Pin<&mut Self>) -> Option> { + self.project().inner.as_pin_mut() } } impl Future for InnerMap where - K: Clone + Unpin, - S: Future + Unpin, + K: Clone, + S: Future, { type Output = (K, Option); - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let Some(st) = self.inner_pin() else { - return Poll::Ready((self.key.clone(), None)); + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let mut this = self.project(); + + let Some(st) = this.inner.as_mut().as_pin_mut() else { + return Poll::Ready((this.key.clone(), None)); }; let output = futures::ready!(st.poll(cx)); - self.inner.take(); - Poll::Ready((self.key.clone(), Some(output))) + this.inner.set(None); + Poll::Ready((this.key.clone(), Some(output))) } } impl Stream for InnerMap where - K: Clone + Unpin, - S: Stream + Unpin, + K: Clone, + S: Stream, { type Item = (K, Option); - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let Some(st) = self.inner_pin() else { + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let mut this = self.project(); + + let Some(st) = this.inner.as_mut().as_pin_mut() else { // Note: While we could panic for any attempts to poll the stream that doesnt exist or have been terminated, // we opt to just return `Poll::Ready(None)` and letting upstream define how to handle a terminated stream // although in the future this could change as we should not be polling any terminated streams or futures. @@ -104,38 +118,39 @@ where match st.poll_next(cx) { Poll::Ready(Some(value)) => { - if self.wake_on_success { + if *this.wake_on_success { // Since we made progress, we should attempt to proceed further by waking up the task // TODO: Find a better way to wake task up without needing to call the waker on every successful result // from stream cx.waker().wake_by_ref(); } - Poll::Ready(Some((self.key.clone(), Some(value)))) + Poll::Ready(Some((this.key.clone(), Some(value)))) } Poll::Ready(None) => { // Note: Although some streams can return a `Poll::Ready(None)`, we will have to assume that the stream is completely finished // and terminated at this point and that we should not attempt to poll again. // In the future, we could probably provide a flag that would allow us to take the inner stream or keep it and attempt on polling it again // without actually terminating it. - self.inner.take(); - Poll::Ready(Some((self.key.clone(), None))) + this.inner.set(None); + Poll::Ready(Some((this.key.clone(), None))) } Poll::Pending => Poll::Pending, } } } -#[cfg(feature = "std")] -pub struct Timed(Timeout); +#[cfg(all(feature = "std", feature = "timeout"))] +#[pin_project::pin_project] +pub struct Timed(#[pin] Timeout); -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] impl Timed { pub(crate) fn into_inner(self) -> F { self.0.into_inner() } } -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] impl core::ops::Deref for Timed { type Target = F; fn deref(&self) -> &Self::Target { @@ -143,38 +158,40 @@ impl core::ops::Deref for Timed { } } -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] impl core::ops::DerefMut for Timed { fn deref_mut(&mut self) -> &mut Self::Target { self.0.deref_mut() } } -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] impl Timed { pub(crate) fn new(item: F, timeout: core::time::Duration) -> Self { Self(item.timeout(timeout)) } } -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] impl Future for Timed where - F: Future + Unpin, + F: Future, { type Output = std::io::Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - Pin::new(&mut self.0).poll(cx) + let this = self.as_mut().project(); + this.0.poll(cx) } } -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] impl Stream for Timed where - F: Stream + Unpin, + F: Stream, { type Item = std::io::Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.0).poll_next(cx) + let this = self.as_mut().project(); + this.0.poll_next(cx) } } diff --git a/src/futures.rs b/src/futures.rs index 7075abb..ec8b005 100644 --- a/src/futures.rs +++ b/src/futures.rs @@ -1,9 +1,9 @@ pub mod optional; pub mod ordered; pub mod set; -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] pub mod timeout_map; -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] pub mod timeout_set; use crate::common::InnerMap; @@ -16,6 +16,7 @@ use futures::{Stream, StreamExt}; pub struct FutureMap { list: FuturesUnordered>, empty: bool, + terminate_on_empty: bool, waker: Option, } @@ -31,15 +32,21 @@ impl FutureMap { Self { list: FuturesUnordered::new(), empty: true, + terminate_on_empty: false, waker: None, } } + + /// Set flag to terminate stream after all futures are completed + pub fn set_terminate_on_empty(&mut self, terminate: bool) { + self.terminate_on_empty = terminate; + } } impl FutureMap where - K: Clone + PartialEq + Unpin, - T: Future + Unpin, + K: Clone + PartialEq, + T: Future, { /// Insert a future into the map with a unique key. /// The function will return true if the map does not have the key present, @@ -64,45 +71,48 @@ where /// Will return false if future does not exist or if value is the same as /// previously set. pub fn set_wake_on_success(&mut self, key: &K, wake_on_success: bool) -> bool { - self.list - .iter_mut() - .find(|st| st.key().eq(key)) - .is_some_and(|st| st.set_wake_on_success(wake_on_success)) + Pin::new(&mut self.list) + .iter_pin_mut() + .find(|st| st.as_ref().key_pin().eq(key)) + .is_some_and(|st| st.set_wake_on_success_pin(wake_on_success)) } /// An iterator visiting all key-value pairs in arbitrary order. pub fn iter(&self) -> impl Iterator { - self.list.iter().filter_map(|st| st.key_value()) - } - - /// An iterator visiting all key-value pairs mutably in arbitrary order. - pub fn iter_mut(&mut self) -> impl Iterator { - self.list.iter_mut().filter_map(|st| st.key_value_mut()) + Pin::new(&self.list) + .iter_pin_ref() + .filter_map(|st| st.key_value_pin_ref()) + .map(|(key, future)| (key, future.get_ref())) } /// An iterator visiting all key-value pairs with a pinned valued in arbitrary order pub fn iter_pin(&mut self) -> impl Iterator)> { - self.list.iter_mut().filter_map(|st| st.key_value_pin()) + Pin::new(&mut self.list) + .iter_pin_mut() + .filter_map(|st| st.key_value_pin()) } /// Returns an iterator visiting all keys in arbitrary order. pub fn keys(&self) -> impl Iterator { - self.list.iter().map(|st| st.key()) + Pin::new(&self.list) + .iter_pin_ref() + .filter_map(|st| st.key_value_pin_ref().map(|(key, _)| key)) } /// An iterator visiting all values in arbitrary order. pub fn values(&self) -> impl Iterator { - self.list.iter().filter_map(|st| st.inner()) - } - - /// An iterator visiting all values mutably in arbitrary order. - pub fn values_mut(&mut self) -> impl Iterator { - self.list.iter_mut().filter_map(|st| st.inner_mut()) + Pin::new(&self.list) + .iter_pin_ref() + .filter_map(|st| st.inner_pin_ref()) + .map(Pin::get_ref) } /// Returns `true` if the map contains a future for the specified key. pub fn contains_key(&self, key: &K) -> bool { - self.list.iter().any(|st| st.key().eq(key)) + Pin::new(&self.list) + .iter_pin_ref() + .filter(|st| st.as_ref().inner_pin_ref().is_some()) + .any(|st| st.key_pin().eq(key)) } /// Clears the map. @@ -112,13 +122,51 @@ where /// Returns a reference to the future corresponding to the key. pub fn get(&self, key: &K) -> Option<&T> { - self.list - .iter() - .find(|st| st.key().eq(key)) - .and_then(|st| st.inner()) + Pin::new(&self.list) + .iter_pin_ref() + .find(|st| st.as_ref().key_pin().eq(key)) + .and_then(|st| st.inner_pin_ref()) + .map(Pin::get_ref) + } + + /// Returns a pinned future corresponding to the key. + pub fn get_pinned(&mut self, key: &K) -> Option> { + Pin::new(&mut self.list) + .iter_pin_mut() + .find(|st| st.as_ref().key_pin().eq(key)) + .and_then(|st| st.inner_pin()) + } + + /// Returns the number of futures in the map. + pub fn len(&self) -> usize { + Pin::new(&self.list) + .iter_pin_ref() + .filter(|st| st.as_ref().inner_pin_ref().is_some()) + .count() + } + + /// Return `true` map contains no elements. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +impl FutureMap +where + K: Clone + PartialEq, + T: Future + Unpin, +{ + /// An iterator visiting all key-value pairs mutably in arbitrary order. + pub fn iter_mut(&mut self) -> impl Iterator { + self.list.iter_mut().filter_map(|st| st.key_value_mut()) + } + + /// An iterator visiting all values mutably in arbitrary order. + pub fn values_mut(&mut self) -> impl Iterator { + self.list.iter_mut().filter_map(|st| st.inner_mut()) } - /// Returns a mutable future to the value corresponding to the key. + /// Returns a mutable future corresponding to the key. pub fn get_mut(&mut self, key: &K) -> Option<&mut T> { self.list .iter_mut() @@ -126,7 +174,7 @@ where .and_then(|st| st.inner_mut()) } - /// Returns a muable future or default value if it does not exist. + /// Returns a mutable future or default value if it does not exist. pub fn get_mut_or_default(&mut self, key: &K) -> &mut T where T: Default, @@ -135,14 +183,6 @@ where self.get_mut(key).expect("valid entry") } - /// Returns a pinned future corresponding to the key. - pub fn get_pinned(&mut self, key: &K) -> Option> { - self.list - .iter_mut() - .find(|st| st.key().eq(key)) - .and_then(|st| st.inner_pin()) - } - /// Removes a key from the map, returning the future. pub fn remove(&mut self, key: &K) -> Option { self.list @@ -150,22 +190,12 @@ where .find(|st| st.key().eq(key)) .and_then(|st| st.take_inner()) } - - /// Returns the number of futures in the map. - pub fn len(&self) -> usize { - self.list.iter().filter(|st| st.inner().is_some()).count() - } - - /// Return `true` map contains no elements. - pub fn is_empty(&self) -> bool { - self.list.is_empty() || self.list.iter().all(|st| st.inner().is_none()) - } } impl FromIterator<(K, T)> for FutureMap where - K: Clone + PartialEq + Unpin, - T: Future + Unpin, + K: Clone + PartialEq, + T: Future, { fn from_iter>(iter: I) -> Self { let mut maps = Self::new(); @@ -178,8 +208,8 @@ where impl Stream for FutureMap where - K: Clone + PartialEq + Unpin, - T: Future + Unpin, + K: Clone, + T: Future, { type Item = (K, T::Output); @@ -188,9 +218,7 @@ where match self.list.poll_next_unpin(cx) { Poll::Ready(Some((key, Some(item)))) => return Poll::Ready(Some((key, item))), // We continue in case there is any progress on the set of streams - Poll::Ready(Some((key, None))) => { - self.remove(&key); - } + Poll::Ready(Some((_key, None))) => continue, Poll::Ready(None) => { // While we could allow the stream to continue to be pending, it would make more sense to notify that the stream // is empty without needing to explicitly check while polling the actual "map" itself @@ -199,6 +227,9 @@ where // We do this so that we are not returning `Poll::Ready(None)` each time the map is polled // as that may be seen as UB and may cause an increase in cpu usage if self.empty { + if self.terminate_on_empty { + return Poll::Ready(None); + } self.waker = Some(cx.waker().clone()); return Poll::Pending; } @@ -222,11 +253,11 @@ where impl FusedStream for FutureMap where - K: Clone + PartialEq + Unpin, - T: Future + Unpin, + K: Clone, + T: Future, { fn is_terminated(&self) -> bool { - self.list.is_terminated() + self.terminate_on_empty && self.list.is_terminated() } } @@ -244,6 +275,16 @@ mod test { assert!(!map.insert(1, pending::<()>())); } + #[test] + fn supports_unboxed_async_future() { + let mut map = FutureMap::new(); + assert!(map.insert(1, async { 42 })); + + futures::executor::block_on(async move { + assert_eq!(map.next().await, Some((1, 42))); + }); + } + #[test] fn poll_multiple_keyed_streams() { let mut map = FutureMap::new(); diff --git a/src/futures/ordered.rs b/src/futures/ordered.rs index 365b0d8..8544d8e 100644 --- a/src/futures/ordered.rs +++ b/src/futures/ordered.rs @@ -6,8 +6,10 @@ use futures::Stream; /// An unbounded queue of futures imposed a FIFO order while polling one future at a time /// and returning the output to stream before popping the next future in queue to be polled. +#[pin_project::pin_project] pub struct OrderedFutureSet { queue: VecDeque, + #[pin] current_future: Option, waker: Option, } @@ -36,6 +38,15 @@ impl OrderedFutureSet { } } + /// Push a future to the back of a pinned queue. + pub fn push_pinned(self: Pin<&mut Self>, fut: F) { + let this = self.project(); + this.queue.push_back(fut); + if let Some(waker) = this.waker.take() { + waker.wake(); + } + } + /// Remove a future from the front of the queue pub fn pop_front(&mut self) -> Option { let fut = self.queue.pop_front(); @@ -45,6 +56,16 @@ impl OrderedFutureSet { fut } + /// Remove a future from the front of a pinned queue. + pub fn pop_front_pinned(self: Pin<&mut Self>) -> Option { + let this = self.project(); + let fut = this.queue.pop_front(); + if let Some(waker) = this.waker.take() { + waker.wake(); + } + fut + } + /// Remove a future from the back of the queue pub fn pop_back(&mut self) -> Option { let fut = self.queue.pop_back(); @@ -53,12 +74,19 @@ impl OrderedFutureSet { } fut } + + /// Remove a future from the back of a pinned queue. + pub fn pop_back_pinned(self: Pin<&mut Self>) -> Option { + let this = self.project(); + let fut = this.queue.pop_back(); + if let Some(waker) = this.waker.take() { + waker.wake(); + } + fut + } } -impl FromIterator for OrderedFutureSet -where - F: Future + Unpin, -{ +impl FromIterator for OrderedFutureSet { fn from_iter>(iter: T) -> Self { let mut ordered = Self::new(); for fut in iter { @@ -70,39 +98,44 @@ where impl Stream for OrderedFutureSet where - F: Future + Unpin, + F: Future, { type Item = F::Output; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll> { - let this = &mut *self; - - loop { - if this.current_future.is_none() { - let Some(fut) = this.queue.pop_front() else { - break; - }; - this.current_future.replace(fut); - } + fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll> { + let mut this = self.project(); + + if this.current_future.as_ref().get_ref().is_none() { + let Some(fut) = this.queue.pop_front() else { + this.waker.replace(cx.waker().clone()); + return Poll::Pending; + }; + this.current_future.set(Some(fut)); + } - match this.current_future.as_mut() { - Some(fut) => { - let output = futures::ready!(Pin::new(fut).poll(cx)); - this.current_future.take(); - cx.waker().wake_by_ref(); - return Poll::Ready(Some(output)); - } - None => { - this.waker.replace(cx.waker().clone()); - } + let fut = this + .current_future + .as_mut() + .as_pin_mut() + .expect("current future was initialized"); + + match fut.poll(cx) { + Poll::Ready(output) => { + this.current_future.set(None); + cx.waker().wake_by_ref(); + Poll::Ready(Some(output)) + } + Poll::Pending => { + this.waker.replace(cx.waker().clone()); + Poll::Pending } } - - this.waker.replace(cx.waker().clone()); - Poll::Pending } fn size_hint(&self) -> (usize, Option) { - (self.queue.len(), None) + ( + self.queue.len() + usize::from(self.current_future.is_some()), + None, + ) } } @@ -167,4 +200,24 @@ mod tests { assert_eq!(items, vec![1, 2, 4]); }) } + + #[test] + fn supports_unboxed_async_futures() { + async fn value(value: u8) -> u8 { + value + } + + futures::executor::block_on(async move { + let mut fifo = OrderedFutureSet::new(); + fifo.push(value(1)); + fifo.push(value(2)); + futures::pin_mut!(fifo); + + assert_eq!(fifo.as_mut().next().await, Some(1)); + assert_eq!(fifo.as_mut().next().await, Some(2)); + + fifo.as_mut().push_pinned(value(3)); + assert_eq!(fifo.as_mut().next().await, Some(3)); + }); + } } diff --git a/src/futures/set.rs b/src/futures/set.rs index b8f95ae..e8f27ad 100644 --- a/src/futures/set.rs +++ b/src/futures/set.rs @@ -13,7 +13,7 @@ pub struct FutureSet { impl Default for FutureSet where - S: Future + Unpin, + S: Future, { fn default() -> Self { Self::new() @@ -22,7 +22,7 @@ where impl FutureSet where - S: Future + Unpin, + S: Future, { /// Creates an empty ['FutureSet`] pub fn new() -> Self { @@ -43,11 +43,6 @@ where self.map.iter().map(|(_, st)| st) } - /// An iterator visiting all futures mutably in arbitrary order. - pub fn iter_mut(&mut self) -> impl Iterator { - self.map.iter_mut().map(|(_, st)| st) - } - /// An iterator visiting all futures pinned valued in arbitrary order pub fn iter_pin(&mut self) -> impl Iterator> { self.map.iter_pin().map(|(_, st)| st) @@ -69,9 +64,19 @@ where } } -impl FromIterator for FutureSet +impl FutureSet where S: Future + Unpin, +{ + /// An iterator visiting all futures mutably in arbitrary order. + pub fn iter_mut(&mut self) -> impl Iterator { + self.map.iter_mut().map(|(_, st)| st) + } +} + +impl FromIterator for FutureSet +where + S: Future, { fn from_iter>(iter: I) -> Self { let mut maps = Self::new(); @@ -84,7 +89,7 @@ where impl Stream for FutureSet where - S: Future + Unpin, + S: Future, { type Item = S::Output; @@ -101,7 +106,7 @@ where impl FusedStream for FutureSet where - S: Future + Unpin, + S: Future, { fn is_terminated(&self) -> bool { self.map.is_terminated() @@ -126,4 +131,14 @@ mod test { assert_eq!(val, Some(1)); }); } + + #[test] + fn supports_unboxed_async_future() { + let mut set = FutureSet::new(); + assert!(set.insert(async { 42 })); + + futures::executor::block_on(async move { + assert_eq!(set.next().await, Some(42)); + }); + } } diff --git a/src/futures/timeout_map.rs b/src/futures/timeout_map.rs index 7fc3a42..0d2e57b 100644 --- a/src/futures/timeout_map.rs +++ b/src/futures/timeout_map.rs @@ -28,8 +28,8 @@ impl DerefMut for TimeoutFutureMap { impl TimeoutFutureMap where - K: Clone + PartialEq + Unpin, - F: Future + Unpin, + K: Clone + PartialEq, + F: Future, { /// Create an empty [`TimeoutFutureMap`] pub fn new(duration: Duration) -> Self { @@ -49,8 +49,8 @@ where impl Stream for TimeoutFutureMap where - K: Clone + PartialEq + Unpin, - F: Future + Unpin, + K: Clone + PartialEq, + F: Future, { type Item = (K, std::io::Result); fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { @@ -64,8 +64,8 @@ where impl FusedStream for TimeoutFutureMap where - K: Clone + PartialEq + Unpin, - F: Future + Unpin, + K: Clone + PartialEq, + F: Future, { fn is_terminated(&self) -> bool { self.map.is_terminated() @@ -107,4 +107,18 @@ mod test { assert_eq!(val, 0); }); } + + #[test] + fn supports_unboxed_async_future() { + let mut map = TimeoutFutureMap::new(Duration::from_secs(1)); + assert!(map.insert(1, async { 42 })); + + futures::executor::block_on(async move { + let Some((1, Ok(val))) = map.next().await else { + unreachable!("result is err"); + }; + + assert_eq!(val, 42); + }); + } } diff --git a/src/futures/timeout_set.rs b/src/futures/timeout_set.rs index fe50a27..9863b12 100644 --- a/src/futures/timeout_set.rs +++ b/src/futures/timeout_set.rs @@ -28,7 +28,7 @@ impl DerefMut for TimeoutFutureSet { impl TimeoutFutureSet where - F: Future + Unpin, + F: Future, { /// Create an empty [`TimeoutFutureSet`] pub fn new(duration: Duration) -> Self { @@ -46,7 +46,7 @@ where impl Stream for TimeoutFutureSet where - F: Future + Unpin, + F: Future, { type Item = std::io::Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { @@ -60,7 +60,7 @@ where impl FusedStream for TimeoutFutureSet where - F: Future + Unpin, + F: Future, { fn is_terminated(&self) -> bool { self.set.is_terminated() diff --git a/src/optional.rs b/src/optional.rs index 11d6ec5..9cdd1b3 100644 --- a/src/optional.rs +++ b/src/optional.rs @@ -1,4 +1,4 @@ -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] pub mod timeout; use core::future::Future; @@ -7,6 +7,7 @@ use core::task::{Context, Poll, Waker}; use futures::future::FusedFuture; use futures::stream::FusedStream; use futures::Stream; +use pin_project::pin_project; /// A reusable future or stream based on `Option`. /// @@ -15,13 +16,13 @@ use futures::Stream; /// is set via [`Optional::replace`], it would then be polled once [`Optional`] /// is polled. Once the future is polled to completion, the results will be returned, with /// [`Optional`] being empty. +#[pin_project] pub struct Optional { + #[pin] task: Option, waker: Option, } -impl Unpin for Optional {} - impl Default for Optional { fn default() -> Self { Self { @@ -112,6 +113,17 @@ impl Optional { fut } + /// Replaces the current future or stream in place without moving the previous value. + pub fn set(self: Pin<&mut Self>, task: T) { + let mut this = self.project(); + + this.task.set(Some(task)); + + if let Some(waker) = this.waker.take() { + waker.wake(); + } + } + /// Returns a constructed `Option>`. pub fn as_pin_mut(&mut self) -> Option> where @@ -120,6 +132,11 @@ impl Optional { self.task.as_mut().map(Pin::new) } + /// Return a constructed `Option>`. + pub fn pinned_as_mut(self: Pin<&mut Self>) -> Option> { + self.project().task.as_pin_mut() + } + /// Returns a constructed `Option>`. pub fn as_pin_ref(&self) -> Option> where @@ -127,27 +144,33 @@ impl Optional { { self.task.as_ref().map(Pin::new) } + + /// Return a constructed `Option>`. + pub fn pinned_as_ref(self: Pin<&Self>) -> Option> { + self.project_ref().task.as_pin_ref() + } } impl Future for Optional where - F: Future + Unpin, + F: Future, { type Output = F::Output; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll { - let Some(future) = self.as_pin_mut() else { - self.waker.replace(cx.waker().clone()); + fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll { + let mut this = self.project(); + let Some(future) = this.task.as_mut().as_pin_mut() else { + this.waker.replace(cx.waker().clone()); return Poll::Pending; }; match future.poll(cx) { Poll::Ready(output) => { - self.task.take(); + this.task.set(None); Poll::Ready(output) } Poll::Pending => { - self.waker.replace(cx.waker().clone()); + this.waker.replace(cx.waker().clone()); Poll::Pending } } @@ -156,7 +179,7 @@ where impl FusedFuture for Optional where - F: Future + Unpin, + F: Future, { fn is_terminated(&self) -> bool { self.task.is_none() @@ -165,24 +188,25 @@ where impl Stream for Optional where - S: Stream + Unpin, + S: Stream, { type Item = S::Item; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let Some(stream) = self.as_pin_mut() else { - self.waker.replace(cx.waker().clone()); + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let mut this = self.project(); + let Some(stream) = this.task.as_mut().as_pin_mut() else { + this.waker.replace(cx.waker().clone()); return Poll::Pending; }; match stream.poll_next(cx) { Poll::Ready(Some(output)) => Poll::Ready(Some(output)), Poll::Ready(None) => { - self.task.take(); + this.task.set(None); Poll::Ready(None) } Poll::Pending => { - self.waker.replace(cx.waker().clone()); + this.waker.replace(cx.waker().clone()); Poll::Pending } } @@ -198,7 +222,7 @@ where impl FusedStream for Optional where - S: Stream + Unpin, + S: Stream, { fn is_terminated(&self) -> bool { self.task.is_none() @@ -239,6 +263,29 @@ mod test { assert!(future.is_none()); } + #[test] + fn reusable_pinned_optional_future() { + async fn set_value(value: i32) -> i32 { + value + } + + let future = Optional::new(set_value(0)); + futures::pin_mut!(future); + assert!(future.is_some()); + let waker = futures::task::noop_waker_ref(); + + let value = future.as_mut().poll(&mut Context::from_waker(waker)); + assert_eq!(value, Poll::Ready(0)); + assert!(future.is_none()); + + future.as_mut().set(set_value(1)); + assert!(future.is_some()); + + let value = future.as_mut().poll(&mut Context::from_waker(waker)); + assert_eq!(value, Poll::Ready(1)); + assert!(future.is_none()); + } + #[test] fn convert_future_to_optional_future() { let fut = futures::future::ready(0); @@ -293,6 +340,37 @@ mod test { assert!(stream.is_none()); } + #[test] + fn reusable_pinned_optional_stream() { + async fn set_val(value: i32) -> i32 { + value + } + + let stream = Optional::new(futures::stream::once(set_val(0))); + futures::pin_mut!(stream); + assert!(stream.is_some()); + let waker = futures::task::noop_waker_ref(); + + let val = stream.as_mut().poll_next(&mut Context::from_waker(waker)); + assert_eq!(val, Poll::Ready(Some(0))); + assert!(stream.is_some()); + + let val = stream.as_mut().poll_next(&mut Context::from_waker(waker)); + assert_eq!(val, Poll::Ready(None)); + assert!(stream.is_none()); + + stream.as_mut().set(futures::stream::once(set_val(1))); + assert!(stream.is_some()); + + let val = stream.as_mut().poll_next(&mut Context::from_waker(waker)); + assert_eq!(val, Poll::Ready(Some(1))); + assert!(stream.is_some()); + + let val = stream.as_mut().poll_next(&mut Context::from_waker(waker)); + assert_eq!(val, Poll::Ready(None)); + assert!(stream.is_none()); + } + #[test] fn convert_stream_to_optional_stream() { let st = futures::stream::once(async { 0 }).boxed(); @@ -310,4 +388,13 @@ mod test { assert_eq!(val, Poll::Ready(None)); assert!(stream.is_none()); } + + #[test] + fn pinned_accessors_support_not_unpin() { + let optional = Optional::new(async { 42 }); + futures::pin_mut!(optional); + + assert!(optional.as_ref().pinned_as_ref().is_some()); + assert!(optional.as_mut().pinned_as_mut().is_some()); + } } diff --git a/src/optional/timeout.rs b/src/optional/timeout.rs index 3c8b56a..32b8e70 100644 --- a/src/optional/timeout.rs +++ b/src/optional/timeout.rs @@ -8,8 +8,10 @@ use core::time::Duration; use futures::Stream; /// A reusable future or stream based on `Option` that will time out after a specific duration as elapse. +#[pin_project::pin_project] pub struct TimeoutOptional { duration: Duration, + #[pin] task: Optional>, } @@ -46,7 +48,7 @@ impl TimeoutOptional { /// Construct a new [`TimeoutOptional`] with an existing [`Future`]. pub fn new_with_future(duration: Duration, task: T) -> Self where - T: Future + Unpin, + T: Future, { Self { duration, @@ -57,7 +59,7 @@ impl TimeoutOptional { /// Construct a new [`TimeoutOptional`] with an existing [`Stream`]. pub fn new_with_stream(duration: Duration, task: T) -> Self where - T: Stream + Unpin, + T: Stream, { Self { duration, @@ -70,28 +72,37 @@ impl TimeoutOptional { let prev = self.task.replace(Timed::new(task, self.duration)); prev.map(|item| item.into_inner()) } + + /// Replaces the current future or stream in place without moving the previous value. + pub fn set(self: Pin<&mut Self>, task: T) { + let this = self.project(); + this.task.set(Timed::new(task, *this.duration)); + } } -impl Future for TimeoutOptional { +impl Future for TimeoutOptional { type Output = std::io::Result; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - Pin::new(&mut self.task).poll(cx) + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let mut this = self.project(); + Pin::new(&mut this.task).poll(cx) } } -impl Stream for TimeoutOptional { +impl Stream for TimeoutOptional { type Item = std::io::Result; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.task).poll_next(cx) + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let mut this = self.project(); + Pin::new(&mut this.task).poll_next(cx) } } #[cfg(test)] mod test { - use super::*; + use crate::optional::timeout::TimeoutOptional; + use core::future::pending; + use core::pin::Pin; + use core::time::Duration; use futures::future::ready; - use std::future::pending; - use std::time::Duration; #[test] fn test_timeout_optional_ready() { @@ -120,4 +131,25 @@ mod test { } }) } + + #[test] + fn reusable_pinned_timeout_optional_future() { + async fn value(value: i32) -> i32 { + value + } + + let task = TimeoutOptional::new_with_future(Duration::from_secs(1), value(0)); + futures::pin_mut!(task); + + futures::executor::block_on(async { + assert_eq!(task.as_mut().await.expect("future should not time out"), 0); + assert!(task.is_none()); + + task.as_mut().set(value(1)); + assert!(task.is_some()); + + assert_eq!(task.as_mut().await.expect("future should not time out"), 1); + assert!(task.is_none()); + }); + } } diff --git a/src/stream.rs b/src/stream.rs index efbc8bf..b2fa755 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1,11 +1,12 @@ pub mod optional; pub mod set; -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] pub mod timeout_map; -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "timeout"))] pub mod timeout_set; use crate::common::InnerMap; +use alloc::boxed::Box; use core::pin::Pin; use core::task::{Context, Poll, Waker}; use futures::stream::{FusedStream, SelectAll}; @@ -13,15 +14,16 @@ use futures::{Stream, StreamExt}; /// Combining multiple streams into one, with each stream having a unique key. pub struct StreamMap { - list: SelectAll>, + list: SelectAll>>>, empty: bool, + terminate_on_empty: bool, waker: Option, } impl Default for StreamMap where - K: Clone + Unpin, - T: Stream + Unpin, + K: Clone, + T: Stream, { fn default() -> Self { Self::new() @@ -30,23 +32,29 @@ where impl StreamMap where - K: Clone + Unpin, - T: Stream + Unpin, + K: Clone, + T: Stream, { /// Creates an empty [`StreamMap`] pub fn new() -> Self { Self { list: SelectAll::new(), empty: true, + terminate_on_empty: false, waker: None, } } + + /// Set flag to terminate stream after all streams are completed + pub fn set_terminate_on_empty(&mut self, terminate: bool) { + self.terminate_on_empty = terminate; + } } impl StreamMap where - K: Clone + PartialEq + Unpin, - T: Stream + Unpin, + K: Clone + PartialEq, + T: Stream, { /// Insert a stream into the map with a unique key. /// The function will return true if the map does not have the key present, @@ -56,7 +64,7 @@ where return false; } - let st = InnerMap::new(key, stream); + let st = Box::pin(InnerMap::new(key, stream)); self.list.push(st); if let Some(waker) = self.waker.take() { @@ -73,43 +81,46 @@ where pub fn set_wake_on_success(&mut self, key: &K, wake_on_success: bool) -> bool { self.list .iter_mut() - .find(|st| st.key().eq(key)) - .is_some_and(|st| st.set_wake_on_success(wake_on_success)) + .find(|st| st.as_ref().key_pin().eq(key)) + .is_some_and(|st| st.as_mut().set_wake_on_success_pin(wake_on_success)) } /// An iterator visiting all key-value pairs in arbitrary order. pub fn iter(&self) -> impl Iterator { - self.list.iter().filter_map(|st| st.key_value()) - } - - /// An iterator visiting all key-value pairs mutably in arbitrary order. - pub fn iter_mut(&mut self) -> impl Iterator { - self.list.iter_mut().filter_map(|st| st.key_value_mut()) + self.list + .iter() + .filter_map(|st| st.as_ref().key_value_pin_ref()) + .map(|(key, stream)| (key, stream.get_ref())) } /// An iterator visiting all key-value pairs with a pinned valued in arbitrary order pub fn iter_pin(&mut self) -> impl Iterator)> { - self.list.iter_mut().filter_map(|st| st.key_value_pin()) + self.list + .iter_mut() + .filter_map(|st| st.as_mut().key_value_pin()) } /// Returns an iterator visiting all keys in arbitrary order. pub fn keys(&self) -> impl Iterator { - self.list.iter().map(|st| st.key()) + self.list + .iter() + .filter_map(|st| st.as_ref().key_value_pin_ref().map(|(key, _)| key)) } /// An iterator visiting all values in arbitrary order. pub fn values(&self) -> impl Iterator { - self.list.iter().filter_map(|st| st.inner()) - } - - /// An iterator visiting all values mutably in arbitrary order. - pub fn values_mut(&mut self) -> impl Iterator { - self.list.iter_mut().filter_map(|st| st.inner_mut()) + self.list + .iter() + .filter_map(|st| st.as_ref().inner_pin_ref()) + .map(Pin::get_ref) } /// Returns `true` if the map contains a stream for the specified key. pub fn contains_key(&self, key: &K) -> bool { - self.list.iter().any(|st| st.key().eq(key)) + self.list + .iter() + .filter(|st| st.as_ref().inner_pin_ref().is_some()) + .any(|st| st.as_ref().key_pin().eq(key)) } /// Clears the map. @@ -121,19 +132,55 @@ where pub fn get(&self, key: &K) -> Option<&T> { self.list .iter() - .find(|st| st.key().eq(key)) - .and_then(|st| st.inner()) + .find(|st| st.as_ref().key_pin().eq(key)) + .and_then(|st| st.as_ref().inner_pin_ref()) + .map(Pin::get_ref) } - /// Returns a mutable stream to the value corresponding to the key. - pub fn get_mut(&mut self, key: &K) -> Option<&mut T> { + /// Returns a pinned stream corresponding to the key. + pub fn get_pinned(&mut self, key: &K) -> Option> { self.list .iter_mut() - .find(|st| st.key().eq(key)) - .and_then(|st| st.inner_mut()) + .find(|st| st.as_ref().key_pin().eq(key)) + .and_then(|st| st.as_mut().inner_pin()) + } + + /// Returns the number of streams in the map. + pub fn len(&self) -> usize { + self.list + .iter() + .filter(|st| st.as_ref().inner_pin_ref().is_some()) + .count() + } + + /// Return `true` map contains no elements. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +impl StreamMap +where + K: Clone + PartialEq, + T: Stream + Unpin, +{ + /// An iterator visiting all key-value pairs mutably in arbitrary order. + pub fn iter_mut(&mut self) -> impl Iterator { + self.iter_pin() + .map(|(key, stream)| (key, Pin::get_mut(stream))) + } + + /// An iterator visiting all values mutably in arbitrary order. + pub fn values_mut(&mut self) -> impl Iterator { + self.iter_pin().map(|(_, stream)| Pin::get_mut(stream)) } - /// Returns a muable stream or default value if it does not exist. + /// Returns a mutable stream corresponding to the key. + pub fn get_mut(&mut self, key: &K) -> Option<&mut T> { + self.get_pinned(key).map(Pin::get_mut) + } + + /// Returns a mutable stream or default value if it does not exist. pub fn get_mut_or_default(&mut self, key: &K) -> &mut T where T: Default, @@ -142,37 +189,19 @@ where self.get_mut(key).expect("valid entry") } - /// Returns a pinned stream corresponding to the key. - pub fn get_pinned(&mut self, key: &K) -> Option> { - self.list - .iter_mut() - .find(|st| st.key().eq(key)) - .and_then(|st| st.inner_pin()) - } - /// Removes a key from the map, returning the stream. pub fn remove(&mut self, key: &K) -> Option { self.list .iter_mut() - .find(|st| st.key().eq(key)) - .and_then(|st| st.take_inner()) - } - - /// Returns the number of streams in the map. - pub fn len(&self) -> usize { - self.list.iter().filter(|st| st.inner().is_some()).count() - } - - /// Return `true` map contains no elements. - pub fn is_empty(&self) -> bool { - self.list.is_empty() || self.list.iter().all(|st| st.inner().is_none()) + .find(|st| st.as_ref().key_pin().eq(key)) + .and_then(|st| st.as_mut().take_inner_pin()) } } impl FromIterator<(K, T)> for StreamMap where - K: Clone + PartialEq + Unpin, - T: Stream + Unpin, + K: Clone + PartialEq, + T: Stream, { fn from_iter>(iter: I) -> Self { let mut maps = Self::new(); @@ -185,13 +214,16 @@ where impl Stream for StreamMap where - K: Clone + PartialEq + Unpin, - T: Stream + Unpin, + K: Clone, + T: Stream, { type Item = (K, T::Item); fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { if self.list.is_empty() { + if self.terminate_on_empty { + return Poll::Ready(None); + } self.waker = Some(cx.waker().clone()); return Poll::Pending; } @@ -200,9 +232,7 @@ where match self.list.poll_next_unpin(cx) { Poll::Ready(Some((key, Some(item)))) => return Poll::Ready(Some((key, item))), // We continue in case there is any progress on the set of streams - Poll::Ready(Some((key, None))) => { - self.remove(&key); - } + Poll::Ready(Some((_key, None))) => continue, Poll::Ready(None) => { // While we could allow the stream to continue to be pending, it would make more sense to notify that the stream // is empty without needing to explicitly check while polling the actual "map" itself @@ -234,11 +264,11 @@ where impl FusedStream for StreamMap where - K: Clone + PartialEq + Unpin, - T: Stream + Unpin, + K: Clone, + T: Stream, { fn is_terminated(&self) -> bool { - self.list.is_terminated() + self.terminate_on_empty && self.list.is_terminated() } } @@ -285,6 +315,22 @@ mod test { assert!(!map.insert(1, empty::<()>())); } + #[test] + fn supports_unboxed_async_stream() { + let mut map = StreamMap::new(); + let stream = futures::stream::unfold(0, |value| async move { + (value < 3).then_some((value, value + 1)) + }); + assert!(map.insert("numbers", stream)); + + futures::executor::block_on(async move { + assert_eq!(map.next().await, Some(("numbers", 0))); + assert_eq!(map.next().await, Some(("numbers", 1))); + assert_eq!(map.next().await, Some(("numbers", 2))); + assert_eq!(map.next().await, None); + }); + } + #[test] fn poll_multiple_keyed_streams() { let mut map = StreamMap::new(); diff --git a/src/stream/set.rs b/src/stream/set.rs index ba3e1f2..bdc798a 100644 --- a/src/stream/set.rs +++ b/src/stream/set.rs @@ -12,7 +12,7 @@ pub struct StreamSet { impl Default for StreamSet where - S: Stream + Unpin, + S: Stream, { fn default() -> Self { Self::new() @@ -21,7 +21,7 @@ where impl StreamSet where - S: Stream + Unpin, + S: Stream, { /// Creates an empty ['StreamSet`] pub fn new() -> Self { @@ -42,11 +42,6 @@ where self.map.iter().map(|(_, st)| st) } - /// An iterator visiting all streams mutably in arbitrary order. - pub fn iter_mut(&mut self) -> impl Iterator { - self.map.iter_mut().map(|(_, st)| st) - } - /// An iterator visiting all streams pinned valued in arbitrary order pub fn iter_pin(&mut self) -> impl Iterator> { self.map.iter_pin().map(|(_, st)| st) @@ -68,9 +63,19 @@ where } } -impl FromIterator for StreamSet +impl StreamSet where S: Stream + Unpin, +{ + /// An iterator visiting all streams mutably in arbitrary order. + pub fn iter_mut(&mut self) -> impl Iterator { + self.map.iter_mut().map(|(_, st)| st) + } +} + +impl FromIterator for StreamSet +where + S: Stream, { fn from_iter>(iter: I) -> Self { let mut maps = Self::new(); @@ -83,7 +88,7 @@ where impl Stream for StreamSet where - S: Stream + Unpin, + S: Stream, { type Item = S::Item; @@ -100,7 +105,7 @@ where impl FusedStream for StreamSet where - S: Stream + Unpin, + S: Stream, { fn is_terminated(&self) -> bool { self.map.is_terminated() @@ -125,4 +130,20 @@ mod test { assert_eq!(val, Some(1)); }); } + + #[test] + fn supports_unboxed_async_stream() { + let mut set = StreamSet::new(); + let stream = futures::stream::unfold(0, |value| async move { + (value < 3).then_some((value, value + 1)) + }); + assert!(set.insert(stream)); + + futures::executor::block_on(async move { + assert_eq!(set.next().await, Some(0)); + assert_eq!(set.next().await, Some(1)); + assert_eq!(set.next().await, Some(2)); + assert_eq!(set.next().await, None); + }); + } } diff --git a/src/stream/timeout_map.rs b/src/stream/timeout_map.rs index 272d7cc..dc2fd45 100644 --- a/src/stream/timeout_map.rs +++ b/src/stream/timeout_map.rs @@ -27,8 +27,8 @@ impl DerefMut for TimeoutStreamMap { impl TimeoutStreamMap where - K: Clone + PartialEq + Unpin, - S: Stream + Unpin, + K: Clone + PartialEq, + S: Stream, { /// Create an empty [`TimeoutStreamMap`] pub fn new(duration: Duration) -> Self { @@ -48,8 +48,8 @@ where impl Stream for TimeoutStreamMap where - K: Clone + PartialEq + Unpin, - S: Stream + Unpin, + K: Clone, + S: Stream, { type Item = (K, std::io::Result); fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { @@ -63,8 +63,8 @@ where impl FusedStream for TimeoutStreamMap where - K: Clone + PartialEq + Unpin, - S: Stream + Unpin, + K: Clone, + S: Stream, { fn is_terminated(&self) -> bool { self.map.is_terminated() diff --git a/src/stream/timeout_set.rs b/src/stream/timeout_set.rs index 4190922..4ebf8c7 100644 --- a/src/stream/timeout_set.rs +++ b/src/stream/timeout_set.rs @@ -27,7 +27,7 @@ impl DerefMut for TimeoutStreamSet { impl TimeoutStreamSet where - S: Stream + Unpin, + S: Stream, { /// Create an empty ['TimeoutStreamSet'] pub fn new(duration: Duration) -> Self { @@ -45,7 +45,7 @@ where impl Stream for TimeoutStreamSet where - S: Stream + Unpin, + S: Stream, { type Item = std::io::Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { @@ -59,7 +59,7 @@ where impl FusedStream for TimeoutStreamSet where - S: Stream + Unpin, + S: Stream, { fn is_terminated(&self) -> bool { self.set.is_terminated()