add tryx_collect and tryx_join_all adapters

Author: sunshowers

Cargo.toml

@@ -14,9 +14,13 @@ futures-core = { version = "0.3.28", default-features = false }
 futures-sink = { version = "0.3.28", default-features = false }
 futures-util = { version = "0.3.28", default-features = false, features = ["sink"] }
 tokio = { version = "1.28.2", optional = true }
+futures = { version = "0.3.28", optional = true }
+pin-project-lite = "0.2.9"
 
 [dev-dependencies]
 anyhow = "1.0.72"
+async-stream = "0.3.5"
+debug-ignore = "1.0.5"
 tempfile = "3.6.0"
 tokio = { version = "1.28.2", features = ["fs", "io-util", "macros", "rt", "rt-multi-thread", "sync", "time"] }
 tokio-test = "0.4.0"
@@ -31,4 +35,4 @@ alloc = ["futures-core/alloc", "futures-sink/alloc", "futures-util/alloc"]
 macros = []
 
 # Internal-only feature, used for documentation and doctests. Not part of the public API.
-internal-docs = ["dep:tokio", "tokio?/io-util", "tokio?/macros", "tokio?/rt"]
+internal-docs = ["dep:futures", "dep:tokio", "tokio?/io-util", "tokio?/macros", "tokio?/rt"]

build.rs

@@ -0,0 +1,41 @@
+// The rustc-cfg listed below are considered public API, but it is *unstable*
+// and outside of the normal semver guarantees:
+//
+// - `futures_no_atomic_cas`
+//      Assume the target does *not* support atomic CAS operations.
+//      This is usually detected automatically by the build script, but you may
+//      need to enable it manually when building for custom targets or using
+//      non-cargo build systems that don't run the build script.
+//
+// With the exceptions mentioned above, the rustc-cfg emitted by the build
+// script are *not* public API.
+
+#![warn(rust_2018_idioms, single_use_lifetimes)]
+
+use std::env;
+
+include!("no_atomic_cas.rs");
+
+fn main() {
+    let target = match env::var("TARGET") {
+        Ok(target) => target,
+        Err(e) => {
+            println!(
+                "cargo:warning={}: unable to get TARGET environment variable: {}",
+                env!("CARGO_PKG_NAME"),
+                e
+            );
+            return;
+        }
+    };
+
+    // Note that this is `no_*`, not `has_*`. This allows treating
+    // `cfg(target_has_atomic = "ptr")` as true when the build script doesn't
+    // run. This is needed for compatibility with non-cargo build systems that
+    // don't run the build script.
+    if NO_ATOMIC_CAS.contains(&&*target) {
+        println!("cargo:rustc-cfg=futures_no_atomic_cas");
+    }
+
+    println!("cargo:rerun-if-changed=no_atomic_cas.rs");
+}

no_atomic_cas.rs

@@ -0,0 +1,17 @@
+// This file is @generated by no-atomic-cas.sh.
+// It is not intended for manual editing.
+
+const NO_ATOMIC_CAS: &[&str] = &[
+    "armv4t-none-eabi",
+    "armv5te-none-eabi",
+    "avr-unknown-gnu-atmega328",
+    "bpfeb-unknown-none",
+    "bpfel-unknown-none",
+    "msp430-none-elf",
+    "riscv32i-unknown-none-elf",
+    "riscv32im-unknown-none-elf",
+    "riscv32imc-unknown-none-elf",
+    "thumbv4t-none-eabi",
+    "thumbv5te-none-eabi",
+    "thumbv6m-none-eabi",
+];

scripts/no-atomic-cas.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+set -euo pipefail
+IFS=$'\n\t'
+cd "$(dirname "$0")"/..
+
+# Update the list of targets that do not support atomic CAS operations.
+#
+# Usage:
+#    ./scripts/no-atomic-cas.sh
+
+file="no_atomic_cas.rs"
+
+while [[ "$#" -gt 0 ]]; do
+    case $1 in
+        -h|--help) echo "Usage: no-atomic-cas.sh"; exit 0 ;;
+        *) echo "Unknown parameter passed: $1"; exit 1 ;;
+    esac
+    shift
+done
+
+# For -Z unstable-options
+export RUSTC_BOOTSTRAP=1
+
+no_atomic_cas=()
+for target in $(rustc --print target-list); do
+    target_spec=$(rustc --print target-spec-json -Z unstable-options --target "${target}")
+    res=$(jq <<<"${target_spec}" -r 'select(."atomic-cas" == false)')
+    [[ -z "${res}" ]] || no_atomic_cas+=("${target}")
+done
+
+cat >"${file}" <<EOF
+// This file is @generated by $(basename "$0").
+// It is not intended for manual editing.
+
+const NO_ATOMIC_CAS: &[&str] = &[
+EOF
+for target in "${no_atomic_cas[@]}"; do
+    echo "    \"${target}\"," >>"${file}"
+done
+cat >>"${file}" <<EOF
+];
+EOF

