additional futures extension utils

Signed-off-by: Jason Volk <jason@zemos.net>

src/core/utils/future/ext_ext.rs

@@ -0,0 +1,34 @@
+//! Extended external extensions to futures::FutureExt
+
+use std::marker::Unpin;
+
+use futures::{future, future::Select, Future};
+
+/// This interface is not necessarily complete; feel free to add as-needed.
+pub trait ExtExt<T>
+where
+	Self: Future<Output = T> + Send,
+{
+	fn until<A, B, F>(self, f: F) -> Select<A, B>
+	where
+		Self: Sized,
+		F: FnOnce() -> B,
+		A: Future<Output = T> + From<Self> + Send + Unpin,
+		B: Future<Output = ()> + Send + Unpin;
+}
+
+impl<T, Fut> ExtExt<T> for Fut
+where
+	Fut: Future<Output = T> + Send,
+{
+	#[inline]
+	fn until<A, B, F>(self, f: F) -> Select<A, B>
+	where
+		Self: Sized,
+		F: FnOnce() -> B,
+		A: Future<Output = T> + From<Self> + Send + Unpin,
+		B: Future<Output = ()> + Send + Unpin,
+	{
+		future::select(self.into(), f())
+	}
+}

src/core/utils/future/mod.rs

@@ -1,5 +1,7 @@
+mod ext_ext;
 mod option_ext;
 mod try_ext_ext;
 
+pub use ext_ext::ExtExt;
 pub use option_ext::OptionExt;
 pub use try_ext_ext::TryExtExt;

src/core/utils/future/try_ext_ext.rs

@@ -4,8 +4,11 @@
 // caller only ever caring about result status while discarding all contents.
 #![allow(clippy::wrong_self_convention)]
 
+use std::marker::Unpin;
+
 use futures::{
-	future::{MapOkOrElse, UnwrapOrElse},
+	future,
+	future::{MapOkOrElse, TrySelect, UnwrapOrElse},
 	TryFuture, TryFutureExt,
 };
 
@@ -46,6 +49,13 @@ where
 	where
 		Self: Sized;
 
+	fn try_until<A, B, F>(self, f: F) -> TrySelect<A, B>
+	where
+		Self: Sized,
+		F: FnOnce() -> B,
+		A: TryFuture<Ok = Self::Ok> + From<Self> + Send + Unpin,
+		B: TryFuture<Ok = (), Error = Self::Error> + Send + Unpin;
+
 	fn unwrap_or(
 		self,
 		default: Self::Ok,
@@ -110,6 +120,17 @@ where
 		self.map_ok_or(None, Some)
 	}
 
+	#[inline]
+	fn try_until<A, B, F>(self, f: F) -> TrySelect<A, B>
+	where
+		Self: Sized,
+		F: FnOnce() -> B,
+		A: TryFuture<Ok = Self::Ok> + From<Self> + Send + Unpin,
+		B: TryFuture<Ok = (), Error = Self::Error> + Send + Unpin,
+	{
+		future::try_select(self.into(), f())
+	}
+
 	#[inline]
 	fn unwrap_or(
 		self,

src/core/utils/stream/mod.rs

@@ -8,6 +8,7 @@ mod ready;
 mod tools;
 mod try_broadband;
 mod try_ready;
+mod try_tools;
 mod wideband;
 
 pub use band::{
@@ -23,4 +24,5 @@ pub use ready::ReadyExt;
 pub use tools::Tools;
 pub use try_broadband::TryBroadbandExt;
 pub use try_ready::TryReadyExt;
+pub use try_tools::TryTools;
 pub use wideband::WidebandExt;

src/core/utils/stream/try_ready.rs

@@ -3,7 +3,7 @@
 
 use futures::{
 	future::{ready, Ready},
-	stream::{AndThen, TryFilterMap, TryFold, TryForEach, TryStream, TryStreamExt},
+	stream::{AndThen, TryFilterMap, TryFold, TryForEach, TryStream, TryStreamExt, TryTakeWhile},
 };
 
 use crate::Result;
@@ -56,6 +56,13 @@ where
 	) -> TryForEach<Self, Ready<Result<(), E>>, impl FnMut(S::Ok) -> Ready<Result<(), E>>>
 	where
 		F: FnMut(S::Ok) -> Result<(), E>;
+
+	fn ready_try_take_while<F>(
+		self,
+		f: F,
+	) -> TryTakeWhile<Self, Ready<Result<bool, E>>, impl FnMut(&S::Ok) -> Ready<Result<bool, E>>>
+	where
+		F: Fn(&S::Ok) -> Result<bool, E>;
 }
 
 impl<T, E, S> TryReadyExt<T, E, S> for S
@@ -122,4 +129,15 @@ where
 	{
 		self.try_for_each(move |t| ready(f(t)))
 	}
+
+	#[inline]
+	fn ready_try_take_while<F>(
+		self,
+		f: F,
+	) -> TryTakeWhile<Self, Ready<Result<bool, E>>, impl FnMut(&S::Ok) -> Ready<Result<bool, E>>>
+	where
+		F: Fn(&S::Ok) -> Result<bool, E>,
+	{
+		self.try_take_while(move |t| ready(f(t)))
+	}
 }

src/core/utils/stream/try_tools.rs

@@ -0,0 +1,44 @@
+//! TryStreamTools for futures::TryStream
+#![allow(clippy::type_complexity)]
+
+use futures::{future, future::Ready, stream::TryTakeWhile, TryStream, TryStreamExt};
+
+use crate::Result;
+
+/// TryStreamTools
+pub trait TryTools<T, E, S>
+where
+	S: TryStream<Ok = T, Error = E, Item = Result<T, E>> + Send + ?Sized,
+	Self: TryStream + Send + Sized,
+{
+	fn try_take(
+		self,
+		n: usize,
+	) -> TryTakeWhile<
+		Self,
+		Ready<Result<bool, S::Error>>,
+		impl FnMut(&S::Ok) -> Ready<Result<bool, S::Error>>,
+	>;
+}
+
+impl<T, E, S> TryTools<T, E, S> for S
+where
+	S: TryStream<Ok = T, Error = E, Item = Result<T, E>> + Send + ?Sized,
+	Self: TryStream + Send + Sized,
+{
+	#[inline]
+	fn try_take(
+		self,
+		mut n: usize,
+	) -> TryTakeWhile<
+		Self,
+		Ready<Result<bool, S::Error>>,
+		impl FnMut(&S::Ok) -> Ready<Result<bool, S::Error>>,
+	> {
+		self.try_take_while(move |_| {
+			let res = future::ok(n > 0);
+			n = n.saturating_sub(1);
+			res
+		})
+	}
+}