src/future/mod.rs

@@ -0,0 +1,9 @@
+//! Alternative adapters for asynchronous values.
+//!
+//! This module contains adapters that are similar to the `try_` adapters in the [`futures::future`]
+//! module, but don't cancel other futures if one of them errors out.
+
+#[cfg(feature = "alloc")]
+mod tryx_join_all;
+#[cfg(feature = "alloc")]
+pub use tryx_join_all::{tryx_join_all, TryxJoinAll};

src/future/tryx_join_all.rs

@@ -0,0 +1,218 @@
+//! Definition of the `TryxJoinAll` adapter, waiting for all of a list of
+//! futures to finish with either success or error.
+
+use crate::{
+    stream::{TryStreamExt, TryxCollect},
+    support::assert_future,
+};
+use alloc::{boxed::Box, vec::Vec};
+use core::{
+    fmt,
+    future::Future,
+    iter::FromIterator,
+    mem,
+    pin::Pin,
+    task::{Context, Poll},
+};
+use futures_core::future::TryFuture;
+use futures_util::{
+    future::{IntoFuture, MaybeDone, TryFutureExt},
+    stream::FuturesOrdered,
+};
+
+#[cfg(not(futures_no_atomic_cas))]
+pub(crate) const SMALL: usize = 30;
+
+/// Future for the [`tryx_join_all`] function.
+#[must_use = "futures do nothing unless you `.await` or poll them"]
+pub struct TryxJoinAll<F>
+where
+    F: TryFuture,
+{
+    kind: TryxJoinAllKind<F>,
+}
+
+enum TryxJoinAllKind<F>
+where
+    F: TryFuture,
+{
+    Small {
+        elems: Pin<Box<[MaybeDone<IntoFuture<F>>]>>,
+    },
+    #[cfg(not(futures_no_atomic_cas))]
+    Big {
+        // The use of FuturesOrdered here ensures that in case of errors, the first future listed in
+        // the iterator that errors out will be returned.
+        fut: TryxCollect<FuturesOrdered<IntoFuture<F>>, Vec<F::Ok>>,
+    },
+}
+
+impl<F> fmt::Debug for TryxJoinAll<F>
+where
+    F: TryFuture + fmt::Debug,
+    F::Ok: fmt::Debug,
+    F::Error: fmt::Debug,
+    F::Output: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self.kind {
+            TryxJoinAllKind::Small { ref elems } => {
+                f.debug_struct("TryxJoinAll").field("elems", elems).finish()
+            }
+            #[cfg(not(futures_no_atomic_cas))]
+            TryxJoinAllKind::Big { ref fut, .. } => fmt::Debug::fmt(fut, f),
+        }
+    }
+}
+
+/// Creates a future which represents either a collection of the results of the futures given or an
+/// error.
+///
+/// The returned future will drive execution for all of its underlying futures, collecting the
+/// results into a destination `Vec<T>` in the same order as they were provided.
+///
+/// Unlike [`futures::future::try_join_all`], if any future returns an error then all other futures
+/// will **not** be canceled. Instead, all other futures will be run to completion.
+///
+/// * If all futures complete successfully, then the returned future will succeed with a
+///   `Vec` of all the successful results.
+/// * If one or more futures fail, then the returned future will error out with the error
+///   for the first future listed that failed.
+///
+/// This function is only available when the `std` or `alloc` feature of this library is activated,
+/// and it is activated by default.
+///
+/// # Why use `tryx_join_all`?
+///
+/// See the documentation for [`tryx_join`](crate::tryx_join) for a discussion of why you might
+/// want to use a `tryx_` adapter.
+///
+/// # See Also
+///
+/// `tryx_join_all` will switch to the more powerful [`FuturesOrdered`] for performance reasons if
+/// the number of futures is large. You may want to look into using it or its counterpart
+/// [`FuturesUnordered`][futures::stream::FuturesUnordered] directly.
+///
+/// Some examples for additional functionality provided by these are:
+///
+///  * Adding new futures to the set even after it has been started.
+///
+///  * Only polling the specific futures that have been woken. In cases where you have a lot of
+///    futures this will result in much more efficient polling.
+///
+/// # Examples
+///
+/// ```
+/// # futures::executor::block_on(async {
+/// use futures::future;
+/// use cancel_safe_futures::future::tryx_join_all;
+///
+/// let futures = vec![
+///     future::ok::<u32, u32>(1),
+///     future::ok::<u32, u32>(2),
+///     future::ok::<u32, u32>(3),
+/// ];
+///
+/// assert_eq!(tryx_join_all(futures).await, Ok(vec![1, 2, 3]));
+///
+/// let futures = vec![
+///     future::ok::<u32, u32>(1),
+///     future::err::<u32, u32>(2),
+///     future::ok::<u32, u32>(3),
+/// ];
+///
+/// assert_eq!(tryx_join_all(futures).await, Err(2));
+/// # });
+/// ```
+pub fn tryx_join_all<I>(iter: I) -> TryxJoinAll<I::Item>
+where
+    I: IntoIterator,
+    I::Item: TryFuture,
+{
+    let iter = iter.into_iter().map(TryFutureExt::into_future);
+
+    #[cfg(futures_no_atomic_cas)]
+    {
+        let kind = TryxJoinAllKind::Small {
+            elems: iter.map(MaybeDone::Future).collect::<Box<[_]>>().into(),
+        };
+
+        assert_future::<Result<Vec<<I::Item as TryFuture>::Ok>, <I::Item as TryFuture>::Error>, _>(
+            TryxJoinAll { kind },
+        )
+    }
+
+    #[cfg(not(futures_no_atomic_cas))]
+    {
+        let kind = match iter.size_hint().1 {
+            Some(max) if max <= SMALL => TryxJoinAllKind::Small {
+                elems: iter.map(MaybeDone::Future).collect::<Box<[_]>>().into(),
+            },
+            _ => TryxJoinAllKind::Big {
+                fut: iter.collect::<FuturesOrdered<_>>().tryx_collect(),
+            },
+        };
+
+        assert_future::<Result<Vec<<I::Item as TryFuture>::Ok>, <I::Item as TryFuture>::Error>, _>(
+            TryxJoinAll { kind },
+        )
+    }
+}
+
+impl<F> Future for TryxJoinAll<F>
+where
+    F: TryFuture,
+{
+    type Output = Result<Vec<F::Ok>, F::Error>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        enum FinalState {
+            Pending,
+            AllDone,
+        }
+
+        match &mut self.kind {
+            TryxJoinAllKind::Small { elems } => {
+                let mut state = FinalState::AllDone;
+
+                for elem in iter_pin_mut(elems.as_mut()) {
+                    match elem.poll(cx) {
+                        Poll::Pending => state = FinalState::Pending,
+                        Poll::Ready(()) => {}
+                    }
+                }
+
+                match state {
+                    FinalState::Pending => Poll::Pending,
+                    FinalState::AllDone => {
+                        let mut elems = mem::replace(elems, Box::pin([]));
+                        let results: Result<Vec<_>, _> = iter_pin_mut(elems.as_mut())
+                            .map(|e| e.take_output().unwrap())
+                            .collect();
+                        Poll::Ready(results)
+                    }
+                }
+            }
+            #[cfg(not(futures_no_atomic_cas))]
+            TryxJoinAllKind::Big { fut } => Pin::new(fut).poll(cx),
+        }
+    }
+}
+
+impl<F> FromIterator<F> for TryxJoinAll<F>
+where
+    F: TryFuture,
+{
+    fn from_iter<T: IntoIterator<Item = F>>(iter: T) -> Self {
+        tryx_join_all(iter)
+    }
+}
+
+pub(crate) fn iter_pin_mut<T>(slice: Pin<&mut [T]>) -> impl Iterator<Item = Pin<&mut T>> {
+    // Safety: `std` _could_ make this unsound if it were to decide Pin's
+    // invariants aren't required to transmit through slices. Otherwise this has
+    // the same safety as a normal field pin projection.
+    unsafe { slice.get_unchecked_mut() }
+        .iter_mut()
+        .map(|t| unsafe { Pin::new_unchecked(t) })
+}

src/lib.rs

@@ -72,6 +72,8 @@
 //! The `tryx` library includes:
 //!
 //! * [`tryx_join`]: similar to [`tokio::try_join`].
+//! * [`future::tryx_join_all`]: similar to [`futures::future::try_join_all`].
+//! * [`TryStreamExt`]: contains extension methods similar to [`futures::stream::TryStreamExt`].
 //!
 //! ### Example
 //!
@@ -87,6 +89,9 @@
 #![warn(missing_docs)]
 #![cfg_attr(doc_cfg, feature(doc_cfg, doc_auto_cfg))]
 
+#[cfg(feature = "alloc")]
+extern crate alloc;
+
 // Includes re-exports used by macros.
 //
 // This module is not intended to be part of the public API. In general, any
@@ -95,7 +100,11 @@
 #[doc(hidden)]
 pub mod macros;
 
+pub mod future;
 pub mod prelude;
 pub mod sink;
+pub mod stream;
+mod support;
 
 pub use sink::SinkExt;
+pub use stream::TryStreamExt;