diff options
184 files changed, 22862 insertions, 0 deletions
diff --git a/Android.bp b/Android.bp new file mode 100644 index 0000000..fb2aea9 --- /dev/null +++ b/Android.bp @@ -0,0 +1,50 @@ +// This file is generated by cargo2android.py. + +rust_library_host_rlib { + name: "libfutures_util", + crate_name: "futures_util", + srcs: ["src/lib.rs"], + edition: "2018", + features: [ + "alloc", + "async-await", + "async-await-macro", + "channel", + "default", + "futures-macro", + "io", + "proc-macro-hack", + "proc-macro-nested", + "sink", + "slab", + "std", + ], + rlibs: [ + "libfutures_channel", + "libfutures_core", + "libfutures_io", + "libfutures_sink", + "libfutures_task", + "libmemchr", + "libpin_utils", + "libproc_macro_nested", + "libslab", + ], + proc_macros: [ + "libfutures_macro", + "libproc_macro_hack", + ], +} + +// dependent_library ["feature_list"] +// futures-core-0.3.4 "alloc,std" +// futures-macro-0.3.4 +// futures-task-0.3.4 "alloc,std" +// pin-utils-0.1.0-alpha.4 +// proc-macro-hack-0.5.11 +// proc-macro-nested-0.1.3 +// proc-macro2-1.0.9 "default,proc-macro" +// quote-1.0.3 "default,proc-macro" +// slab-0.4.2 +// syn-1.0.16 "clone-impls,default,derive,full,parsing,printing,proc-macro,quote" +// unicode-xid-0.2.0 "default" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..125ee8c --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,97 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies +# +# If you believe there's an error in this file please file an +# issue against the rust-lang/cargo repository. If you're +# editing this file be aware that the upstream Cargo.toml +# will likely look very different (and much more reasonable) + +[package] +edition = "2018" +name = "futures-util" +version = "0.3.4" +authors = ["Alex Crichton <alex@alexcrichton.com>"] +description = "Common utilities and extension traits for the futures-rs library.\n" +homepage = "https://rust-lang.github.io/futures-rs" +documentation = "https://docs.rs/futures-util/0.3.0" +license = "MIT OR Apache-2.0" +repository = "https://github.com/rust-lang/futures-rs" +[package.metadata.docs.rs] +all-features = true +[dependencies.futures-channel] +version = "0.3.4" +features = ["std"] +optional = true +default-features = false + +[dependencies.futures-core] +version = "0.3.4" +default-features = false + +[dependencies.futures-io] +version = "0.3.4" +features = ["std"] +optional = true +default-features = false + +[dependencies.futures-macro] +version = "0.3.4" +optional = true +default-features = false + +[dependencies.futures-sink] +version = "0.3.4" +optional = true +default-features = false + +[dependencies.futures-task] +version = "0.3.4" +default-features = false + +[dependencies.futures_01] +version = "0.1.25" +optional = true +package = "futures" + +[dependencies.memchr] +version = "2.2" +optional = true + +[dependencies.pin-utils] +version = "0.1.0-alpha.4" + +[dependencies.proc-macro-hack] +version = "0.5.9" +optional = true + +[dependencies.proc-macro-nested] +version = "0.1.2" +optional = true + +[dependencies.slab] +version = "0.4" +optional = true + +[dependencies.tokio-io] +version = "0.1.9" +optional = true + +[features] +alloc = ["futures-core/alloc", "futures-task/alloc"] +async-await = [] +async-await-macro = ["async-await", "futures-macro", "proc-macro-hack", "proc-macro-nested"] +bilock = [] +cfg-target-has-atomic = ["futures-core/cfg-target-has-atomic", "futures-task/cfg-target-has-atomic"] +channel = ["std", "futures-channel"] +compat = ["std", "futures_01"] +default = ["std", "async-await", "async-await-macro"] +io = ["std", "futures-io", "memchr"] +io-compat = ["io", "compat", "tokio-io"] +read-initializer = ["io", "futures-io/read-initializer", "futures-io/unstable"] +sink = ["futures-sink"] +std = ["alloc", "futures-core/std", "futures-task/std", "slab"] +unstable = ["futures-core/unstable", "futures-task/unstable"] diff --git a/Cargo.toml.orig b/Cargo.toml.orig new file mode 100644 index 0000000..7d47baf --- /dev/null +++ b/Cargo.toml.orig @@ -0,0 +1,50 @@ +[package] +name = "futures-util" +edition = "2018" +version = "0.3.4" +authors = ["Alex Crichton <alex@alexcrichton.com>"] +license = "MIT OR Apache-2.0" +repository = "https://github.com/rust-lang/futures-rs" +homepage = "https://rust-lang.github.io/futures-rs" +documentation = "https://docs.rs/futures-util/0.3.0" +description = """ +Common utilities and extension traits for the futures-rs library. +""" + +[features] +default = ["std", "async-await", "async-await-macro"] +std = ["alloc", "futures-core/std", "futures-task/std", "slab"] +alloc = ["futures-core/alloc", "futures-task/alloc"] +async-await = [] +async-await-macro = ["async-await", "futures-macro", "proc-macro-hack", "proc-macro-nested"] +compat = ["std", "futures_01"] +io-compat = ["io", "compat", "tokio-io"] +sink = ["futures-sink"] +io = ["std", "futures-io", "memchr"] +channel = ["std", "futures-channel"] + +# Unstable features +# These features are outside of the normal semver guarantees and require the +# `unstable` feature as an explicit opt-in to unstable API. +unstable = ["futures-core/unstable", "futures-task/unstable"] +cfg-target-has-atomic = ["futures-core/cfg-target-has-atomic", "futures-task/cfg-target-has-atomic"] +bilock = [] +read-initializer = ["io", "futures-io/read-initializer", "futures-io/unstable"] + +[dependencies] +futures-core = { path = "../futures-core", version = "0.3.4", default-features = false } +futures-task = { path = "../futures-task", version = "0.3.4", default-features = false } +futures-channel = { path = "../futures-channel", version = "0.3.4", default-features = false, features = ["std"], optional = true } +futures-io = { path = "../futures-io", version = "0.3.4", default-features = false, features = ["std"], optional = true } +futures-sink = { path = "../futures-sink", version = "0.3.4", default-features = false, optional = true } +futures-macro = { path = "../futures-macro", version = "0.3.4", default-features = false, optional = true } +proc-macro-hack = { version = "0.5.9", optional = true } +proc-macro-nested = { version = "0.1.2", optional = true } +slab = { version = "0.4", optional = true } +memchr = { version = "2.2", optional = true } +futures_01 = { version = "0.1.25", optional = true, package = "futures" } +tokio-io = { version = "0.1.9", optional = true } +pin-utils = "0.1.0-alpha.4" + +[package.metadata.docs.rs] +all-features = true @@ -0,0 +1 @@ +LICENSE-APACHE
\ No newline at end of file diff --git a/LICENSE-APACHE b/LICENSE-APACHE new file mode 100644 index 0000000..9eb0b09 --- /dev/null +++ b/LICENSE-APACHE @@ -0,0 +1,202 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright (c) 2016 Alex Crichton +Copyright (c) 2017 The Tokio Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/LICENSE-MIT b/LICENSE-MIT new file mode 100644 index 0000000..8ad082e --- /dev/null +++ b/LICENSE-MIT @@ -0,0 +1,26 @@ +Copyright (c) 2016 Alex Crichton +Copyright (c) 2017 The Tokio Authors + +Permission is hereby granted, free of charge, to any +person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without +limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software +is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice +shall be included in all copies or substantial portions +of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. diff --git a/METADATA b/METADATA new file mode 100644 index 0000000..4a196e3 --- /dev/null +++ b/METADATA @@ -0,0 +1,17 @@ +name: "futures-util" +description: + "Common utilities and extension traits for the futures-rs library." + +third_party { + url { + type: HOMEPAGE + value: "https://crates.io/crates/futures-util" + } + url { + type: GIT + value: "https://github.com/rust-lang/futures-rs" + } + version: "0.3.4" + last_upgrade_date { year: 2020 month: 3 day: 17 } + license_type: NOTICE +} diff --git a/MODULE_LICENSE_APACHE2 b/MODULE_LICENSE_APACHE2 new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/MODULE_LICENSE_APACHE2 @@ -0,0 +1 @@ +LICENSE
\ No newline at end of file diff --git a/benches/futures_unordered.rs b/benches/futures_unordered.rs new file mode 100644 index 0000000..05e9ead --- /dev/null +++ b/benches/futures_unordered.rs @@ -0,0 +1,43 @@ +#![feature(test)] + +extern crate test; +use crate::test::Bencher; + +use futures::channel::oneshot; +use futures::executor::block_on; +use futures::future; +use futures::stream::{StreamExt, FuturesUnordered}; +use futures::task::Poll; +use std::collections::VecDeque; +use std::thread; + +#[bench] +fn oneshots(b: &mut Bencher) { + const NUM: usize = 10_000; + + b.iter(|| { + let mut txs = VecDeque::with_capacity(NUM); + let mut rxs = FuturesUnordered::new(); + + for _ in 0..NUM { + let (tx, rx) = oneshot::channel(); + txs.push_back(tx); + rxs.push(rx); + } + + thread::spawn(move || { + while let Some(tx) = txs.pop_front() { + let _ = tx.send("hello"); + } + }); + + block_on(future::poll_fn(move |cx| { + loop { + if let Poll::Ready(None) = rxs.poll_next_unpin(cx) { + break + } + } + Poll::Ready(()) + })) + }); +} diff --git a/benches_disabled/bilock.rs b/benches_disabled/bilock.rs new file mode 100644 index 0000000..78b5edb --- /dev/null +++ b/benches_disabled/bilock.rs @@ -0,0 +1,126 @@ +#![feature(test)] + +#[cfg(feature = "bilock")] +mod bench { +use futures::task::{Context, Waker}; +use futures::executor::LocalPool; +use futures_util::lock::BiLock; +use futures_util::lock::BiLockAcquire; +use futures_util::lock::BiLockAcquired; +use futures_util::task::ArcWake; + +use std::sync::Arc; +use test::Bencher; + +fn notify_noop() -> Waker { + struct Noop; + + impl ArcWake for Noop { + fn wake(_: &Arc<Self>) {} + } + + ArcWake::into_waker(Arc::new(Noop)) +} + + +/// Pseudo-stream which simply calls `lock.poll()` on `poll` +struct LockStream { + lock: BiLockAcquire<u32>, +} + +impl LockStream { + fn new(lock: BiLock<u32>) -> LockStream { + LockStream { + lock: lock.lock() + } + } + + /// Release a lock after it was acquired in `poll`, + /// so `poll` could be called again. + fn release_lock(&mut self, guard: BiLockAcquired<u32>) { + self.lock = guard.unlock().lock() + } +} + +impl Stream for LockStream { + type Item = BiLockAcquired<u32>; + type Error = (); + + fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::Item>, Self::Error> { + self.lock.poll(cx).map(|a| a.map(Some)) + } +} + + +#[bench] +fn contended(b: &mut Bencher) { + let pool = LocalPool::new(); + let mut exec = pool.executor(); + let waker = notify_noop(); + let mut map = task::LocalMap::new(); + let mut waker = task::Context::new(&mut map, &waker, &mut exec); + + b.iter(|| { + let (x, y) = BiLock::new(1); + + let mut x = LockStream::new(x); + let mut y = LockStream::new(y); + + for _ in 0..1000 { + let x_guard = match x.poll_next(&mut waker) { + Ok(Poll::Ready(Some(guard))) => guard, + _ => panic!(), + }; + + // Try poll second lock while first lock still holds the lock + match y.poll_next(&mut waker) { + Ok(Poll::Pending) => (), + _ => panic!(), + }; + + x.release_lock(x_guard); + + let y_guard = match y.poll_next(&mut waker) { + Ok(Poll::Ready(Some(guard))) => guard, + _ => panic!(), + }; + + y.release_lock(y_guard); + } + (x, y) + }); +} + +#[bench] +fn lock_unlock(b: &mut Bencher) { + let pool = LocalPool::new(); + let mut exec = pool.executor(); + let waker = notify_noop(); + let mut map = task::LocalMap::new(); + let mut waker = task::Context::new(&mut map, &waker, &mut exec); + + b.iter(|| { + let (x, y) = BiLock::new(1); + + let mut x = LockStream::new(x); + let mut y = LockStream::new(y); + + for _ in 0..1000 { + let x_guard = match x.poll_next(&mut waker) { + Ok(Poll::Ready(Some(guard))) => guard, + _ => panic!(), + }; + + x.release_lock(x_guard); + + let y_guard = match y.poll_next(&mut waker) { + Ok(Poll::Ready(Some(guard))) => guard, + _ => panic!(), + }; + + y.release_lock(y_guard); + } + (x, y) + }) +} +} diff --git a/src/async_await/join_mod.rs b/src/async_await/join_mod.rs new file mode 100644 index 0000000..2e27f21 --- /dev/null +++ b/src/async_await/join_mod.rs @@ -0,0 +1,82 @@ +//! The `join` macro. + +use proc_macro_hack::proc_macro_hack; + +#[doc(hidden)] +#[macro_export] +macro_rules! document_join_macro { + ($join:item $try_join:item) => { + /// Polls multiple futures simultaneously, returning a tuple + /// of all results once complete. + /// + /// While `join!(a, b)` is similar to `(a.await, b.await)`, + /// `join!` polls both futures concurrently and therefore is more efficent. + /// + /// This macro is only usable inside of async functions, closures, and blocks. + /// It is also gated behind the `async-await` feature of this library, which is + /// activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::join; + /// + /// let a = async { 1 }; + /// let b = async { 2 }; + /// + /// assert_eq!(join!(a, b), (1, 2)); + /// # }); + /// ``` + $join + + /// Polls multiple futures simultaneously, resolving to a [`Result`] containing + /// either a tuple of the successful outputs or an error. + /// + /// `try_join!` is similar to [`join!`], but completes immediately if any of + /// the futures return an error. + /// + /// This macro is only usable inside of async functions, closures, and blocks. + /// It is also gated behind the `async-await` feature of this library, which is + /// activated by default. + /// + /// # Examples + /// + /// When used on multiple futures that return `Ok`, `try_join!` will return + /// `Ok` of a tuple of the values: + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::try_join; + /// + /// let a = async { Ok::<i32, i32>(1) }; + /// let b = async { Ok::<u64, i32>(2) }; + /// + /// assert_eq!(try_join!(a, b), Ok((1, 2))); + /// # }); + /// ``` + /// + /// If one of the futures resolves to an error, `try_join!` will return + /// that error: + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::try_join; + /// + /// let a = async { Ok::<i32, i32>(1) }; + /// let b = async { Err::<u64, i32>(2) }; + /// + /// assert_eq!(try_join!(a, b), Err(2)); + /// # }); + /// ``` + $try_join + } +} + +document_join_macro! { + #[proc_macro_hack(support_nested)] + pub use futures_macro::join; + + #[proc_macro_hack(support_nested)] + pub use futures_macro::try_join; +} diff --git a/src/async_await/mod.rs b/src/async_await/mod.rs new file mode 100644 index 0000000..69cae13 --- /dev/null +++ b/src/async_await/mod.rs @@ -0,0 +1,51 @@ +//! Await +//! +//! This module contains a number of functions and combinators for working +//! with `async`/`await` code. + +use futures_core::future::Future; +use futures_core::stream::Stream; + +#[doc(hidden)] +pub use futures_core::future::FusedFuture; +#[doc(hidden)] +pub use futures_core::stream::FusedStream; + +#[macro_use] +mod poll; +pub use self::poll::*; + +#[macro_use] +mod pending; +pub use self::pending::*; + +// Primary export is a macro +#[cfg(feature = "async-await-macro")] +mod join_mod; +#[cfg(feature = "async-await-macro")] +pub use self::join_mod::*; + +// Primary export is a macro +#[cfg(feature = "async-await-macro")] +mod select_mod; +#[cfg(feature = "async-await-macro")] +pub use self::select_mod::*; + +#[cfg(feature = "std")] +#[cfg(feature = "async-await-macro")] +mod random; +#[cfg(feature = "std")] +#[cfg(feature = "async-await-macro")] +pub use self::random::*; + +#[doc(hidden)] +#[inline(always)] +pub fn assert_unpin<T: Unpin>(_: &T) {} + +#[doc(hidden)] +#[inline(always)] +pub fn assert_fused_future<T: Future + FusedFuture>(_: &T) {} + +#[doc(hidden)] +#[inline(always)] +pub fn assert_fused_stream<T: Stream + FusedStream>(_: &T) {} diff --git a/src/async_await/pending.rs b/src/async_await/pending.rs new file mode 100644 index 0000000..b143869 --- /dev/null +++ b/src/async_await/pending.rs @@ -0,0 +1,43 @@ +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; + +/// A macro which yields to the event loop once. +/// +/// This is equivalent to returning [`Poll::Pending`](futures_core::task::Poll) +/// from a [`Future::poll`](futures_core::future::Future::poll) implementation. +/// Similarly, when using this macro, it must be ensured that [`wake`](std::task::Waker::wake) +/// is called somewhere when further progress can be made. +/// +/// This macro is only usable inside of async functions, closures, and blocks. +/// It is also gated behind the `async-await` feature of this library, which is +/// activated by default. +#[macro_export] +macro_rules! pending { + () => { + $crate::async_await::pending_once().await + } +} + +#[doc(hidden)] +pub fn pending_once() -> PendingOnce { + PendingOnce { is_ready: false } +} + +#[allow(missing_debug_implementations)] +#[doc(hidden)] +pub struct PendingOnce { + is_ready: bool, +} + +impl Future for PendingOnce { + type Output = (); + fn poll(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> { + if self.is_ready { + Poll::Ready(()) + } else { + self.is_ready = true; + Poll::Pending + } + } +} diff --git a/src/async_await/poll.rs b/src/async_await/poll.rs new file mode 100644 index 0000000..dffa94b --- /dev/null +++ b/src/async_await/poll.rs @@ -0,0 +1,35 @@ +use crate::future::FutureExt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; + +/// A macro which returns the result of polling a future once within the +/// current `async` context. +/// +/// This macro is only usable inside of `async` functions, closures, and blocks. +/// It is also gated behind the `async-await` feature of this library, which is +/// activated by default. +#[macro_export] +macro_rules! poll { + ($x:expr $(,)?) => { + $crate::async_await::poll($x).await + } +} + +#[doc(hidden)] +pub fn poll<F: Future + Unpin>(future: F) -> PollOnce<F> { + PollOnce { future } +} + +#[allow(missing_debug_implementations)] +#[doc(hidden)] +pub struct PollOnce<F: Future + Unpin> { + future: F, +} + +impl<F: Future + Unpin> Future for PollOnce<F> { + type Output = Poll<F::Output>; + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + Poll::Ready(self.future.poll_unpin(cx)) + } +} diff --git a/src/async_await/random.rs b/src/async_await/random.rs new file mode 100644 index 0000000..4f8c725 --- /dev/null +++ b/src/async_await/random.rs @@ -0,0 +1,54 @@ +use std::{ + cell::Cell, + collections::hash_map::DefaultHasher, + hash::Hasher, + num::Wrapping, + sync::atomic::{AtomicUsize, Ordering}, +}; + +// Based on [Fisher–Yates shuffle]. +// +// [Fisher–Yates shuffle]: https://en.wikipedia.org/wiki/Fisher–Yates_shuffle +#[doc(hidden)] +pub fn shuffle<T>(slice: &mut [T]) { + for i in (1..slice.len()).rev() { + slice.swap(i, gen_index(i + 1)); + } +} + +/// Return a value from `0..n`. +fn gen_index(n: usize) -> usize { + (random() % n as u64) as usize +} + +/// Pseudorandom number generator based on [xorshift*]. +/// +/// [xorshift*]: https://en.wikipedia.org/wiki/Xorshift#xorshift* +fn random() -> u64 { + thread_local! { + static RNG: Cell<Wrapping<u64>> = Cell::new(Wrapping(prng_seed())); + } + + fn prng_seed() -> u64 { + static COUNTER: AtomicUsize = AtomicUsize::new(0); + + // Any non-zero seed will do + let mut seed = 0; + while seed == 0 { + let mut hasher = DefaultHasher::new(); + hasher.write_usize(COUNTER.fetch_add(1, Ordering::Relaxed)); + seed = hasher.finish(); + } + seed + } + + RNG.with(|rng| { + let mut x = rng.get(); + debug_assert_ne!(x.0, 0); + x ^= x >> 12; + x ^= x << 25; + x ^= x >> 27; + rng.set(x); + x.0.wrapping_mul(0x2545_f491_4f6c_dd1d) + }) +} diff --git a/src/async_await/select_mod.rs b/src/async_await/select_mod.rs new file mode 100644 index 0000000..38153c7 --- /dev/null +++ b/src/async_await/select_mod.rs @@ -0,0 +1,320 @@ +//! The `select` macro. + +use proc_macro_hack::proc_macro_hack; + +#[doc(hidden)] +#[macro_export] +macro_rules! document_select_macro { + // This branch is required for `futures 0.3.1`, from before select_biased was introduced + ($select:item) => { + /// Polls multiple futures and streams simultaneously, executing the branch + /// for the future that finishes first. If multiple futures are ready, + /// one will be pseudo-randomly selected at runtime. Futures directly + /// passed to `select!` must be `Unpin` and implement `FusedFuture`. + /// + /// If an expression which yields a `Future` is passed to `select!` + /// (e.g. an `async fn` call) instead of a `Future` by name the `Unpin` + /// requirement is relaxed, since the macro will pin the resulting `Future` + /// on the stack. However the `Future` returned by the expression must + /// still implement `FusedFuture`. This difference is presented + /// + /// Futures and streams which are not already fused can be fused using the + /// `.fuse()` method. Note, though, that fusing a future or stream directly + /// in the call to `select!` will not be enough to prevent it from being + /// polled after completion if the `select!` call is in a loop, so when + /// `select!`ing in a loop, users should take care to `fuse()` outside of + /// the loop. + /// + /// `select!` can be used as an expression and will return the return + /// value of the selected branch. For this reason the return type of every + /// branch in a `select!` must be the same. + /// + /// This macro is only usable inside of async functions, closures, and blocks. + /// It is also gated behind the `async-await` feature of this library, which is + /// activated by default. + /// + /// Note that `select!` relies on `proc-macro-hack`, and may require to set the + /// compiler's recursion limit very high, e.g. `#![recursion_limit="1024"]`. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::select; + /// let mut a = future::ready(4); + /// let mut b = future::pending::<()>(); + /// + /// let res = select! { + /// a_res = a => a_res + 1, + /// _ = b => 0, + /// }; + /// assert_eq!(res, 5); + /// # }); + /// ``` + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// use futures::select; + /// let mut st = stream::iter(vec![2]).fuse(); + /// let mut fut = future::pending::<()>(); + /// + /// select! { + /// x = st.next() => assert_eq!(Some(2), x), + /// _ = fut => panic!(), + /// }; + /// # }); + /// ``` + /// + /// As described earlier, `select` can directly select on expressions + /// which return `Future`s - even if those do not implement `Unpin`: + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// use futures::select; + /// + /// // Calling the following async fn returns a Future which does not + /// // implement Unpin + /// async fn async_identity_fn(arg: usize) -> usize { + /// arg + /// } + /// + /// let res = select! { + /// a_res = async_identity_fn(62).fuse() => a_res + 1, + /// b_res = async_identity_fn(13).fuse() => b_res, + /// }; + /// assert!(res == 63 || res == 12); + /// # }); + /// ``` + /// + /// If a similar async function is called outside of `select` to produce + /// a `Future`, the `Future` must be pinned in order to be able to pass + /// it to `select`. This can be achieved via `Box::pin` for pinning a + /// `Future` on the heap or the `pin_mut!` macro for pinning a `Future` + /// on the stack. + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// use futures::select; + /// use futures::pin_mut; + /// + /// // Calling the following async fn returns a Future which does not + /// // implement Unpin + /// async fn async_identity_fn(arg: usize) -> usize { + /// arg + /// } + /// + /// let fut_1 = async_identity_fn(1).fuse(); + /// let fut_2 = async_identity_fn(2).fuse(); + /// let mut fut_1 = Box::pin(fut_1); // Pins the Future on the heap + /// pin_mut!(fut_2); // Pins the Future on the stack + /// + /// let res = select! { + /// a_res = fut_1 => a_res, + /// b_res = fut_2 => b_res, + /// }; + /// assert!(res == 1 || res == 2); + /// # }); + /// ``` + /// + /// `select` also accepts a `complete` branch and a `default` branch. + /// `complete` will run if all futures and streams have already been + /// exhausted. `default` will run if no futures or streams are + /// immediately ready. `complete` takes priority over `default` in + /// the case where all futures have completed. + /// A motivating use-case for passing `Future`s by name as well as for + /// `complete` blocks is to call `select!` in a loop, which is + /// demonstrated in the following example: + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::select; + /// let mut a_fut = future::ready(4); + /// let mut b_fut = future::ready(6); + /// let mut total = 0; + /// + /// loop { + /// select! { + /// a = a_fut => total += a, + /// b = b_fut => total += b, + /// complete => break, + /// default => panic!(), // never runs (futures run first, then complete) + /// }; + /// } + /// assert_eq!(total, 10); + /// # }); + /// ``` + /// + /// Note that the futures that have been matched over can still be mutated + /// from inside the `select!` block's branches. This can be used to implement + /// more complex behavior such as timer resets or writing into the head of + /// a stream. + $select + }; + + ($select:item $select_biased:item) => { + $crate::document_select_macro!($select); + + /// Polls multiple futures and streams simultaneously, executing the branch + /// for the future that finishes first. Unlike [`select!`], if multiple futures are ready, + /// one will be selected in order of declaration. Futures directly + /// passed to `select_biased!` must be `Unpin` and implement `FusedFuture`. + /// + /// If an expression which yields a `Future` is passed to `select_biased!` + /// (e.g. an `async fn` call) instead of a `Future` by name the `Unpin` + /// requirement is relaxed, since the macro will pin the resulting `Future` + /// on the stack. However the `Future` returned by the expression must + /// still implement `FusedFuture`. This difference is presented + /// + /// Futures and streams which are not already fused can be fused using the + /// `.fuse()` method. Note, though, that fusing a future or stream directly + /// in the call to `select_biased!` will not be enough to prevent it from being + /// polled after completion if the `select_biased!` call is in a loop, so when + /// `select_biased!`ing in a loop, users should take care to `fuse()` outside of + /// the loop. + /// + /// `select_biased!` can be used as an expression and will return the return + /// value of the selected branch. For this reason the return type of every + /// branch in a `select_biased!` must be the same. + /// + /// This macro is only usable inside of async functions, closures, and blocks. + /// It is also gated behind the `async-await` feature of this library, which is + /// activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::select_biased; + /// let mut a = future::ready(4); + /// let mut b = future::pending::<()>(); + /// + /// let res = select_biased! { + /// a_res = a => a_res + 1, + /// _ = b => 0, + /// }; + /// assert_eq!(res, 5); + /// # }); + /// ``` + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// use futures::select_biased; + /// let mut st = stream::iter(vec![2]).fuse(); + /// let mut fut = future::pending::<()>(); + /// + /// select_biased! { + /// x = st.next() => assert_eq!(Some(2), x), + /// _ = fut => panic!(), + /// }; + /// # }); + /// ``` + /// + /// As described earlier, `select_biased` can directly select on expressions + /// which return `Future`s - even if those do not implement `Unpin`: + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// use futures::select_biased; + /// + /// // Calling the following async fn returns a Future which does not + /// // implement Unpin + /// async fn async_identity_fn(arg: usize) -> usize { + /// arg + /// } + /// + /// let res = select_biased! { + /// a_res = async_identity_fn(62).fuse() => a_res + 1, + /// b_res = async_identity_fn(13).fuse() => b_res, + /// }; + /// assert!(res == 63 || res == 12); + /// # }); + /// ``` + /// + /// If a similar async function is called outside of `select_biased` to produce + /// a `Future`, the `Future` must be pinned in order to be able to pass + /// it to `select_biased`. This can be achieved via `Box::pin` for pinning a + /// `Future` on the heap or the `pin_mut!` macro for pinning a `Future` + /// on the stack. + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// use futures::select_biased; + /// use futures::pin_mut; + /// + /// // Calling the following async fn returns a Future which does not + /// // implement Unpin + /// async fn async_identity_fn(arg: usize) -> usize { + /// arg + /// } + /// + /// let fut_1 = async_identity_fn(1).fuse(); + /// let fut_2 = async_identity_fn(2).fuse(); + /// let mut fut_1 = Box::pin(fut_1); // Pins the Future on the heap + /// pin_mut!(fut_2); // Pins the Future on the stack + /// + /// let res = select_biased! { + /// a_res = fut_1 => a_res, + /// b_res = fut_2 => b_res, + /// }; + /// assert!(res == 1 || res == 2); + /// # }); + /// ``` + /// + /// `select_biased` also accepts a `complete` branch and a `default` branch. + /// `complete` will run if all futures and streams have already been + /// exhausted. `default` will run if no futures or streams are + /// immediately ready. `complete` takes priority over `default` in + /// the case where all futures have completed. + /// A motivating use-case for passing `Future`s by name as well as for + /// `complete` blocks is to call `select_biased!` in a loop, which is + /// demonstrated in the following example: + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::select_biased; + /// let mut a_fut = future::ready(4); + /// let mut b_fut = future::ready(6); + /// let mut total = 0; + /// + /// loop { + /// select_biased! { + /// a = a_fut => total += a, + /// b = b_fut => total += b, + /// complete => break, + /// default => panic!(), // never runs (futures run first, then complete) + /// }; + /// } + /// assert_eq!(total, 10); + /// # }); + /// ``` + /// + /// Note that the futures that have been matched over can still be mutated + /// from inside the `select_biased!` block's branches. This can be used to implement + /// more complex behavior such as timer resets or writing into the head of + /// a stream. + /// + /// [`select!`]: macro.select.html + $select_biased + }; +} + +document_select_macro! { + #[cfg(feature = "std")] + #[proc_macro_hack(support_nested)] + pub use futures_macro::select; + + #[proc_macro_hack(support_nested)] + pub use futures_macro::select_biased; +} diff --git a/src/compat/compat01as03.rs b/src/compat/compat01as03.rs new file mode 100644 index 0000000..9bb00bf --- /dev/null +++ b/src/compat/compat01as03.rs @@ -0,0 +1,472 @@ +use futures_01::executor::{ + spawn as spawn01, Notify as Notify01, NotifyHandle as NotifyHandle01, + Spawn as Spawn01, UnsafeNotify as UnsafeNotify01, +}; +use futures_01::{ + Async as Async01, Future as Future01, + Stream as Stream01, +}; +#[cfg(feature = "sink")] +use futures_01::{AsyncSink as AsyncSink01, Sink as Sink01}; +use futures_core::{task as task03, future::Future as Future03, stream::Stream as Stream03}; +use std::pin::Pin; +use std::task::Context; +#[cfg(feature = "sink")] +use futures_sink::Sink as Sink03; + +#[cfg(feature = "io-compat")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use io::{AsyncRead01CompatExt, AsyncWrite01CompatExt}; + +/// Converts a futures 0.1 Future, Stream, AsyncRead, or AsyncWrite +/// object to a futures 0.3-compatible version, +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Compat01As03<T> { + pub(crate) inner: Spawn01<T>, +} + +impl<T> Unpin for Compat01As03<T> {} + +impl<T> Compat01As03<T> { + /// Wraps a futures 0.1 Future, Stream, AsyncRead, or AsyncWrite + /// object in a futures 0.3-compatible wrapper. + pub fn new(object: T) -> Compat01As03<T> { + Compat01As03 { + inner: spawn01(object), + } + } + + fn in_notify<R>(&mut self, cx: &mut Context<'_>, f: impl FnOnce(&mut T) -> R) -> R { + let notify = &WakerToHandle(cx.waker()); + self.inner.poll_fn_notify(notify, 0, f) + } + + /// Get a reference to 0.1 Future, Stream, AsyncRead, or AsyncWrite object contained within. + pub fn get_ref(&self) -> &T { + self.inner.get_ref() + } + + /// Get a mutable reference to 0.1 Future, Stream, AsyncRead or AsyncWrite object contained + /// within. + pub fn get_mut(&mut self) -> &mut T { + self.inner.get_mut() + } + + /// Consume this wrapper to return the underlying 0.1 Future, Stream, AsyncRead, or + /// AsyncWrite object. + pub fn into_inner(self) -> T { + self.inner.into_inner() + } +} + +/// Extension trait for futures 0.1 [`Future`](futures_01::future::Future) +pub trait Future01CompatExt: Future01 { + /// Converts a futures 0.1 + /// [`Future<Item = T, Error = E>`](futures_01::future::Future) + /// into a futures 0.3 + /// [`Future<Output = Result<T, E>>`](futures_core::future::Future). + /// + /// ``` + /// # futures::executor::block_on(async { + /// # // TODO: These should be all using `futures::compat`, but that runs up against Cargo + /// # // feature issues + /// use futures_util::compat::Future01CompatExt; + /// + /// let future = futures_01::future::ok::<u32, ()>(1); + /// assert_eq!(future.compat().await, Ok(1)); + /// # }); + /// ``` + fn compat(self) -> Compat01As03<Self> + where + Self: Sized, + { + Compat01As03::new(self) + } +} +impl<Fut: Future01> Future01CompatExt for Fut {} + +/// Extension trait for futures 0.1 [`Stream`](futures_01::stream::Stream) +pub trait Stream01CompatExt: Stream01 { + /// Converts a futures 0.1 + /// [`Stream<Item = T, Error = E>`](futures_01::stream::Stream) + /// into a futures 0.3 + /// [`Stream<Item = Result<T, E>>`](futures_core::stream::Stream). + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::StreamExt; + /// use futures_util::compat::Stream01CompatExt; + /// + /// let stream = futures_01::stream::once::<u32, ()>(Ok(1)); + /// let mut stream = stream.compat(); + /// assert_eq!(stream.next().await, Some(Ok(1))); + /// assert_eq!(stream.next().await, None); + /// # }); + /// ``` + fn compat(self) -> Compat01As03<Self> + where + Self: Sized, + { + Compat01As03::new(self) + } +} +impl<St: Stream01> Stream01CompatExt for St {} + +/// Extension trait for futures 0.1 [`Sink`](futures_01::sink::Sink) +#[cfg(feature = "sink")] +pub trait Sink01CompatExt: Sink01 { + /// Converts a futures 0.1 + /// [`Sink<SinkItem = T, SinkError = E>`](futures_01::sink::Sink) + /// into a futures 0.3 + /// [`Sink<T, Error = E>`](futures_sink::Sink). + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::{sink::SinkExt, stream::StreamExt}; + /// use futures_util::compat::{Stream01CompatExt, Sink01CompatExt}; + /// + /// let (tx, rx) = futures_01::unsync::mpsc::channel(1); + /// let (mut tx, mut rx) = (tx.sink_compat(), rx.compat()); + /// + /// tx.send(1).await.unwrap(); + /// drop(tx); + /// assert_eq!(rx.next().await, Some(Ok(1))); + /// assert_eq!(rx.next().await, None); + /// # }); + /// ``` + fn sink_compat(self) -> Compat01As03Sink<Self, Self::SinkItem> + where + Self: Sized, + { + Compat01As03Sink::new(self) + } +} +#[cfg(feature = "sink")] +impl<Si: Sink01> Sink01CompatExt for Si {} + +fn poll_01_to_03<T, E>(x: Result<Async01<T>, E>) -> task03::Poll<Result<T, E>> { + match x? { + Async01::Ready(t) => task03::Poll::Ready(Ok(t)), + Async01::NotReady => task03::Poll::Pending, + } +} + +impl<Fut: Future01> Future03 for Compat01As03<Fut> { + type Output = Result<Fut::Item, Fut::Error>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> task03::Poll<Self::Output> { + poll_01_to_03(self.in_notify(cx, Future01::poll)) + } +} + +impl<St: Stream01> Stream03 for Compat01As03<St> { + type Item = Result<St::Item, St::Error>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> task03::Poll<Option<Self::Item>> { + match self.in_notify(cx, Stream01::poll)? { + Async01::Ready(Some(t)) => task03::Poll::Ready(Some(Ok(t))), + Async01::Ready(None) => task03::Poll::Ready(None), + Async01::NotReady => task03::Poll::Pending, + } + } +} + +/// Converts a futures 0.1 Sink object to a futures 0.3-compatible version +#[cfg(feature = "sink")] +#[derive(Debug)] +#[must_use = "sinks do nothing unless polled"] +pub struct Compat01As03Sink<S, SinkItem> { + pub(crate) inner: Spawn01<S>, + pub(crate) buffer: Option<SinkItem>, + pub(crate) close_started: bool, +} + +#[cfg(feature = "sink")] +impl<S, SinkItem> Unpin for Compat01As03Sink<S, SinkItem> {} + +#[cfg(feature = "sink")] +impl<S, SinkItem> Compat01As03Sink<S, SinkItem> { + /// Wraps a futures 0.1 Sink object in a futures 0.3-compatible wrapper. + pub fn new(inner: S) -> Compat01As03Sink<S, SinkItem> { + Compat01As03Sink { + inner: spawn01(inner), + buffer: None, + close_started: false + } + } + + fn in_notify<R>( + &mut self, + cx: &mut Context<'_>, + f: impl FnOnce(&mut S) -> R, + ) -> R { + let notify = &WakerToHandle(cx.waker()); + self.inner.poll_fn_notify(notify, 0, f) + } + + /// Get a reference to 0.1 Sink object contained within. + pub fn get_ref(&self) -> &S { + self.inner.get_ref() + } + + /// Get a mutable reference to 0.1 Sink contained within. + pub fn get_mut(&mut self) -> &mut S { + self.inner.get_mut() + } + + /// Consume this wrapper to return the underlying 0.1 Sink. + pub fn into_inner(self) -> S { + self.inner.into_inner() + } +} + +#[cfg(feature = "sink")] +impl<S, SinkItem> Stream03 for Compat01As03Sink<S, SinkItem> +where + S: Stream01, +{ + type Item = Result<S::Item, S::Error>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> task03::Poll<Option<Self::Item>> { + match self.in_notify(cx, Stream01::poll)? { + Async01::Ready(Some(t)) => task03::Poll::Ready(Some(Ok(t))), + Async01::Ready(None) => task03::Poll::Ready(None), + Async01::NotReady => task03::Poll::Pending, + } + } +} + +#[cfg(feature = "sink")] +impl<S, SinkItem> Sink03<SinkItem> for Compat01As03Sink<S, SinkItem> +where + S: Sink01<SinkItem = SinkItem>, +{ + type Error = S::SinkError; + + fn start_send( + mut self: Pin<&mut Self>, + item: SinkItem, + ) -> Result<(), Self::Error> { + debug_assert!(self.buffer.is_none()); + self.buffer = Some(item); + Ok(()) + } + + fn poll_ready( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> task03::Poll<Result<(), Self::Error>> { + match self.buffer.take() { + Some(item) => match self.in_notify(cx, |f| f.start_send(item))? { + AsyncSink01::Ready => task03::Poll::Ready(Ok(())), + AsyncSink01::NotReady(i) => { + self.buffer = Some(i); + task03::Poll::Pending + } + }, + None => task03::Poll::Ready(Ok(())), + } + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> task03::Poll<Result<(), Self::Error>> { + let item = self.buffer.take(); + match self.in_notify(cx, |f| match item { + Some(i) => match f.start_send(i)? { + AsyncSink01::Ready => f.poll_complete().map(|i| (i, None)), + AsyncSink01::NotReady(t) => { + Ok((Async01::NotReady, Some(t))) + } + }, + None => f.poll_complete().map(|i| (i, None)), + })? { + (Async01::Ready(_), _) => task03::Poll::Ready(Ok(())), + (Async01::NotReady, item) => { + self.buffer = item; + task03::Poll::Pending + } + } + } + + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> task03::Poll<Result<(), Self::Error>> { + let item = self.buffer.take(); + let close_started = self.close_started; + + let result = self.in_notify(cx, |f| { + if !close_started { + if let Some(item) = item { + if let AsyncSink01::NotReady(item) = f.start_send(item)? { + return Ok((Async01::NotReady, Some(item), false)); + } + } + + if let Async01::NotReady = f.poll_complete()? { + return Ok((Async01::NotReady, None, false)); + } + } + + Ok((<S as Sink01>::close(f)?, None, true)) + }); + + match result? { + (Async01::Ready(_), _, _) => task03::Poll::Ready(Ok(())), + (Async01::NotReady, item, close_started) => { + self.buffer = item; + self.close_started = close_started; + task03::Poll::Pending + } + } + } +} + +struct NotifyWaker(task03::Waker); + +#[allow(missing_debug_implementations)] // false positive: this is private type +#[derive(Clone)] +struct WakerToHandle<'a>(&'a task03::Waker); + +impl From<WakerToHandle<'_>> for NotifyHandle01 { + fn from(handle: WakerToHandle<'_>) -> NotifyHandle01 { + let ptr = Box::new(NotifyWaker(handle.0.clone())); + + unsafe { NotifyHandle01::new(Box::into_raw(ptr)) } + } +} + +impl Notify01 for NotifyWaker { + fn notify(&self, _: usize) { + self.0.wake_by_ref(); + } +} + +unsafe impl UnsafeNotify01 for NotifyWaker { + unsafe fn clone_raw(&self) -> NotifyHandle01 { + WakerToHandle(&self.0).into() + } + + unsafe fn drop_raw(&self) { + let ptr: *const dyn UnsafeNotify01 = self; + drop(Box::from_raw(ptr as *mut dyn UnsafeNotify01)); + } +} + +#[cfg(feature = "io-compat")] +mod io { + use super::*; + #[cfg(feature = "read-initializer")] + use futures_io::Initializer; + use futures_io::{AsyncRead as AsyncRead03, AsyncWrite as AsyncWrite03}; + use std::io::Error; + use tokio_io::{AsyncRead as AsyncRead01, AsyncWrite as AsyncWrite01}; + + /// Extension trait for tokio-io [`AsyncRead`](tokio_io::AsyncRead) + pub trait AsyncRead01CompatExt: AsyncRead01 { + /// Converts a tokio-io [`AsyncRead`](tokio_io::AsyncRead) into a futures-io 0.3 + /// [`AsyncRead`](futures_io::AsyncRead). + /// + /// ``` + /// #![feature(impl_trait_in_bindings)] + /// # #![allow(incomplete_features)] + /// # futures::executor::block_on(async { + /// use futures::io::AsyncReadExt; + /// use futures_util::compat::AsyncRead01CompatExt; + /// + /// let input = b"Hello World!"; + /// let reader: impl tokio_io::AsyncRead = std::io::Cursor::new(input); + /// let mut reader: impl futures::io::AsyncRead + Unpin = reader.compat(); + /// + /// let mut output = Vec::with_capacity(12); + /// reader.read_to_end(&mut output).await.unwrap(); + /// assert_eq!(output, input); + /// # }); + /// ``` + fn compat(self) -> Compat01As03<Self> + where + Self: Sized, + { + Compat01As03::new(self) + } + } + impl<R: AsyncRead01> AsyncRead01CompatExt for R {} + + /// Extension trait for tokio-io [`AsyncWrite`](tokio_io::AsyncWrite) + pub trait AsyncWrite01CompatExt: AsyncWrite01 { + /// Converts a tokio-io [`AsyncWrite`](tokio_io::AsyncWrite) into a futures-io 0.3 + /// [`AsyncWrite`](futures_io::AsyncWrite). + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::AsyncWriteExt; + /// use futures_util::compat::AsyncWrite01CompatExt; + /// + /// let input = b"Hello World!"; + /// let mut cursor = std::io::Cursor::new(Vec::with_capacity(12)); + /// + /// let mut writer = (&mut cursor).compat(); + /// writer.write_all(input).await.unwrap(); + /// + /// assert_eq!(cursor.into_inner(), input); + /// # }); + /// ``` + fn compat(self) -> Compat01As03<Self> + where + Self: Sized, + { + Compat01As03::new(self) + } + } + impl<W: AsyncWrite01> AsyncWrite01CompatExt for W {} + + impl<R: AsyncRead01> AsyncRead03 for Compat01As03<R> { + #[cfg(feature = "read-initializer")] + unsafe fn initializer(&self) -> Initializer { + // check if `prepare_uninitialized_buffer` needs zeroing + if self.inner.get_ref().prepare_uninitialized_buffer(&mut [1]) { + Initializer::zeroing() + } else { + Initializer::nop() + } + } + + fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) + -> task03::Poll<Result<usize, Error>> + { + poll_01_to_03(self.in_notify(cx, |x| x.poll_read(buf))) + } + } + + impl<W: AsyncWrite01> AsyncWrite03 for Compat01As03<W> { + fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) + -> task03::Poll<Result<usize, Error>> + { + poll_01_to_03(self.in_notify(cx, |x| x.poll_write(buf))) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) + -> task03::Poll<Result<(), Error>> + { + poll_01_to_03(self.in_notify(cx, AsyncWrite01::poll_flush)) + } + + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) + -> task03::Poll<Result<(), Error>> + { + poll_01_to_03(self.in_notify(cx, AsyncWrite01::shutdown)) + } + } +} diff --git a/src/compat/compat03as01.rs b/src/compat/compat03as01.rs new file mode 100644 index 0000000..3fd2ae0 --- /dev/null +++ b/src/compat/compat03as01.rs @@ -0,0 +1,299 @@ +use futures_01::{ + task as task01, Async as Async01, Future as Future01, Poll as Poll01, + Stream as Stream01, +}; +#[cfg(feature = "sink")] +use futures_01::{ + AsyncSink as AsyncSink01, Sink as Sink01, StartSend as StartSend01, +}; +use futures_core::{ + task::{RawWaker, RawWakerVTable}, + future::TryFuture as TryFuture03, + stream::TryStream as TryStream03, +}; +#[cfg(feature = "sink")] +use futures_sink::Sink as Sink03; +use crate::task::{ + self as task03, + ArcWake as ArcWake03, + WakerRef, +}; +#[cfg(feature = "sink")] +use std::marker::PhantomData; +use std::{ + mem, + pin::Pin, + sync::Arc, + task::Context, +}; + +/// Converts a futures 0.3 [`TryFuture`](futures_core::future::TryFuture) or +/// [`TryStream`](futures_core::stream::TryStream) into a futures 0.1 +/// [`Future`](futures_01::future::Future) or +/// [`Stream`](futures_01::stream::Stream). +#[derive(Debug, Clone, Copy)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Compat<T> { + pub(crate) inner: T, +} + +/// Converts a futures 0.3 [`Sink`](futures_sink::Sink) into a futures 0.1 +/// [`Sink`](futures_01::sink::Sink). +#[cfg(feature = "sink")] +#[derive(Debug)] +#[must_use = "sinks do nothing unless polled"] +pub struct CompatSink<T, Item> { + inner: T, + _phantom: PhantomData<fn(Item)>, +} + +impl<T> Compat<T> { + /// Creates a new [`Compat`]. + /// + /// For types which implement appropriate futures `0.3` + /// traits, the result will be a type which implements + /// the corresponding futures 0.1 type. + pub fn new(inner: T) -> Compat<T> { + Compat { inner } + } + + /// Get a reference to 0.3 Future, Stream, AsyncRead, or AsyncWrite object + /// contained within. + pub fn get_ref(&self) -> &T { + &self.inner + } + + /// Get a mutable reference to 0.3 Future, Stream, AsyncRead, or AsyncWrite object + /// contained within. + pub fn get_mut(&mut self) -> &mut T { + &mut self.inner + } + + /// Returns the inner item. + pub fn into_inner(self) -> T { + self.inner + } +} + +#[cfg(feature = "sink")] +impl<T, Item> CompatSink<T, Item> { + /// Creates a new [`CompatSink`]. + pub fn new(inner: T) -> Self { + CompatSink { + inner, + _phantom: PhantomData, + } + } + + /// Get a reference to 0.3 Sink contained within. + pub fn get_ref(&self) -> &T { + &self.inner + } + + /// Get a mutable reference to 0.3 Sink contained within. + pub fn get_mut(&mut self) -> &mut T { + &mut self.inner + } + + /// Returns the inner item. + pub fn into_inner(self) -> T { + self.inner + } +} + +fn poll_03_to_01<T, E>(x: task03::Poll<Result<T, E>>) + -> Result<Async01<T>, E> +{ + match x? { + task03::Poll::Ready(t) => Ok(Async01::Ready(t)), + task03::Poll::Pending => Ok(Async01::NotReady), + } +} + +impl<Fut> Future01 for Compat<Fut> +where + Fut: TryFuture03 + Unpin, +{ + type Item = Fut::Ok; + type Error = Fut::Error; + + fn poll(&mut self) -> Poll01<Self::Item, Self::Error> { + with_context(self, |inner, cx| poll_03_to_01(inner.try_poll(cx))) + } +} + +impl<St> Stream01 for Compat<St> +where + St: TryStream03 + Unpin, +{ + type Item = St::Ok; + type Error = St::Error; + + fn poll(&mut self) -> Poll01<Option<Self::Item>, Self::Error> { + with_context(self, |inner, cx| match inner.try_poll_next(cx)? { + task03::Poll::Ready(None) => Ok(Async01::Ready(None)), + task03::Poll::Ready(Some(t)) => Ok(Async01::Ready(Some(t))), + task03::Poll::Pending => Ok(Async01::NotReady), + }) + } +} + +#[cfg(feature = "sink")] +impl<T, Item> Sink01 for CompatSink<T, Item> +where + T: Sink03<Item> + Unpin, +{ + type SinkItem = Item; + type SinkError = T::Error; + + fn start_send( + &mut self, + item: Self::SinkItem, + ) -> StartSend01<Self::SinkItem, Self::SinkError> { + with_sink_context(self, |mut inner, cx| { + match inner.as_mut().poll_ready(cx)? { + task03::Poll::Ready(()) => { + inner.start_send(item).map(|()| AsyncSink01::Ready) + } + task03::Poll::Pending => Ok(AsyncSink01::NotReady(item)), + } + }) + } + + fn poll_complete(&mut self) -> Poll01<(), Self::SinkError> { + with_sink_context(self, |inner, cx| poll_03_to_01(inner.poll_flush(cx))) + } + + fn close(&mut self) -> Poll01<(), Self::SinkError> { + with_sink_context(self, |inner, cx| poll_03_to_01(inner.poll_close(cx))) + } +} + +#[derive(Clone)] +struct Current(task01::Task); + +impl Current { + fn new() -> Current { + Current(task01::current()) + } + + fn as_waker(&self) -> WakerRef<'_> { + unsafe fn ptr_to_current<'a>(ptr: *const ()) -> &'a Current { + &*(ptr as *const Current) + } + fn current_to_ptr(current: &Current) -> *const () { + current as *const Current as *const () + } + + unsafe fn clone(ptr: *const ()) -> RawWaker { + // Lazily create the `Arc` only when the waker is actually cloned. + // FIXME: remove `transmute` when a `Waker` -> `RawWaker` conversion + // function is landed in `core`. + mem::transmute::<task03::Waker, RawWaker>( + task03::waker(Arc::new(ptr_to_current(ptr).clone())) + ) + } + unsafe fn drop(_: *const ()) {} + unsafe fn wake(ptr: *const ()) { + ptr_to_current(ptr).0.notify() + } + + let ptr = current_to_ptr(self); + let vtable = &RawWakerVTable::new(clone, wake, wake, drop); + WakerRef::new_unowned(std::mem::ManuallyDrop::new(unsafe { + task03::Waker::from_raw(RawWaker::new(ptr, vtable)) + })) + } +} + +impl ArcWake03 for Current { + fn wake_by_ref(arc_self: &Arc<Self>) { + arc_self.0.notify(); + } +} + +fn with_context<T, R, F>(compat: &mut Compat<T>, f: F) -> R +where + T: Unpin, + F: FnOnce(Pin<&mut T>, &mut Context<'_>) -> R, +{ + let current = Current::new(); + let waker = current.as_waker(); + let mut cx = Context::from_waker(&waker); + f(Pin::new(&mut compat.inner), &mut cx) +} + +#[cfg(feature = "sink")] +fn with_sink_context<T, Item, R, F>(compat: &mut CompatSink<T, Item>, f: F) -> R +where + T: Unpin, + F: FnOnce(Pin<&mut T>, &mut Context<'_>) -> R, +{ + let current = Current::new(); + let waker = current.as_waker(); + let mut cx = Context::from_waker(&waker); + f(Pin::new(&mut compat.inner), &mut cx) +} + +#[cfg(feature = "io-compat")] +mod io { + use super::*; + use futures_io::{AsyncRead as AsyncRead03, AsyncWrite as AsyncWrite03}; + use tokio_io::{AsyncRead as AsyncRead01, AsyncWrite as AsyncWrite01}; + + fn poll_03_to_io<T>(x: task03::Poll<Result<T, std::io::Error>>) + -> Result<T, std::io::Error> + { + match x { + task03::Poll::Ready(Ok(t)) => Ok(t), + task03::Poll::Pending => Err(std::io::ErrorKind::WouldBlock.into()), + task03::Poll::Ready(Err(e)) => Err(e), + } + } + + impl<R: AsyncRead03 + Unpin> std::io::Read for Compat<R> { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> { + let current = Current::new(); + let waker = current.as_waker(); + let mut cx = Context::from_waker(&waker); + poll_03_to_io(Pin::new(&mut self.inner).poll_read(&mut cx, buf)) + } + } + + impl<R: AsyncRead03 + Unpin> AsyncRead01 for Compat<R> { + #[cfg(feature = "read-initializer")] + unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool { + let initializer = self.inner.initializer(); + let does_init = initializer.should_initialize(); + if does_init { + initializer.initialize(buf); + } + does_init + } + } + + impl<W: AsyncWrite03 + Unpin> std::io::Write for Compat<W> { + fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> { + let current = Current::new(); + let waker = current.as_waker(); + let mut cx = Context::from_waker(&waker); + poll_03_to_io(Pin::new(&mut self.inner).poll_write(&mut cx, buf)) + } + + fn flush(&mut self) -> std::io::Result<()> { + let current = Current::new(); + let waker = current.as_waker(); + let mut cx = Context::from_waker(&waker); + poll_03_to_io(Pin::new(&mut self.inner).poll_flush(&mut cx)) + } + } + + impl<W: AsyncWrite03 + Unpin> AsyncWrite01 for Compat<W> { + fn shutdown(&mut self) -> std::io::Result<Async01<()>> { + let current = Current::new(); + let waker = current.as_waker(); + let mut cx = Context::from_waker(&waker); + poll_03_to_01(Pin::new(&mut self.inner).poll_close(&mut cx)) + } + } +} diff --git a/src/compat/executor.rs b/src/compat/executor.rs new file mode 100644 index 0000000..82cb496 --- /dev/null +++ b/src/compat/executor.rs @@ -0,0 +1,87 @@ +use super::{Compat, Future01CompatExt}; +use crate::{ + future::{FutureExt, TryFutureExt, UnitError}, + task::SpawnExt, +}; +use futures_01::future::{ExecuteError as ExecuteError01, Executor as Executor01}; +use futures_01::Future as Future01; +use futures_task::{FutureObj, Spawn as Spawn03, SpawnError as SpawnError03}; + +/// A future that can run on a futures 0.1 +/// [`Executor`](futures_01::future::Executor). +pub type Executor01Future = Compat<UnitError<FutureObj<'static, ()>>>; + +/// Extension trait for futures 0.1 [`Executor`](futures_01::future::Executor). +pub trait Executor01CompatExt: Executor01<Executor01Future> + Clone + Send + 'static { + /// Converts a futures 0.1 [`Executor`](futures_01::future::Executor) into a + /// futures 0.3 [`Spawn`](futures_task::Spawn). + /// + /// ``` + /// use futures::task::SpawnExt; + /// use futures::future::{FutureExt, TryFutureExt}; + /// use futures_util::compat::Executor01CompatExt; + /// use tokio::executor::DefaultExecutor; + /// + /// # let (tx, rx) = futures::channel::oneshot::channel(); + /// + /// let spawner = DefaultExecutor::current().compat(); + /// let future03 = async move { + /// println!("Running on the pool"); + /// spawner.spawn(async { + /// println!("Spawned!"); + /// # tx.send(42).unwrap(); + /// }).unwrap(); + /// }; + /// + /// let future01 = future03.unit_error().boxed().compat(); + /// + /// tokio::run(future01); + /// # futures::executor::block_on(rx).unwrap(); + /// ``` + fn compat(self) -> Executor01As03<Self> + where + Self: Sized; +} + +impl<Ex> Executor01CompatExt for Ex +where + Ex: Executor01<Executor01Future> + Clone + Send + 'static, +{ + fn compat(self) -> Executor01As03<Self> { + Executor01As03 { executor01: self } + } +} + +/// Converts a futures 0.1 [`Executor`](futures_01::future::Executor) into a +/// futures 0.3 [`Spawn`](futures_task::Spawn). +#[derive(Debug, Clone)] +pub struct Executor01As03<Ex> { + executor01: Ex, +} + +impl<Ex> Spawn03 for Executor01As03<Ex> +where + Ex: Executor01<Executor01Future> + Clone + Send + 'static, +{ + fn spawn_obj(&self, future: FutureObj<'static, ()>) -> Result<(), SpawnError03> { + let future = future.unit_error().compat(); + + self.executor01 + .execute(future) + .map_err(|_| SpawnError03::shutdown()) + } +} + +#[allow(single_use_lifetimes)] // https://github.com/rust-lang/rust/issues/55058 +impl<Sp, Fut> Executor01<Fut> for Compat<Sp> +where + for<'a> &'a Sp: Spawn03, + Fut: Future01<Item = (), Error = ()> + Send + 'static, +{ + fn execute(&self, future: Fut) -> Result<(), ExecuteError01<Fut>> { + (&self.inner) + .spawn(future.compat().map(|_| ())) + .expect("unable to spawn future from Compat executor"); + Ok(()) + } +} diff --git a/src/compat/mod.rs b/src/compat/mod.rs new file mode 100644 index 0000000..1826836 --- /dev/null +++ b/src/compat/mod.rs @@ -0,0 +1,19 @@ +//! Futures 0.1 / 0.3 shims +//! +//! This module is only available when the `compat` feature of this +//! library is activated. + +mod executor; +pub use self::executor::{Executor01CompatExt, Executor01Future, Executor01As03}; + +mod compat01as03; +pub use self::compat01as03::{Compat01As03, Future01CompatExt, Stream01CompatExt}; +#[cfg(feature = "sink")] +pub use self::compat01as03::{Compat01As03Sink, Sink01CompatExt}; +#[cfg(feature = "io-compat")] +pub use self::compat01as03::{AsyncRead01CompatExt, AsyncWrite01CompatExt}; + +mod compat03as01; +pub use self::compat03as01::Compat; +#[cfg(feature = "sink")] +pub use self::compat03as01::CompatSink; diff --git a/src/future/abortable.rs b/src/future/abortable.rs new file mode 100644 index 0000000..281cf6b --- /dev/null +++ b/src/future/abortable.rs @@ -0,0 +1,177 @@ +use crate::task::AtomicWaker; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; +use core::fmt; +use core::pin::Pin; +use core::sync::atomic::{AtomicBool, Ordering}; +use alloc::sync::Arc; + +/// A future which can be remotely short-circuited using an `AbortHandle`. +#[derive(Debug, Clone)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Abortable<Fut> { + future: Fut, + inner: Arc<AbortInner>, +} + +impl<Fut: Unpin> Unpin for Abortable<Fut> {} + +impl<Fut> Abortable<Fut> where Fut: Future { + unsafe_pinned!(future: Fut); + + /// Creates a new `Abortable` future using an existing `AbortRegistration`. + /// `AbortRegistration`s can be acquired through `AbortHandle::new`. + /// + /// When `abort` is called on the handle tied to `reg` or if `abort` has + /// already been called, the future will complete immediately without making + /// any further progress. + /// + /// Example: + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::{Abortable, AbortHandle, Aborted}; + /// + /// let (abort_handle, abort_registration) = AbortHandle::new_pair(); + /// let future = Abortable::new(async { 2 }, abort_registration); + /// abort_handle.abort(); + /// assert_eq!(future.await, Err(Aborted)); + /// # }); + /// ``` + pub fn new(future: Fut, reg: AbortRegistration) -> Self { + Abortable { + future, + inner: reg.inner, + } + } +} + +/// A registration handle for a `Abortable` future. +/// Values of this type can be acquired from `AbortHandle::new` and are used +/// in calls to `Abortable::new`. +#[derive(Debug)] +pub struct AbortRegistration { + inner: Arc<AbortInner>, +} + +/// A handle to a `Abortable` future. +#[derive(Debug, Clone)] +pub struct AbortHandle { + inner: Arc<AbortInner>, +} + +impl AbortHandle { + /// Creates an (`AbortHandle`, `AbortRegistration`) pair which can be used + /// to abort a running future. + /// + /// This function is usually paired with a call to `Abortable::new`. + /// + /// Example: + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::{Abortable, AbortHandle, Aborted}; + /// + /// let (abort_handle, abort_registration) = AbortHandle::new_pair(); + /// let future = Abortable::new(async { 2 }, abort_registration); + /// abort_handle.abort(); + /// assert_eq!(future.await, Err(Aborted)); + /// # }); + /// ``` + pub fn new_pair() -> (Self, AbortRegistration) { + let inner = Arc::new(AbortInner { + waker: AtomicWaker::new(), + cancel: AtomicBool::new(false), + }); + + ( + AbortHandle { + inner: inner.clone(), + }, + AbortRegistration { + inner, + }, + ) + } +} + +// Inner type storing the waker to awaken and a bool indicating that it +// should be cancelled. +#[derive(Debug)] +struct AbortInner { + waker: AtomicWaker, + cancel: AtomicBool, +} + +/// Creates a new `Abortable` future and a `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<Fut>(future: Fut) -> (Abortable<Fut>, AbortHandle) + where Fut: Future +{ + let (handle, reg) = AbortHandle::new_pair(); + ( + Abortable::new(future, reg), + handle, + ) +} + +/// Indicator that the `Abortable` future was aborted. +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub struct Aborted; + +impl fmt::Display for Aborted { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "`Abortable` future has been aborted") + } +} + +#[cfg(feature = "std")] +impl std::error::Error for Aborted {} + +impl<Fut> Future for Abortable<Fut> where Fut: Future { + type Output = Result<Fut::Output, Aborted>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + // Check if the future has been aborted + if self.inner.cancel.load(Ordering::Relaxed) { + return Poll::Ready(Err(Aborted)) + } + + // attempt to complete the future + if let Poll::Ready(x) = self.as_mut().future().poll(cx) { + return Poll::Ready(Ok(x)) + } + + // Register to receive a wakeup if the future is aborted in the... future + self.inner.waker.register(cx.waker()); + + // Check to see if the future was aborted between the first check and + // registration. + // Checking with `Relaxed` is sufficient because `register` introduces an + // `AcqRel` barrier. + if self.inner.cancel.load(Ordering::Relaxed) { + return Poll::Ready(Err(Aborted)) + } + + Poll::Pending + } +} + +impl AbortHandle { + /// Abort the `Abortable` future associated with this handle. + /// + /// Notifies the Abortable future associated with this handle that it + /// should abort. Note that if the future is currently being polled on + /// another thread, it will not immediately stop running. Instead, it will + /// continue to run until its poll method returns. + pub fn abort(&self) { + self.inner.cancel.store(true, Ordering::Relaxed); + self.inner.waker.wake(); + } +} diff --git a/src/future/either.rs b/src/future/either.rs new file mode 100644 index 0000000..24fbbe7 --- /dev/null +++ b/src/future/either.rs @@ -0,0 +1,304 @@ +use core::pin::Pin; +use core::task::{Context, Poll}; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::{FusedStream, Stream}; +#[cfg(feature = "sink")] +use futures_sink::Sink; + +/// Combines two different futures, streams, or sinks having the same associated types into a single +/// type. +#[derive(Debug, Clone)] +pub enum Either<A, B> { + /// First branch of the type + Left(A), + /// Second branch of the type + Right(B), +} + +impl<A, B, T> Either<(T, A), (T, B)> { + /// Factor out a homogeneous type from an either of pairs. + /// + /// Here, the homogeneous type is the first element of the pairs. + pub fn factor_first(self) -> (T, Either<A, B>) { + match self { + Either::Left((x, a)) => (x, Either::Left(a)), + Either::Right((x, b)) => (x, Either::Right(b)), + } + } +} + +impl<A, B, T> Either<(A, T), (B, T)> { + /// Factor out a homogeneous type from an either of pairs. + /// + /// Here, the homogeneous type is the second element of the pairs. + pub fn factor_second(self) -> (Either<A, B>, T) { + match self { + Either::Left((a, x)) => (Either::Left(a), x), + Either::Right((b, x)) => (Either::Right(b), x), + } + } +} + +impl<T> Either<T, T> { + /// Extract the value of an either over two equivalent types. + pub fn into_inner(self) -> T { + match self { + Either::Left(x) => x, + Either::Right(x) => x, + } + } +} + +impl<A, B> Future for Either<A, B> +where + A: Future, + B: Future<Output = A::Output>, +{ + type Output = A::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<A::Output> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll(cx), + Either::Right(x) => Pin::new_unchecked(x).poll(cx), + } + } + } +} + +impl<A, B> FusedFuture for Either<A, B> +where + A: FusedFuture, + B: FusedFuture<Output = A::Output>, +{ + fn is_terminated(&self) -> bool { + match self { + Either::Left(x) => x.is_terminated(), + Either::Right(x) => x.is_terminated(), + } + } +} + +impl<A, B> Stream for Either<A, B> +where + A: Stream, + B: Stream<Item = A::Item>, +{ + type Item = A::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<A::Item>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_next(cx), + Either::Right(x) => Pin::new_unchecked(x).poll_next(cx), + } + } + } +} + +impl<A, B> FusedStream for Either<A, B> +where + A: FusedStream, + B: FusedStream<Item = A::Item>, +{ + fn is_terminated(&self) -> bool { + match self { + Either::Left(x) => x.is_terminated(), + Either::Right(x) => x.is_terminated(), + } + } +} + +#[cfg(feature = "sink")] +impl<A, B, Item> Sink<Item> for Either<A, B> +where + A: Sink<Item>, + B: Sink<Item, Error = A::Error>, +{ + type Error = A::Error; + + fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_ready(cx), + Either::Right(x) => Pin::new_unchecked(x).poll_ready(cx), + } + } + } + + fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).start_send(item), + Either::Right(x) => Pin::new_unchecked(x).start_send(item), + } + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_flush(cx), + Either::Right(x) => Pin::new_unchecked(x).poll_flush(cx), + } + } + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_close(cx), + Either::Right(x) => Pin::new_unchecked(x).poll_close(cx), + } + } + } +} + +#[cfg(feature = "io")] +#[cfg(feature = "std")] +mod if_std { + use super::Either; + use core::pin::Pin; + use core::task::{Context, Poll}; + #[cfg(feature = "read-initializer")] + use futures_io::Initializer; + use futures_io::{ + AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, IoSlice, IoSliceMut, Result, SeekFrom, + }; + + impl<A, B> AsyncRead for Either<A, B> + where + A: AsyncRead, + B: AsyncRead, + { + #[cfg(feature = "read-initializer")] + unsafe fn initializer(&self) -> Initializer { + match self { + Either::Left(x) => x.initializer(), + Either::Right(x) => x.initializer(), + } + } + + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll<Result<usize>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_read(cx, buf), + Either::Right(x) => Pin::new_unchecked(x).poll_read(cx, buf), + } + } + } + + fn poll_read_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &mut [IoSliceMut<'_>], + ) -> Poll<Result<usize>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_read_vectored(cx, bufs), + Either::Right(x) => Pin::new_unchecked(x).poll_read_vectored(cx, bufs), + } + } + } + } + + impl<A, B> AsyncWrite for Either<A, B> + where + A: AsyncWrite, + B: AsyncWrite, + { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<Result<usize>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_write(cx, buf), + Either::Right(x) => Pin::new_unchecked(x).poll_write(cx, buf), + } + } + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<Result<usize>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_write_vectored(cx, bufs), + Either::Right(x) => Pin::new_unchecked(x).poll_write_vectored(cx, bufs), + } + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_flush(cx), + Either::Right(x) => Pin::new_unchecked(x).poll_flush(cx), + } + } + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_close(cx), + Either::Right(x) => Pin::new_unchecked(x).poll_close(cx), + } + } + } + } + + impl<A, B> AsyncSeek for Either<A, B> + where + A: AsyncSeek, + B: AsyncSeek, + { + fn poll_seek( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + pos: SeekFrom, + ) -> Poll<Result<u64>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_seek(cx, pos), + Either::Right(x) => Pin::new_unchecked(x).poll_seek(cx, pos), + } + } + } + } + + impl<A, B> AsyncBufRead for Either<A, B> + where + A: AsyncBufRead, + B: AsyncBufRead, + { + fn poll_fill_buf( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<&[u8]>> { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).poll_fill_buf(cx), + Either::Right(x) => Pin::new_unchecked(x).poll_fill_buf(cx), + } + } + } + + fn consume(self: Pin<&mut Self>, amt: usize) { + unsafe { + match self.get_unchecked_mut() { + Either::Left(x) => Pin::new_unchecked(x).consume(amt), + Either::Right(x) => Pin::new_unchecked(x).consume(amt), + } + } + } + } +} diff --git a/src/future/future/catch_unwind.rs b/src/future/future/catch_unwind.rs new file mode 100644 index 0000000..e88cce7 --- /dev/null +++ b/src/future/future/catch_unwind.rs @@ -0,0 +1,31 @@ +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; +use std::any::Any; +use std::pin::Pin; +use std::panic::{catch_unwind, UnwindSafe, AssertUnwindSafe}; + +/// Future for the [`catch_unwind`](super::FutureExt::catch_unwind) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct CatchUnwind<Fut> { + future: Fut, +} + +impl<Fut> CatchUnwind<Fut> where Fut: Future + UnwindSafe { + unsafe_pinned!(future: Fut); + + pub(super) fn new(future: Fut) -> CatchUnwind<Fut> { + CatchUnwind { future } + } +} + +impl<Fut> Future for CatchUnwind<Fut> + where Fut: Future + UnwindSafe, +{ + type Output = Result<Fut::Output, Box<dyn Any + Send>>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + catch_unwind(AssertUnwindSafe(|| self.future().poll(cx)))?.map(Ok) + } +} diff --git a/src/future/future/chain.rs b/src/future/future/chain.rs new file mode 100644 index 0000000..3f248e8 --- /dev/null +++ b/src/future/future/chain.rs @@ -0,0 +1,58 @@ +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; + +#[must_use = "futures do nothing unless you `.await` or poll them"] +#[derive(Debug)] +pub(crate) enum Chain<Fut1, Fut2, Data> { + First(Fut1, Option<Data>), + Second(Fut2), + Empty, +} + +impl<Fut1: Unpin, Fut2: Unpin, Data> Unpin for Chain<Fut1, Fut2, Data> {} + +impl<Fut1, Fut2, Data> Chain<Fut1, Fut2, Data> { + pub(crate)fn is_terminated(&self) -> bool { + if let Chain::Empty = *self { true } else { false } + } +} + +impl<Fut1, Fut2, Data> Chain<Fut1, Fut2, Data> + where Fut1: Future, + Fut2: Future, +{ + pub(crate) fn new(fut1: Fut1, data: Data) -> Chain<Fut1, Fut2, Data> { + Chain::First(fut1, Some(data)) + } + + pub(crate) fn poll<F>( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + f: F, + ) -> Poll<Fut2::Output> + where F: FnOnce(Fut1::Output, Data) -> Fut2, + { + let mut f = Some(f); + + // Safe to call `get_unchecked_mut` because we won't move the futures. + let this = unsafe { self.get_unchecked_mut() }; + + loop { + let (output, data) = match this { + Chain::First(fut1, data) => { + let output = ready!(unsafe { Pin::new_unchecked(fut1) }.poll(cx)); + (output, data.take().unwrap()) + } + Chain::Second(fut2) => { + return unsafe { Pin::new_unchecked(fut2) }.poll(cx); + } + Chain::Empty => unreachable!() + }; + + *this = Chain::Empty; // Drop fut1 + let fut2 = (f.take().unwrap())(output, data); + *this = Chain::Second(fut2) + } + } +} diff --git a/src/future/future/flatten.rs b/src/future/future/flatten.rs new file mode 100644 index 0000000..16b3a19 --- /dev/null +++ b/src/future/future/flatten.rs @@ -0,0 +1,56 @@ +use super::chain::Chain; +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Future for the [`flatten`](super::FutureExt::flatten) method. +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Flatten<Fut> + where Fut: Future, +{ + state: Chain<Fut, Fut::Output, ()>, +} + +impl<Fut> Flatten<Fut> + where Fut: Future, + Fut::Output: Future, +{ + unsafe_pinned!(state: Chain<Fut, Fut::Output, ()>); + + pub(super) fn new(future: Fut) -> Flatten<Fut> { + Flatten { + state: Chain::new(future, ()), + } + } +} + +impl<Fut> fmt::Debug for Flatten<Fut> + where Fut: Future + fmt::Debug, + Fut::Output: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Flatten") + .field("state", &self.state) + .finish() + } +} + +impl<Fut> FusedFuture for Flatten<Fut> + where Fut: Future, + Fut::Output: Future, +{ + fn is_terminated(&self) -> bool { self.state.is_terminated() } +} + +impl<Fut> Future for Flatten<Fut> + where Fut: Future, + Fut::Output: Future, +{ + type Output = <Fut::Output as Future>::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + self.state().poll(cx, |a, ()| a) + } +} diff --git a/src/future/future/flatten_stream.rs b/src/future/future/flatten_stream.rs new file mode 100644 index 0000000..d110886 --- /dev/null +++ b/src/future/future/flatten_stream.rs @@ -0,0 +1,89 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Stream for the [`flatten_stream`](super::FutureExt::flatten_stream) method. +#[must_use = "streams do nothing unless polled"] +pub struct FlattenStream<Fut: Future> { + state: State<Fut, Fut::Output>, +} + +impl<Fut: Future> FlattenStream<Fut> { + unsafe_pinned!(state: State<Fut, Fut::Output>); + + pub(super) fn new(future: Fut) -> FlattenStream<Fut> { + FlattenStream { + state: State::Future(future) + } + } +} + +impl<Fut> fmt::Debug for FlattenStream<Fut> + where Fut: Future + fmt::Debug, + Fut::Output: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FlattenStream") + .field("state", &self.state) + .finish() + } +} + +#[derive(Debug)] +enum State<Fut, St> { + // future is not yet called or called and not ready + Future(Fut), + // future resolved to Stream + Stream(St), +} + +impl<Fut, St> State<Fut, St> { + fn get_pin_mut(self: Pin<&mut Self>) -> State<Pin<&mut Fut>, Pin<&mut St>> { + // safety: data is never moved via the resulting &mut reference + match unsafe { self.get_unchecked_mut() } { + // safety: the future we're re-pinning here will never be moved; + // it will just be polled, then dropped in place + State::Future(f) => State::Future(unsafe { Pin::new_unchecked(f) }), + // safety: the stream we're repinning here will never be moved; + // it will just be polled, then dropped in place + State::Stream(s) => State::Stream(unsafe { Pin::new_unchecked(s) }), + } + } +} + +impl<Fut> FusedStream for FlattenStream<Fut> + where Fut: Future, + Fut::Output: Stream + FusedStream, +{ + fn is_terminated(&self) -> bool { + match &self.state { + State::Future(_) => false, + State::Stream(stream) => stream.is_terminated(), + } + } +} + +impl<Fut> Stream for FlattenStream<Fut> + where Fut: Future, + Fut::Output: Stream, +{ + type Item = <Fut::Output as Stream>::Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + loop { + match self.as_mut().state().get_pin_mut() { + State::Future(f) => { + let stream = ready!(f.poll(cx)); + // Future resolved to stream. + // We do not return, but poll that + // stream in the next loop iteration. + self.as_mut().state().set(State::Stream(stream)); + } + State::Stream(s) => return s.poll_next(cx), + } + } + } +} diff --git a/src/future/future/fuse.rs b/src/future/future/fuse.rs new file mode 100644 index 0000000..b5ef913 --- /dev/null +++ b/src/future/future/fuse.rs @@ -0,0 +1,90 @@ +use core::pin::Pin; +use futures_core::future::{Future, FusedFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Future for the [`fuse`](super::FutureExt::fuse) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Fuse<Fut> { + future: Option<Fut>, +} + +impl<Fut: Future> Fuse<Fut> { + unsafe_pinned!(future: Option<Fut>); + + pub(super) fn new(f: Fut) -> Fuse<Fut> { + Fuse { + future: Some(f), + } + } + + /// Creates a new `Fuse`-wrapped future which is already terminated. + /// + /// This can be useful in combination with looping and the `select!` + /// macro, which bypasses terminated futures. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::future::{Fuse, FusedFuture, FutureExt}; + /// use futures::select; + /// use futures::stream::StreamExt; + /// use futures::pin_mut; + /// + /// let (sender, mut stream) = mpsc::unbounded(); + /// + /// // Send a few messages into the stream + /// sender.unbounded_send(()).unwrap(); + /// sender.unbounded_send(()).unwrap(); + /// drop(sender); + /// + /// // Use `Fuse::termianted()` to create an already-terminated future + /// // which may be instantiated later. + /// let foo_printer = Fuse::terminated(); + /// pin_mut!(foo_printer); + /// + /// loop { + /// select! { + /// _ = foo_printer => {}, + /// () = stream.select_next_some() => { + /// if !foo_printer.is_terminated() { + /// println!("Foo is already being printed!"); + /// } else { + /// foo_printer.set(async { + /// // do some other async operations + /// println!("Printing foo from `foo_printer` future"); + /// }.fuse()); + /// } + /// }, + /// complete => break, // `foo_printer` is terminated and the stream is done + /// } + /// } + /// # }); + /// ``` + pub fn terminated() -> Fuse<Fut> { + Fuse { future: None } + } +} + +impl<Fut: Future> FusedFuture for Fuse<Fut> { + fn is_terminated(&self) -> bool { + self.future.is_none() + } +} + +impl<Fut: Future> Future for Fuse<Fut> { + type Output = Fut::Output; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Fut::Output> { + let v = match self.as_mut().future().as_pin_mut() { + Some(fut) => ready!(fut.poll(cx)), + None => return Poll::Pending, + }; + + self.as_mut().future().set(None); + Poll::Ready(v) + } +} diff --git a/src/future/future/inspect.rs b/src/future/future/inspect.rs new file mode 100644 index 0000000..d67455a --- /dev/null +++ b/src/future/future/inspect.rs @@ -0,0 +1,47 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Future for the [`inspect`](super::FutureExt::inspect) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Inspect<Fut, F> { + future: Fut, + f: Option<F>, +} + +impl<Fut: Future, F: FnOnce(&Fut::Output)> Inspect<Fut, F> { + unsafe_pinned!(future: Fut); + unsafe_unpinned!(f: Option<F>); + + pub(super) fn new(future: Fut, f: F) -> Inspect<Fut, F> { + Inspect { + future, + f: Some(f), + } + } +} + +impl<Fut: Future + Unpin, F> Unpin for Inspect<Fut, F> {} + +impl<Fut, F> FusedFuture for Inspect<Fut, F> + where Fut: FusedFuture, + F: FnOnce(&Fut::Output), +{ + fn is_terminated(&self) -> bool { self.future.is_terminated() } +} + +impl<Fut, F> Future for Inspect<Fut, F> + where Fut: Future, + F: FnOnce(&Fut::Output), +{ + type Output = Fut::Output; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Fut::Output> { + let e = ready!(self.as_mut().future().poll(cx)); + let f = self.as_mut().f().take().expect("cannot poll Inspect twice"); + f(&e); + Poll::Ready(e) + } +} diff --git a/src/future/future/into_stream.rs b/src/future/future/into_stream.rs new file mode 100644 index 0000000..616c4cb --- /dev/null +++ b/src/future/future/into_stream.rs @@ -0,0 +1,43 @@ +use crate::stream::{self, Once}; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Stream for the [`into_stream`](super::FutureExt::into_stream) method. +#[must_use = "streams do nothing unless polled"] +#[derive(Debug)] +pub struct IntoStream<Fut> { + inner: Once<Fut> +} + +impl<Fut: Future> IntoStream<Fut> { + unsafe_pinned!(inner: Once<Fut>); + + pub(super) fn new(future: Fut) -> IntoStream<Fut> { + IntoStream { + inner: stream::once(future) + } + } +} + +impl<Fut: Future> Stream for IntoStream<Fut> { + type Item = Fut::Output; + + #[inline] + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + self.inner().poll_next(cx) + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + self.inner.size_hint() + } +} + +impl<Fut: Future> FusedStream for IntoStream<Fut> { + fn is_terminated(&self) -> bool { + self.inner.is_terminated() + } +} diff --git a/src/future/future/map.rs b/src/future/future/map.rs new file mode 100644 index 0000000..b5fbfb1 --- /dev/null +++ b/src/future/future/map.rs @@ -0,0 +1,49 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Future for the [`map`](super::FutureExt::map) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Map<Fut, F> { + future: Fut, + f: Option<F>, +} + +impl<Fut, F> Map<Fut, F> { + unsafe_pinned!(future: Fut); + unsafe_unpinned!(f: Option<F>); + + /// Creates a new Map. + pub(super) fn new(future: Fut, f: F) -> Map<Fut, F> { + Map { future, f: Some(f) } + } +} + +impl<Fut: Unpin, F> Unpin for Map<Fut, F> {} + +impl<Fut, F, T> FusedFuture for Map<Fut, F> + where Fut: Future, + F: FnOnce(Fut::Output) -> T, +{ + fn is_terminated(&self) -> bool { self.f.is_none() } +} + +impl<Fut, F, T> Future for Map<Fut, F> + where Fut: Future, + F: FnOnce(Fut::Output) -> T, +{ + type Output = T; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { + self.as_mut() + .future() + .poll(cx) + .map(|output| { + let f = self.f().take() + .expect("Map must not be polled after it returned `Poll::Ready`"); + f(output) + }) + } +} diff --git a/src/future/future/mod.rs b/src/future/future/mod.rs new file mode 100644 index 0000000..e58cafc --- /dev/null +++ b/src/future/future/mod.rs @@ -0,0 +1,558 @@ +//! Futures +//! +//! This module contains a number of functions for working with `Future`s, +//! including the `FutureExt` trait which adds methods to `Future` types. + +use super::{assert_future, Either}; +#[cfg(feature = "alloc")] +use alloc::boxed::Box; +use core::pin::Pin; +#[cfg(feature = "alloc")] +use futures_core::future::{BoxFuture, LocalBoxFuture}; +use futures_core::{ + future::Future, + stream::Stream, + task::{Context, Poll}, +}; + +// Combinators + +mod flatten; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::flatten::Flatten; + +mod flatten_stream; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::flatten_stream::FlattenStream; + +mod fuse; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::fuse::Fuse; + +mod into_stream; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::into_stream::IntoStream; + +mod map; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::map::Map; + +mod then; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::then::Then; + +mod inspect; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::inspect::Inspect; + +mod unit_error; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::unit_error::UnitError; + +mod never_error; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::never_error::NeverError; + +#[cfg(feature = "std")] +mod catch_unwind; +#[cfg(feature = "std")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::catch_unwind::CatchUnwind; + +#[cfg(feature = "channel")] +#[cfg(feature = "std")] +mod remote_handle; +#[cfg(feature = "channel")] +#[cfg(feature = "std")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::remote_handle::{Remote, RemoteHandle}; + +#[cfg(feature = "std")] +mod shared; +#[cfg(feature = "std")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::shared::Shared; + +// Implementation details + +mod chain; +pub(crate) use self::chain::Chain; + +impl<T: ?Sized> FutureExt for T where T: Future {} + +/// An extension trait for `Future`s that provides a variety of convenient +/// adapters. +pub trait FutureExt: Future { + /// Map this future's output to a different type, returning a new future of + /// the resulting type. + /// + /// This function is similar to the `Option::map` or `Iterator::map` where + /// it will change the type of the underlying future. This is useful to + /// chain along a computation once a future has been resolved. + /// + /// Note that this function consumes the receiving future and returns a + /// wrapped version of it, similar to the existing `map` methods in the + /// standard library. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// + /// let future = async { 1 }; + /// let new_future = future.map(|x| x + 3); + /// assert_eq!(new_future.await, 4); + /// # }); + /// ``` + fn map<U, F>(self, f: F) -> Map<Self, F> + where + F: FnOnce(Self::Output) -> U, + Self: Sized, + { + assert_future::<U, _>(Map::new(self, f)) + } + + /// Chain on a computation for when a future finished, passing the result of + /// the future to the provided closure `f`. + /// + /// The returned value of the closure must implement the `Future` trait + /// and can represent some more work to be done before the composed future + /// is finished. + /// + /// The closure `f` is only run *after* successful completion of the `self` + /// future. + /// + /// Note that this function consumes the receiving future and returns a + /// wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// + /// let future_of_1 = async { 1 }; + /// let future_of_4 = future_of_1.then(|x| async move { x + 3 }); + /// assert_eq!(future_of_4.await, 4); + /// # }); + /// ``` + fn then<Fut, F>(self, f: F) -> Then<Self, Fut, F> + where + F: FnOnce(Self::Output) -> Fut, + Fut: Future, + Self: Sized, + { + assert_future::<Fut::Output, _>(Then::new(self, f)) + } + + /// Wrap this future in an `Either` future, making it the left-hand variant + /// of that `Either`. + /// + /// This can be used in combination with the `right_future` method to write `if` + /// statements that evaluate to different futures in different branches. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// + /// let x = 6; + /// let future = if x < 10 { + /// async { true }.left_future() + /// } else { + /// async { false }.right_future() + /// }; + /// + /// assert_eq!(future.await, true); + /// # }); + /// ``` + fn left_future<B>(self) -> Either<Self, B> + where + B: Future<Output = Self::Output>, + Self: Sized, + { + Either::Left(self) + } + + /// Wrap this future in an `Either` future, making it the right-hand variant + /// of that `Either`. + /// + /// This can be used in combination with the `left_future` method to write `if` + /// statements that evaluate to different futures in different branches. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// + /// let x = 6; + /// let future = if x > 10 { + /// async { true }.left_future() + /// } else { + /// async { false }.right_future() + /// }; + /// + /// assert_eq!(future.await, false); + /// # }); + /// ``` + fn right_future<A>(self) -> Either<A, Self> + where + A: Future<Output = Self::Output>, + Self: Sized, + { + Either::Right(self) + } + + /// Convert this future into a single element stream. + /// + /// The returned stream contains single success if this future resolves to + /// success or single error if this future resolves into error. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// use futures::stream::StreamExt; + /// + /// let future = async { 17 }; + /// let stream = future.into_stream(); + /// let collected: Vec<_> = stream.collect().await; + /// assert_eq!(collected, vec![17]); + /// # }); + /// ``` + fn into_stream(self) -> IntoStream<Self> + where + Self: Sized, + { + IntoStream::new(self) + } + + /// Flatten the execution of this future when the output of this + /// future is itself another future. + /// + /// This can be useful when combining futures together to flatten the + /// computation out the final result. + /// + /// This method is roughly equivalent to `self.then(|x| x)`. + /// + /// Note that this function consumes the receiving future and returns a + /// wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// + /// let nested_future = async { async { 1 } }; + /// let future = nested_future.flatten(); + /// assert_eq!(future.await, 1); + /// # }); + /// ``` + fn flatten(self) -> Flatten<Self> + where + Self::Output: Future, + Self: Sized, + { + let f = Flatten::new(self); + assert_future::<<<Self as Future>::Output as Future>::Output, _>(f) + } + + /// Flatten the execution of this future when the successful result of this + /// future is a stream. + /// + /// This can be useful when stream initialization is deferred, and it is + /// convenient to work with that stream as if stream was available at the + /// call site. + /// + /// Note that this function consumes this future and returns a wrapped + /// version of it. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// use futures::stream::{self, StreamExt}; + /// + /// let stream_items = vec![17, 18, 19]; + /// let future_of_a_stream = async { stream::iter(stream_items) }; + /// + /// let stream = future_of_a_stream.flatten_stream(); + /// let list: Vec<_> = stream.collect().await; + /// assert_eq!(list, vec![17, 18, 19]); + /// # }); + /// ``` + fn flatten_stream(self) -> FlattenStream<Self> + where + Self::Output: Stream, + Self: Sized, + { + FlattenStream::new(self) + } + + /// Fuse a future such that `poll` will never again be called once it has + /// completed. This method can be used to turn any `Future` into a + /// `FusedFuture`. + /// + /// Normally, once a future has returned `Poll::Ready` from `poll`, + /// any further calls could exhibit bad behavior such as blocking + /// forever, panicking, never returning, etc. If it is known that `poll` + /// may be called too often then this method can be used to ensure that it + /// has defined semantics. + /// + /// If a `fuse`d future is `poll`ed after having returned `Poll::Ready` + /// previously, it will return `Poll::Pending`, from `poll` again (and will + /// continue to do so for all future calls to `poll`). + /// + /// This combinator will drop the underlying future as soon as it has been + /// completed to ensure resources are reclaimed as soon as possible. + fn fuse(self) -> Fuse<Self> + where + Self: Sized, + { + let f = Fuse::new(self); + assert_future::<Self::Output, _>(f) + } + + /// Do something with the output of a future before passing it on. + /// + /// When using futures, you'll often chain several of them together. While + /// working on such code, you might want to check out what's happening at + /// various parts in the pipeline, without consuming the intermediate + /// value. To do that, insert a call to `inspect`. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// + /// let future = async { 1 }; + /// let new_future = future.inspect(|&x| println!("about to resolve: {}", x)); + /// assert_eq!(new_future.await, 1); + /// # }); + /// ``` + fn inspect<F>(self, f: F) -> Inspect<Self, F> + where + F: FnOnce(&Self::Output), + Self: Sized, + { + assert_future::<Self::Output, _>(Inspect::new(self, f)) + } + + /// Catches unwinding panics while polling the future. + /// + /// In general, panics within a future can propagate all the way out to the + /// task level. This combinator makes it possible to halt unwinding within + /// the future itself. It's most commonly used within task executors. It's + /// not recommended to use this for error handling. + /// + /// Note that this method requires the `UnwindSafe` bound from the standard + /// library. This isn't always applied automatically, and the standard + /// library provides an `AssertUnwindSafe` wrapper type to apply it + /// after-the fact. To assist using this method, the `Future` trait is also + /// implemented for `AssertUnwindSafe<F>` where `F` implements `Future`. + /// + /// This method is only available when the `std` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::{self, FutureExt, Ready}; + /// + /// let future = future::ready(2); + /// assert!(future.catch_unwind().await.is_ok()); + /// + /// let future = future::lazy(|_| -> Ready<i32> { + /// unimplemented!() + /// }); + /// assert!(future.catch_unwind().await.is_err()); + /// # }); + /// ``` + #[cfg(feature = "std")] + fn catch_unwind(self) -> CatchUnwind<Self> + where + Self: Sized + ::std::panic::UnwindSafe, + { + CatchUnwind::new(self) + } + + /// Create a cloneable handle to this future where all handles will resolve + /// to the same result. + /// + /// The `shared` combinator method provides a method to convert any future + /// into a cloneable future. It enables a future to be polled by multiple + /// threads. + /// + /// This method is only available when the `std` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// + /// let future = async { 6 }; + /// let shared1 = future.shared(); + /// let shared2 = shared1.clone(); + /// + /// assert_eq!(6, shared1.await); + /// assert_eq!(6, shared2.await); + /// # }); + /// ``` + /// + /// ``` + /// // Note, unlike most examples this is written in the context of a + /// // synchronous function to better illustrate the cross-thread aspect of + /// // the `shared` combinator. + /// + /// # futures::executor::block_on(async { + /// use futures::future::FutureExt; + /// use futures::executor::block_on; + /// use std::thread; + /// + /// let future = async { 6 }; + /// let shared1 = future.shared(); + /// let shared2 = shared1.clone(); + /// let join_handle = thread::spawn(move || { + /// assert_eq!(6, block_on(shared2)); + /// }); + /// assert_eq!(6, shared1.await); + /// join_handle.join().unwrap(); + /// # }); + /// ``` + #[cfg(feature = "std")] + fn shared(self) -> Shared<Self> + where + Self: Sized, + Self::Output: Clone, + { + Shared::new(self) + } + + /// Turn this future into a future that yields `()` on completion and sends + /// its output to another future on a separate task. + /// + /// This can be used with spawning executors to easily retrieve the result + /// of a future executing on a separate task or thread. + /// + /// This method is only available when the `std` feature of this + /// library is activated, and it is activated by default. + #[cfg(feature = "channel")] + #[cfg(feature = "std")] + fn remote_handle(self) -> (Remote<Self>, RemoteHandle<Self::Output>) + where + Self: Sized, + { + remote_handle::remote_handle(self) + } + + /// Wrap the future in a Box, pinning it. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg(feature = "alloc")] + fn boxed<'a>(self) -> BoxFuture<'a, Self::Output> + where + Self: Sized + Send + 'a, + { + Box::pin(self) + } + + /// Wrap the future in a Box, pinning it. + /// + /// Similar to `boxed`, but without the `Send` requirement. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg(feature = "alloc")] + fn boxed_local<'a>(self) -> LocalBoxFuture<'a, Self::Output> + where + Self: Sized + 'a, + { + Box::pin(self) + } + + /// Turns a [`Future<Output = T>`](Future) into a + /// [`TryFuture<Ok = T, Error = ()`>](futures_core::future::TryFuture). + fn unit_error(self) -> UnitError<Self> + where + Self: Sized, + { + UnitError::new(self) + } + + /// Turns a [`Future<Output = T>`](Future) into a + /// [`TryFuture<Ok = T, Error = Never`>](futures_core::future::TryFuture). + fn never_error(self) -> NeverError<Self> + where + Self: Sized, + { + NeverError::new(self) + } + + /// A convenience for calling `Future::poll` on `Unpin` future types. + fn poll_unpin(&mut self, cx: &mut Context<'_>) -> Poll<Self::Output> + where + Self: Unpin, + { + Pin::new(self).poll(cx) + } + + /// Evaluates and consumes the future, returning the resulting output if + /// the future is ready after the first call to `Future::poll`. + /// + /// If `poll` instead returns `Poll::Pending`, `None` is returned. + /// + /// This method is useful in cases where immediacy is more important than + /// waiting for a result. It is also convenient for quickly obtaining + /// the value of a future that is known to always resolve immediately. + /// + /// # Examples + /// + /// ``` + /// # use futures::prelude::*; + /// use futures::{future::ready, future::pending}; + /// let future_ready = ready("foobar"); + /// let future_pending = pending::<&'static str>(); + /// + /// assert_eq!(future_ready.now_or_never(), Some("foobar")); + /// assert_eq!(future_pending.now_or_never(), None); + /// ``` + /// + /// In cases where it is absolutely known that a future should always + /// resolve immediately and never return `Poll::Pending`, this method can + /// be combined with `expect()`: + /// + /// ``` + /// # use futures::{prelude::*, future::ready}; + /// let future_ready = ready("foobar"); + /// + /// assert_eq!(future_ready.now_or_never().expect("Future not ready"), "foobar"); + /// ``` + fn now_or_never(mut self) -> Option<Self::Output> + where + Self: Sized, + { + let noop_waker = crate::task::noop_waker(); + let mut cx = Context::from_waker(&noop_waker); + + // SAFETY: This is safe because this method consumes the future, so `poll` is + // only going to be called once. Thus it doesn't matter to us if the + // future is `Unpin` or not. + let pinned = unsafe { Pin::new_unchecked(&mut self) }; + + match pinned.poll(&mut cx) { + Poll::Ready(x) => Some(x), + _ => None, + } + } +} diff --git a/src/future/future/never_error.rs b/src/future/future/never_error.rs new file mode 100644 index 0000000..5a68e6f --- /dev/null +++ b/src/future/future/never_error.rs @@ -0,0 +1,36 @@ +use crate::never::Never; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{self, Poll}; +use pin_utils::unsafe_pinned; + +/// Future for the [`never_error`](super::FutureExt::never_error) combinator. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct NeverError<Fut> { + future: Fut, +} + +impl<Fut> NeverError<Fut> { + unsafe_pinned!(future: Fut); + + pub(super) fn new(future: Fut) -> NeverError<Fut> { + NeverError { future } + } +} + +impl<Fut: Unpin> Unpin for NeverError<Fut> {} + +impl<Fut: FusedFuture> FusedFuture for NeverError<Fut> { + fn is_terminated(&self) -> bool { self.future.is_terminated() } +} + +impl<Fut, T> Future for NeverError<Fut> + where Fut: Future<Output = T>, +{ + type Output = Result<T, Never>; + + fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { + self.future().poll(cx).map(Ok) + } +} diff --git a/src/future/future/remote_handle.rs b/src/future/future/remote_handle.rs new file mode 100644 index 0000000..11b2a65 --- /dev/null +++ b/src/future/future/remote_handle.rs @@ -0,0 +1,114 @@ +use { + crate::future::{CatchUnwind, FutureExt}, + futures_channel::oneshot::{self, Sender, Receiver}, + futures_core::{ + future::Future, + task::{Context, Poll}, + }, + pin_utils::{unsafe_pinned, unsafe_unpinned}, + std::{ + any::Any, + fmt, + panic::{self, AssertUnwindSafe}, + pin::Pin, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + thread, + }, +}; + +/// The handle to a remote future returned by +/// [`remote_handle`](crate::future::FutureExt::remote_handle). When you drop this, +/// the remote future will be woken up to be dropped by the executor. +#[must_use = "futures do nothing unless you `.await` or poll them"] +#[derive(Debug)] +pub struct RemoteHandle<T> { + rx: Receiver<thread::Result<T>>, + keep_running: Arc<AtomicBool>, +} + +impl<T> RemoteHandle<T> { + /// Drops this handle *without* canceling the underlying future. + /// + /// This method can be used if you want to drop the handle, but let the + /// execution continue. + pub fn forget(self) { + self.keep_running.store(true, Ordering::SeqCst); + } +} + +impl<T: Send + 'static> Future for RemoteHandle<T> { + type Output = T; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { + match ready!(self.rx.poll_unpin(cx)) { + Ok(Ok(output)) => Poll::Ready(output), + Ok(Err(e)) => panic::resume_unwind(e), + Err(e) => panic::resume_unwind(Box::new(e)), + } + } +} + +type SendMsg<Fut> = Result<<Fut as Future>::Output, Box<(dyn Any + Send + 'static)>>; + +/// A future which sends its output to the corresponding `RemoteHandle`. +/// Created by [`remote_handle`](crate::future::FutureExt::remote_handle). +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Remote<Fut: Future> { + tx: Option<Sender<SendMsg<Fut>>>, + keep_running: Arc<AtomicBool>, + future: CatchUnwind<AssertUnwindSafe<Fut>>, +} + +impl<Fut: Future + fmt::Debug> fmt::Debug for Remote<Fut> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("Remote") + .field(&self.future) + .finish() + } +} + +impl<Fut: Future + Unpin> Unpin for Remote<Fut> {} + +impl<Fut: Future> Remote<Fut> { + unsafe_pinned!(future: CatchUnwind<AssertUnwindSafe<Fut>>); + unsafe_unpinned!(tx: Option<Sender<SendMsg<Fut>>>); +} + +impl<Fut: Future> Future for Remote<Fut> { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + if let Poll::Ready(_) = self.as_mut().tx().as_mut().unwrap().poll_canceled(cx) { + if !self.keep_running.load(Ordering::SeqCst) { + // Cancelled, bail out + return Poll::Ready(()) + } + } + + let output = ready!(self.as_mut().future().poll(cx)); + + // if the receiving end has gone away then that's ok, we just ignore the + // send error here. + drop(self.as_mut().tx().take().unwrap().send(output)); + Poll::Ready(()) + } +} + +pub(super) fn remote_handle<Fut: Future>(future: Fut) -> (Remote<Fut>, RemoteHandle<Fut::Output>) { + let (tx, rx) = oneshot::channel(); + let keep_running = Arc::new(AtomicBool::new(false)); + + // AssertUnwindSafe is used here because `Send + 'static` is basically + // an alias for an implementation of the `UnwindSafe` trait but we can't + // express that in the standard library right now. + let wrapped = Remote { + future: AssertUnwindSafe(future).catch_unwind(), + tx: Some(tx), + keep_running: keep_running.clone(), + }; + + (wrapped, RemoteHandle { rx, keep_running }) +} diff --git a/src/future/future/shared.rs b/src/future/future/shared.rs new file mode 100644 index 0000000..816f5dd --- /dev/null +++ b/src/future/future/shared.rs @@ -0,0 +1,335 @@ +use crate::task::{ArcWake, waker_ref}; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll, Waker}; +use slab::Slab; +use std::cell::UnsafeCell; +use std::fmt; +use std::pin::Pin; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering::SeqCst; +use std::sync::{Arc, Mutex}; + +/// Future for the [`shared`](super::FutureExt::shared) method. +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Shared<Fut: Future> { + inner: Option<Arc<Inner<Fut>>>, + waker_key: usize, +} + +struct Inner<Fut: Future> { + future_or_output: UnsafeCell<FutureOrOutput<Fut>>, + notifier: Arc<Notifier>, +} + +struct Notifier { + state: AtomicUsize, + wakers: Mutex<Option<Slab<Option<Waker>>>>, +} + +// The future itself is polled behind the `Arc`, so it won't be moved +// when `Shared` is moved. +impl<Fut: Future> Unpin for Shared<Fut> {} + +impl<Fut: Future> fmt::Debug for Shared<Fut> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Shared") + .field("inner", &self.inner) + .field("waker_key", &self.waker_key) + .finish() + } +} + +impl<Fut: Future> fmt::Debug for Inner<Fut> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Inner").finish() + } +} + +enum FutureOrOutput<Fut: Future> { + Future(Fut), + Output(Fut::Output), +} + +unsafe impl<Fut> Send for Inner<Fut> +where + Fut: Future + Send, + Fut::Output: Send + Sync, +{} + +unsafe impl<Fut> Sync for Inner<Fut> +where + Fut: Future + Send, + Fut::Output: Send + Sync, +{} + +const IDLE: usize = 0; +const POLLING: usize = 1; +const REPOLL: usize = 2; +const COMPLETE: usize = 3; +const POISONED: usize = 4; + +const NULL_WAKER_KEY: usize = usize::max_value(); + +impl<Fut: Future> Shared<Fut> { + pub(super) fn new(future: Fut) -> Shared<Fut> { + let inner = Inner { + future_or_output: UnsafeCell::new(FutureOrOutput::Future(future)), + notifier: Arc::new(Notifier { + state: AtomicUsize::new(IDLE), + wakers: Mutex::new(Some(Slab::new())), + }), + }; + + Shared { + inner: Some(Arc::new(inner)), + waker_key: NULL_WAKER_KEY, + } + } +} + +impl<Fut> Shared<Fut> +where + Fut: Future, + Fut::Output: Clone, +{ + /// Returns [`Some`] containing a reference to this [`Shared`]'s output if + /// it has already been computed by a clone or [`None`] if it hasn't been + /// computed yet or this [`Shared`] already returned its output from + /// [`poll`](Future::poll). + pub fn peek(&self) -> Option<&Fut::Output> { + if let Some(inner) = self.inner.as_ref() { + match inner.notifier.state.load(SeqCst) { + COMPLETE => unsafe { return Some(inner.output()) }, + POISONED => panic!("inner future panicked during poll"), + _ => {} + } + } + None + } + + /// Registers the current task to receive a wakeup when `Inner` is awoken. + fn set_waker(&mut self, cx: &mut Context<'_>) { + // Acquire the lock first before checking COMPLETE to ensure there + // isn't a race. + let mut wakers_guard = if let Some(inner) = self.inner.as_ref() { + inner.notifier.wakers.lock().unwrap() + } else { + return; + }; + + let wakers = if let Some(wakers) = wakers_guard.as_mut() { + wakers + } else { + return; + }; + + if self.waker_key == NULL_WAKER_KEY { + self.waker_key = wakers.insert(Some(cx.waker().clone())); + } else { + let waker_slot = &mut wakers[self.waker_key]; + let needs_replacement = if let Some(_old_waker) = waker_slot { + // If there's still an unwoken waker in the slot, only replace + // if the current one wouldn't wake the same task. + // TODO: This API is currently not available, so replace always + // !waker.will_wake_nonlocal(old_waker) + true + } else { + true + }; + if needs_replacement { + *waker_slot = Some(cx.waker().clone()); + } + } + debug_assert!(self.waker_key != NULL_WAKER_KEY); + } + + /// Safety: callers must first ensure that `self.inner.state` + /// is `COMPLETE` + unsafe fn take_or_clone_output(&mut self) -> Fut::Output { + let inner = self.inner.take().unwrap(); + + match Arc::try_unwrap(inner) { + Ok(inner) => match inner.future_or_output.into_inner() { + FutureOrOutput::Output(item) => item, + FutureOrOutput::Future(_) => unreachable!(), + }, + Err(inner) => inner.output().clone(), + } + } +} + +impl<Fut> Inner<Fut> +where + Fut: Future, + Fut::Output: Clone, +{ + /// Safety: callers must first ensure that `self.inner.state` + /// is `COMPLETE` + unsafe fn output(&self) -> &Fut::Output { + match &*self.future_or_output.get() { + FutureOrOutput::Output(ref item) => &item, + FutureOrOutput::Future(_) => unreachable!(), + } + } +} + +impl<Fut> FusedFuture for Shared<Fut> +where + Fut: Future, + Fut::Output: Clone, +{ + fn is_terminated(&self) -> bool { + self.inner.is_none() + } +} + +impl<Fut> Future for Shared<Fut> +where + Fut: Future, + Fut::Output: Clone, +{ + type Output = Fut::Output; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = &mut *self; + + this.set_waker(cx); + + let inner = if let Some(inner) = this.inner.as_ref() { + inner + } else { + panic!("Shared future polled again after completion"); + }; + + match inner.notifier.state.compare_and_swap(IDLE, POLLING, SeqCst) { + IDLE => { + // Lock acquired, fall through + } + POLLING | REPOLL => { + // Another task is currently polling, at this point we just want + // to ensure that the waker for this task is registered + + return Poll::Pending; + } + COMPLETE => { + // Safety: We're in the COMPLETE state + return unsafe { Poll::Ready(this.take_or_clone_output()) }; + } + POISONED => panic!("inner future panicked during poll"), + _ => unreachable!(), + } + + let waker = waker_ref(&inner.notifier); + let mut cx = Context::from_waker(&waker); + + struct Reset<'a>(&'a AtomicUsize); + + impl Drop for Reset<'_> { + fn drop(&mut self) { + use std::thread; + + if thread::panicking() { + self.0.store(POISONED, SeqCst); + } + } + } + + let _reset = Reset(&inner.notifier.state); + + let output = loop { + let future = unsafe { + match &mut *inner.future_or_output.get() { + FutureOrOutput::Future(fut) => Pin::new_unchecked(fut), + _ => unreachable!(), + } + }; + + let poll = future.poll(&mut cx); + + match poll { + Poll::Pending => { + let state = &inner.notifier.state; + match state.compare_and_swap(POLLING, IDLE, SeqCst) { + POLLING => { + // Success + return Poll::Pending; + } + REPOLL => { + // Was woken since: Gotta poll again! + let prev = state.swap(POLLING, SeqCst); + assert_eq!(prev, REPOLL); + } + _ => unreachable!(), + } + } + Poll::Ready(output) => break output, + } + }; + + unsafe { + *inner.future_or_output.get() = + FutureOrOutput::Output(output); + } + + inner.notifier.state.store(COMPLETE, SeqCst); + + // Wake all tasks and drop the slab + let mut wakers_guard = inner.notifier.wakers.lock().unwrap(); + let wakers = &mut wakers_guard.take().unwrap(); + for (_key, opt_waker) in wakers { + if let Some(waker) = opt_waker.take() { + waker.wake(); + } + } + + drop(_reset); // Make borrow checker happy + drop(wakers_guard); + + // Safety: We're in the COMPLETE state + unsafe { Poll::Ready(this.take_or_clone_output()) } + } +} + +impl<Fut> Clone for Shared<Fut> +where + Fut: Future, +{ + fn clone(&self) -> Self { + Shared { + inner: self.inner.clone(), + waker_key: NULL_WAKER_KEY, + } + } +} + +impl<Fut> Drop for Shared<Fut> +where + Fut: Future, +{ + fn drop(&mut self) { + if self.waker_key != NULL_WAKER_KEY { + if let Some(ref inner) = self.inner { + if let Ok(mut wakers) = inner.notifier.wakers.lock() { + if let Some(wakers) = wakers.as_mut() { + wakers.remove(self.waker_key); + } + } + } + } + } +} + +impl ArcWake for Notifier { + fn wake_by_ref(arc_self: &Arc<Self>) { + arc_self.state.compare_and_swap(POLLING, REPOLL, SeqCst); + + let wakers = &mut *arc_self.wakers.lock().unwrap(); + if let Some(wakers) = wakers.as_mut() { + for (_key, opt_waker) in wakers { + if let Some(waker) = opt_waker.take() { + waker.wake(); + } + } + } + } +} diff --git a/src/future/future/then.rs b/src/future/future/then.rs new file mode 100644 index 0000000..9f30f09 --- /dev/null +++ b/src/future/future/then.rs @@ -0,0 +1,46 @@ +use super::Chain; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Future for the [`then`](super::FutureExt::then) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Then<Fut1, Fut2, F> { + chain: Chain<Fut1, Fut2, F>, +} + +impl<Fut1, Fut2, F> Then<Fut1, Fut2, F> + where Fut1: Future, + Fut2: Future, +{ + unsafe_pinned!(chain: Chain<Fut1, Fut2, F>); + + /// Creates a new `Then`. + pub(super) fn new(future: Fut1, f: F) -> Then<Fut1, Fut2, F> { + Then { + chain: Chain::new(future, f), + } + } +} + +impl<Fut1, Fut2, F> FusedFuture for Then<Fut1, Fut2, F> + where Fut1: Future, + Fut2: Future, + F: FnOnce(Fut1::Output) -> Fut2, +{ + fn is_terminated(&self) -> bool { self.chain.is_terminated() } +} + +impl<Fut1, Fut2, F> Future for Then<Fut1, Fut2, F> + where Fut1: Future, + Fut2: Future, + F: FnOnce(Fut1::Output) -> Fut2, +{ + type Output = Fut2::Output; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Fut2::Output> { + self.as_mut().chain().poll(cx, |output, f| f(output)) + } +} diff --git a/src/future/future/unit_error.rs b/src/future/future/unit_error.rs new file mode 100644 index 0000000..679e988 --- /dev/null +++ b/src/future/future/unit_error.rs @@ -0,0 +1,35 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Future for the [`unit_error`](super::FutureExt::unit_error) combinator. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct UnitError<Fut> { + future: Fut, +} + +impl<Fut> UnitError<Fut> { + unsafe_pinned!(future: Fut); + + pub(super) fn new(future: Fut) -> UnitError<Fut> { + UnitError { future } + } +} + +impl<Fut: Unpin> Unpin for UnitError<Fut> {} + +impl<Fut: FusedFuture> FusedFuture for UnitError<Fut> { + fn is_terminated(&self) -> bool { self.future.is_terminated() } +} + +impl<Fut, T> Future for UnitError<Fut> + where Fut: Future<Output = T>, +{ + type Output = Result<T, ()>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<T, ()>> { + self.future().poll(cx).map(Ok) + } +} diff --git a/src/future/join.rs b/src/future/join.rs new file mode 100644 index 0000000..5af5b40 --- /dev/null +++ b/src/future/join.rs @@ -0,0 +1,214 @@ +#![allow(non_snake_case)] + +use crate::future::{MaybeDone, maybe_done}; +use core::fmt; +use core::pin::Pin; +use futures_core::future::{Future, FusedFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; +use super::assert_future; + +macro_rules! generate { + ($( + $(#[$doc:meta])* + ($Join:ident, <$($Fut:ident),*>), + )*) => ($( + $(#[$doc])* + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct $Join<$($Fut: Future),*> { + $($Fut: MaybeDone<$Fut>,)* + } + + impl<$($Fut),*> fmt::Debug for $Join<$($Fut),*> + where + $( + $Fut: Future + fmt::Debug, + $Fut::Output: fmt::Debug, + )* + { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct(stringify!($Join)) + $(.field(stringify!($Fut), &self.$Fut))* + .finish() + } + } + + impl<$($Fut: Future),*> $Join<$($Fut),*> { + fn new($($Fut: $Fut),*) -> $Join<$($Fut),*> { + $Join { + $($Fut: maybe_done($Fut)),* + } + } + $( + unsafe_pinned!($Fut: MaybeDone<$Fut>); + )* + } + + impl<$($Fut: Future),*> Future for $Join<$($Fut),*> { + type Output = ($($Fut::Output),*); + + fn poll( + mut self: Pin<&mut Self>, cx: &mut Context<'_> + ) -> Poll<Self::Output> { + let mut all_done = true; + $( + all_done &= self.as_mut().$Fut().poll(cx).is_ready(); + )* + + if all_done { + Poll::Ready(($(self.as_mut().$Fut().take_output().unwrap()), *)) + } else { + Poll::Pending + } + } + } + + impl<$($Fut: FusedFuture),*> FusedFuture for $Join<$($Fut),*> { + fn is_terminated(&self) -> bool { + $( + self.$Fut.is_terminated() + ) && * + } + } + )*) +} + +generate! { + /// Future for the [`join`](join()) function. + (Join, <Fut1, Fut2>), + + /// Future for the [`join3`] function. + (Join3, <Fut1, Fut2, Fut3>), + + /// Future for the [`join4`] function. + (Join4, <Fut1, Fut2, Fut3, Fut4>), + + /// Future for the [`join5`] function. + (Join5, <Fut1, Fut2, Fut3, Fut4, Fut5>), +} + +/// Joins the result of two futures, waiting for them both to complete. +/// +/// This function will return a new future which awaits both futures to +/// complete. The returned future will finish with a tuple of both results. +/// +/// Note that this function consumes the passed futures and returns a +/// wrapped version of it. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = async { 1 }; +/// let b = async { 2 }; +/// let pair = future::join(a, b); +/// +/// assert_eq!(pair.await, (1, 2)); +/// # }); +/// ``` +pub fn join<Fut1, Fut2>(future1: Fut1, future2: Fut2) -> Join<Fut1, Fut2> +where + Fut1: Future, + Fut2: Future, +{ + let f = Join::new(future1, future2); + assert_future::<(Fut1::Output, Fut2::Output), _>(f) +} + +/// Same as [`join`](join()), but with more futures. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = async { 1 }; +/// let b = async { 2 }; +/// let c = async { 3 }; +/// let tuple = future::join3(a, b, c); +/// +/// assert_eq!(tuple.await, (1, 2, 3)); +/// # }); +/// ``` +pub fn join3<Fut1, Fut2, Fut3>( + future1: Fut1, + future2: Fut2, + future3: Fut3, +) -> Join3<Fut1, Fut2, Fut3> +where + Fut1: Future, + Fut2: Future, + Fut3: Future, +{ + Join3::new(future1, future2, future3) +} + +/// Same as [`join`](join()), but with more futures. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = async { 1 }; +/// let b = async { 2 }; +/// let c = async { 3 }; +/// let d = async { 4 }; +/// let tuple = future::join4(a, b, c, d); +/// +/// assert_eq!(tuple.await, (1, 2, 3, 4)); +/// # }); +/// ``` +pub fn join4<Fut1, Fut2, Fut3, Fut4>( + future1: Fut1, + future2: Fut2, + future3: Fut3, + future4: Fut4, +) -> Join4<Fut1, Fut2, Fut3, Fut4> +where + Fut1: Future, + Fut2: Future, + Fut3: Future, + Fut4: Future, +{ + Join4::new(future1, future2, future3, future4) +} + +/// Same as [`join`](join()), but with more futures. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = async { 1 }; +/// let b = async { 2 }; +/// let c = async { 3 }; +/// let d = async { 4 }; +/// let e = async { 5 }; +/// let tuple = future::join5(a, b, c, d, e); +/// +/// assert_eq!(tuple.await, (1, 2, 3, 4, 5)); +/// # }); +/// ``` +pub fn join5<Fut1, Fut2, Fut3, Fut4, Fut5>( + future1: Fut1, + future2: Fut2, + future3: Fut3, + future4: Fut4, + future5: Fut5, +) -> Join5<Fut1, Fut2, Fut3, Fut4, Fut5> +where + Fut1: Future, + Fut2: Future, + Fut3: Future, + Fut4: Future, + Fut5: Future, +{ + Join5::new(future1, future2, future3, future4, future5) +} diff --git a/src/future/join_all.rs b/src/future/join_all.rs new file mode 100644 index 0000000..0740885 --- /dev/null +++ b/src/future/join_all.rs @@ -0,0 +1,159 @@ +//! Definition of the `JoinAll` combinator, waiting for all of a list of futures +//! to finish. + +use core::fmt; +use core::future::Future; +use core::iter::FromIterator; +use core::mem; +use core::pin::Pin; +use core::task::{Context, Poll}; +use alloc::boxed::Box; +use alloc::vec::Vec; + +#[derive(Debug)] +enum ElemState<F> +where + F: Future, +{ + Pending(F), + Done(Option<F::Output>), +} + +impl<F> ElemState<F> +where + F: Future, +{ + fn pending_pin_mut(self: Pin<&mut Self>) -> Option<Pin<&mut F>> { + // Safety: Basic enum pin projection, no drop + optionally Unpin based + // on the type of this variant + match unsafe { self.get_unchecked_mut() } { + ElemState::Pending(f) => Some(unsafe { Pin::new_unchecked(f) }), + ElemState::Done(_) => None, + } + } + + fn take_done(self: Pin<&mut Self>) -> Option<F::Output> { + // Safety: Going from pin to a variant we never pin-project + match unsafe { self.get_unchecked_mut() } { + ElemState::Pending(_) => None, + ElemState::Done(output) => output.take(), + } + } +} + +impl<F> Unpin for ElemState<F> +where + F: Future + Unpin, +{ +} + +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) }) +} + +/// Future for the [`join_all`] function. +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct JoinAll<F> +where + F: Future, +{ + elems: Pin<Box<[ElemState<F>]>>, +} + +impl<F> fmt::Debug for JoinAll<F> +where + F: Future + fmt::Debug, + F::Output: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("JoinAll") + .field("elems", &self.elems) + .finish() + } +} + +/// Creates a future which represents a collection of the outputs of the futures +/// given. +/// +/// 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. +/// +/// This function is only available when the `std` or `alloc` feature of this +/// library is activated, and it is activated by default. +/// +/// # See Also +/// +/// This is purposefully a very simple API for basic use-cases. In a lot of +/// cases you will want to use the more powerful +/// [`FuturesUnordered`][crate::stream::FuturesUnordered] APIs, some +/// examples of additional functionality that provides: +/// +/// * 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::join_all; +/// +/// async fn foo(i: u32) -> u32 { i } +/// +/// let futures = vec![foo(1), foo(2), foo(3)]; +/// +/// assert_eq!(join_all(futures).await, [1, 2, 3]); +/// # }); +/// ``` +pub fn join_all<I>(i: I) -> JoinAll<I::Item> +where + I: IntoIterator, + I::Item: Future, +{ + let elems: Box<[_]> = i.into_iter().map(ElemState::Pending).collect(); + JoinAll { elems: elems.into() } +} + +impl<F> Future for JoinAll<F> +where + F: Future, +{ + type Output = Vec<F::Output>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mut all_done = true; + + for mut elem in iter_pin_mut(self.elems.as_mut()) { + if let Some(pending) = elem.as_mut().pending_pin_mut() { + if let Poll::Ready(output) = pending.poll(cx) { + elem.set(ElemState::Done(Some(output))); + } else { + all_done = false; + } + } + } + + if all_done { + let mut elems = mem::replace(&mut self.elems, Box::pin([])); + let result = iter_pin_mut(elems.as_mut()) + .map(|e| e.take_done().unwrap()) + .collect(); + Poll::Ready(result) + } else { + Poll::Pending + } + } +} + +impl<F: Future> FromIterator<F> for JoinAll<F> { + fn from_iter<T: IntoIterator<Item = F>>(iter: T) -> Self { + join_all(iter) + } +} diff --git a/src/future/lazy.rs b/src/future/lazy.rs new file mode 100644 index 0000000..5e72218 --- /dev/null +++ b/src/future/lazy.rs @@ -0,0 +1,54 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll}; + +/// Future for the [`lazy`] function. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Lazy<F> { + f: Option<F> +} + +// safe because we never generate `Pin<&mut F>` +impl<F> Unpin for Lazy<F> {} + +/// Creates a new future that allows delayed execution of a closure. +/// +/// The provided closure is only run once the future is polled. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = future::lazy(|_| 1); +/// assert_eq!(a.await, 1); +/// +/// let b = future::lazy(|_| -> i32 { +/// panic!("oh no!") +/// }); +/// drop(b); // closure is never run +/// # }); +/// ``` +pub fn lazy<F, R>(f: F) -> Lazy<F> + where F: FnOnce(&mut Context<'_>) -> R, +{ + Lazy { f: Some(f) } +} + +impl<F, R> FusedFuture for Lazy<F> + where F: FnOnce(&mut Context<'_>) -> R, +{ + fn is_terminated(&self) -> bool { self.f.is_none() } +} + +impl<F, R> Future for Lazy<F> + where F: FnOnce(&mut Context<'_>) -> R, +{ + type Output = R; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<R> { + Poll::Ready((self.f.take().unwrap())(cx)) + } +} diff --git a/src/future/maybe_done.rs b/src/future/maybe_done.rs new file mode 100644 index 0000000..f16f889 --- /dev/null +++ b/src/future/maybe_done.rs @@ -0,0 +1,104 @@ +//! Definition of the MaybeDone combinator + +use core::mem; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll}; + +/// A future that may have completed. +/// +/// This is created by the [`maybe_done()`] function. +#[derive(Debug)] +pub enum MaybeDone<Fut: Future> { + /// A not-yet-completed future + Future(Fut), + /// The output of the completed future + Done(Fut::Output), + /// The empty variant after the result of a [`MaybeDone`] has been + /// taken using the [`take_output`](MaybeDone::take_output) method. + Gone, +} + +// Safe because we never generate `Pin<&mut Fut::Output>` +impl<Fut: Future + Unpin> Unpin for MaybeDone<Fut> {} + +/// Wraps a future into a `MaybeDone` +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// use futures::pin_mut; +/// +/// let future = future::maybe_done(async { 5 }); +/// pin_mut!(future); +/// assert_eq!(future.as_mut().take_output(), None); +/// let () = future.as_mut().await; +/// assert_eq!(future.as_mut().take_output(), Some(5)); +/// assert_eq!(future.as_mut().take_output(), None); +/// # }); +/// ``` +pub fn maybe_done<Fut: Future>(future: Fut) -> MaybeDone<Fut> { + MaybeDone::Future(future) +} + +impl<Fut: Future> MaybeDone<Fut> { + /// Returns an [`Option`] containing a mutable reference to the output of the future. + /// The output of this method will be [`Some`] if and only if the inner + /// future has been completed and [`take_output`](MaybeDone::take_output) + /// has not yet been called. + #[inline] + pub fn output_mut(self: Pin<&mut Self>) -> Option<&mut Fut::Output> { + unsafe { + let this = self.get_unchecked_mut(); + match this { + MaybeDone::Done(res) => Some(res), + _ => None, + } + } + } + + /// Attempt to take the output of a `MaybeDone` without driving it + /// towards completion. + #[inline] + pub fn take_output(self: Pin<&mut Self>) -> Option<Fut::Output> { + unsafe { + let this = self.get_unchecked_mut(); + match this { + MaybeDone::Done(_) => {}, + MaybeDone::Future(_) | MaybeDone::Gone => return None, + }; + if let MaybeDone::Done(output) = mem::replace(this, MaybeDone::Gone) { + Some(output) + } else { + unreachable!() + } + } + } +} + +impl<Fut: Future> FusedFuture for MaybeDone<Fut> { + fn is_terminated(&self) -> bool { + match self { + MaybeDone::Future(_) => false, + MaybeDone::Done(_) | MaybeDone::Gone => true, + } + } +} + +impl<Fut: Future> Future for MaybeDone<Fut> { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let res = unsafe { + match self.as_mut().get_unchecked_mut() { + MaybeDone::Future(a) => ready!(Pin::new_unchecked(a).poll(cx)), + MaybeDone::Done(_) => return Poll::Ready(()), + MaybeDone::Gone => panic!("MaybeDone polled after value taken"), + } + }; + self.set(MaybeDone::Done(res)); + Poll::Ready(()) + } +} diff --git a/src/future/mod.rs b/src/future/mod.rs new file mode 100644 index 0000000..3f4bb01 --- /dev/null +++ b/src/future/mod.rs @@ -0,0 +1,110 @@ +//! Futures +//! +//! This module contains a number of functions for working with `Future`s, +//! including the [`FutureExt`] trait and the [`TryFutureExt`] trait which add +//! methods to `Future` types. + +#[cfg(feature = "alloc")] +pub use futures_core::future::{BoxFuture, LocalBoxFuture}; +pub use futures_core::future::{FusedFuture, Future, TryFuture}; +pub use futures_task::{FutureObj, LocalFutureObj, UnsafeFutureObj}; + +// Extension traits and combinators + +#[allow(clippy::module_inception)] +mod future; +pub use self::future::{ + Flatten, FlattenStream, Fuse, FutureExt, Inspect, IntoStream, Map, NeverError, Then, UnitError, +}; + +#[cfg(feature = "std")] +pub use self::future::CatchUnwind; + +#[cfg(feature = "channel")] +#[cfg(feature = "std")] +pub use self::future::{Remote, RemoteHandle}; + +#[cfg(feature = "std")] +pub use self::future::Shared; + +mod try_future; +pub use self::try_future::{ + AndThen, ErrInto, InspectErr, InspectOk, IntoFuture, MapErr, MapOk, OrElse, TryFlattenStream, + TryFutureExt, UnwrapOrElse, +}; + +#[cfg(feature = "sink")] +pub use self::try_future::FlattenSink; + +// Primitive futures + +mod lazy; +pub use self::lazy::{lazy, Lazy}; + +mod pending; +pub use self::pending::{pending, Pending}; + +mod maybe_done; +pub use self::maybe_done::{maybe_done, MaybeDone}; + +mod option; +pub use self::option::OptionFuture; + +mod poll_fn; +pub use self::poll_fn::{poll_fn, PollFn}; + +mod ready; +pub use self::ready::{err, ok, ready, Ready}; + +mod join; +pub use self::join::{join, join3, join4, join5, Join, Join3, Join4, Join5}; + +#[cfg(feature = "alloc")] +mod join_all; +#[cfg(feature = "alloc")] +pub use self::join_all::{join_all, JoinAll}; + +mod select; +pub use self::select::{select, Select}; + +#[cfg(feature = "alloc")] +mod select_all; +#[cfg(feature = "alloc")] +pub use self::select_all::{select_all, SelectAll}; + +mod try_join; +pub use self::try_join::{ + try_join, try_join3, try_join4, try_join5, TryJoin, TryJoin3, TryJoin4, TryJoin5, +}; + +#[cfg(feature = "alloc")] +mod try_join_all; +#[cfg(feature = "alloc")] +pub use self::try_join_all::{try_join_all, TryJoinAll}; + +mod try_select; +pub use self::try_select::{try_select, TrySelect}; + +#[cfg(feature = "alloc")] +mod select_ok; +#[cfg(feature = "alloc")] +pub use self::select_ok::{select_ok, SelectOk}; + +mod either; +pub use self::either::Either; + +cfg_target_has_atomic! { + #[cfg(feature = "alloc")] + mod abortable; + #[cfg(feature = "alloc")] + pub use self::abortable::{abortable, Abortable, AbortHandle, AbortRegistration, Aborted}; +} + +// Just a helper function to ensure the futures we're returning all have the +// right implementations. +fn assert_future<T, F>(future: F) -> F +where + F: Future<Output = T>, +{ + future +} diff --git a/src/future/option.rs b/src/future/option.rs new file mode 100644 index 0000000..2141352 --- /dev/null +++ b/src/future/option.rs @@ -0,0 +1,62 @@ +//! Definition of the `Option` (optional step) combinator + +use core::pin::Pin; +use futures_core::future::{Future, FusedFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// A future representing a value which may or may not be present. +/// +/// Created by the [`From`] implementation for [`Option`](std::option::Option). +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future::OptionFuture; +/// +/// let mut a: OptionFuture<_> = Some(async { 123 }).into(); +/// assert_eq!(a.await, Some(123)); +/// +/// a = None.into(); +/// assert_eq!(a.await, None); +/// # }); +/// ``` +#[derive(Debug, Clone)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct OptionFuture<F> { + option: Option<F>, +} + +impl<F> OptionFuture<F> { + unsafe_pinned!(option: Option<F>); +} + +impl<F: Future> Future for OptionFuture<F> { + type Output = Option<F::Output>; + + fn poll( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + match self.option().as_pin_mut() { + Some(x) => x.poll(cx).map(Some), + None => Poll::Ready(None), + } + } +} + +impl<F: FusedFuture> FusedFuture for OptionFuture<F> { + fn is_terminated(&self) -> bool { + match &self.option { + Some(x) => x.is_terminated(), + None => true, + } + } +} + +impl<T> From<Option<T>> for OptionFuture<T> { + fn from(option: Option<T>) -> Self { + OptionFuture { option } + } +} diff --git a/src/future/pending.rs b/src/future/pending.rs new file mode 100644 index 0000000..5a7bbb8 --- /dev/null +++ b/src/future/pending.rs @@ -0,0 +1,56 @@ +use core::marker; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll}; + +/// Future for the [`pending()`] function. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Pending<T> { + _data: marker::PhantomData<T>, +} + +impl<T> FusedFuture for Pending<T> { + fn is_terminated(&self) -> bool { + true + } +} + +/// Creates a future which never resolves, representing a computation that never +/// finishes. +/// +/// The returned future will forever return [`Poll::Pending`]. +/// +/// # Examples +/// +/// ```ignore +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let future = future::pending(); +/// let () = future.await; +/// unreachable!(); +/// # }); +/// ``` +pub fn pending<T>() -> Pending<T> { + Pending { + _data: marker::PhantomData, + } +} + +impl<T> Future for Pending<T> { + type Output = T; + + fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<T> { + Poll::Pending + } +} + +impl<T> Unpin for Pending<T> { +} + +impl<T> Clone for Pending<T> { + fn clone(&self) -> Self { + pending() + } +} diff --git a/src/future/poll_fn.rs b/src/future/poll_fn.rs new file mode 100644 index 0000000..b7b10be --- /dev/null +++ b/src/future/poll_fn.rs @@ -0,0 +1,56 @@ +//! Definition of the `PollFn` adapter combinator + +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; + +/// Future for the [`poll_fn`] function. +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct PollFn<F> { + f: F, +} + +impl<F> Unpin for PollFn<F> {} + +/// Creates a new future wrapping around a function returning [`Poll`]. +/// +/// Polling the returned future delegates to the wrapped function. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future::poll_fn; +/// use futures::task::{Context, Poll}; +/// +/// fn read_line(_cx: &mut Context<'_>) -> Poll<String> { +/// Poll::Ready("Hello, World!".into()) +/// } +/// +/// let read_future = poll_fn(read_line); +/// assert_eq!(read_future.await, "Hello, World!".to_owned()); +/// # }); +/// ``` +pub fn poll_fn<T, F>(f: F) -> PollFn<F> +where + F: FnMut(&mut Context<'_>) -> Poll<T> +{ + PollFn { f } +} + +impl<F> fmt::Debug for PollFn<F> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PollFn").finish() + } +} + +impl<T, F> Future for PollFn<F> + where F: FnMut(&mut Context<'_>) -> Poll<T>, +{ + type Output = T; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { + (&mut self.f)(cx) + } +} diff --git a/src/future/ready.rs b/src/future/ready.rs new file mode 100644 index 0000000..48661b3 --- /dev/null +++ b/src/future/ready.rs @@ -0,0 +1,81 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll}; + +/// Future for the [`ready`](ready()) function. +#[derive(Debug, Clone)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Ready<T>(Option<T>); + +impl<T> Ready<T> { + /// Unwraps the value from this immediately ready future. + #[inline] + pub fn into_inner(mut self) -> T { + self.0.take().unwrap() + } +} + +impl<T> Unpin for Ready<T> {} + +impl<T> FusedFuture for Ready<T> { + fn is_terminated(&self) -> bool { + self.0.is_none() + } +} + +impl<T> Future for Ready<T> { + type Output = T; + + #[inline] + fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<T> { + Poll::Ready(self.0.take().unwrap()) + } +} + +/// Creates a future that is immediately ready with a value. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = future::ready(1); +/// assert_eq!(a.await, 1); +/// # }); +/// ``` +pub fn ready<T>(t: T) -> Ready<T> { + Ready(Some(t)) +} + +/// Create a future that is immediately ready with a success value. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = future::ok::<i32, i32>(1); +/// assert_eq!(a.await, Ok(1)); +/// # }); +/// ``` +pub fn ok<T, E>(t: T) -> Ready<Result<T, E>> { + Ready(Some(Ok(t))) +} + +/// Create a future that is immediately ready with an error value. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = future::err::<i32, i32>(1); +/// assert_eq!(a.await, Err(1)); +/// # }); +/// ``` +pub fn err<T, E>(err: E) -> Ready<Result<T, E>> { + Ready(Some(Err(err))) +} diff --git a/src/future/select.rs b/src/future/select.rs new file mode 100644 index 0000000..91b467d --- /dev/null +++ b/src/future/select.rs @@ -0,0 +1,83 @@ +use core::pin::Pin; +use futures_core::future::{Future, FusedFuture}; +use futures_core::task::{Context, Poll}; +use crate::future::{Either, FutureExt}; + +/// Future for the [`select()`] function. +#[must_use = "futures do nothing unless you `.await` or poll them"] +#[derive(Debug)] +pub struct Select<A, B> { + inner: Option<(A, B)>, +} + +impl<A: Unpin, B: Unpin> Unpin for Select<A, B> {} + +/// Waits for either one of two differently-typed futures to complete. +/// +/// This function will return a new future which awaits for either one of both +/// futures to complete. The returned future will finish with both the value +/// resolved and a future representing the completion of the other work. +/// +/// Note that this function consumes the receiving futures and returns a +/// wrapped version of them. +/// +/// Also note that if both this and the second future have the same +/// output type you can use the `Either::factor_first` method to +/// conveniently extract out the value at the end. +/// +/// # Examples +/// +/// ``` +/// use futures::future::{self, Either, Future, FutureExt}; +/// +/// // A poor-man's join implemented on top of select +/// +/// fn join<A, B, E>(a: A, b: B) -> impl Future<Output=(A::Output, B::Output)> +/// where A: Future + Unpin, +/// B: Future + Unpin, +/// { +/// future::select(a, b).then(|either| { +/// match either { +/// Either::Left((x, b)) => b.map(move |y| (x, y)).left_future(), +/// Either::Right((y, a)) => a.map(move |x| (x, y)).right_future(), +/// } +/// }) +/// } +/// ``` +pub fn select<A, B>(future1: A, future2: B) -> Select<A, B> + where A: Future + Unpin, B: Future + Unpin +{ + Select { inner: Some((future1, future2)) } +} + +impl<A, B> Future for Select<A, B> +where + A: Future + Unpin, + B: Future + Unpin, +{ + type Output = Either<(A::Output, B), (B::Output, A)>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let (mut a, mut b) = self.inner.take().expect("cannot poll Select twice"); + match a.poll_unpin(cx) { + Poll::Ready(x) => Poll::Ready(Either::Left((x, b))), + Poll::Pending => match b.poll_unpin(cx) { + Poll::Ready(x) => Poll::Ready(Either::Right((x, a))), + Poll::Pending => { + self.inner = Some((a, b)); + Poll::Pending + } + } + } + } +} + +impl<A, B> FusedFuture for Select<A, B> +where + A: Future + Unpin, + B: Future + Unpin, +{ + fn is_terminated(&self) -> bool { + self.inner.is_none() + } +} diff --git a/src/future/select_all.rs b/src/future/select_all.rs new file mode 100644 index 0000000..9f7fb24 --- /dev/null +++ b/src/future/select_all.rs @@ -0,0 +1,69 @@ +use crate::future::FutureExt; +use core::iter::FromIterator; +use core::mem; +use core::pin::Pin; +use alloc::vec::Vec; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; + +/// Future for the [`select_all`] function. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct SelectAll<Fut> { + inner: Vec<Fut>, +} + +impl<Fut: Unpin> Unpin for SelectAll<Fut> {} + +/// Creates a new future which will select over a list of futures. +/// +/// The returned future will wait for any future within `iter` to be ready. Upon +/// completion the item resolved will be returned, along with the index of the +/// future that was ready and the list of all the remaining futures. +/// +/// There are no guarantees provided on the order of the list with the remaining +/// futures. They might be swapped around, reversed, or completely random. +/// +/// This function is only available when the `std` or `alloc` feature of this +/// library is activated, and it is activated by default. +/// +/// # Panics +/// +/// This function will panic if the iterator specified contains no items. +pub fn select_all<I>(iter: I) -> SelectAll<I::Item> + where I: IntoIterator, + I::Item: Future + Unpin, +{ + let ret = SelectAll { + inner: iter.into_iter().collect() + }; + assert!(!ret.inner.is_empty()); + ret +} + +impl<Fut: Future + Unpin> Future for SelectAll<Fut> { + type Output = (Fut::Output, usize, Vec<Fut>); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let item = self.inner.iter_mut().enumerate().find_map(|(i, f)| { + match f.poll_unpin(cx) { + Poll::Pending => None, + Poll::Ready(e) => Some((i, e)), + } + }); + match item { + Some((idx, res)) => { + let _ = self.inner.swap_remove(idx); + let rest = mem::replace(&mut self.inner, Vec::new()); + Poll::Ready((res, idx, rest)) + } + None => Poll::Pending, + } + } +} + +impl<Fut: Future + Unpin> FromIterator<Fut> for SelectAll<Fut> { + fn from_iter<T: IntoIterator<Item = Fut>>(iter: T) -> Self { + select_all(iter) + } +} diff --git a/src/future/select_ok.rs b/src/future/select_ok.rs new file mode 100644 index 0000000..7f4f4d6 --- /dev/null +++ b/src/future/select_ok.rs @@ -0,0 +1,83 @@ +use crate::future::TryFutureExt; +use core::iter::FromIterator; +use core::mem; +use core::pin::Pin; +use alloc::vec::Vec; +use futures_core::future::{Future, TryFuture}; +use futures_core::task::{Context, Poll}; + +/// Future for the [`select_ok`] function. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct SelectOk<Fut> { + inner: Vec<Fut>, +} + +impl<Fut: Unpin> Unpin for SelectOk<Fut> {} + +/// Creates a new future which will select the first successful future over a list of futures. +/// +/// The returned future will wait for any future within `iter` to be ready and Ok. Unlike +/// `select_all`, this will only return the first successful completion, or the last +/// failure. This is useful in contexts where any success is desired and failures +/// are ignored, unless all the futures fail. +/// +/// This function is only available when the `std` or `alloc` feature of this +/// library is activated, and it is activated by default. +/// +/// # Panics +/// +/// This function will panic if the iterator specified contains no items. +pub fn select_ok<I>(iter: I) -> SelectOk<I::Item> + where I: IntoIterator, + I::Item: TryFuture + Unpin, +{ + let ret = SelectOk { + inner: iter.into_iter().collect() + }; + assert!(!ret.inner.is_empty(), "iterator provided to select_ok was empty"); + ret +} + +impl<Fut: TryFuture + Unpin> Future for SelectOk<Fut> { + type Output = Result<(Fut::Ok, Vec<Fut>), Fut::Error>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + // loop until we've either exhausted all errors, a success was hit, or nothing is ready + loop { + let item = self.inner.iter_mut().enumerate().find_map(|(i, f)| { + match f.try_poll_unpin(cx) { + Poll::Pending => None, + Poll::Ready(e) => Some((i, e)), + } + }); + match item { + Some((idx, res)) => { + // always remove Ok or Err, if it's not the last Err continue looping + drop(self.inner.remove(idx)); + match res { + Ok(e) => { + let rest = mem::replace(&mut self.inner, Vec::new()); + return Poll::Ready(Ok((e, rest))) + } + Err(e) => { + if self.inner.is_empty() { + return Poll::Ready(Err(e)) + } + } + } + } + None => { + // based on the filter above, nothing is ready, return + return Poll::Pending + } + } + } + } +} + +impl<Fut: TryFuture + Unpin> FromIterator<Fut> for SelectOk<Fut> { + fn from_iter<T: IntoIterator<Item = Fut>>(iter: T) -> Self { + select_ok(iter) + } +} diff --git a/src/future/try_future/and_then.rs b/src/future/try_future/and_then.rs new file mode 100644 index 0000000..37333e0 --- /dev/null +++ b/src/future/try_future/and_then.rs @@ -0,0 +1,53 @@ +use super::{TryChain, TryChainAction}; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Future for the [`and_then`](super::TryFutureExt::and_then) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct AndThen<Fut1, Fut2, F> { + try_chain: TryChain<Fut1, Fut2, F>, +} + +impl<Fut1, Fut2, F> AndThen<Fut1, Fut2, F> + where Fut1: TryFuture, + Fut2: TryFuture, +{ + unsafe_pinned!(try_chain: TryChain<Fut1, Fut2, F>); + + /// Creates a new `Then`. + pub(super) fn new(future: Fut1, f: F) -> AndThen<Fut1, Fut2, F> { + AndThen { + try_chain: TryChain::new(future, f), + } + } +} + +impl<Fut1, Fut2, F> FusedFuture for AndThen<Fut1, Fut2, F> + where Fut1: TryFuture, + Fut2: TryFuture<Error = Fut1::Error>, + F: FnOnce(Fut1::Ok) -> Fut2, +{ + fn is_terminated(&self) -> bool { + self.try_chain.is_terminated() + } +} + +impl<Fut1, Fut2, F> Future for AndThen<Fut1, Fut2, F> + where Fut1: TryFuture, + Fut2: TryFuture<Error = Fut1::Error>, + F: FnOnce(Fut1::Ok) -> Fut2, +{ + type Output = Result<Fut2::Ok, Fut2::Error>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + self.try_chain().poll(cx, |result, async_op| { + match result { + Ok(ok) => TryChainAction::Future(async_op(ok)), + Err(err) => TryChainAction::Output(Err(err)), + } + }) + } +} diff --git a/src/future/try_future/err_into.rs b/src/future/try_future/err_into.rs new file mode 100644 index 0000000..731fcae --- /dev/null +++ b/src/future/try_future/err_into.rs @@ -0,0 +1,48 @@ +use core::marker::PhantomData; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Future for the [`err_into`](super::TryFutureExt::err_into) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct ErrInto<Fut, E> { + future: Fut, + _marker: PhantomData<E>, +} + +impl<Fut: Unpin, E> Unpin for ErrInto<Fut, E> {} + +impl<Fut, E> ErrInto<Fut, E> { + unsafe_pinned!(future: Fut); + + pub(super) fn new(future: Fut) -> ErrInto<Fut, E> { + ErrInto { + future, + _marker: PhantomData, + } + } +} + +impl<Fut, E> FusedFuture for ErrInto<Fut, E> + where Fut: TryFuture + FusedFuture, + Fut::Error: Into<E>, +{ + fn is_terminated(&self) -> bool { self.future.is_terminated() } +} + +impl<Fut, E> Future for ErrInto<Fut, E> + where Fut: TryFuture, + Fut::Error: Into<E>, +{ + type Output = Result<Fut::Ok, E>; + + fn poll( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + self.future().try_poll(cx) + .map(|res| res.map_err(Into::into)) + } +} diff --git a/src/future/try_future/flatten_sink.rs b/src/future/try_future/flatten_sink.rs new file mode 100644 index 0000000..d6863dd --- /dev/null +++ b/src/future/try_future/flatten_sink.rs @@ -0,0 +1,76 @@ +use super::FlattenStreamSink; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; +use pin_utils::unsafe_pinned; + +/// Sink for the [`flatten_sink`](super::TryFutureExt::flatten_sink) method. +#[derive(Debug)] +#[must_use = "sinks do nothing unless polled"] +pub struct FlattenSink<Fut, Si> +where + Fut: TryFuture<Ok = Si>, +{ + inner: FlattenStreamSink<Fut>, +} + +impl<Fut, Si> FlattenSink<Fut, Si> +where + Fut: TryFuture<Ok = Si>, +{ + unsafe_pinned!(inner: FlattenStreamSink<Fut>); + + pub(super) fn new(future: Fut) -> Self { + Self { + inner: FlattenStreamSink::new(future), + } + } +} + +impl<Fut, S> FusedStream for FlattenSink<Fut, S> +where + Fut: TryFuture<Ok = S>, + S: TryStream<Error = Fut::Error> + FusedStream, +{ + fn is_terminated(&self) -> bool { + self.inner.is_terminated() + } +} + +impl<Fut, S> Stream for FlattenSink<Fut, S> +where + Fut: TryFuture<Ok = S>, + S: TryStream<Error = Fut::Error>, +{ + type Item = Result<S::Ok, Fut::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + self.inner().poll_next(cx) + } +} + +impl<Fut, Si, Item> Sink<Item> for FlattenSink<Fut, Si> +where + Fut: TryFuture<Ok = Si>, + Si: Sink<Item, Error = Fut::Error>, +{ + type Error = Fut::Error; + + fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { + self.inner().poll_ready(cx) + } + + fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> { + self.inner().start_send(item) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { + self.inner().poll_flush(cx) + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { + self.inner().poll_close(cx) + } +} diff --git a/src/future/try_future/flatten_stream_sink.rs b/src/future/try_future/flatten_stream_sink.rs new file mode 100644 index 0000000..5a56bf7 --- /dev/null +++ b/src/future/try_future/flatten_stream_sink.rs @@ -0,0 +1,181 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::unsafe_pinned; + +#[must_use = "streams do nothing unless polled"] +pub(crate) struct FlattenStreamSink<Fut> +where + Fut: TryFuture, +{ + state: State<Fut, Fut::Ok>, +} + +impl<Fut> Unpin for FlattenStreamSink<Fut> +where + Fut: TryFuture + Unpin, + Fut::Ok: Unpin, +{ +} + +impl<Fut> fmt::Debug for FlattenStreamSink<Fut> +where + Fut: TryFuture + fmt::Debug, + Fut::Ok: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FlattenStreamSink") + .field("state", &self.state) + .finish() + } +} + +impl<Fut> FlattenStreamSink<Fut> +where + Fut: TryFuture, +{ + unsafe_pinned!(state: State<Fut, Fut::Ok>); + + pub(crate) fn new(future: Fut) -> Self { + Self { + state: State::Future(future), + } + } +} + +#[derive(Debug)] +enum State<Fut, S> { + // future is not yet called or called and not ready + Future(Fut), + // future resolved to Stream or Sink + StreamOrSink(S), + // future resolved to error + Done, +} + +impl<Fut, S> State<Fut, S> { + fn get_pin_mut(self: Pin<&mut Self>) -> State<Pin<&mut Fut>, Pin<&mut S>> { + // safety: data is never moved via the resulting &mut reference + match unsafe { self.get_unchecked_mut() } { + // safety: the future we're re-pinning here will never be moved; + // it will just be polled, then dropped in place + State::Future(f) => State::Future(unsafe { Pin::new_unchecked(f) }), + // safety: the stream we're repinning here will never be moved; + // it will just be polled, then dropped in place + State::StreamOrSink(s) => State::StreamOrSink(unsafe { Pin::new_unchecked(s) }), + State::Done => State::Done, + } + } +} + +impl<Fut> State<Fut, Fut::Ok> +where + Fut: TryFuture, +{ + fn poll_future(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Fut::Error>> { + if let State::Future(f) = self.as_mut().get_pin_mut() { + match ready!(f.try_poll(cx)) { + Ok(s) => { + // Future resolved to stream. + // We do not return, but poll that + // stream in the next loop iteration. + self.set(State::StreamOrSink(s)); + } + Err(e) => { + // Future resolved to error. + // We have neither a pollable stream nor a future. + self.set(State::Done); + return Poll::Ready(Err(e)); + } + } + } + Poll::Ready(Ok(())) + } +} + +impl<Fut> FusedStream for FlattenStreamSink<Fut> +where + Fut: TryFuture, + Fut::Ok: TryStream<Error = Fut::Error> + FusedStream, +{ + fn is_terminated(&self) -> bool { + match &self.state { + State::Future(_) => false, + State::StreamOrSink(stream) => stream.is_terminated(), + State::Done => true, + } + } +} + +impl<Fut> Stream for FlattenStreamSink<Fut> +where + Fut: TryFuture, + Fut::Ok: TryStream<Error = Fut::Error>, +{ + type Item = Result<<Fut::Ok as TryStream>::Ok, Fut::Error>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + ready!(self.as_mut().state().poll_future(cx)?); + match self.as_mut().state().get_pin_mut() { + State::StreamOrSink(s) => s.try_poll_next(cx), + State::Done => Poll::Ready(None), + State::Future(_) => unreachable!(), + } + } +} + +#[cfg(feature = "sink")] +impl<Fut, Item> Sink<Item> for FlattenStreamSink<Fut> +where + Fut: TryFuture, + Fut::Ok: Sink<Item, Error = Fut::Error>, +{ + type Error = Fut::Error; + + fn poll_ready( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + ready!(self.as_mut().state().poll_future(cx)?); + match self.as_mut().state().get_pin_mut() { + State::StreamOrSink(s) => s.poll_ready(cx), + State::Done => panic!("poll_ready called after eof"), + State::Future(_) => unreachable!(), + } + } + + fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> { + match self.state().get_pin_mut() { + State::StreamOrSink(s) => s.start_send(item), + State::Future(_) => panic!("poll_ready not called first"), + State::Done => panic!("start_send called after eof"), + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { + match self.state().get_pin_mut() { + State::StreamOrSink(s) => s.poll_flush(cx), + // if sink not yet resolved, nothing written ==> everything flushed + State::Future(_) => Poll::Ready(Ok(())), + State::Done => panic!("poll_flush called after eof"), + } + } + + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + let res = match self.as_mut().state().get_pin_mut() { + State::StreamOrSink(s) => s.poll_close(cx), + State::Future(_) | State::Done => Poll::Ready(Ok(())), + }; + if res.is_ready() { + self.as_mut().state().set(State::Done); + } + res + } +} diff --git a/src/future/try_future/inspect_err.rs b/src/future/try_future/inspect_err.rs new file mode 100644 index 0000000..8700337 --- /dev/null +++ b/src/future/try_future/inspect_err.rs @@ -0,0 +1,53 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Future for the [`inspect_err`](super::TryFutureExt::inspect_err) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct InspectErr<Fut, F> { + future: Fut, + f: Option<F>, +} + +impl<Fut: Unpin, F> Unpin for InspectErr<Fut, F> {} + +impl<Fut, F> InspectErr<Fut, F> +where + Fut: TryFuture, + F: FnOnce(&Fut::Error), +{ + unsafe_pinned!(future: Fut); + unsafe_unpinned!(f: Option<F>); + + pub(super) fn new(future: Fut, f: F) -> Self { + Self { future, f: Some(f) } + } +} + +impl<Fut, F> FusedFuture for InspectErr<Fut, F> +where + Fut: TryFuture + FusedFuture, + F: FnOnce(&Fut::Error), +{ + fn is_terminated(&self) -> bool { + self.future.is_terminated() + } +} + +impl<Fut, F> Future for InspectErr<Fut, F> +where + Fut: TryFuture, + F: FnOnce(&Fut::Error), +{ + type Output = Result<Fut::Ok, Fut::Error>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let e = ready!(self.as_mut().future().try_poll(cx)); + if let Err(e) = &e { + self.as_mut().f().take().expect("cannot poll InspectErr twice")(e); + } + Poll::Ready(e) + } +} diff --git a/src/future/try_future/inspect_ok.rs b/src/future/try_future/inspect_ok.rs new file mode 100644 index 0000000..3d0a972 --- /dev/null +++ b/src/future/try_future/inspect_ok.rs @@ -0,0 +1,53 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Future for the [`inspect_ok`](super::TryFutureExt::inspect_ok) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct InspectOk<Fut, F> { + future: Fut, + f: Option<F>, +} + +impl<Fut: Unpin, F> Unpin for InspectOk<Fut, F> {} + +impl<Fut, F> InspectOk<Fut, F> +where + Fut: TryFuture, + F: FnOnce(&Fut::Ok), +{ + unsafe_pinned!(future: Fut); + unsafe_unpinned!(f: Option<F>); + + pub(super) fn new(future: Fut, f: F) -> Self { + Self { future, f: Some(f) } + } +} + +impl<Fut, F> FusedFuture for InspectOk<Fut, F> +where + Fut: TryFuture + FusedFuture, + F: FnOnce(&Fut::Ok), +{ + fn is_terminated(&self) -> bool { + self.future.is_terminated() + } +} + +impl<Fut, F> Future for InspectOk<Fut, F> +where + Fut: TryFuture, + F: FnOnce(&Fut::Ok), +{ + type Output = Result<Fut::Ok, Fut::Error>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let e = ready!(self.as_mut().future().try_poll(cx)); + if let Ok(e) = &e { + self.as_mut().f().take().expect("cannot poll InspectOk twice")(e); + } + Poll::Ready(e) + } +} diff --git a/src/future/try_future/into_future.rs b/src/future/try_future/into_future.rs new file mode 100644 index 0000000..a766d5b --- /dev/null +++ b/src/future/try_future/into_future.rs @@ -0,0 +1,36 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Future for the [`into_future`](super::TryFutureExt::into_future) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct IntoFuture<Fut> { + future: Fut, +} + +impl<Fut> IntoFuture<Fut> { + unsafe_pinned!(future: Fut); + + #[inline] + pub(super) fn new(future: Fut) -> IntoFuture<Fut> { + IntoFuture { future } + } +} + +impl<Fut: TryFuture + FusedFuture> FusedFuture for IntoFuture<Fut> { + fn is_terminated(&self) -> bool { self.future.is_terminated() } +} + +impl<Fut: TryFuture> Future for IntoFuture<Fut> { + type Output = Result<Fut::Ok, Fut::Error>; + + #[inline] + fn poll( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + self.future().try_poll(cx) + } +} diff --git a/src/future/try_future/map_err.rs b/src/future/try_future/map_err.rs new file mode 100644 index 0000000..8edebad --- /dev/null +++ b/src/future/try_future/map_err.rs @@ -0,0 +1,52 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Future for the [`map_err`](super::TryFutureExt::map_err) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct MapErr<Fut, F> { + future: Fut, + f: Option<F>, +} + +impl<Fut, F> MapErr<Fut, F> { + unsafe_pinned!(future: Fut); + unsafe_unpinned!(f: Option<F>); + + /// Creates a new MapErr. + pub(super) fn new(future: Fut, f: F) -> MapErr<Fut, F> { + MapErr { future, f: Some(f) } + } +} + +impl<Fut: Unpin, F> Unpin for MapErr<Fut, F> {} + +impl<Fut, F, E> FusedFuture for MapErr<Fut, F> + where Fut: TryFuture, + F: FnOnce(Fut::Error) -> E, +{ + fn is_terminated(&self) -> bool { self.f.is_none() } +} + +impl<Fut, F, E> Future for MapErr<Fut, F> + where Fut: TryFuture, + F: FnOnce(Fut::Error) -> E, +{ + type Output = Result<Fut::Ok, E>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + self.as_mut() + .future() + .try_poll(cx) + .map(|result| { + let f = self.as_mut().f().take() + .expect("MapErr must not be polled after it returned `Poll::Ready`"); + result.map_err(f) + }) + } +} diff --git a/src/future/try_future/map_ok.rs b/src/future/try_future/map_ok.rs new file mode 100644 index 0000000..ab28f14 --- /dev/null +++ b/src/future/try_future/map_ok.rs @@ -0,0 +1,54 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Future for the [`map_ok`](super::TryFutureExt::map_ok) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct MapOk<Fut, F> { + future: Fut, + f: Option<F>, +} + +impl<Fut, F> MapOk<Fut, F> { + unsafe_pinned!(future: Fut); + unsafe_unpinned!(f: Option<F>); + + /// Creates a new MapOk. + pub(super) fn new(future: Fut, f: F) -> MapOk<Fut, F> { + MapOk { future, f: Some(f) } + } +} + +impl<Fut: Unpin, F> Unpin for MapOk<Fut, F> {} + +impl<Fut, F, T> FusedFuture for MapOk<Fut, F> + where Fut: TryFuture, + F: FnOnce(Fut::Ok) -> T, +{ + fn is_terminated(&self) -> bool { + self.f.is_none() + } +} + +impl<Fut, F, T> Future for MapOk<Fut, F> + where Fut: TryFuture, + F: FnOnce(Fut::Ok) -> T, +{ + type Output = Result<T, Fut::Error>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + self.as_mut() + .future() + .try_poll(cx) + .map(|result| { + let op = self.as_mut().f().take() + .expect("MapOk must not be polled after it returned `Poll::Ready`"); + result.map(op) + }) + } +} diff --git a/src/future/try_future/map_ok_or_else.rs b/src/future/try_future/map_ok_or_else.rs new file mode 100644 index 0000000..730b679 --- /dev/null +++ b/src/future/try_future/map_ok_or_else.rs @@ -0,0 +1,59 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Future for the [`map_ok_or_else`](super::TryFutureExt::map_ok_or_else) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct MapOkOrElse<Fut, F, E> { + future: Fut, + f: Option<F>, + e: Option<E>, +} + +impl<Fut, F, E> MapOkOrElse<Fut, F, E> { + unsafe_pinned!(future: Fut); + unsafe_unpinned!(f: Option<F>); + unsafe_unpinned!(e: Option<E>); + + /// Creates a new MapOkOrElse. + pub(super) fn new(future: Fut, e: E, f: F) -> Self { + Self { future, f: Some(f), e: Some(e) } + } +} + +impl<Fut: Unpin, F, E> Unpin for MapOkOrElse<Fut, F, E> {} + +impl<Fut, F, E, T> FusedFuture for MapOkOrElse<Fut, F, E> + where Fut: TryFuture, + F: FnOnce(Fut::Ok) -> T, + E: FnOnce(Fut::Error) -> T, +{ + fn is_terminated(&self) -> bool { + self.f.is_none() || self.e.is_none() + } +} + +impl<Fut, F, E, T> Future for MapOkOrElse<Fut, F, E> + where Fut: TryFuture, + F: FnOnce(Fut::Ok) -> T, + E: FnOnce(Fut::Error) -> T, +{ + type Output = T; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + self.as_mut() + .future() + .try_poll(cx) + .map(|result| { + match result { + Ok(i) => (self.as_mut().f().take().expect("MapOkOrElse must not be polled after it returned `Poll::Ready`"))(i), + Err(e) => (self.as_mut().e().take().expect("MapOkOrElse must not be polled after it returned `Poll::Ready`"))(e), + } + }) + } +} diff --git a/src/future/try_future/mod.rs b/src/future/try_future/mod.rs new file mode 100644 index 0000000..e8e059e --- /dev/null +++ b/src/future/try_future/mod.rs @@ -0,0 +1,536 @@ +//! Futures +//! +//! This module contains a number of functions for working with `Future`s, +//! including the `FutureExt` trait which adds methods to `Future` types. + +#[cfg(feature = "compat")] +use crate::compat::Compat; +use core::pin::Pin; +use futures_core::{ + future::TryFuture, + stream::TryStream, + task::{Context, Poll}, +}; +#[cfg(feature = "sink")] +use futures_sink::Sink; + +// Combinators + +mod and_then; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::and_then::AndThen; + +mod err_into; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::err_into::ErrInto; + +#[cfg(feature = "sink")] +mod flatten_sink; +#[cfg(feature = "sink")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::flatten_sink::FlattenSink; + +mod inspect_ok; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::inspect_ok::InspectOk; + +mod inspect_err; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::inspect_err::InspectErr; + +mod into_future; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::into_future::IntoFuture; + +mod map_err; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::map_err::MapErr; + +mod map_ok; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::map_ok::MapOk; + +mod map_ok_or_else; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::map_ok_or_else::MapOkOrElse; + +mod or_else; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::or_else::OrElse; + +mod try_flatten_stream; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_flatten_stream::TryFlattenStream; + +mod unwrap_or_else; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::unwrap_or_else::UnwrapOrElse; + +// Implementation details + +mod flatten_stream_sink; +pub(crate) use self::flatten_stream_sink::FlattenStreamSink; + +mod try_chain; +pub(crate) use self::try_chain::{TryChain, TryChainAction}; + +impl<Fut: ?Sized + TryFuture> TryFutureExt for Fut {} + +/// Adapters specific to [`Result`]-returning futures +pub trait TryFutureExt: TryFuture { + /// Flattens the execution of this future when the successful result of this + /// future is a [`Sink`]. + /// + /// This can be useful when sink initialization is deferred, and it is + /// convenient to work with that sink as if the sink was available at the + /// call site. + /// + /// Note that this function consumes this future and returns a wrapped + /// version of it. + /// + /// # Examples + /// + /// ``` + /// use futures::future::{Future, TryFutureExt}; + /// use futures::sink::Sink; + /// # use futures::channel::mpsc::{self, SendError}; + /// # type T = i32; + /// # type E = SendError; + /// + /// fn make_sink_async() -> impl Future<Output = Result< + /// impl Sink<T, Error = E>, + /// E, + /// >> { // ... } + /// # let (tx, _rx) = mpsc::unbounded::<i32>(); + /// # futures::future::ready(Ok(tx)) + /// # } + /// fn take_sink(sink: impl Sink<T, Error = E>) { /* ... */ } + /// + /// let fut = make_sink_async(); + /// take_sink(fut.flatten_sink()) + /// ``` + #[cfg(feature = "sink")] + fn flatten_sink<Item>(self) -> FlattenSink<Self, Self::Ok> + where + Self::Ok: Sink<Item, Error = Self::Error>, + Self: Sized, + { + FlattenSink::new(self) + } + + /// Maps this future's success value to a different value. + /// + /// This method can be used to change the [`Ok`](TryFuture::Ok) type of the + /// future into a different type. It is similar to the [`Result::map`] + /// method. You can use this method to chain along a computation once the + /// future has been resolved. + /// + /// The provided closure `f` will only be called if this future is resolved + /// to an [`Ok`]. If it resolves to an [`Err`], panics, or is dropped, then + /// the provided closure will never be invoked. + /// + /// Note that this method consumes the future it is called on and returns a + /// wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future = async { Ok::<i32, i32>(1) }; + /// let future = future.map_ok(|x| x + 3); + /// assert_eq!(future.await, Ok(4)); + /// # }); + /// ``` + /// + /// Calling [`map_ok`](TryFutureExt::map_ok) on an errored future has no + /// effect: + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future = async { Err::<i32, i32>(1) }; + /// let future = future.map_ok(|x| x + 3); + /// assert_eq!(future.await, Err(1)); + /// # }); + /// ``` + fn map_ok<T, F>(self, f: F) -> MapOk<Self, F> + where + F: FnOnce(Self::Ok) -> T, + Self: Sized, + { + MapOk::new(self, f) + } + + /// Maps this future's success value to a different value, and permits for error handling resulting in the same type. + /// + /// This method can be used to coalesce your [`Ok`](TryFuture::Ok) type and [`Error`](TryFuture::Error) into another type, + /// where that type is the same for both outcomes. + /// + /// The provided closure `f` will only be called if this future is resolved + /// to an [`Ok`]. If it resolves to an [`Err`], panics, or is dropped, then + /// the provided closure will never be invoked. + /// + /// The provided closure `e` will only be called if this future is resolved + /// to an [`Err`]. If it resolves to an [`Ok`], panics, or is dropped, then + /// the provided closure will never be invoked. + /// + /// Note that this method consumes the future it is called on and returns a + /// wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future = async { Ok::<i32, i32>(5) }; + /// let future = future.map_ok_or_else(|x| x * 2, |x| x + 3); + /// assert_eq!(future.await, 8); + /// + /// let future = async { Err::<i32, i32>(5) }; + /// let future = future.map_ok_or_else(|x| x * 2, |x| x + 3); + /// assert_eq!(future.await, 10); + /// # }); + /// ``` + /// + fn map_ok_or_else<T, E, F>(self, e: E, f: F) -> MapOkOrElse<Self, F, E> + where + F: FnOnce(Self::Ok) -> T, + E: FnOnce(Self::Error) -> T, + Self: Sized, + { + MapOkOrElse::new(self, e, f) + } + + /// Maps this future's error value to a different value. + /// + /// This method can be used to change the [`Error`](TryFuture::Error) type + /// of the future into a different type. It is similar to the + /// [`Result::map_err`] method. You can use this method for example to + /// ensure that futures have the same [`Error`](TryFuture::Error) type when + /// using [`select!`] or [`join!`]. + /// + /// The provided closure `f` will only be called if this future is resolved + /// to an [`Err`]. If it resolves to an [`Ok`], panics, or is dropped, then + /// the provided closure will never be invoked. + /// + /// Note that this method consumes the future it is called on and returns a + /// wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future = async { Err::<i32, i32>(1) }; + /// let future = future.map_err(|x| x + 3); + /// assert_eq!(future.await, Err(4)); + /// # }); + /// ``` + /// + /// Calling [`map_err`](TryFutureExt::map_err) on a successful future has + /// no effect: + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future = async { Ok::<i32, i32>(1) }; + /// let future = future.map_err(|x| x + 3); + /// assert_eq!(future.await, Ok(1)); + /// # }); + /// ``` + fn map_err<E, F>(self, f: F) -> MapErr<Self, F> + where + F: FnOnce(Self::Error) -> E, + Self: Sized, + { + MapErr::new(self, f) + } + + /// Maps this future's [`Error`](TryFuture::Error) to a new error type + /// using the [`Into`](std::convert::Into) trait. + /// + /// This method does for futures what the `?`-operator does for + /// [`Result`]: It lets the compiler infer the type of the resulting + /// error. Just as [`map_err`](TryFutureExt::map_err), this is useful for + /// example to ensure that futures have the same [`Error`](TryFuture::Error) + /// type when using [`select!`] or [`join!`]. + /// + /// Note that this method consumes the future it is called on and returns a + /// wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future_err_u8 = async { Err::<(), u8>(1) }; + /// let future_err_i32 = future_err_u8.err_into::<i32>(); + /// # }); + /// ``` + fn err_into<E>(self) -> ErrInto<Self, E> + where + Self: Sized, + Self::Error: Into<E>, + { + ErrInto::new(self) + } + + /// Executes another future after this one resolves successfully. The + /// success value is passed to a closure to create this subsequent future. + /// + /// The provided closure `f` will only be called if this future is resolved + /// to an [`Ok`]. If this future resolves to an [`Err`], panics, or is + /// dropped, then the provided closure will never be invoked. The + /// [`Error`](TryFuture::Error) type of this future and the future + /// returned by `f` have to match. + /// + /// Note that this method consumes the future it is called on and returns a + /// wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future = async { Ok::<i32, i32>(1) }; + /// let future = future.and_then(|x| async move { Ok::<i32, i32>(x + 3) }); + /// assert_eq!(future.await, Ok(4)); + /// # }); + /// ``` + /// + /// Calling [`and_then`](TryFutureExt::and_then) on an errored future has no + /// effect: + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future = async { Err::<i32, i32>(1) }; + /// let future = future.and_then(|x| async move { Err::<i32, i32>(x + 3) }); + /// assert_eq!(future.await, Err(1)); + /// # }); + /// ``` + fn and_then<Fut, F>(self, f: F) -> AndThen<Self, Fut, F> + where + F: FnOnce(Self::Ok) -> Fut, + Fut: TryFuture<Error = Self::Error>, + Self: Sized, + { + AndThen::new(self, f) + } + + /// Executes another future if this one resolves to an error. The + /// error value is passed to a closure to create this subsequent future. + /// + /// The provided closure `f` will only be called if this future is resolved + /// to an [`Err`]. If this future resolves to an [`Ok`], panics, or is + /// dropped, then the provided closure will never be invoked. The + /// [`Ok`](TryFuture::Ok) type of this future and the future returned by `f` + /// have to match. + /// + /// Note that this method consumes the future it is called on and returns a + /// wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future = async { Err::<i32, i32>(1) }; + /// let future = future.or_else(|x| async move { Err::<i32, i32>(x + 3) }); + /// assert_eq!(future.await, Err(4)); + /// # }); + /// ``` + /// + /// Calling [`or_else`](TryFutureExt::or_else) on a successful future has + /// no effect: + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future = async { Ok::<i32, i32>(1) }; + /// let future = future.or_else(|x| async move { Ok::<i32, i32>(x + 3) }); + /// assert_eq!(future.await, Ok(1)); + /// # }); + /// ``` + fn or_else<Fut, F>(self, f: F) -> OrElse<Self, Fut, F> + where + F: FnOnce(Self::Error) -> Fut, + Fut: TryFuture<Ok = Self::Ok>, + Self: Sized, + { + OrElse::new(self, f) + } + + /// Do something with the success value of a future before passing it on. + /// + /// When using futures, you'll often chain several of them together. While + /// working on such code, you might want to check out what's happening at + /// various parts in the pipeline, without consuming the intermediate + /// value. To do that, insert a call to `inspect_ok`. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::TryFutureExt; + /// + /// let future = async { Ok::<_, ()>(1) }; + /// let new_future = future.inspect_ok(|&x| println!("about to resolve: {}", x)); + /// assert_eq!(new_future.await, Ok(1)); + /// # }); + /// ``` + fn inspect_ok<F>(self, f: F) -> InspectOk<Self, F> + where + F: FnOnce(&Self::Ok), + Self: Sized, + { + InspectOk::new(self, f) + } + + /// Do something with the error value of a future before passing it on. + /// + /// When using futures, you'll often chain several of them together. While + /// working on such code, you might want to check out what's happening at + /// various parts in the pipeline, without consuming the intermediate + /// value. To do that, insert a call to `inspect_err`. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::TryFutureExt; + /// + /// let future = async { Err::<(), _>(1) }; + /// let new_future = future.inspect_err(|&x| println!("about to error: {}", x)); + /// assert_eq!(new_future.await, Err(1)); + /// # }); + /// ``` + fn inspect_err<F>(self, f: F) -> InspectErr<Self, F> + where + F: FnOnce(&Self::Error), + Self: Sized, + { + InspectErr::new(self, f) + } + + /// Flatten the execution of this future when the successful result of this + /// future is a stream. + /// + /// This can be useful when stream initialization is deferred, and it is + /// convenient to work with that stream as if stream was available at the + /// call site. + /// + /// Note that this function consumes this future and returns a wrapped + /// version of it. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future::TryFutureExt; + /// use futures::stream::{self, TryStreamExt}; + /// + /// let stream_items = vec![17, 18, 19].into_iter().map(Ok); + /// let future_of_a_stream = async { Ok::<_, ()>(stream::iter(stream_items)) }; + /// + /// let stream = future_of_a_stream.try_flatten_stream(); + /// let list = stream.try_collect::<Vec<_>>().await; + /// assert_eq!(list, Ok(vec![17, 18, 19])); + /// # }); + /// ``` + fn try_flatten_stream(self) -> TryFlattenStream<Self> + where + Self::Ok: TryStream<Error = Self::Error>, + Self: Sized, + { + TryFlattenStream::new(self) + } + + /// Unwraps this future's ouput, producing a future with this future's + /// [`Ok`](TryFuture::Ok) type as its + /// [`Output`](std::future::Future::Output) type. + /// + /// If this future is resolved successfully, the returned future will + /// contain the original future's success value as output. Otherwise, the + /// closure `f` is called with the error value to produce an alternate + /// success value. + /// + /// This method is similar to the [`Result::unwrap_or_else`] method. + /// + /// # Examples + /// + /// ``` + /// use futures::future::TryFutureExt; + /// + /// # futures::executor::block_on(async { + /// let future = async { Err::<(), &str>("Boom!") }; + /// let future = future.unwrap_or_else(|_| ()); + /// assert_eq!(future.await, ()); + /// # }); + /// ``` + fn unwrap_or_else<F>(self, f: F) -> UnwrapOrElse<Self, F> + where + Self: Sized, + F: FnOnce(Self::Error) -> Self::Ok, + { + UnwrapOrElse::new(self, f) + } + + /// Wraps a [`TryFuture`] into a future compatable with libraries using + /// futures 0.1 future definitons. Requires the `compat` feature to enable. + #[cfg(feature = "compat")] + fn compat(self) -> Compat<Self> + where + Self: Sized + Unpin, + { + Compat::new(self) + } + + /// Wraps a [`TryFuture`] into a type that implements + /// [`Future`](std::future::Future). + /// + /// [`TryFuture`]s currently do not implement the + /// [`Future`](std::future::Future) trait due to limitations of the + /// compiler. + /// + /// # Examples + /// + /// ``` + /// use futures::future::{Future, TryFuture, TryFutureExt}; + /// + /// # type T = i32; + /// # type E = (); + /// fn make_try_future() -> impl TryFuture<Ok = T, Error = E> { // ... } + /// # async { Ok::<i32, ()>(1) } + /// # } + /// fn take_future(future: impl Future<Output = Result<T, E>>) { /* ... */ } + /// + /// take_future(make_try_future().into_future()); + /// ``` + fn into_future(self) -> IntoFuture<Self> + where + Self: Sized, + { + IntoFuture::new(self) + } + + /// A convenience method for calling [`TryFuture::try_poll`] on [`Unpin`] + /// future types. + fn try_poll_unpin(&mut self, cx: &mut Context<'_>) -> Poll<Result<Self::Ok, Self::Error>> + where + Self: Unpin, + { + Pin::new(self).try_poll(cx) + } +} diff --git a/src/future/try_future/or_else.rs b/src/future/try_future/or_else.rs new file mode 100644 index 0000000..a9c006f --- /dev/null +++ b/src/future/try_future/or_else.rs @@ -0,0 +1,56 @@ +use super::{TryChain, TryChainAction}; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Future for the [`or_else`](super::TryFutureExt::or_else) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct OrElse<Fut1, Fut2, F> { + try_chain: TryChain<Fut1, Fut2, F>, +} + +impl<Fut1, Fut2, F> OrElse<Fut1, Fut2, F> + where Fut1: TryFuture, + Fut2: TryFuture, +{ + unsafe_pinned!(try_chain: TryChain<Fut1, Fut2, F>); + + /// Creates a new `Then`. + pub(super) fn new(future: Fut1, f: F) -> OrElse<Fut1, Fut2, F> { + OrElse { + try_chain: TryChain::new(future, f), + } + } +} + +impl<Fut1, Fut2, F> FusedFuture for OrElse<Fut1, Fut2, F> + where Fut1: TryFuture, + Fut2: TryFuture<Ok = Fut1::Ok>, + F: FnOnce(Fut1::Error) -> Fut2, +{ + fn is_terminated(&self) -> bool { + self.try_chain.is_terminated() + } +} + +impl<Fut1, Fut2, F> Future for OrElse<Fut1, Fut2, F> + where Fut1: TryFuture, + Fut2: TryFuture<Ok = Fut1::Ok>, + F: FnOnce(Fut1::Error) -> Fut2, +{ + type Output = Result<Fut2::Ok, Fut2::Error>; + + fn poll( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + self.try_chain().poll(cx, |result, async_op| { + match result { + Ok(ok) => TryChainAction::Output(Ok(ok)), + Err(err) => TryChainAction::Future(async_op(err)), + } + }) + } +} diff --git a/src/future/try_future/try_chain.rs b/src/future/try_future/try_chain.rs new file mode 100644 index 0000000..662bdf2 --- /dev/null +++ b/src/future/try_future/try_chain.rs @@ -0,0 +1,108 @@ +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::task::{Context, Poll}; + +#[must_use = "futures do nothing unless you `.await` or poll them"] +#[derive(Debug)] +pub(crate) enum TryChain<Fut1, Fut2, Data> { + First(Fut1, Option<Data>), + Second(Fut2), + Empty, +} + +impl<Fut1: Unpin, Fut2: Unpin, Data> Unpin for TryChain<Fut1, Fut2, Data> {} + +pub(crate) enum TryChainAction<Fut2> + where Fut2: TryFuture, +{ + Future(Fut2), + Output(Result<Fut2::Ok, Fut2::Error>), +} + +impl<Fut1, Fut2, Data> TryChain<Fut1, Fut2, Data> + where Fut1: TryFuture, + Fut2: TryFuture, +{ + pub(crate) fn new(fut1: Fut1, data: Data) -> TryChain<Fut1, Fut2, Data> { + TryChain::First(fut1, Some(data)) + } + + pub(crate) fn is_terminated(&self) -> bool { + match self { + TryChain::First(..) | TryChain::Second(_) => false, + TryChain::Empty => true, + } + } + + pub(crate) fn poll<F>( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + f: F, + ) -> Poll<Result<Fut2::Ok, Fut2::Error>> + where F: FnOnce(Result<Fut1::Ok, Fut1::Error>, Data) -> TryChainAction<Fut2>, + { + let mut f = Some(f); + + // Safe to call `get_unchecked_mut` because we won't move the futures. + let this = unsafe { self.get_unchecked_mut() }; + + loop { + let (output, data) = match this { + TryChain::First(fut1, data) => { + // Poll the first future + let output = ready!(unsafe { Pin::new_unchecked(fut1) }.try_poll(cx)); + (output, data.take().unwrap()) + } + TryChain::Second(fut2) => { + // Poll the second future + return unsafe { Pin::new_unchecked(fut2) } + .try_poll(cx) + .map(|res| { + *this = TryChain::Empty; // Drop fut2. + res + }); + } + TryChain::Empty => { + panic!("future must not be polled after it returned `Poll::Ready`"); + } + }; + + *this = TryChain::Empty; // Drop fut1 + let f = f.take().unwrap(); + match f(output, data) { + TryChainAction::Future(fut2) => *this = TryChain::Second(fut2), + TryChainAction::Output(output) => return Poll::Ready(output), + } + } + } +} + +#[cfg(test)] +mod tests { + use std::pin::Pin; + use std::task::Poll; + + use futures_test::task::noop_context; + + use crate::future::ready; + + use super::{TryChain, TryChainAction}; + + #[test] + fn try_chain_is_terminated() { + let mut cx = noop_context(); + + let mut future = TryChain::new(ready(Ok(1)), ()); + assert!(!future.is_terminated()); + + let res = Pin::new(&mut future).poll( + &mut cx, + |res: Result<usize, ()>, ()| { + assert!(res.is_ok()); + TryChainAction::Future(ready(Ok(2))) + }, + ); + assert_eq!(res, Poll::Ready::<Result<usize, ()>>(Ok(2))); + assert!(future.is_terminated()); + } +} diff --git a/src/future/try_future/try_flatten_stream.rs b/src/future/try_future/try_flatten_stream.rs new file mode 100644 index 0000000..2462431 --- /dev/null +++ b/src/future/try_future/try_flatten_stream.rs @@ -0,0 +1,91 @@ +use super::FlattenStreamSink; +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::stream::{FusedStream, Stream, TryStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::unsafe_pinned; + +/// Stream for the [`try_flatten_stream`](super::TryFutureExt::try_flatten_stream) method. +#[must_use = "streams do nothing unless polled"] +pub struct TryFlattenStream<Fut> +where + Fut: TryFuture, +{ + inner: FlattenStreamSink<Fut>, +} + +impl<Fut: TryFuture> TryFlattenStream<Fut> +where + Fut: TryFuture, + Fut::Ok: TryStream<Error = Fut::Error>, +{ + unsafe_pinned!(inner: FlattenStreamSink<Fut>); + + pub(super) fn new(future: Fut) -> Self { + Self { + inner: FlattenStreamSink::new(future), + } + } +} + +impl<Fut> fmt::Debug for TryFlattenStream<Fut> +where + Fut: TryFuture + fmt::Debug, + Fut::Ok: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TryFlattenStream") + .field("inner", &self.inner) + .finish() + } +} + +impl<Fut> FusedStream for TryFlattenStream<Fut> +where + Fut: TryFuture, + Fut::Ok: TryStream<Error = Fut::Error> + FusedStream, +{ + fn is_terminated(&self) -> bool { + self.inner.is_terminated() + } +} + +impl<Fut> Stream for TryFlattenStream<Fut> +where + Fut: TryFuture, + Fut::Ok: TryStream<Error = Fut::Error>, +{ + type Item = Result<<Fut::Ok as TryStream>::Ok, Fut::Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + self.inner().poll_next(cx) + } +} + +#[cfg(feature = "sink")] +impl<Fut, Item> Sink<Item> for TryFlattenStream<Fut> +where + Fut: TryFuture, + Fut::Ok: TryStream<Error = Fut::Error> + Sink<Item, Error = Fut::Error>, +{ + type Error = Fut::Error; + + fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { + self.inner().poll_ready(cx) + } + + fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> { + self.inner().start_send(item) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { + self.inner().poll_flush(cx) + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { + self.inner().poll_close(cx) + } +} diff --git a/src/future/try_future/unwrap_or_else.rs b/src/future/try_future/unwrap_or_else.rs new file mode 100644 index 0000000..286cc00 --- /dev/null +++ b/src/future/try_future/unwrap_or_else.rs @@ -0,0 +1,55 @@ +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Future for the [`unwrap_or_else`](super::TryFutureExt::unwrap_or_else) +/// method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct UnwrapOrElse<Fut, F> { + future: Fut, + f: Option<F>, +} + +impl<Fut, F> UnwrapOrElse<Fut, F> { + unsafe_pinned!(future: Fut); + unsafe_unpinned!(f: Option<F>); + + /// Creates a new UnwrapOrElse. + pub(super) fn new(future: Fut, f: F) -> UnwrapOrElse<Fut, F> { + UnwrapOrElse { future, f: Some(f) } + } +} + +impl<Fut: Unpin, F> Unpin for UnwrapOrElse<Fut, F> {} + +impl<Fut, F> FusedFuture for UnwrapOrElse<Fut, F> + where Fut: TryFuture, + F: FnOnce(Fut::Error) -> Fut::Ok, +{ + fn is_terminated(&self) -> bool { + self.f.is_none() + } +} + +impl<Fut, F> Future for UnwrapOrElse<Fut, F> + where Fut: TryFuture, + F: FnOnce(Fut::Error) -> Fut::Ok, +{ + type Output = Fut::Ok; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + self.as_mut() + .future() + .try_poll(cx) + .map(|result| { + let op = self.as_mut().f().take() + .expect("UnwrapOrElse already returned `Poll::Ready` before"); + result.unwrap_or_else(op) + }) + } +} diff --git a/src/future/try_join.rs b/src/future/try_join.rs new file mode 100644 index 0000000..da85eff --- /dev/null +++ b/src/future/try_join.rs @@ -0,0 +1,262 @@ +#![allow(non_snake_case)] + +use crate::future::{MaybeDone, maybe_done, TryFutureExt, IntoFuture}; +use core::fmt; +use core::pin::Pin; +use futures_core::future::{Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +macro_rules! generate { + ($( + $(#[$doc:meta])* + ($Join:ident, <Fut1, $($Fut:ident),*>), + )*) => ($( + $(#[$doc])* + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct $Join<Fut1: TryFuture, $($Fut: TryFuture),*> { + Fut1: MaybeDone<IntoFuture<Fut1>>, + $($Fut: MaybeDone<IntoFuture<$Fut>>,)* + } + + impl<Fut1, $($Fut),*> fmt::Debug for $Join<Fut1, $($Fut),*> + where + Fut1: TryFuture + fmt::Debug, + Fut1::Ok: fmt::Debug, + Fut1::Error: fmt::Debug, + $( + $Fut: TryFuture + fmt::Debug, + $Fut::Ok: fmt::Debug, + $Fut::Error: fmt::Debug, + )* + { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct(stringify!($Join)) + .field("Fut1", &self.Fut1) + $(.field(stringify!($Fut), &self.$Fut))* + .finish() + } + } + + impl<Fut1, $($Fut),*> $Join<Fut1, $($Fut),*> + where + Fut1: TryFuture, + $( + $Fut: TryFuture<Error=Fut1::Error> + ),* + { + fn new(Fut1: Fut1, $($Fut: $Fut),*) -> $Join<Fut1, $($Fut),*> { + $Join { + Fut1: maybe_done(TryFutureExt::into_future(Fut1)), + $($Fut: maybe_done(TryFutureExt::into_future($Fut))),* + } + } + + unsafe_pinned!(Fut1: MaybeDone<IntoFuture<Fut1>>); + $( + unsafe_pinned!($Fut: MaybeDone<IntoFuture<$Fut>>); + )* + } + + impl<Fut1, $($Fut),*> Future for $Join<Fut1, $($Fut),*> + where + Fut1: TryFuture, + $( + $Fut: TryFuture<Error=Fut1::Error> + ),* + { + type Output = Result<(Fut1::Ok, $($Fut::Ok),*), Fut1::Error>; + + fn poll( + mut self: Pin<&mut Self>, cx: &mut Context<'_> + ) -> Poll<Self::Output> { + let mut all_done = true; + if self.as_mut().Fut1().poll(cx).is_pending() { + all_done = false; + } else if self.as_mut().Fut1().output_mut().unwrap().is_err() { + return Poll::Ready(Err( + self.as_mut().Fut1().take_output().unwrap().err().unwrap())); + } + $( + if self.as_mut().$Fut().poll(cx).is_pending() { + all_done = false; + } else if self.as_mut().$Fut().output_mut().unwrap().is_err() { + return Poll::Ready(Err( + self.as_mut().$Fut().take_output().unwrap().err().unwrap())); + } + )* + + if all_done { + Poll::Ready(Ok(( + self.as_mut().Fut1().take_output().unwrap().ok().unwrap(), + $( + self.as_mut().$Fut().take_output().unwrap().ok().unwrap() + ),* + ))) + } else { + Poll::Pending + } + } + } + )*) +} + +generate! { + /// Future for the [`try_join`](try_join()) function. + (TryJoin, <Fut1, Fut2>), + + /// Future for the [`try_join3`] function. + (TryJoin3, <Fut1, Fut2, Fut3>), + + /// Future for the [`try_join4`] function. + (TryJoin4, <Fut1, Fut2, Fut3, Fut4>), + + /// Future for the [`try_join5`] function. + (TryJoin5, <Fut1, Fut2, Fut3, Fut4, Fut5>), +} + +/// Joins the result of two futures, waiting for them both to complete or +/// for one to produce an error. +/// +/// This function will return a new future which awaits both futures to +/// complete. If successful, the returned future will finish with a tuple of +/// both results. If unsuccesful, it will complete with the first error +/// encountered. +/// +/// Note that this function consumes the passed futures and returns a +/// wrapped version of it. +/// +/// # Examples +/// +/// When used on multiple futures that return [`Ok`], `try_join` will return +/// [`Ok`] of a tuple of the values: +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = future::ready(Ok::<i32, i32>(1)); +/// let b = future::ready(Ok::<i32, i32>(2)); +/// let pair = future::try_join(a, b); +/// +/// assert_eq!(pair.await, Ok((1, 2))); +/// # }); +/// ``` +/// +/// If one of the futures resolves to an error, `try_join` will return +/// that error: +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = future::ready(Ok::<i32, i32>(1)); +/// let b = future::ready(Err::<i32, i32>(2)); +/// let pair = future::try_join(a, b); +/// +/// assert_eq!(pair.await, Err(2)); +/// # }); +/// ``` +pub fn try_join<Fut1, Fut2>(future1: Fut1, future2: Fut2) -> TryJoin<Fut1, Fut2> +where + Fut1: TryFuture, + Fut2: TryFuture<Error = Fut1::Error>, +{ + TryJoin::new(future1, future2) +} + +/// Same as [`try_join`](try_join()), but with more futures. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = future::ready(Ok::<i32, i32>(1)); +/// let b = future::ready(Ok::<i32, i32>(2)); +/// let c = future::ready(Ok::<i32, i32>(3)); +/// let tuple = future::try_join3(a, b, c); +/// +/// assert_eq!(tuple.await, Ok((1, 2, 3))); +/// # }); +/// ``` +pub fn try_join3<Fut1, Fut2, Fut3>( + future1: Fut1, + future2: Fut2, + future3: Fut3, +) -> TryJoin3<Fut1, Fut2, Fut3> +where + Fut1: TryFuture, + Fut2: TryFuture<Error = Fut1::Error>, + Fut3: TryFuture<Error = Fut1::Error>, +{ + TryJoin3::new(future1, future2, future3) +} + +/// Same as [`try_join`](try_join()), but with more futures. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = future::ready(Ok::<i32, i32>(1)); +/// let b = future::ready(Ok::<i32, i32>(2)); +/// let c = future::ready(Ok::<i32, i32>(3)); +/// let d = future::ready(Ok::<i32, i32>(4)); +/// let tuple = future::try_join4(a, b, c, d); +/// +/// assert_eq!(tuple.await, Ok((1, 2, 3, 4))); +/// # }); +/// ``` +pub fn try_join4<Fut1, Fut2, Fut3, Fut4>( + future1: Fut1, + future2: Fut2, + future3: Fut3, + future4: Fut4, +) -> TryJoin4<Fut1, Fut2, Fut3, Fut4> +where + Fut1: TryFuture, + Fut2: TryFuture<Error = Fut1::Error>, + Fut3: TryFuture<Error = Fut1::Error>, + Fut4: TryFuture<Error = Fut1::Error>, +{ + TryJoin4::new(future1, future2, future3, future4) +} + +/// Same as [`try_join`](try_join()), but with more futures. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future; +/// +/// let a = future::ready(Ok::<i32, i32>(1)); +/// let b = future::ready(Ok::<i32, i32>(2)); +/// let c = future::ready(Ok::<i32, i32>(3)); +/// let d = future::ready(Ok::<i32, i32>(4)); +/// let e = future::ready(Ok::<i32, i32>(5)); +/// let tuple = future::try_join5(a, b, c, d, e); +/// +/// assert_eq!(tuple.await, Ok((1, 2, 3, 4, 5))); +/// # }); +/// ``` +pub fn try_join5<Fut1, Fut2, Fut3, Fut4, Fut5>( + future1: Fut1, + future2: Fut2, + future3: Fut3, + future4: Fut4, + future5: Fut5, +) -> TryJoin5<Fut1, Fut2, Fut3, Fut4, Fut5> +where + Fut1: TryFuture, + Fut2: TryFuture<Error = Fut1::Error>, + Fut3: TryFuture<Error = Fut1::Error>, + Fut4: TryFuture<Error = Fut1::Error>, + Fut5: TryFuture<Error = Fut1::Error>, +{ + TryJoin5::new(future1, future2, future3, future4, future5) +} diff --git a/src/future/try_join_all.rs b/src/future/try_join_all.rs new file mode 100644 index 0000000..30300e4 --- /dev/null +++ b/src/future/try_join_all.rs @@ -0,0 +1,179 @@ +//! Definition of the `TryJoinAll` combinator, waiting for all of a list of +//! futures to finish with either success or error. + +use core::fmt; +use core::future::Future; +use core::iter::FromIterator; +use core::mem; +use core::pin::Pin; +use core::task::{Context, Poll}; +use alloc::boxed::Box; +use alloc::vec::Vec; + +use super::TryFuture; + +#[derive(Debug)] +enum ElemState<F> +where + F: TryFuture, +{ + Pending(F), + Done(Option<F::Ok>), +} + +impl<F> ElemState<F> +where + F: TryFuture, +{ + fn pending_pin_mut(self: Pin<&mut Self>) -> Option<Pin<&mut F>> { + // Safety: Basic enum pin projection, no drop + optionally Unpin based + // on the type of this variant + match unsafe { self.get_unchecked_mut() } { + ElemState::Pending(f) => Some(unsafe { Pin::new_unchecked(f) }), + ElemState::Done(_) => None, + } + } + + fn take_done(self: Pin<&mut Self>) -> Option<F::Ok> { + // Safety: Going from pin to a variant we never pin-project + match unsafe { self.get_unchecked_mut() } { + ElemState::Pending(_) => None, + ElemState::Done(output) => output.take(), + } + } +} + +impl<F> Unpin for ElemState<F> where F: TryFuture + Unpin {} + +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) }) +} + +enum FinalState<E = ()> { + Pending, + AllDone, + Error(E) +} + +/// Future for the [`try_join_all`] function. +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct TryJoinAll<F> +where + F: TryFuture, +{ + elems: Pin<Box<[ElemState<F>]>>, +} + +impl<F> fmt::Debug for TryJoinAll<F> +where + F: TryFuture + fmt::Debug, + F::Ok: fmt::Debug, + F::Error: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TryJoinAll") + .field("elems", &self.elems) + .finish() + } +} + +/// 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. +/// +/// If any future returns an error then all other futures will be canceled and +/// an error will be returned immediately. If all futures complete successfully, +/// however, then the returned future will succeed with a `Vec` of all the +/// successful results. +/// +/// This function is only available when the `std` or `alloc` feature of this +/// library is activated, and it is activated by default. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future::{self, try_join_all}; +/// +/// let futures = vec![ +/// future::ok::<u32, u32>(1), +/// future::ok::<u32, u32>(2), +/// future::ok::<u32, u32>(3), +/// ]; +/// +/// assert_eq!(try_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!(try_join_all(futures).await, Err(2)); +/// # }); +/// ``` +pub fn try_join_all<I>(i: I) -> TryJoinAll<I::Item> +where + I: IntoIterator, + I::Item: TryFuture, +{ + let elems: Box<[_]> = i.into_iter().map(ElemState::Pending).collect(); + TryJoinAll { + elems: elems.into(), + } +} + +impl<F> Future for TryJoinAll<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> { + let mut state = FinalState::AllDone; + + for mut elem in iter_pin_mut(self.elems.as_mut()) { + if let Some(pending) = elem.as_mut().pending_pin_mut() { + match pending.try_poll(cx) { + Poll::Pending => state = FinalState::Pending, + Poll::Ready(output) => match output { + Ok(item) => elem.set(ElemState::Done(Some(item))), + Err(e) => { + state = FinalState::Error(e); + break; + } + } + } + } + } + + match state { + FinalState::Pending => Poll::Pending, + FinalState::AllDone => { + let mut elems = mem::replace(&mut self.elems, Box::pin([])); + let results = iter_pin_mut(elems.as_mut()) + .map(|e| e.take_done().unwrap()) + .collect(); + Poll::Ready(Ok(results)) + }, + FinalState::Error(e) => { + let _ = mem::replace(&mut self.elems, Box::pin([])); + Poll::Ready(Err(e)) + }, + } + } +} + +impl<F: TryFuture> FromIterator<F> for TryJoinAll<F> { + fn from_iter<T: IntoIterator<Item = F>>(iter: T) -> Self { + try_join_all(iter) + } +} diff --git a/src/future/try_select.rs b/src/future/try_select.rs new file mode 100644 index 0000000..56564f5 --- /dev/null +++ b/src/future/try_select.rs @@ -0,0 +1,80 @@ +use core::pin::Pin; +use futures_core::future::{Future, TryFuture}; +use futures_core::task::{Context, Poll}; +use crate::future::{Either, TryFutureExt}; + +/// Future for the [`try_select()`] function. +#[must_use = "futures do nothing unless you `.await` or poll them"] +#[derive(Debug)] +pub struct TrySelect<A, B> { + inner: Option<(A, B)>, +} + +impl<A: Unpin, B: Unpin> Unpin for TrySelect<A, B> {} + +/// Waits for either one of two differently-typed futures to complete. +/// +/// This function will return a new future which awaits for either one of both +/// futures to complete. The returned future will finish with both the value +/// resolved and a future representing the completion of the other work. +/// +/// Note that this function consumes the receiving futures and returns a +/// wrapped version of them. +/// +/// Also note that if both this and the second future have the same +/// success/error type you can use the `Either::factor_first` method to +/// conveniently extract out the value at the end. +/// +/// # Examples +/// +/// ``` +/// use futures::future::{self, Either, Future, FutureExt, TryFuture, TryFutureExt}; +/// +/// // A poor-man's try_join implemented on top of select +/// +/// fn try_join<A, B, E>(a: A, b: B) -> impl TryFuture<Ok=(A::Ok, B::Ok), Error=E> +/// where A: TryFuture<Error = E> + Unpin + 'static, +/// B: TryFuture<Error = E> + Unpin + 'static, +/// E: 'static, +/// { +/// future::try_select(a, b).then(|res| -> Box<dyn Future<Output = Result<_, _>> + Unpin> { +/// match res { +/// Ok(Either::Left((x, b))) => Box::new(b.map_ok(move |y| (x, y))), +/// Ok(Either::Right((y, a))) => Box::new(a.map_ok(move |x| (x, y))), +/// Err(Either::Left((e, _))) => Box::new(future::err(e)), +/// Err(Either::Right((e, _))) => Box::new(future::err(e)), +/// } +/// }) +/// } +/// ``` +pub fn try_select<A, B>(future1: A, future2: B) -> TrySelect<A, B> + where A: TryFuture + Unpin, B: TryFuture + Unpin +{ + TrySelect { inner: Some((future1, future2)) } +} + +impl<A: Unpin, B: Unpin> Future for TrySelect<A, B> + where A: TryFuture, B: TryFuture +{ + #[allow(clippy::type_complexity)] + type Output = Result< + Either<(A::Ok, B), (B::Ok, A)>, + Either<(A::Error, B), (B::Error, A)>, + >; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let (mut a, mut b) = self.inner.take().expect("cannot poll Select twice"); + match a.try_poll_unpin(cx) { + Poll::Ready(Err(x)) => Poll::Ready(Err(Either::Left((x, b)))), + Poll::Ready(Ok(x)) => Poll::Ready(Ok(Either::Left((x, b)))), + Poll::Pending => match b.try_poll_unpin(cx) { + Poll::Ready(Err(x)) => Poll::Ready(Err(Either::Right((x, a)))), + Poll::Ready(Ok(x)) => Poll::Ready(Ok(Either::Right((x, a)))), + Poll::Pending => { + self.inner = Some((a, b)); + Poll::Pending + } + } + } + } +} diff --git a/src/io/allow_std.rs b/src/io/allow_std.rs new file mode 100644 index 0000000..346e9ba --- /dev/null +++ b/src/io/allow_std.rs @@ -0,0 +1,179 @@ +use futures_core::task::{Context, Poll}; +#[cfg(feature = "read-initializer")] +use futures_io::Initializer; +use futures_io::{AsyncRead, AsyncWrite, AsyncSeek, AsyncBufRead, IoSlice, IoSliceMut, SeekFrom}; +use std::{fmt, io}; +use std::pin::Pin; + +/// A simple wrapper type which allows types which implement only +/// implement `std::io::Read` or `std::io::Write` +/// to be used in contexts which expect an `AsyncRead` or `AsyncWrite`. +/// +/// If these types issue an error with the kind `io::ErrorKind::WouldBlock`, +/// it is expected that they will notify the current task on readiness. +/// Synchronous `std` types should not issue errors of this kind and +/// are safe to use in this context. However, using these types with +/// `AllowStdIo` will cause the event loop to block, so they should be used +/// with care. +#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)] +pub struct AllowStdIo<T>(T); + +impl<T> Unpin for AllowStdIo<T> {} + +macro_rules! try_with_interrupt { + ($e:expr) => { + loop { + match $e { + Ok(e) => { + break e; + } + Err(ref e) if e.kind() == ::std::io::ErrorKind::Interrupted => { + continue; + } + Err(e) => { + return Poll::Ready(Err(e)); + } + } + } + } +} + +impl<T> AllowStdIo<T> { + /// Creates a new `AllowStdIo` from an existing IO object. + pub fn new(io: T) -> Self { + AllowStdIo(io) + } + + /// Returns a reference to the contained IO object. + pub fn get_ref(&self) -> &T { + &self.0 + } + + /// Returns a mutable reference to the contained IO object. + pub fn get_mut(&mut self) -> &mut T { + &mut self.0 + } + + /// Consumes self and returns the contained IO object. + pub fn into_inner(self) -> T { + self.0 + } +} + +impl<T> io::Write for AllowStdIo<T> where T: io::Write { + fn write(&mut self, buf: &[u8]) -> io::Result<usize> { + self.0.write(buf) + } + fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> { + self.0.write_vectored(bufs) + } + fn flush(&mut self) -> io::Result<()> { + self.0.flush() + } + fn write_all(&mut self, buf: &[u8]) -> io::Result<()> { + self.0.write_all(buf) + } + fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> { + self.0.write_fmt(fmt) + } +} + +impl<T> AsyncWrite for AllowStdIo<T> where T: io::Write { + fn poll_write(mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &[u8]) + -> Poll<io::Result<usize>> + { + Poll::Ready(Ok(try_with_interrupt!(self.0.write(buf)))) + } + + fn poll_write_vectored(mut self: Pin<&mut Self>, _: &mut Context<'_>, bufs: &[IoSlice<'_>]) + -> Poll<io::Result<usize>> + { + Poll::Ready(Ok(try_with_interrupt!(self.0.write_vectored(bufs)))) + } + + fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + try_with_interrupt!(self.0.flush()); + Poll::Ready(Ok(())) + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.poll_flush(cx) + } +} + +impl<T> io::Read for AllowStdIo<T> where T: io::Read { + fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { + self.0.read(buf) + } + fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> { + self.0.read_vectored(bufs) + } + #[cfg(feature = "read-initializer")] + unsafe fn initializer(&self) -> Initializer { + self.0.initializer() + } + fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> { + self.0.read_to_end(buf) + } + fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> { + self.0.read_to_string(buf) + } + fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> { + self.0.read_exact(buf) + } +} + +impl<T> AsyncRead for AllowStdIo<T> where T: io::Read { + fn poll_read(mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &mut [u8]) + -> Poll<io::Result<usize>> + { + Poll::Ready(Ok(try_with_interrupt!(self.0.read(buf)))) + } + + fn poll_read_vectored(mut self: Pin<&mut Self>, _: &mut Context<'_>, bufs: &mut [IoSliceMut<'_>]) + -> Poll<io::Result<usize>> + { + Poll::Ready(Ok(try_with_interrupt!(self.0.read_vectored(bufs)))) + } + + #[cfg(feature = "read-initializer")] + unsafe fn initializer(&self) -> Initializer { + self.0.initializer() + } +} + +impl<T> io::Seek for AllowStdIo<T> where T: io::Seek { + fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> { + self.0.seek(pos) + } +} + +impl<T> AsyncSeek for AllowStdIo<T> where T: io::Seek { + fn poll_seek(mut self: Pin<&mut Self>, _: &mut Context<'_>, pos: SeekFrom) + -> Poll<io::Result<u64>> + { + Poll::Ready(Ok(try_with_interrupt!(self.0.seek(pos)))) + } +} + +impl<T> io::BufRead for AllowStdIo<T> where T: io::BufRead { + fn fill_buf(&mut self) -> io::Result<&[u8]> { + self.0.fill_buf() + } + fn consume(&mut self, amt: usize) { + self.0.consume(amt) + } +} + +impl<T> AsyncBufRead for AllowStdIo<T> where T: io::BufRead { + fn poll_fill_buf(mut self: Pin<&mut Self>, _: &mut Context<'_>) + -> Poll<io::Result<&[u8]>> + { + let this: *mut Self = &mut *self as *mut _; + Poll::Ready(Ok(try_with_interrupt!(unsafe { &mut *this }.0.fill_buf()))) + } + + fn consume(mut self: Pin<&mut Self>, amt: usize) { + self.0.consume(amt) + } +} diff --git a/src/io/buf_reader.rs b/src/io/buf_reader.rs new file mode 100644 index 0000000..96d3f28 --- /dev/null +++ b/src/io/buf_reader.rs @@ -0,0 +1,259 @@ +use futures_core::task::{Context, Poll}; +#[cfg(feature = "read-initializer")] +use futures_io::Initializer; +use futures_io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, IoSlice, IoSliceMut, SeekFrom}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; +use std::io::{self, Read}; +use std::pin::Pin; +use std::{cmp, fmt}; +use super::DEFAULT_BUF_SIZE; + +/// The `BufReader` struct adds buffering to any reader. +/// +/// It can be excessively inefficient to work directly with a [`AsyncRead`] +/// instance. A `BufReader` performs large, infrequent reads on the underlying +/// [`AsyncRead`] and maintains an in-memory buffer of the results. +/// +/// `BufReader` can improve the speed of programs that make *small* and +/// *repeated* read calls to the same file or network socket. It does not +/// help when reading very large amounts at once, or reading just one or a few +/// times. It also provides no advantage when reading from a source that is +/// already in memory, like a `Vec<u8>`. +/// +/// When the `BufReader` is dropped, the contents of its buffer will be +/// discarded. Creating multiple instances of a `BufReader` on the same +/// stream can cause data loss. +/// +/// [`AsyncRead`]: futures_io::AsyncRead +/// +// TODO: Examples +pub struct BufReader<R> { + inner: R, + buf: Box<[u8]>, + pos: usize, + cap: usize, +} + +impl<R> BufReader<R> { + unsafe_pinned!(inner: R); + unsafe_unpinned!(pos: usize); + unsafe_unpinned!(cap: usize); +} + +impl<R: AsyncRead> BufReader<R> { + /// Creates a new `BufReader` with a default buffer capacity. The default is currently 8 KB, + /// but may change in the future. + pub fn new(inner: R) -> Self { + Self::with_capacity(DEFAULT_BUF_SIZE, inner) + } + + /// Creates a new `BufReader` with the specified buffer capacity. + pub fn with_capacity(capacity: usize, inner: R) -> Self { + unsafe { + let mut buffer = Vec::with_capacity(capacity); + buffer.set_len(capacity); + super::initialize(&inner, &mut buffer); + Self { + inner, + buf: buffer.into_boxed_slice(), + pos: 0, + cap: 0, + } + } + } + + /// Gets a reference to the underlying reader. + /// + /// It is inadvisable to directly read from the underlying reader. + pub fn get_ref(&self) -> &R { + &self.inner + } + + /// Gets a mutable reference to the underlying reader. + /// + /// It is inadvisable to directly read from the underlying reader. + pub fn get_mut(&mut self) -> &mut R { + &mut self.inner + } + + /// Gets a pinned mutable reference to the underlying reader. + /// + /// It is inadvisable to directly read from the underlying reader. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> { + self.inner() + } + + /// Consumes this `BufWriter`, returning the underlying reader. + /// + /// Note that any leftover data in the internal buffer is lost. + pub fn into_inner(self) -> R { + self.inner + } + + /// Returns a reference to the internally buffered data. + /// + /// Unlike `fill_buf`, this will not attempt to fill the buffer if it is empty. + pub fn buffer(&self) -> &[u8] { + &self.buf[self.pos..self.cap] + } + + /// Invalidates all data in the internal buffer. + #[inline] + fn discard_buffer(mut self: Pin<&mut Self>) { + *self.as_mut().pos() = 0; + *self.cap() = 0; + } +} + +impl<R: AsyncRead> AsyncRead for BufReader<R> { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll<io::Result<usize>> { + // If we don't have any buffered data and we're doing a massive read + // (larger than our internal buffer), bypass our internal buffer + // entirely. + if self.pos == self.cap && buf.len() >= self.buf.len() { + let res = ready!(self.as_mut().inner().poll_read(cx, buf)); + self.discard_buffer(); + return Poll::Ready(res); + } + let mut rem = ready!(self.as_mut().poll_fill_buf(cx))?; + let nread = rem.read(buf)?; + self.consume(nread); + Poll::Ready(Ok(nread)) + } + + fn poll_read_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &mut [IoSliceMut<'_>], + ) -> Poll<io::Result<usize>> { + let total_len = bufs.iter().map(|b| b.len()).sum::<usize>(); + if self.pos == self.cap && total_len >= self.buf.len() { + let res = ready!(self.as_mut().inner().poll_read_vectored(cx, bufs)); + self.discard_buffer(); + return Poll::Ready(res); + } + let mut rem = ready!(self.as_mut().poll_fill_buf(cx))?; + let nread = rem.read_vectored(bufs)?; + self.consume(nread); + Poll::Ready(Ok(nread)) + } + + // we can't skip unconditionally because of the large buffer case in read. + #[cfg(feature = "read-initializer")] + unsafe fn initializer(&self) -> Initializer { + self.inner.initializer() + } +} + +impl<R: AsyncRead> AsyncBufRead for BufReader<R> { + fn poll_fill_buf( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<io::Result<&[u8]>> { + let Self { inner, buf, cap, pos } = unsafe { self.get_unchecked_mut() }; + let mut inner = unsafe { Pin::new_unchecked(inner) }; + + // If we've reached the end of our internal buffer then we need to fetch + // some more data from the underlying reader. + // Branch using `>=` instead of the more correct `==` + // to tell the compiler that the pos..cap slice is always valid. + if *pos >= *cap { + debug_assert!(*pos == *cap); + *cap = ready!(inner.as_mut().poll_read(cx, buf))?; + *pos = 0; + } + Poll::Ready(Ok(&buf[*pos..*cap])) + } + + fn consume(mut self: Pin<&mut Self>, amt: usize) { + *self.as_mut().pos() = cmp::min(self.pos + amt, self.cap); + } +} + +impl<R: AsyncWrite> AsyncWrite for BufReader<R> { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.inner().poll_write(cx, buf) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + self.inner().poll_write_vectored(cx, bufs) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.inner().poll_flush(cx) + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.inner().poll_close(cx) + } +} + +impl<R: fmt::Debug> fmt::Debug for BufReader<R> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BufReader") + .field("reader", &self.inner) + .field("buffer", &format_args!("{}/{}", self.cap - self.pos, self.buf.len())) + .finish() + } +} + +impl<R: AsyncRead + AsyncSeek> AsyncSeek for BufReader<R> { + /// Seek to an offset, in bytes, in the underlying reader. + /// + /// The position used for seeking with `SeekFrom::Current(_)` is the + /// position the underlying reader would be at if the `BufReader` had no + /// internal buffer. + /// + /// Seeking always discards the internal buffer, even if the seek position + /// would otherwise fall within it. This guarantees that calling + /// `.into_inner()` immediately after a seek yields the underlying reader + /// at the same position. + /// + /// See [`AsyncSeek`](futures_io::AsyncSeek) for more details. + /// + /// Note: In the edge case where you're seeking with `SeekFrom::Current(n)` + /// where `n` minus the internal buffer length overflows an `i64`, two + /// seeks will be performed instead of one. If the second seek returns + /// `Err`, the underlying reader will be left at the same position it would + /// have if you called `seek` with `SeekFrom::Current(0)`. + fn poll_seek( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + pos: SeekFrom, + ) -> Poll<io::Result<u64>> { + let result: u64; + if let SeekFrom::Current(n) = pos { + let remainder = (self.cap - self.pos) as i64; + // it should be safe to assume that remainder fits within an i64 as the alternative + // means we managed to allocate 8 exbibytes and that's absurd. + // But it's not out of the realm of possibility for some weird underlying reader to + // support seeking by i64::min_value() so we need to handle underflow when subtracting + // remainder. + if let Some(offset) = n.checked_sub(remainder) { + result = ready!(self.as_mut().inner().poll_seek(cx, SeekFrom::Current(offset)))?; + } else { + // seek backwards by our remainder, and then by the offset + ready!(self.as_mut().inner().poll_seek(cx, SeekFrom::Current(-remainder)))?; + self.as_mut().discard_buffer(); + result = ready!(self.as_mut().inner().poll_seek(cx, SeekFrom::Current(n)))?; + } + } else { + // Seeking with Start/End doesn't care about our buffer length. + result = ready!(self.as_mut().inner().poll_seek(cx, pos))?; + } + self.discard_buffer(); + Poll::Ready(Ok(result)) + } +} diff --git a/src/io/buf_writer.rs b/src/io/buf_writer.rs new file mode 100644 index 0000000..b0afbd8 --- /dev/null +++ b/src/io/buf_writer.rs @@ -0,0 +1,222 @@ +use futures_core::task::{Context, Poll}; +#[cfg(feature = "read-initializer")] +use futures_io::Initializer; +use futures_io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, IoSlice, IoSliceMut, SeekFrom}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; +use std::fmt; +use std::io::{self, Write}; +use std::pin::Pin; +use super::DEFAULT_BUF_SIZE; + +/// Wraps a writer and buffers its output. +/// +/// It can be excessively inefficient to work directly with something that +/// implements [`AsyncWrite`]. A `BufWriter` keeps an in-memory buffer of data and +/// writes it to an underlying writer in large, infrequent batches. +/// +/// `BufWriter` can improve the speed of programs that make *small* and +/// *repeated* write calls to the same file or network socket. It does not +/// help when writing very large amounts at once, or writing just one or a few +/// times. It also provides no advantage when writing to a destination that is +/// in memory, like a `Vec<u8>`. +/// +/// When the `BufWriter` is dropped, the contents of its buffer will be +/// discarded. Creating multiple instances of a `BufWriter` on the same +/// stream can cause data loss. If you need to write out the contents of its +/// buffer, you must manually call flush before the writer is dropped. +/// +/// [`AsyncWrite`]: futures_io::AsyncWrite +/// [`flush`]: super::AsyncWriteExt::flush +/// +// TODO: Examples +pub struct BufWriter<W> { + inner: W, + buf: Vec<u8>, + written: usize, +} + +impl<W> BufWriter<W> { + unsafe_pinned!(inner: W); + unsafe_unpinned!(buf: Vec<u8>); +} + +impl<W: AsyncWrite> BufWriter<W> { + /// Creates a new `BufWriter` with a default buffer capacity. The default is currently 8 KB, + /// but may change in the future. + pub fn new(inner: W) -> Self { + Self::with_capacity(DEFAULT_BUF_SIZE, inner) + } + + /// Creates a new `BufWriter` with the specified buffer capacity. + pub fn with_capacity(cap: usize, inner: W) -> Self { + Self { + inner, + buf: Vec::with_capacity(cap), + written: 0, + } + } + + fn flush_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + let Self { inner, buf, written } = unsafe { self.get_unchecked_mut() }; + let mut inner = unsafe { Pin::new_unchecked(inner) }; + + let len = buf.len(); + let mut ret = Ok(()); + while *written < len { + match ready!(inner.as_mut().poll_write(cx, &buf[*written..])) { + Ok(0) => { + ret = Err(io::Error::new( + io::ErrorKind::WriteZero, + "failed to write the buffered data", + )); + break; + } + Ok(n) => *written += n, + Err(e) => { + ret = Err(e); + break; + } + } + } + if *written > 0 { + buf.drain(..*written); + } + *written = 0; + Poll::Ready(ret) + } + + /// Gets a reference to the underlying writer. + pub fn get_ref(&self) -> &W { + &self.inner + } + + /// Gets a mutable reference to the underlying writer. + /// + /// It is inadvisable to directly write to the underlying writer. + pub fn get_mut(&mut self) -> &mut W { + &mut self.inner + } + + /// Gets a pinned mutable reference to the underlying writer. + /// + /// It is inadvisable to directly write to the underlying writer. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut W> { + self.inner() + } + + /// Consumes this `BufWriter`, returning the underlying writer. + /// + /// Note that any leftover data in the internal buffer is lost. + pub fn into_inner(self) -> W { + self.inner + } + + /// Returns a reference to the internally buffered data. + pub fn buffer(&self) -> &[u8] { + &self.buf + } +} + +impl<W: AsyncWrite> AsyncWrite for BufWriter<W> { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + if self.buf.len() + buf.len() > self.buf.capacity() { + ready!(self.as_mut().flush_buf(cx))?; + } + if buf.len() >= self.buf.capacity() { + self.inner().poll_write(cx, buf) + } else { + Poll::Ready(self.buf().write(buf)) + } + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + let total_len = bufs.iter().map(|b| b.len()).sum::<usize>(); + if self.buf.len() + total_len > self.buf.capacity() { + ready!(self.as_mut().flush_buf(cx))?; + } + if total_len >= self.buf.capacity() { + self.inner().poll_write_vectored(cx, bufs) + } else { + Poll::Ready(self.buf().write_vectored(bufs)) + } + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + ready!(self.as_mut().flush_buf(cx))?; + self.inner().poll_flush(cx) + } + + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + ready!(self.as_mut().flush_buf(cx))?; + self.inner().poll_close(cx) + } +} + +impl<W: AsyncRead> AsyncRead for BufWriter<W> { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll<io::Result<usize>> { + self.inner().poll_read(cx, buf) + } + + fn poll_read_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &mut [IoSliceMut<'_>], + ) -> Poll<io::Result<usize>> { + self.inner().poll_read_vectored(cx, bufs) + } + + // we can't skip unconditionally because of the large buffer case in read. + #[cfg(feature = "read-initializer")] + unsafe fn initializer(&self) -> Initializer { + self.inner.initializer() + } +} + +impl<W: AsyncBufRead> AsyncBufRead for BufWriter<W> { + fn poll_fill_buf( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<io::Result<&[u8]>> { + self.inner().poll_fill_buf(cx) + } + + fn consume(self: Pin<&mut Self>, amt: usize) { + self.inner().consume(amt) + } +} + +impl<W: fmt::Debug> fmt::Debug for BufWriter<W> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BufWriter") + .field("writer", &self.inner) + .field("buffer", &format_args!("{}/{}", self.buf.len(), self.buf.capacity())) + .field("written", &self.written) + .finish() + } +} + +impl<W: AsyncWrite + AsyncSeek> AsyncSeek for BufWriter<W> { + /// Seek to the offset, in bytes, in the underlying writer. + /// + /// Seeking always writes out the internal buffer before seeking. + fn poll_seek( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + pos: SeekFrom, + ) -> Poll<io::Result<u64>> { + ready!(self.as_mut().flush_buf(cx))?; + self.inner().poll_seek(cx, pos) + } +} diff --git a/src/io/chain.rs b/src/io/chain.rs new file mode 100644 index 0000000..64bbdec --- /dev/null +++ b/src/io/chain.rs @@ -0,0 +1,166 @@ +use futures_core::task::{Context, Poll}; +#[cfg(feature = "read-initializer")] +use futures_io::Initializer; +use futures_io::{AsyncBufRead, AsyncRead, IoSliceMut}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; +use std::fmt; +use std::io; +use std::pin::Pin; + +/// Reader for the [`chain`](super::AsyncReadExt::chain) method. +#[must_use = "readers do nothing unless polled"] +pub struct Chain<T, U> { + first: T, + second: U, + done_first: bool, +} + +impl<T, U> Unpin for Chain<T, U> +where + T: Unpin, + U: Unpin, +{ +} + +impl<T, U> Chain<T, U> +where + T: AsyncRead, + U: AsyncRead, +{ + unsafe_pinned!(first: T); + unsafe_pinned!(second: U); + unsafe_unpinned!(done_first: bool); + + pub(super) fn new(first: T, second: U) -> Self { + Self { + first, + second, + done_first: false, + } + } + + /// Gets references to the underlying readers in this `Chain`. + pub fn get_ref(&self) -> (&T, &U) { + (&self.first, &self.second) + } + + /// Gets mutable references to the underlying readers in this `Chain`. + /// + /// Care should be taken to avoid modifying the internal I/O state of the + /// underlying readers as doing so may corrupt the internal state of this + /// `Chain`. + pub fn get_mut(&mut self) -> (&mut T, &mut U) { + (&mut self.first, &mut self.second) + } + + /// Gets pinned mutable references to the underlying readers in this `Chain`. + /// + /// Care should be taken to avoid modifying the internal I/O state of the + /// underlying readers as doing so may corrupt the internal state of this + /// `Chain`. + pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut T>, Pin<&mut U>) { + unsafe { + let Self { first, second, .. } = self.get_unchecked_mut(); + (Pin::new_unchecked(first), Pin::new_unchecked(second)) + } + } + + /// Consumes the `Chain`, returning the wrapped readers. + pub fn into_inner(self) -> (T, U) { + (self.first, self.second) + } +} + +impl<T, U> fmt::Debug for Chain<T, U> +where + T: fmt::Debug, + U: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Chain") + .field("t", &self.first) + .field("u", &self.second) + .field("done_first", &self.done_first) + .finish() + } +} + +impl<T, U> AsyncRead for Chain<T, U> +where + T: AsyncRead, + U: AsyncRead, +{ + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll<io::Result<usize>> { + if !self.done_first { + match ready!(self.as_mut().first().poll_read(cx, buf)?) { + 0 if !buf.is_empty() => *self.as_mut().done_first() = true, + n => return Poll::Ready(Ok(n)), + } + } + self.second().poll_read(cx, buf) + } + + fn poll_read_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &mut [IoSliceMut<'_>], + ) -> Poll<io::Result<usize>> { + if !self.done_first { + let n = ready!(self.as_mut().first().poll_read_vectored(cx, bufs)?); + if n == 0 && bufs.iter().any(|b| !b.is_empty()) { + *self.as_mut().done_first() = true + } else { + return Poll::Ready(Ok(n)); + } + } + self.second().poll_read_vectored(cx, bufs) + } + + #[cfg(feature = "read-initializer")] + unsafe fn initializer(&self) -> Initializer { + let initializer = self.first.initializer(); + if initializer.should_initialize() { + initializer + } else { + self.second.initializer() + } + } +} + +impl<T, U> AsyncBufRead for Chain<T, U> +where + T: AsyncBufRead, + U: AsyncBufRead, +{ + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + let Self { + first, + second, + done_first, + } = unsafe { self.get_unchecked_mut() }; + let first = unsafe { Pin::new_unchecked(first) }; + let second = unsafe { Pin::new_unchecked(second) }; + + if !*done_first { + match ready!(first.poll_fill_buf(cx)?) { + buf if buf.is_empty() => { + *done_first = true; + } + buf => return Poll::Ready(Ok(buf)), + } + } + second.poll_fill_buf(cx) + } + + fn consume(self: Pin<&mut Self>, amt: usize) { + if !self.done_first { + self.first().consume(amt) + } else { + self.second().consume(amt) + } + } +} diff --git a/src/io/close.rs b/src/io/close.rs new file mode 100644 index 0000000..4d56696 --- /dev/null +++ b/src/io/close.rs @@ -0,0 +1,28 @@ +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_io::AsyncWrite; +use std::io; +use std::pin::Pin; + +/// Future for the [`close`](super::AsyncWriteExt::close) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Close<'a, W: ?Sized> { + writer: &'a mut W, +} + +impl<W: ?Sized + Unpin> Unpin for Close<'_, W> {} + +impl<'a, W: AsyncWrite + ?Sized + Unpin> Close<'a, W> { + pub(super) fn new(writer: &'a mut W) -> Self { + Close { writer } + } +} + +impl<W: AsyncWrite + ?Sized + Unpin> Future for Close<'_, W> { + type Output = io::Result<()>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + Pin::new(&mut *self.writer).poll_close(cx) + } +} diff --git a/src/io/copy.rs b/src/io/copy.rs new file mode 100644 index 0000000..9531aab --- /dev/null +++ b/src/io/copy.rs @@ -0,0 +1,63 @@ +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_io::{AsyncRead, AsyncWrite}; +use std::io; +use std::pin::Pin; +use super::{BufReader, copy_buf, CopyBuf}; +use pin_utils::unsafe_pinned; + +/// Creates a future which copies all the bytes from one object to another. +/// +/// The returned future will copy all the bytes read from this `AsyncRead` into the +/// `writer` specified. This future will only complete once the `reader` has hit +/// EOF and all bytes have been written to and flushed from the `writer` +/// provided. +/// +/// On success the number of bytes is returned. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::io::{self, AsyncWriteExt, Cursor}; +/// +/// let reader = Cursor::new([1, 2, 3, 4]); +/// let mut writer = Cursor::new(vec![0u8; 5]); +/// +/// let bytes = io::copy(reader, &mut writer).await?; +/// writer.close().await?; +/// +/// assert_eq!(bytes, 4); +/// assert_eq!(writer.into_inner(), [1, 2, 3, 4, 0]); +/// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); +/// ``` +pub fn copy<R, W>(reader: R, writer: &mut W) -> Copy<'_, R, W> +where + R: AsyncRead, + W: AsyncWrite + Unpin + ?Sized, +{ + Copy { + inner: copy_buf(BufReader::new(reader), writer), + } +} + +/// Future for the [`copy()`] function. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Copy<'a, R, W: ?Sized> { + inner: CopyBuf<'a, BufReader<R>, W>, +} + +impl<'a, R: AsyncRead, W: ?Sized> Unpin for Copy<'a, R, W> where CopyBuf<'a, BufReader<R>, W>: Unpin {} + +impl<'a, R: AsyncRead, W: ?Sized> Copy<'a, R, W> { + unsafe_pinned!(inner: CopyBuf<'a, BufReader<R>, W>); +} + +impl<R: AsyncRead, W: AsyncWrite + Unpin + ?Sized> Future for Copy<'_, R, W> { + type Output = io::Result<u64>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + self.inner().poll(cx) + } +} diff --git a/src/io/copy_buf.rs b/src/io/copy_buf.rs new file mode 100644 index 0000000..9881182 --- /dev/null +++ b/src/io/copy_buf.rs @@ -0,0 +1,87 @@ +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_io::{AsyncBufRead, AsyncWrite}; +use std::io; +use std::pin::Pin; + +/// Creates a future which copies all the bytes from one object to another. +/// +/// The returned future will copy all the bytes read from this `AsyncBufRead` into the +/// `writer` specified. This future will only complete once the `reader` has hit +/// EOF and all bytes have been written to and flushed from the `writer` +/// provided. +/// +/// On success the number of bytes is returned. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::io::{self, AsyncWriteExt, Cursor}; +/// +/// let reader = Cursor::new([1, 2, 3, 4]); +/// let mut writer = Cursor::new(vec![0u8; 5]); +/// +/// let bytes = io::copy_buf(reader, &mut writer).await?; +/// writer.close().await?; +/// +/// assert_eq!(bytes, 4); +/// assert_eq!(writer.into_inner(), [1, 2, 3, 4, 0]); +/// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); +/// ``` +pub fn copy_buf<R, W>(reader: R, writer: &mut W) -> CopyBuf<'_, R, W> +where + R: AsyncBufRead, + W: AsyncWrite + Unpin + ?Sized, +{ + CopyBuf { + reader, + writer, + amt: 0, + } +} + +/// Future for the [`copy_buf()`] function. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct CopyBuf<'a, R, W: ?Sized> { + reader: R, + writer: &'a mut W, + amt: u64, +} + +impl<R: Unpin, W: ?Sized> Unpin for CopyBuf<'_, R, W> {} + +impl<R, W: Unpin + ?Sized> CopyBuf<'_, R, W> { + fn project(self: Pin<&mut Self>) -> (Pin<&mut R>, Pin<&mut W>, &mut u64) { + unsafe { + let this = self.get_unchecked_mut(); + (Pin::new_unchecked(&mut this.reader), Pin::new(&mut *this.writer), &mut this.amt) + } + } +} + +impl<R, W> Future for CopyBuf<'_, R, W> + where R: AsyncBufRead, + W: AsyncWrite + Unpin + ?Sized, +{ + type Output = io::Result<u64>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let (mut reader, mut writer, amt) = self.project(); + loop { + let buffer = ready!(reader.as_mut().poll_fill_buf(cx))?; + if buffer.is_empty() { + ready!(writer.as_mut().poll_flush(cx))?; + return Poll::Ready(Ok(*amt)); + } + + let i = ready!(writer.as_mut().poll_write(cx, buffer))?; + if i == 0 { + return Poll::Ready(Err(io::ErrorKind::WriteZero.into())) + } + *amt += i as u64; + reader.as_mut().consume(i); + } + } +} diff --git a/src/io/cursor.rs b/src/io/cursor.rs new file mode 100644 index 0000000..d135923 --- /dev/null +++ b/src/io/cursor.rs @@ -0,0 +1,238 @@ +use futures_core::task::{Context, Poll}; +#[cfg(feature = "read_initializer")] +use futures_io::Initializer; +use futures_io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, IoSlice, IoSliceMut, SeekFrom}; +use std::io; +use std::pin::Pin; + +/// A `Cursor` wraps an in-memory buffer and provides it with a +/// [`AsyncSeek`] implementation. +/// +/// `Cursor`s are used with in-memory buffers, anything implementing +/// `AsRef<[u8]>`, to allow them to implement [`AsyncRead`] and/or [`AsyncWrite`], +/// allowing these buffers to be used anywhere you might use a reader or writer +/// that does actual I/O. +/// +/// The standard library implements some I/O traits on various types which +/// are commonly used as a buffer, like `Cursor<`[`Vec`]`<u8>>` and +/// `Cursor<`[`&[u8]`][bytes]`>`. +/// +/// [`AsyncSeek`]: trait.AsyncSeek.html +/// [`AsyncRead`]: trait.AsyncRead.html +/// [`AsyncWrite`]: trait.AsyncWrite.html +/// [bytes]: https://doc.rust-lang.org/std/primitive.slice.html +#[derive(Clone, Debug, Default)] +pub struct Cursor<T> { + inner: io::Cursor<T>, +} + +impl<T> Cursor<T> { + /// Creates a new cursor wrapping the provided underlying in-memory buffer. + /// + /// Cursor initial position is `0` even if underlying buffer (e.g., `Vec`) + /// is not empty. So writing to cursor starts with overwriting `Vec` + /// content, not with appending to it. + /// + /// # Examples + /// + /// ``` + /// use futures::io::Cursor; + /// + /// let buff = Cursor::new(Vec::new()); + /// # fn force_inference(_: &Cursor<Vec<u8>>) {} + /// # force_inference(&buff); + /// ``` + pub fn new(inner: T) -> Cursor<T> { + Cursor { + inner: io::Cursor::new(inner), + } + } + + /// Consumes this cursor, returning the underlying value. + /// + /// # Examples + /// + /// ``` + /// use futures::io::Cursor; + /// + /// let buff = Cursor::new(Vec::new()); + /// # fn force_inference(_: &Cursor<Vec<u8>>) {} + /// # force_inference(&buff); + /// + /// let vec = buff.into_inner(); + /// ``` + pub fn into_inner(self) -> T { + self.inner.into_inner() + } + + /// Gets a reference to the underlying value in this cursor. + /// + /// # Examples + /// + /// ``` + /// use futures::io::Cursor; + /// + /// let buff = Cursor::new(Vec::new()); + /// # fn force_inference(_: &Cursor<Vec<u8>>) {} + /// # force_inference(&buff); + /// + /// let reference = buff.get_ref(); + /// ``` + pub fn get_ref(&self) -> &T { + self.inner.get_ref() + } + + /// Gets a mutable reference to the underlying value in this cursor. + /// + /// Care should be taken to avoid modifying the internal I/O state of the + /// underlying value as it may corrupt this cursor's position. + /// + /// # Examples + /// + /// ``` + /// use futures::io::Cursor; + /// + /// let mut buff = Cursor::new(Vec::new()); + /// # fn force_inference(_: &Cursor<Vec<u8>>) {} + /// # force_inference(&buff); + /// + /// let reference = buff.get_mut(); + /// ``` + pub fn get_mut(&mut self) -> &mut T { + self.inner.get_mut() + } + + /// Returns the current position of this cursor. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncSeekExt, Cursor, SeekFrom}; + /// + /// let mut buff = Cursor::new(vec![1, 2, 3, 4, 5]); + /// + /// assert_eq!(buff.position(), 0); + /// + /// buff.seek(SeekFrom::Current(2)).await?; + /// assert_eq!(buff.position(), 2); + /// + /// buff.seek(SeekFrom::Current(-1)).await?; + /// assert_eq!(buff.position(), 1); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + pub fn position(&self) -> u64 { + self.inner.position() + } + + /// Sets the position of this cursor. + /// + /// # Examples + /// + /// ``` + /// use futures::io::Cursor; + /// + /// let mut buff = Cursor::new(vec![1, 2, 3, 4, 5]); + /// + /// assert_eq!(buff.position(), 0); + /// + /// buff.set_position(2); + /// assert_eq!(buff.position(), 2); + /// + /// buff.set_position(4); + /// assert_eq!(buff.position(), 4); + /// ``` + pub fn set_position(&mut self, pos: u64) { + self.inner.set_position(pos) + } +} + +impl<T> AsyncSeek for Cursor<T> +where + T: AsRef<[u8]> + Unpin, +{ + fn poll_seek( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + pos: SeekFrom, + ) -> Poll<io::Result<u64>> { + Poll::Ready(io::Seek::seek(&mut self.inner, pos)) + } +} + +impl<T: AsRef<[u8]> + Unpin> AsyncRead for Cursor<T> { + #[cfg(feature = "read_initializer")] + #[inline] + unsafe fn initializer(&self) -> Initializer { + io::Read::initializer(&self.inner) + } + + fn poll_read( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Read::read(&mut self.inner, buf)) + } + + fn poll_read_vectored( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + bufs: &mut [IoSliceMut<'_>], + ) -> Poll<io::Result<usize>> { + Poll::Ready(io::Read::read_vectored(&mut self.inner, bufs)) + } +} + +impl<T> AsyncBufRead for Cursor<T> +where + T: AsRef<[u8]> + Unpin, +{ + fn poll_fill_buf(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + Poll::Ready(io::BufRead::fill_buf(&mut self.get_mut().inner)) + } + + fn consume(mut self: Pin<&mut Self>, amt: usize) { + io::BufRead::consume(&mut self.inner, amt) + } +} + +macro_rules! delegate_async_write_to_stdio { + () => { + fn poll_write(mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &[u8]) + -> Poll<io::Result<usize>> + { + Poll::Ready(io::Write::write(&mut self.inner, buf)) + } + + fn poll_write_vectored(mut self: Pin<&mut Self>, _: &mut Context<'_>, bufs: &[IoSlice<'_>]) + -> Poll<io::Result<usize>> + { + Poll::Ready(io::Write::write_vectored(&mut self.inner, bufs)) + } + + fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(io::Write::flush(&mut self.inner)) + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.poll_flush(cx) + } + } +} + +impl AsyncWrite for Cursor<&mut [u8]> { + delegate_async_write_to_stdio!(); +} + +impl AsyncWrite for Cursor<&mut Vec<u8>> { + delegate_async_write_to_stdio!(); +} + +impl AsyncWrite for Cursor<Vec<u8>> { + delegate_async_write_to_stdio!(); +} + +impl AsyncWrite for Cursor<Box<[u8]>> { + delegate_async_write_to_stdio!(); +} diff --git a/src/io/empty.rs b/src/io/empty.rs new file mode 100644 index 0000000..ab2395a --- /dev/null +++ b/src/io/empty.rs @@ -0,0 +1,67 @@ +use futures_core::task::{Context, Poll}; +#[cfg(feature = "read-initializer")] +use futures_io::Initializer; +use futures_io::{AsyncBufRead, AsyncRead}; +use std::fmt; +use std::io; +use std::pin::Pin; + +/// Reader for the [`empty()`] function. +#[must_use = "readers do nothing unless polled"] +pub struct Empty { + _priv: (), +} + +/// Constructs a new handle to an empty reader. +/// +/// All reads from the returned reader will return `Poll::Ready(Ok(0))`. +/// +/// # Examples +/// +/// A slightly sad example of not reading anything into a buffer: +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::io::{self, AsyncReadExt}; +/// +/// let mut buffer = String::new(); +/// let mut reader = io::empty(); +/// reader.read_to_string(&mut buffer).await?; +/// assert!(buffer.is_empty()); +/// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); +/// ``` +pub fn empty() -> Empty { + Empty { _priv: () } +} + +impl AsyncRead for Empty { + #[inline] + fn poll_read( + self: Pin<&mut Self>, + _: &mut Context<'_>, + _: &mut [u8], + ) -> Poll<io::Result<usize>> { + Poll::Ready(Ok(0)) + } + + #[cfg(feature = "read-initializer")] + #[inline] + unsafe fn initializer(&self) -> Initializer { + Initializer::nop() + } +} + +impl AsyncBufRead for Empty { + #[inline] + fn poll_fill_buf(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + Poll::Ready(Ok(&[])) + } + #[inline] + fn consume(self: Pin<&mut Self>, _: usize) {} +} + +impl fmt::Debug for Empty { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.pad("Empty { .. }") + } +} diff --git a/src/io/flush.rs b/src/io/flush.rs new file mode 100644 index 0000000..70b867a --- /dev/null +++ b/src/io/flush.rs @@ -0,0 +1,30 @@ +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_io::AsyncWrite; +use std::io; +use std::pin::Pin; + +/// Future for the [`flush`](super::AsyncWriteExt::flush) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Flush<'a, W: ?Sized> { + writer: &'a mut W, +} + +impl<W: ?Sized + Unpin> Unpin for Flush<'_, W> {} + +impl<'a, W: AsyncWrite + ?Sized + Unpin> Flush<'a, W> { + pub(super) fn new(writer: &'a mut W) -> Self { + Flush { writer } + } +} + +impl<W> Future for Flush<'_, W> + where W: AsyncWrite + ?Sized + Unpin, +{ + type Output = io::Result<()>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + Pin::new(&mut *self.writer).poll_flush(cx) + } +} diff --git a/src/io/into_sink.rs b/src/io/into_sink.rs new file mode 100644 index 0000000..bdc6b34 --- /dev/null +++ b/src/io/into_sink.rs @@ -0,0 +1,108 @@ +use futures_core::task::{Context, Poll}; +use futures_io::AsyncWrite; +use futures_sink::Sink; +use std::io; +use std::pin::Pin; +use std::marker::Unpin; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +#[derive(Debug)] +struct Block<Item> { + offset: usize, + bytes: Item, +} + +/// Sink for the [`into_sink`](super::AsyncWriteExt::into_sink) method. +#[must_use = "sinks do nothing unless polled"] +#[derive(Debug)] +pub struct IntoSink<W, Item> { + writer: W, + /// An outstanding block for us to push into the underlying writer, along with an offset of how + /// far into this block we have written already. + buffer: Option<Block<Item>>, +} + +impl<W: Unpin, Item> Unpin for IntoSink<W, Item> {} + +impl<W: AsyncWrite, Item: AsRef<[u8]>> IntoSink<W, Item> { + unsafe_pinned!(writer: W); + unsafe_unpinned!(buffer: Option<Block<Item>>); + + pub(super) fn new(writer: W) -> Self { + IntoSink { writer, buffer: None } + } + + fn project(self: Pin<&mut Self>) -> (Pin<&mut W>, &mut Option<Block<Item>>) { + unsafe { + let this = self.get_unchecked_mut(); + (Pin::new_unchecked(&mut this.writer), &mut this.buffer) + } + } + + /// If we have an outstanding block in `buffer` attempt to push it into the writer, does _not_ + /// flush the writer after it succeeds in pushing the block into it. + fn poll_flush_buffer( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), io::Error>> + { + let (mut writer, buffer) = self.project(); + if let Some(buffer) = buffer { + loop { + let bytes = buffer.bytes.as_ref(); + let written = ready!(writer.as_mut().poll_write(cx, &bytes[buffer.offset..]))?; + buffer.offset += written; + if buffer.offset == bytes.len() { + break; + } + } + } + *buffer = None; + Poll::Ready(Ok(())) + } + +} + +impl<W: AsyncWrite, Item: AsRef<[u8]>> Sink<Item> for IntoSink<W, Item> { + type Error = io::Error; + + fn poll_ready( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> + { + ready!(self.as_mut().poll_flush_buffer(cx))?; + Poll::Ready(Ok(())) + } + + #[allow(clippy::debug_assert_with_mut_call)] + fn start_send( + mut self: Pin<&mut Self>, + item: Item, + ) -> Result<(), Self::Error> + { + debug_assert!(self.as_mut().buffer().is_none()); + *self.as_mut().buffer() = Some(Block { offset: 0, bytes: item }); + Ok(()) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> + { + ready!(self.as_mut().poll_flush_buffer(cx))?; + ready!(self.as_mut().writer().poll_flush(cx))?; + Poll::Ready(Ok(())) + } + + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> + { + ready!(self.as_mut().poll_flush_buffer(cx))?; + ready!(self.as_mut().writer().poll_close(cx))?; + Poll::Ready(Ok(())) + } +} diff --git a/src/io/lines.rs b/src/io/lines.rs new file mode 100644 index 0000000..2e12616 --- /dev/null +++ b/src/io/lines.rs @@ -0,0 +1,50 @@ +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use futures_io::AsyncBufRead; +use std::io; +use std::mem; +use std::pin::Pin; +use super::read_line::read_line_internal; + +/// Stream for the [`lines`](super::AsyncBufReadExt::lines) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Lines<R> { + reader: R, + buf: String, + bytes: Vec<u8>, + read: usize, +} + +impl<R: Unpin> Unpin for Lines<R> {} + +impl<R: AsyncBufRead> Lines<R> { + pub(super) fn new(reader: R) -> Self { + Self { + reader, + buf: String::new(), + bytes: Vec::new(), + read: 0, + } + } +} + +impl<R: AsyncBufRead> Stream for Lines<R> { + type Item = io::Result<String>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let Self { reader, buf, bytes, read } = unsafe { self.get_unchecked_mut() }; + let reader = unsafe { Pin::new_unchecked(reader) }; + let n = ready!(read_line_internal(reader, cx, buf, bytes, read))?; + if n == 0 && buf.is_empty() { + return Poll::Ready(None) + } + if buf.ends_with('\n') { + buf.pop(); + if buf.ends_with('\r') { + buf.pop(); + } + } + Poll::Ready(Some(Ok(mem::replace(buf, String::new())))) + } +} diff --git a/src/io/mod.rs b/src/io/mod.rs new file mode 100644 index 0000000..43f183f --- /dev/null +++ b/src/io/mod.rs @@ -0,0 +1,675 @@ +//! IO +//! +//! This module contains a number of functions for working with +//! `AsyncRead`, `AsyncWrite`, `AsyncSeek`, and `AsyncBufRead` types, including +//! the `AsyncReadExt`, `AsyncWriteExt`, `AsyncSeekExt`, and `AsyncBufReadExt` +//! traits which add methods to the `AsyncRead`, `AsyncWrite`, `AsyncSeek`, +//! and `AsyncBufRead` types. +//! +//! This module is only available when the `io` and `std` features of this +//! library is activated, and it is activated by default. + +#[cfg(feature = "io-compat")] +use crate::compat::Compat; +use std::ptr; + +pub use futures_io::{ + AsyncRead, AsyncWrite, AsyncSeek, AsyncBufRead, Error, ErrorKind, + IoSlice, IoSliceMut, Result, SeekFrom, +}; +#[cfg(feature = "read-initializer")] +pub use futures_io::Initializer; + +// used by `BufReader` and `BufWriter` +// https://github.com/rust-lang/rust/blob/master/src/libstd/sys_common/io.rs#L1 +const DEFAULT_BUF_SIZE: usize = 8 * 1024; + +/// Initializes a buffer if necessary. +/// +/// A buffer is always initialized if `read-initializer` feature is disabled. +#[inline] +unsafe fn initialize<R: AsyncRead>(_reader: &R, buf: &mut [u8]) { + #[cfg(feature = "read-initializer")] + { + if !_reader.initializer().should_initialize() { + return; + } + } + ptr::write_bytes(buf.as_mut_ptr(), 0, buf.len()) +} + +mod allow_std; +pub use self::allow_std::AllowStdIo; + +mod buf_reader; +pub use self::buf_reader::BufReader; + +mod buf_writer; +pub use self::buf_writer::BufWriter; + +mod chain; +pub use self::chain::Chain; + +mod close; +pub use self::close::Close; + +mod copy; +pub use self::copy::{copy, Copy}; + +mod copy_buf; +pub use self::copy_buf::{copy_buf, CopyBuf}; + +mod cursor; +pub use self::cursor::Cursor; + +mod empty; +pub use self::empty::{empty, Empty}; + +mod flush; +pub use self::flush::Flush; + +#[cfg(feature = "sink")] +mod into_sink; +#[cfg(feature = "sink")] +pub use self::into_sink::IntoSink; + +mod lines; +pub use self::lines::Lines; + +mod read; +pub use self::read::Read; + +mod read_vectored; +pub use self::read_vectored::ReadVectored; + +mod read_exact; +pub use self::read_exact::ReadExact; + +mod read_line; +pub use self::read_line::ReadLine; + +mod read_to_end; +pub use self::read_to_end::ReadToEnd; + +mod read_to_string; +pub use self::read_to_string::ReadToString; + +mod read_until; +pub use self::read_until::ReadUntil; + +mod repeat; +pub use self::repeat::{repeat, Repeat}; + +mod seek; +pub use self::seek::Seek; + +mod sink; +pub use self::sink::{sink, Sink}; + +mod split; +pub use self::split::{ReadHalf, WriteHalf}; + +mod take; +pub use self::take::Take; + +mod window; +pub use self::window::Window; + +mod write; +pub use self::write::Write; + +mod write_vectored; +pub use self::write_vectored::WriteVectored; + +mod write_all; +pub use self::write_all::WriteAll; + +/// An extension trait which adds utility methods to `AsyncRead` types. +pub trait AsyncReadExt: AsyncRead { + /// Creates an adaptor which will chain this stream with another. + /// + /// The returned `AsyncRead` instance will first read all bytes from this object + /// until EOF is encountered. Afterwards the output is equivalent to the + /// output of `next`. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let reader1 = Cursor::new([1, 2, 3, 4]); + /// let reader2 = Cursor::new([5, 6, 7, 8]); + /// + /// let mut reader = reader1.chain(reader2); + /// let mut buffer = Vec::new(); + /// + /// // read the value into a Vec. + /// reader.read_to_end(&mut buffer).await?; + /// assert_eq!(buffer, [1, 2, 3, 4, 5, 6, 7, 8]); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn chain<R>(self, next: R) -> Chain<Self, R> + where + Self: Sized, + R: AsyncRead, + { + Chain::new(self, next) + } + + /// Tries to read some bytes directly into the given `buf` in asynchronous + /// manner, returning a future type. + /// + /// The returned future will resolve to the number of bytes read once the read + /// operation is completed. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let mut reader = Cursor::new([1, 2, 3, 4]); + /// let mut output = [0u8; 5]; + /// + /// let bytes = reader.read(&mut output[..]).await?; + /// + /// // This is only guaranteed to be 4 because `&[u8]` is a synchronous + /// // reader. In a real system you could get anywhere from 1 to + /// // `output.len()` bytes in a single read. + /// assert_eq!(bytes, 4); + /// assert_eq!(output, [1, 2, 3, 4, 0]); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Read<'a, Self> + where Self: Unpin, + { + Read::new(self, buf) + } + + /// Creates a future which will read from the `AsyncRead` into `bufs` using vectored + /// IO operations. + /// + /// The returned future will resolve to the number of bytes read once the read + /// operation is completed. + fn read_vectored<'a>(&'a mut self, bufs: &'a mut [IoSliceMut<'a>]) -> ReadVectored<'a, Self> + where Self: Unpin, + { + ReadVectored::new(self, bufs) + } + + /// Creates a future which will read exactly enough bytes to fill `buf`, + /// returning an error if end of file (EOF) is hit sooner. + /// + /// The returned future will resolve once the read operation is completed. + /// + /// In the case of an error the buffer and the object will be discarded, with + /// the error yielded. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let mut reader = Cursor::new([1, 2, 3, 4]); + /// let mut output = [0u8; 4]; + /// + /// reader.read_exact(&mut output).await?; + /// + /// assert_eq!(output, [1, 2, 3, 4]); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + /// + /// ## EOF is hit before `buf` is filled + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{self, AsyncReadExt, Cursor}; + /// + /// let mut reader = Cursor::new([1, 2, 3, 4]); + /// let mut output = [0u8; 5]; + /// + /// let result = reader.read_exact(&mut output).await; + /// + /// assert_eq!(result.unwrap_err().kind(), io::ErrorKind::UnexpectedEof); + /// # }); + /// ``` + fn read_exact<'a>( + &'a mut self, + buf: &'a mut [u8], + ) -> ReadExact<'a, Self> + where Self: Unpin, + { + ReadExact::new(self, buf) + } + + /// Creates a future which will read all the bytes from this `AsyncRead`. + /// + /// On success the total number of bytes read is returned. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let mut reader = Cursor::new([1, 2, 3, 4]); + /// let mut output = Vec::with_capacity(4); + /// + /// let bytes = reader.read_to_end(&mut output).await?; + /// + /// assert_eq!(bytes, 4); + /// assert_eq!(output, vec![1, 2, 3, 4]); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn read_to_end<'a>( + &'a mut self, + buf: &'a mut Vec<u8>, + ) -> ReadToEnd<'a, Self> + where Self: Unpin, + { + ReadToEnd::new(self, buf) + } + + /// Creates a future which will read all the bytes from this `AsyncRead`. + /// + /// On success the total number of bytes read is returned. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let mut reader = Cursor::new(&b"1234"[..]); + /// let mut buffer = String::with_capacity(4); + /// + /// let bytes = reader.read_to_string(&mut buffer).await?; + /// + /// assert_eq!(bytes, 4); + /// assert_eq!(buffer, String::from("1234")); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn read_to_string<'a>( + &'a mut self, + buf: &'a mut String, + ) -> ReadToString<'a, Self> + where Self: Unpin, + { + ReadToString::new(self, buf) + } + + /// Helper method for splitting this read/write object into two halves. + /// + /// The two halves returned implement the `AsyncRead` and `AsyncWrite` + /// traits, respectively. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{self, AsyncReadExt, Cursor}; + /// + /// // Note that for `Cursor` the read and write halves share a single + /// // seek position. This may or may not be true for other types that + /// // implement both `AsyncRead` and `AsyncWrite`. + /// + /// let reader = Cursor::new([1, 2, 3, 4]); + /// let mut buffer = Cursor::new(vec![0, 0, 0, 0, 5, 6, 7, 8]); + /// let mut writer = Cursor::new(vec![0u8; 5]); + /// + /// { + /// let (buffer_reader, mut buffer_writer) = (&mut buffer).split(); + /// io::copy(reader, &mut buffer_writer).await?; + /// io::copy(buffer_reader, &mut writer).await?; + /// } + /// + /// assert_eq!(buffer.into_inner(), [1, 2, 3, 4, 5, 6, 7, 8]); + /// assert_eq!(writer.into_inner(), [5, 6, 7, 8, 0]); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn split(self) -> (ReadHalf<Self>, WriteHalf<Self>) + where Self: AsyncWrite + Sized, + { + split::split(self) + } + + /// Creates an AsyncRead adapter which will read at most `limit` bytes + /// from the underlying reader. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let reader = Cursor::new(&b"12345678"[..]); + /// let mut buffer = [0; 5]; + /// + /// let mut take = reader.take(4); + /// let n = take.read(&mut buffer).await?; + /// + /// assert_eq!(n, 4); + /// assert_eq!(&buffer, b"1234\0"); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn take(self, limit: u64) -> Take<Self> + where Self: Sized + { + Take::new(self, limit) + } + + /// Wraps an [`AsyncRead`] in a compatibility wrapper that allows it to be + /// used as a futures 0.1 / tokio-io 0.1 `AsyncRead`. If the wrapped type + /// implements [`AsyncWrite`] as well, the result will also implement the + /// futures 0.1 / tokio 0.1 `AsyncWrite` trait. + /// + /// Requires the `io-compat` feature to enable. + #[cfg(feature = "io-compat")] + fn compat(self) -> Compat<Self> + where Self: Sized + Unpin, + { + Compat::new(self) + } +} + +impl<R: AsyncRead + ?Sized> AsyncReadExt for R {} + +/// An extension trait which adds utility methods to `AsyncWrite` types. +pub trait AsyncWriteExt: AsyncWrite { + /// Creates a future which will entirely flush this `AsyncWrite`. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AllowStdIo, AsyncWriteExt}; + /// use std::io::{BufWriter, Cursor}; + /// + /// let mut output = vec![0u8; 5]; + /// + /// { + /// let writer = Cursor::new(&mut output); + /// let mut buffered = AllowStdIo::new(BufWriter::new(writer)); + /// buffered.write_all(&[1, 2]).await?; + /// buffered.write_all(&[3, 4]).await?; + /// buffered.flush().await?; + /// } + /// + /// assert_eq!(output, [1, 2, 3, 4, 0]); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn flush(&mut self) -> Flush<'_, Self> + where Self: Unpin, + { + Flush::new(self) + } + + /// Creates a future which will entirely close this `AsyncWrite`. + fn close(&mut self) -> Close<'_, Self> + where Self: Unpin, + { + Close::new(self) + } + + /// Creates a future which will write bytes from `buf` into the object. + /// + /// The returned future will resolve to the number of bytes written once the write + /// operation is completed. + fn write<'a>(&'a mut self, buf: &'a [u8]) -> Write<'a, Self> + where Self: Unpin, + { + Write::new(self, buf) + } + + /// Creates a future which will write bytes from `bufs` into the object using vectored + /// IO operations. + /// + /// The returned future will resolve to the number of bytes written once the write + /// operation is completed. + fn write_vectored<'a>(&'a mut self, bufs: &'a [IoSlice<'a>]) -> WriteVectored<'a, Self> + where Self: Unpin, + { + WriteVectored::new(self, bufs) + } + + /// Write data into this object. + /// + /// Creates a future that will write the entire contents of the buffer `buf` into + /// this `AsyncWrite`. + /// + /// The returned future will not complete until all the data has been written. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncWriteExt, Cursor}; + /// + /// let mut writer = Cursor::new(vec![0u8; 5]); + /// + /// writer.write_all(&[1, 2, 3, 4]).await?; + /// + /// assert_eq!(writer.into_inner(), [1, 2, 3, 4, 0]); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn write_all<'a>(&'a mut self, buf: &'a [u8]) -> WriteAll<'a, Self> + where Self: Unpin, + { + WriteAll::new(self, buf) + } + + /// Wraps an [`AsyncWrite`] in a compatibility wrapper that allows it to be + /// used as a futures 0.1 / tokio-io 0.1 `AsyncWrite`. + /// Requires the `io-compat` feature to enable. + #[cfg(feature = "io-compat")] + fn compat_write(self) -> Compat<Self> + where Self: Sized + Unpin, + { + Compat::new(self) + } + + + /// Allow using an [`AsyncWrite`] as a [`Sink`](futures_sink::Sink)`<Item: AsRef<[u8]>>`. + /// + /// This adapter produces a sink that will write each value passed to it + /// into the underlying writer. + /// + /// Note that this function consumes the given writer, returning a wrapped + /// version. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::AsyncWriteExt; + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(vec![Ok([1, 2, 3]), Ok([4, 5, 6])]); + /// + /// let mut writer = vec![]; + /// + /// stream.forward((&mut writer).into_sink()).await?; + /// + /// assert_eq!(writer, vec![1, 2, 3, 4, 5, 6]); + /// # Ok::<(), Box<dyn std::error::Error>>(()) + /// # })?; + /// # Ok::<(), Box<dyn std::error::Error>>(()) + /// ``` + #[cfg(feature = "sink")] + fn into_sink<Item: AsRef<[u8]>>(self) -> IntoSink<Self, Item> + where Self: Sized, + { + IntoSink::new(self) + } +} + +impl<W: AsyncWrite + ?Sized> AsyncWriteExt for W {} + +/// An extension trait which adds utility methods to `AsyncSeek` types. +pub trait AsyncSeekExt: AsyncSeek { + /// Creates a future which will seek an IO object, and then yield the + /// new position in the object and the object itself. + /// + /// In the case of an error the buffer and the object will be discarded, with + /// the error yielded. + fn seek(&mut self, pos: SeekFrom) -> Seek<'_, Self> + where Self: Unpin, + { + Seek::new(self, pos) + } +} + +impl<S: AsyncSeek + ?Sized> AsyncSeekExt for S {} + +/// An extension trait which adds utility methods to `AsyncBufRead` types. +pub trait AsyncBufReadExt: AsyncBufRead { + /// Creates a future which will read all the bytes associated with this I/O + /// object into `buf` until the delimiter `byte` or EOF is reached. + /// This method is the async equivalent to [`BufRead::read_until`](std::io::BufRead::read_until). + /// + /// This function will read bytes from the underlying stream until the + /// delimiter or EOF is found. Once found, all bytes up to, and including, + /// the delimiter (if found) will be appended to `buf`. + /// + /// The returned future will resolve to the number of bytes read once the read + /// operation is completed. + /// + /// In the case of an error the buffer and the object will be discarded, with + /// the error yielded. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncBufReadExt, Cursor}; + /// + /// let mut cursor = Cursor::new(b"lorem-ipsum"); + /// let mut buf = vec![]; + /// + /// // cursor is at 'l' + /// let num_bytes = cursor.read_until(b'-', &mut buf).await?; + /// assert_eq!(num_bytes, 6); + /// assert_eq!(buf, b"lorem-"); + /// buf.clear(); + /// + /// // cursor is at 'i' + /// let num_bytes = cursor.read_until(b'-', &mut buf).await?; + /// assert_eq!(num_bytes, 5); + /// assert_eq!(buf, b"ipsum"); + /// buf.clear(); + /// + /// // cursor is at EOF + /// let num_bytes = cursor.read_until(b'-', &mut buf).await?; + /// assert_eq!(num_bytes, 0); + /// assert_eq!(buf, b""); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn read_until<'a>( + &'a mut self, + byte: u8, + buf: &'a mut Vec<u8>, + ) -> ReadUntil<'a, Self> + where Self: Unpin, + { + ReadUntil::new(self, byte, buf) + } + + /// Creates a future which will read all the bytes associated with this I/O + /// object into `buf` until a newline (the 0xA byte) or EOF is reached, + /// This method is the async equivalent to [`BufRead::read_line`](std::io::BufRead::read_line). + /// + /// This function will read bytes from the underlying stream until the + /// newline delimiter (the 0xA byte) or EOF is found. Once found, all bytes + /// up to, and including, the delimiter (if found) will be appended to + /// `buf`. + /// + /// The returned future will resolve to the number of bytes read once the read + /// operation is completed. + /// + /// In the case of an error the buffer and the object will be discarded, with + /// the error yielded. + /// + /// # Errors + /// + /// This function has the same error semantics as [`read_until`] and will + /// also return an error if the read bytes are not valid UTF-8. If an I/O + /// error is encountered then `buf` may contain some bytes already read in + /// the event that all data read so far was valid UTF-8. + /// + /// [`read_until`]: AsyncBufReadExt::read_until + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncBufReadExt, Cursor}; + /// + /// let mut cursor = Cursor::new(b"foo\nbar"); + /// let mut buf = String::new(); + /// + /// // cursor is at 'f' + /// let num_bytes = cursor.read_line(&mut buf).await?; + /// assert_eq!(num_bytes, 4); + /// assert_eq!(buf, "foo\n"); + /// buf.clear(); + /// + /// // cursor is at 'b' + /// let num_bytes = cursor.read_line(&mut buf).await?; + /// assert_eq!(num_bytes, 3); + /// assert_eq!(buf, "bar"); + /// buf.clear(); + /// + /// // cursor is at EOF + /// let num_bytes = cursor.read_line(&mut buf).await?; + /// assert_eq!(num_bytes, 0); + /// assert_eq!(buf, ""); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn read_line<'a>(&'a mut self, buf: &'a mut String) -> ReadLine<'a, Self> + where Self: Unpin, + { + ReadLine::new(self, buf) + } + + /// Returns a stream over the lines of this reader. + /// This method is the async equivalent to [`BufRead::lines`](std::io::BufRead::lines). + /// + /// The stream returned from this function will yield instances of + /// [`io::Result`]`<`[`String`]`>`. Each string returned will *not* have a newline + /// byte (the 0xA byte) or CRLF (0xD, 0xA bytes) at the end. + /// + /// [`io::Result`]: std::io::Result + /// [`String`]: String + /// + /// # Errors + /// + /// Each line of the stream has the same error semantics as [`AsyncBufReadExt::read_line`]. + /// + /// [`AsyncBufReadExt::read_line`]: AsyncBufReadExt::read_line + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncBufReadExt, Cursor}; + /// use futures::stream::StreamExt; + /// + /// let cursor = Cursor::new(b"lorem\nipsum\r\ndolor"); + /// + /// let mut lines_stream = cursor.lines().map(|l| l.unwrap()); + /// assert_eq!(lines_stream.next().await, Some(String::from("lorem"))); + /// assert_eq!(lines_stream.next().await, Some(String::from("ipsum"))); + /// assert_eq!(lines_stream.next().await, Some(String::from("dolor"))); + /// assert_eq!(lines_stream.next().await, None); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + fn lines(self) -> Lines<Self> + where Self: Sized, + { + Lines::new(self) + } +} + +impl<R: AsyncBufRead + ?Sized> AsyncBufReadExt for R {} diff --git a/src/io/read.rs b/src/io/read.rs new file mode 100644 index 0000000..ea25959 --- /dev/null +++ b/src/io/read.rs @@ -0,0 +1,30 @@ +use crate::io::AsyncRead; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use std::io; +use std::pin::Pin; + +/// Future for the [`read`](super::AsyncReadExt::read) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Read<'a, R: ?Sized> { + reader: &'a mut R, + buf: &'a mut [u8], +} + +impl<R: ?Sized + Unpin> Unpin for Read<'_, R> {} + +impl<'a, R: AsyncRead + ?Sized + Unpin> Read<'a, R> { + pub(super) fn new(reader: &'a mut R, buf: &'a mut [u8]) -> Self { + Read { reader, buf } + } +} + +impl<R: AsyncRead + ?Sized + Unpin> Future for Read<'_, R> { + type Output = io::Result<usize>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = &mut *self; + Pin::new(&mut this.reader).poll_read(cx, this.buf) + } +} diff --git a/src/io/read_exact.rs b/src/io/read_exact.rs new file mode 100644 index 0000000..a2bbd40 --- /dev/null +++ b/src/io/read_exact.rs @@ -0,0 +1,41 @@ +use crate::io::AsyncRead; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use std::io; +use std::mem; +use std::pin::Pin; + +/// Future for the [`read_exact`](super::AsyncReadExt::read_exact) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct ReadExact<'a, R: ?Sized> { + reader: &'a mut R, + buf: &'a mut [u8], +} + +impl<R: ?Sized + Unpin> Unpin for ReadExact<'_, R> {} + +impl<'a, R: AsyncRead + ?Sized + Unpin> ReadExact<'a, R> { + pub(super) fn new(reader: &'a mut R, buf: &'a mut [u8]) -> Self { + ReadExact { reader, buf } + } +} + +impl<R: AsyncRead + ?Sized + Unpin> Future for ReadExact<'_, R> { + type Output = io::Result<()>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = &mut *self; + while !this.buf.is_empty() { + let n = ready!(Pin::new(&mut this.reader).poll_read(cx, this.buf))?; + { + let (_, rest) = mem::replace(&mut this.buf, &mut []).split_at_mut(n); + this.buf = rest; + } + if n == 0 { + return Poll::Ready(Err(io::ErrorKind::UnexpectedEof.into())) + } + } + Poll::Ready(Ok(())) + } +} diff --git a/src/io/read_line.rs b/src/io/read_line.rs new file mode 100644 index 0000000..d830514 --- /dev/null +++ b/src/io/read_line.rs @@ -0,0 +1,61 @@ +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_io::AsyncBufRead; +use std::io; +use std::mem; +use std::pin::Pin; +use std::str; +use super::read_until::read_until_internal; + +/// Future for the [`read_line`](super::AsyncBufReadExt::read_line) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct ReadLine<'a, R: ?Sized> { + reader: &'a mut R, + buf: &'a mut String, + bytes: Vec<u8>, + read: usize, +} + +impl<R: ?Sized + Unpin> Unpin for ReadLine<'_, R> {} + +impl<'a, R: AsyncBufRead + ?Sized + Unpin> ReadLine<'a, R> { + pub(super) fn new(reader: &'a mut R, buf: &'a mut String) -> Self { + Self { + reader, + bytes: unsafe { mem::replace(buf.as_mut_vec(), Vec::new()) }, + buf, + read: 0, + } + } +} + +pub(super) fn read_line_internal<R: AsyncBufRead + ?Sized>( + reader: Pin<&mut R>, + cx: &mut Context<'_>, + buf: &mut String, + bytes: &mut Vec<u8>, + read: &mut usize, +) -> Poll<io::Result<usize>> { + let ret = ready!(read_until_internal(reader, cx, b'\n', bytes, read)); + if str::from_utf8(&bytes).is_err() { + Poll::Ready(ret.and_then(|_| { + Err(io::Error::new(io::ErrorKind::InvalidData, "stream did not contain valid UTF-8")) + })) + } else { + debug_assert!(buf.is_empty()); + debug_assert_eq!(*read, 0); + // Safety: `bytes` is a valid UTF-8 because `str::from_utf8` returned `Ok`. + mem::swap(unsafe { buf.as_mut_vec() }, bytes); + Poll::Ready(ret) + } +} + +impl<R: AsyncBufRead + ?Sized + Unpin> Future for ReadLine<'_, R> { + type Output = io::Result<usize>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let Self { reader, buf, bytes, read } = &mut *self; + read_line_internal(Pin::new(reader), cx, buf, bytes, read) + } +} diff --git a/src/io/read_to_end.rs b/src/io/read_to_end.rs new file mode 100644 index 0000000..70b0578 --- /dev/null +++ b/src/io/read_to_end.rs @@ -0,0 +1,90 @@ +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_io::AsyncRead; +use std::io; +use std::pin::Pin; +use std::vec::Vec; + +/// Future for the [`read_to_end`](super::AsyncReadExt::read_to_end) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct ReadToEnd<'a, R: ?Sized> { + reader: &'a mut R, + buf: &'a mut Vec<u8>, + start_len: usize, +} + +impl<R: ?Sized + Unpin> Unpin for ReadToEnd<'_, R> {} + +impl<'a, R: AsyncRead + ?Sized + Unpin> ReadToEnd<'a, R> { + pub(super) fn new(reader: &'a mut R, buf: &'a mut Vec<u8>) -> Self { + let start_len = buf.len(); + Self { + reader, + buf, + start_len, + } + } +} + +struct Guard<'a> { buf: &'a mut Vec<u8>, len: usize } + +impl Drop for Guard<'_> { + fn drop(&mut self) { + unsafe { self.buf.set_len(self.len); } + } +} + +// This uses an adaptive system to extend the vector when it fills. We want to +// avoid paying to allocate and zero a huge chunk of memory if the reader only +// has 4 bytes while still making large reads if the reader does have a ton +// of data to return. Simply tacking on an extra DEFAULT_BUF_SIZE space every +// time is 4,500 times (!) slower than this if the reader has a very small +// amount of data to return. +// +// Because we're extending the buffer with uninitialized data for trusted +// readers, we need to make sure to truncate that if any of this panics. +pub(super) fn read_to_end_internal<R: AsyncRead + ?Sized>( + mut rd: Pin<&mut R>, + cx: &mut Context<'_>, + buf: &mut Vec<u8>, + start_len: usize, +) -> Poll<io::Result<usize>> { + let mut g = Guard { len: buf.len(), buf }; + let ret; + loop { + if g.len == g.buf.len() { + unsafe { + g.buf.reserve(32); + let capacity = g.buf.capacity(); + g.buf.set_len(capacity); + super::initialize(&rd, &mut g.buf[g.len..]); + } + } + + match ready!(rd.as_mut().poll_read(cx, &mut g.buf[g.len..])) { + Ok(0) => { + ret = Poll::Ready(Ok(g.len - start_len)); + break; + } + Ok(n) => g.len += n, + Err(e) => { + ret = Poll::Ready(Err(e)); + break; + } + } + } + + ret +} + +impl<A> Future for ReadToEnd<'_, A> + where A: AsyncRead + ?Sized + Unpin, +{ + type Output = io::Result<usize>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = &mut *self; + read_to_end_internal(Pin::new(&mut this.reader), cx, this.buf, this.start_len) + } +} diff --git a/src/io/read_to_string.rs b/src/io/read_to_string.rs new file mode 100644 index 0000000..56c95ce --- /dev/null +++ b/src/io/read_to_string.rs @@ -0,0 +1,66 @@ +use super::read_to_end::read_to_end_internal; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_io::AsyncRead; +use std::pin::Pin; +use std::vec::Vec; +use std::{io, mem, str}; + +/// Future for the [`read_to_string`](super::AsyncReadExt::read_to_string) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct ReadToString<'a, R: ?Sized> { + reader: &'a mut R, + buf: &'a mut String, + bytes: Vec<u8>, + start_len: usize, +} + +impl<R: ?Sized + Unpin> Unpin for ReadToString<'_, R> {} + +impl<'a, R: AsyncRead + ?Sized + Unpin> ReadToString<'a, R> { + pub(super) fn new(reader: &'a mut R, buf: &'a mut String) -> Self { + let start_len = buf.len(); + Self { + reader, + bytes: unsafe { mem::replace(buf.as_mut_vec(), Vec::new()) }, + buf, + start_len, + } + } +} + +fn read_to_string_internal<R: AsyncRead + ?Sized>( + reader: Pin<&mut R>, + cx: &mut Context<'_>, + buf: &mut String, + bytes: &mut Vec<u8>, + start_len: usize, +) -> Poll<io::Result<usize>> { + let ret = ready!(read_to_end_internal(reader, cx, bytes, start_len)); + if str::from_utf8(&bytes).is_err() { + Poll::Ready(ret.and_then(|_| { + Err(io::Error::new( + io::ErrorKind::InvalidData, + "stream did not contain valid UTF-8", + )) + })) + } else { + debug_assert!(buf.is_empty()); + // Safety: `bytes` is a valid UTF-8 because `str::from_utf8` returned `Ok`. + mem::swap(unsafe { buf.as_mut_vec() }, bytes); + Poll::Ready(ret) + } +} + +impl<A> Future for ReadToString<'_, A> +where + A: AsyncRead + ?Sized + Unpin, +{ + type Output = io::Result<usize>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let Self { reader, buf, bytes, start_len } = &mut *self; + read_to_string_internal(Pin::new(reader), cx, buf, bytes, *start_len) + } +} diff --git a/src/io/read_until.rs b/src/io/read_until.rs new file mode 100644 index 0000000..95c47e0 --- /dev/null +++ b/src/io/read_until.rs @@ -0,0 +1,59 @@ +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_io::AsyncBufRead; +use std::io; +use std::mem; +use std::pin::Pin; + +/// Future for the [`read_until`](super::AsyncBufReadExt::read_until) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct ReadUntil<'a, R: ?Sized> { + reader: &'a mut R, + byte: u8, + buf: &'a mut Vec<u8>, + read: usize, +} + +impl<R: ?Sized + Unpin> Unpin for ReadUntil<'_, R> {} + +impl<'a, R: AsyncBufRead + ?Sized + Unpin> ReadUntil<'a, R> { + pub(super) fn new(reader: &'a mut R, byte: u8, buf: &'a mut Vec<u8>) -> Self { + Self { reader, byte, buf, read: 0 } + } +} + +pub(super) fn read_until_internal<R: AsyncBufRead + ?Sized>( + mut reader: Pin<&mut R>, + cx: &mut Context<'_>, + byte: u8, + buf: &mut Vec<u8>, + read: &mut usize, +) -> Poll<io::Result<usize>> { + loop { + let (done, used) = { + let available = ready!(reader.as_mut().poll_fill_buf(cx))?; + if let Some(i) = memchr::memchr(byte, available) { + buf.extend_from_slice(&available[..=i]); + (true, i + 1) + } else { + buf.extend_from_slice(available); + (false, available.len()) + } + }; + reader.as_mut().consume(used); + *read += used; + if done || used == 0 { + return Poll::Ready(Ok(mem::replace(read, 0))); + } + } +} + +impl<R: AsyncBufRead + ?Sized + Unpin> Future for ReadUntil<'_, R> { + type Output = io::Result<usize>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let Self { reader, byte, buf, read } = &mut *self; + read_until_internal(Pin::new(reader), cx, *byte, buf, read) + } +} diff --git a/src/io/read_vectored.rs b/src/io/read_vectored.rs new file mode 100644 index 0000000..4e22df5 --- /dev/null +++ b/src/io/read_vectored.rs @@ -0,0 +1,30 @@ +use crate::io::AsyncRead; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use std::io::{self, IoSliceMut}; +use std::pin::Pin; + +/// Future for the [`read_vectored`](super::AsyncReadExt::read_vectored) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct ReadVectored<'a, R: ?Sized> { + reader: &'a mut R, + bufs: &'a mut [IoSliceMut<'a>], +} + +impl<R: ?Sized + Unpin> Unpin for ReadVectored<'_, R> {} + +impl<'a, R: AsyncRead + ?Sized + Unpin> ReadVectored<'a, R> { + pub(super) fn new(reader: &'a mut R, bufs: &'a mut [IoSliceMut<'a>]) -> Self { + Self { reader, bufs } + } +} + +impl<R: AsyncRead + ?Sized + Unpin> Future for ReadVectored<'_, R> { + type Output = io::Result<usize>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = &mut *self; + Pin::new(&mut this.reader).poll_read_vectored(cx, this.bufs) + } +} diff --git a/src/io/repeat.rs b/src/io/repeat.rs new file mode 100644 index 0000000..84abd7f --- /dev/null +++ b/src/io/repeat.rs @@ -0,0 +1,73 @@ +use futures_core::task::{Context, Poll}; +#[cfg(feature = "read-initializer")] +use futures_io::Initializer; +use futures_io::{AsyncRead, IoSliceMut}; +use std::fmt; +use std::io; +use std::pin::Pin; + +/// Reader for the [`repeat()`] function. +#[must_use = "readers do nothing unless polled"] +pub struct Repeat { + byte: u8, +} + +/// Creates an instance of a reader that infinitely repeats one byte. +/// +/// All reads from this reader will succeed by filling the specified buffer with +/// the given byte. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::io::{self, AsyncReadExt}; +/// +/// let mut buffer = [0; 3]; +/// let mut reader = io::repeat(0b101); +/// reader.read_exact(&mut buffer).await.unwrap(); +/// assert_eq!(buffer, [0b101, 0b101, 0b101]); +/// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); +/// ``` +pub fn repeat(byte: u8) -> Repeat { + Repeat { byte } +} + +impl AsyncRead for Repeat { + #[inline] + fn poll_read( + self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll<io::Result<usize>> { + for slot in &mut *buf { + *slot = self.byte; + } + Poll::Ready(Ok(buf.len())) + } + + #[inline] + fn poll_read_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &mut [IoSliceMut<'_>], + ) -> Poll<io::Result<usize>> { + let mut nwritten = 0; + for buf in bufs { + nwritten += ready!(self.as_mut().poll_read(cx, buf))?; + } + Poll::Ready(Ok(nwritten)) + } + + #[cfg(feature = "read-initializer")] + #[inline] + unsafe fn initializer(&self) -> Initializer { + Initializer::nop() + } +} + +impl fmt::Debug for Repeat { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.pad("Repeat { .. }") + } +} diff --git a/src/io/seek.rs b/src/io/seek.rs new file mode 100644 index 0000000..0aa2371 --- /dev/null +++ b/src/io/seek.rs @@ -0,0 +1,30 @@ +use crate::io::{AsyncSeek, SeekFrom}; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use std::io; +use std::pin::Pin; + +/// Future for the [`seek`](crate::io::AsyncSeekExt::seek) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Seek<'a, S: ?Sized> { + seek: &'a mut S, + pos: SeekFrom, +} + +impl<S: ?Sized + Unpin> Unpin for Seek<'_, S> {} + +impl<'a, S: AsyncSeek + ?Sized + Unpin> Seek<'a, S> { + pub(super) fn new(seek: &'a mut S, pos: SeekFrom) -> Self { + Self { seek, pos } + } +} + +impl<S: AsyncSeek + ?Sized + Unpin> Future for Seek<'_, S> { + type Output = io::Result<u64>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = &mut *self; + Pin::new(&mut this.seek).poll_seek(cx, this.pos) + } +} diff --git a/src/io/sink.rs b/src/io/sink.rs new file mode 100644 index 0000000..4a32ca7 --- /dev/null +++ b/src/io/sink.rs @@ -0,0 +1,67 @@ +use futures_core::task::{Context, Poll}; +use futures_io::{AsyncWrite, IoSlice}; +use std::fmt; +use std::io; +use std::pin::Pin; + +/// Writer for the [`sink()`] function. +#[must_use = "writers do nothing unless polled"] +pub struct Sink { + _priv: (), +} + +/// Creates an instance of a writer which will successfully consume all data. +/// +/// All calls to `poll_write` on the returned instance will return `Poll::Ready(Ok(buf.len()))` +/// and the contents of the buffer will not be inspected. +/// +/// # Examples +/// +/// ```rust +/// # futures::executor::block_on(async { +/// use futures::io::{self, AsyncWriteExt}; +/// +/// let buffer = vec![1, 2, 3, 5, 8]; +/// let mut writer = io::sink(); +/// let num_bytes = writer.write(&buffer).await?; +/// assert_eq!(num_bytes, 5); +/// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); +/// ``` +pub fn sink() -> Sink { + Sink { _priv: () } +} + +impl AsyncWrite for Sink { + #[inline] + fn poll_write( + self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + Poll::Ready(Ok(buf.len())) + } + + #[inline] + fn poll_write_vectored( + self: Pin<&mut Self>, + _: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + Poll::Ready(Ok(bufs.iter().map(|b| b.len()).sum())) + } + + #[inline] + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(Ok(())) + } + #[inline] + fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(Ok(())) + } +} + +impl fmt::Debug for Sink { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.pad("Sink { .. }") + } +} diff --git a/src/io/split.rs b/src/io/split.rs new file mode 100644 index 0000000..d12df86 --- /dev/null +++ b/src/io/split.rs @@ -0,0 +1,69 @@ +use crate::lock::BiLock; +use futures_core::task::{Context, Poll}; +use futures_io::{AsyncRead, AsyncWrite, IoSlice, IoSliceMut}; +use std::io; +use std::pin::Pin; + +/// The readable half of an object returned from `AsyncRead::split`. +#[derive(Debug)] +pub struct ReadHalf<T> { + handle: BiLock<T>, +} + +/// The writable half of an object returned from `AsyncRead::split`. +#[derive(Debug)] +pub struct WriteHalf<T> { + handle: BiLock<T>, +} + +fn lock_and_then<T, U, E, F>( + lock: &BiLock<T>, + cx: &mut Context<'_>, + f: F +) -> Poll<Result<U, E>> + where F: FnOnce(Pin<&mut T>, &mut Context<'_>) -> Poll<Result<U, E>> +{ + let mut l = ready!(lock.poll_lock(cx)); + f(l.as_pin_mut(), cx) +} + +pub(super) fn split<T: AsyncRead + AsyncWrite>(t: T) -> (ReadHalf<T>, WriteHalf<T>) { + let (a, b) = BiLock::new(t); + (ReadHalf { handle: a }, WriteHalf { handle: b }) +} + +impl<R: AsyncRead> AsyncRead for ReadHalf<R> { + fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) + -> Poll<io::Result<usize>> + { + lock_and_then(&self.handle, cx, |l, cx| l.poll_read(cx, buf)) + } + + fn poll_read_vectored(self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &mut [IoSliceMut<'_>]) + -> Poll<io::Result<usize>> + { + lock_and_then(&self.handle, cx, |l, cx| l.poll_read_vectored(cx, bufs)) + } +} + +impl<W: AsyncWrite> AsyncWrite for WriteHalf<W> { + fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) + -> Poll<io::Result<usize>> + { + lock_and_then(&self.handle, cx, |l, cx| l.poll_write(cx, buf)) + } + + fn poll_write_vectored(self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>]) + -> Poll<io::Result<usize>> + { + lock_and_then(&self.handle, cx, |l, cx| l.poll_write_vectored(cx, bufs)) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + lock_and_then(&self.handle, cx, |l, cx| l.poll_flush(cx)) + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + lock_and_then(&self.handle, cx, |l, cx| l.poll_close(cx)) + } +} diff --git a/src/io/take.rs b/src/io/take.rs new file mode 100644 index 0000000..b1f33fa --- /dev/null +++ b/src/io/take.rs @@ -0,0 +1,210 @@ +use futures_core::task::{Context, Poll}; +#[cfg(feature = "read-initializer")] +use futures_io::Initializer; +use futures_io::{AsyncRead, AsyncBufRead}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; +use std::{cmp, io}; +use std::pin::Pin; + +/// Reader for the [`take`](super::AsyncReadExt::take) method. +#[derive(Debug)] +#[must_use = "readers do nothing unless you `.await` or poll them"] +pub struct Take<R> { + inner: R, + // Add '_' to avoid conflicts with `limit` method. + limit_: u64, +} + +impl<R: Unpin> Unpin for Take<R> { } + +impl<R: AsyncRead> Take<R> { + unsafe_pinned!(inner: R); + unsafe_unpinned!(limit_: u64); + + pub(super) fn new(inner: R, limit: u64) -> Self { + Self { inner, limit_: limit } + } + + /// Returns the remaining number of bytes that can be + /// read before this instance will return EOF. + /// + /// # Note + /// + /// This instance may reach `EOF` after reading fewer bytes than indicated by + /// this method if the underlying [`AsyncRead`] instance reaches EOF. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let reader = Cursor::new(&b"12345678"[..]); + /// let mut buffer = [0; 2]; + /// + /// let mut take = reader.take(4); + /// let n = take.read(&mut buffer).await?; + /// + /// assert_eq!(take.limit(), 2); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + pub fn limit(&self) -> u64 { + self.limit_ + } + + /// Sets the number of bytes that can be read before this instance will + /// return EOF. This is the same as constructing a new `Take` instance, so + /// the amount of bytes read and the previous limit value don't matter when + /// calling this method. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let reader = Cursor::new(&b"12345678"[..]); + /// let mut buffer = [0; 4]; + /// + /// let mut take = reader.take(4); + /// let n = take.read(&mut buffer).await?; + /// + /// assert_eq!(n, 4); + /// assert_eq!(take.limit(), 0); + /// + /// take.set_limit(10); + /// let n = take.read(&mut buffer).await?; + /// assert_eq!(n, 4); + /// + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + pub fn set_limit(&mut self, limit: u64) { + self.limit_ = limit + } + + /// Gets a reference to the underlying reader. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let reader = Cursor::new(&b"12345678"[..]); + /// let mut buffer = [0; 4]; + /// + /// let mut take = reader.take(4); + /// let n = take.read(&mut buffer).await?; + /// + /// let cursor_ref = take.get_ref(); + /// assert_eq!(cursor_ref.position(), 4); + /// + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + pub fn get_ref(&self) -> &R { + &self.inner + } + + /// Gets a mutable reference to the underlying reader. + /// + /// Care should be taken to avoid modifying the internal I/O state of the + /// underlying reader as doing so may corrupt the internal limit of this + /// `Take`. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let reader = Cursor::new(&b"12345678"[..]); + /// let mut buffer = [0; 4]; + /// + /// let mut take = reader.take(4); + /// let n = take.read(&mut buffer).await?; + /// + /// let cursor_mut = take.get_mut(); + /// + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + pub fn get_mut(&mut self) -> &mut R { + &mut self.inner + } + + /// Gets a pinned mutable reference to the underlying reader. + /// + /// Care should be taken to avoid modifying the internal I/O state of the + /// underlying reader as doing so may corrupt the internal limit of this + /// `Take`. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> { + self.inner() + } + + /// Consumes the `Take`, returning the wrapped reader. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::io::{AsyncReadExt, Cursor}; + /// + /// let reader = Cursor::new(&b"12345678"[..]); + /// let mut buffer = [0; 4]; + /// + /// let mut take = reader.take(4); + /// let n = take.read(&mut buffer).await?; + /// + /// let cursor = take.into_inner(); + /// assert_eq!(cursor.position(), 4); + /// + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + pub fn into_inner(self) -> R { + self.inner + } +} + +impl<R: AsyncRead> AsyncRead for Take<R> { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll<Result<usize, io::Error>> { + if self.limit_ == 0 { + return Poll::Ready(Ok(0)); + } + + let max = std::cmp::min(buf.len() as u64, self.limit_) as usize; + let n = ready!(self.as_mut().inner().poll_read(cx, &mut buf[..max]))?; + *self.as_mut().limit_() -= n as u64; + Poll::Ready(Ok(n)) + } + + #[cfg(feature = "read-initializer")] + unsafe fn initializer(&self) -> Initializer { + self.inner.initializer() + } +} + +impl<R: AsyncBufRead> AsyncBufRead for Take<R> { + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + let Self { inner, limit_ } = unsafe { self.get_unchecked_mut() }; + let inner = unsafe { Pin::new_unchecked(inner) }; + + // Don't call into inner reader at all at EOF because it may still block + if *limit_ == 0 { + return Poll::Ready(Ok(&[])); + } + + let buf = ready!(inner.poll_fill_buf(cx)?); + let cap = cmp::min(buf.len() as u64, *limit_) as usize; + Poll::Ready(Ok(&buf[..cap])) + } + + fn consume(mut self: Pin<&mut Self>, amt: usize) { + // Don't let callers reset the limit by passing an overlarge value + let amt = cmp::min(amt as u64, self.limit_) as usize; + *self.as_mut().limit_() -= amt as u64; + self.inner().consume(amt); + } +} diff --git a/src/io/window.rs b/src/io/window.rs new file mode 100644 index 0000000..3424197 --- /dev/null +++ b/src/io/window.rs @@ -0,0 +1,107 @@ +use std::ops::{Bound, Range, RangeBounds}; + +/// A owned window around an underlying buffer. +/// +/// Normally slices work great for considering sub-portions of a buffer, but +/// unfortunately a slice is a *borrowed* type in Rust which has an associated +/// lifetime. When working with future and async I/O these lifetimes are not +/// always appropriate, and are sometimes difficult to store in tasks. This +/// type strives to fill this gap by providing an "owned slice" around an +/// underlying buffer of bytes. +/// +/// A `Window<T>` wraps an underlying buffer, `T`, and has configurable +/// start/end indexes to alter the behavior of the `AsRef<[u8]>` implementation +/// that this type carries. +/// +/// This type can be particularly useful when working with the `write_all` +/// combinator in this crate. Data can be sliced via `Window`, consumed by +/// `write_all`, and then earned back once the write operation finishes through +/// the `into_inner` method on this type. +#[derive(Debug)] +pub struct Window<T> { + inner: T, + range: Range<usize>, +} + +impl<T: AsRef<[u8]>> Window<T> { + /// Creates a new window around the buffer `t` defaulting to the entire + /// slice. + /// + /// Further methods can be called on the returned `Window<T>` to alter the + /// window into the data provided. + pub fn new(t: T) -> Self { + Self { + range: 0..t.as_ref().len(), + inner: t, + } + } + + /// Gets a shared reference to the underlying buffer inside of this + /// `Window`. + pub fn get_ref(&self) -> &T { + &self.inner + } + + /// Gets a mutable reference to the underlying buffer inside of this + /// `Window`. + pub fn get_mut(&mut self) -> &mut T { + &mut self.inner + } + + /// Consumes this `Window`, returning the underlying buffer. + pub fn into_inner(self) -> T { + self.inner + } + + /// Returns the starting index of this window into the underlying buffer + /// `T`. + pub fn start(&self) -> usize { + self.range.start + } + + /// Returns the end index of this window into the underlying buffer + /// `T`. + pub fn end(&self) -> usize { + self.range.end + } + + /// Changes the range of this window to the range specified. + /// + /// # Panics + /// + /// This method will panic if `range` is out of bounds for the underlying + /// slice or if [`start_bound()`] of `range` comes after the [`end_bound()`]. + /// + /// [`start_bound()`]: std::ops::RangeBounds::start_bound + /// [`end_bound()`]: std::ops::RangeBounds::end_bound + pub fn set<R: RangeBounds<usize>>(&mut self, range: R) { + let start = match range.start_bound() { + Bound::Included(n) => *n, + Bound::Excluded(n) => *n + 1, + Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + Bound::Included(n) => *n + 1, + Bound::Excluded(n) => *n, + Bound::Unbounded => self.inner.as_ref().len(), + }; + + assert!(end <= self.inner.as_ref().len()); + assert!(start <= end); + + self.range.start = start; + self.range.end = end; + } +} + +impl<T: AsRef<[u8]>> AsRef<[u8]> for Window<T> { + fn as_ref(&self) -> &[u8] { + &self.inner.as_ref()[self.range.start..self.range.end] + } +} + +impl<T: AsMut<[u8]>> AsMut<[u8]> for Window<T> { + fn as_mut(&mut self) -> &mut [u8] { + &mut self.inner.as_mut()[self.range.start..self.range.end] + } +} diff --git a/src/io/write.rs b/src/io/write.rs new file mode 100644 index 0000000..c47ef9e --- /dev/null +++ b/src/io/write.rs @@ -0,0 +1,30 @@ +use crate::io::AsyncWrite; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use std::io; +use std::pin::Pin; + +/// Future for the [`write`](super::AsyncWriteExt::write) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Write<'a, W: ?Sized> { + writer: &'a mut W, + buf: &'a [u8], +} + +impl<W: ?Sized + Unpin> Unpin for Write<'_, W> {} + +impl<'a, W: AsyncWrite + ?Sized + Unpin> Write<'a, W> { + pub(super) fn new(writer: &'a mut W, buf: &'a [u8]) -> Self { + Self { writer, buf } + } +} + +impl<W: AsyncWrite + ?Sized + Unpin> Future for Write<'_, W> { + type Output = io::Result<usize>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = &mut *self; + Pin::new(&mut this.writer).poll_write(cx, this.buf) + } +} diff --git a/src/io/write_all.rs b/src/io/write_all.rs new file mode 100644 index 0000000..57f1400 --- /dev/null +++ b/src/io/write_all.rs @@ -0,0 +1,42 @@ +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_io::AsyncWrite; +use std::io; +use std::mem; +use std::pin::Pin; + +/// Future for the [`write_all`](super::AsyncWriteExt::write_all) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct WriteAll<'a, W: ?Sized> { + writer: &'a mut W, + buf: &'a [u8], +} + +impl<W: ?Sized + Unpin> Unpin for WriteAll<'_, W> {} + +impl<'a, W: AsyncWrite + ?Sized + Unpin> WriteAll<'a, W> { + pub(super) fn new(writer: &'a mut W, buf: &'a [u8]) -> Self { + WriteAll { writer, buf } + } +} + +impl<W: AsyncWrite + ?Sized + Unpin> Future for WriteAll<'_, W> { + type Output = io::Result<()>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + let this = &mut *self; + while !this.buf.is_empty() { + let n = ready!(Pin::new(&mut this.writer).poll_write(cx, this.buf))?; + { + let (_, rest) = mem::replace(&mut this.buf, &[]).split_at(n); + this.buf = rest; + } + if n == 0 { + return Poll::Ready(Err(io::ErrorKind::WriteZero.into())) + } + } + + Poll::Ready(Ok(())) + } +} diff --git a/src/io/write_vectored.rs b/src/io/write_vectored.rs new file mode 100644 index 0000000..14a01d7 --- /dev/null +++ b/src/io/write_vectored.rs @@ -0,0 +1,30 @@ +use crate::io::AsyncWrite; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use std::io::{self, IoSlice}; +use std::pin::Pin; + +/// Future for the [`write_vectored`](super::AsyncWriteExt::write_vectored) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct WriteVectored<'a, W: ?Sized> { + writer: &'a mut W, + bufs: &'a [IoSlice<'a>], +} + +impl<W: ?Sized + Unpin> Unpin for WriteVectored<'_, W> {} + +impl<'a, W: AsyncWrite + ?Sized + Unpin> WriteVectored<'a, W> { + pub(super) fn new(writer: &'a mut W, bufs: &'a [IoSlice<'a>]) -> Self { + Self { writer, bufs } + } +} + +impl<W: AsyncWrite + ?Sized + Unpin> Future for WriteVectored<'_, W> { + type Output = io::Result<usize>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = &mut *self; + Pin::new(&mut this.writer).poll_write_vectored(cx, this.bufs) + } +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..21645b1 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,122 @@ +//! Combinators and utilities for working with `Future`s, `Stream`s, `Sink`s, +//! and the `AsyncRead` and `AsyncWrite` traits. + +#![cfg_attr(feature = "cfg-target-has-atomic", feature(cfg_target_has_atomic))] +#![cfg_attr(feature = "read-initializer", feature(read_initializer))] + +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs, missing_debug_implementations, rust_2018_idioms, unreachable_pub)] +// It cannot be included in the published code because this lints have false positives in the minimum required version. +#![cfg_attr(test, warn(single_use_lifetimes))] +#![warn(clippy::all)] + +// The solution for this lint is not available on 1.39 which is the current minimum supported version. +// Can be removed as of minimum supported 1.40 or if https://github.com/rust-lang/rust-clippy/issues/3941 +// get's implemented. +#![allow(clippy::mem_replace_with_default)] + +#![doc(test(attr(deny(warnings), allow(dead_code, unused_assignments, unused_variables))))] + +#![doc(html_root_url = "https://docs.rs/futures-util/0.3.0")] + +#[cfg(all(feature = "cfg-target-has-atomic", not(feature = "unstable")))] +compile_error!("The `cfg-target-has-atomic` feature requires the `unstable` feature as an explicit opt-in to unstable features"); + +#[cfg(all(feature = "bilock", not(feature = "unstable")))] +compile_error!("The `bilock` feature requires the `unstable` feature as an explicit opt-in to unstable features"); + +#[cfg(all(feature = "read-initializer", not(feature = "unstable")))] +compile_error!("The `read-initializer` feature requires the `unstable` feature as an explicit opt-in to unstable features"); + +#[cfg(feature = "alloc")] +extern crate alloc; + +#[macro_use(ready)] +extern crate futures_core; + +// Macro re-exports +pub use futures_core::ready; +pub use pin_utils::pin_mut; + +// Not public API. +#[cfg(feature = "async-await")] +#[macro_use] +#[doc(hidden)] +pub mod async_await; +#[cfg(feature = "async-await")] +#[doc(hidden)] +pub use self::async_await::*; + +// Not public API. +#[doc(hidden)] +pub use futures_core::core_reexport; + +macro_rules! cfg_target_has_atomic { + ($($item:item)*) => {$( + #[cfg_attr(feature = "cfg-target-has-atomic", cfg(target_has_atomic = "ptr"))] + $item + )*}; +} + +#[cfg(feature = "sink")] +macro_rules! delegate_sink { + ($field:ident, $item:ty) => { + fn poll_ready( + self: Pin<&mut Self>, + cx: &mut $crate::core_reexport::task::Context<'_>, + ) -> $crate::core_reexport::task::Poll<Result<(), Self::Error>> { + self.$field().poll_ready(cx) + } + + fn start_send( + self: Pin<&mut Self>, + item: $item, + ) -> Result<(), Self::Error> { + self.$field().start_send(item) + } + + fn poll_flush( + self: Pin<&mut Self>, + cx: &mut $crate::core_reexport::task::Context<'_>, + ) -> $crate::core_reexport::task::Poll<Result<(), Self::Error>> { + self.$field().poll_flush(cx) + } + + fn poll_close( + self: Pin<&mut Self>, + cx: &mut $crate::core_reexport::task::Context<'_>, + ) -> $crate::core_reexport::task::Poll<Result<(), Self::Error>> { + self.$field().poll_close(cx) + } + } +} + +pub mod future; +#[doc(hidden)] pub use crate::future::{FutureExt, TryFutureExt}; + +pub mod stream; +#[doc(hidden)] pub use crate::stream::{StreamExt, TryStreamExt}; + +#[cfg(feature = "sink")] +pub mod sink; +#[cfg(feature = "sink")] +#[doc(hidden)] pub use crate::sink::SinkExt; + +pub mod task; + +pub mod never; + +#[cfg(feature = "compat")] +pub mod compat; + +#[cfg(feature = "io")] +#[cfg(feature = "std")] +pub mod io; +#[cfg(feature = "io")] +#[cfg(feature = "std")] +#[doc(hidden)] pub use crate::io::{AsyncReadExt, AsyncWriteExt, AsyncSeekExt, AsyncBufReadExt}; + +cfg_target_has_atomic! { + #[cfg(feature = "alloc")] + pub mod lock; +} diff --git a/src/lock/bilock.rs b/src/lock/bilock.rs new file mode 100644 index 0000000..374ea50 --- /dev/null +++ b/src/lock/bilock.rs @@ -0,0 +1,282 @@ +//! Futures-powered synchronization primitives. + +#[cfg(feature = "bilock")] +use futures_core::future::Future; +use futures_core::task::{Context, Poll, Waker}; +use core::cell::UnsafeCell; +#[cfg(any(feature = "bilock", feature = "sink"))] +use core::fmt; +use core::ops::{Deref, DerefMut}; +use core::pin::Pin; +use core::sync::atomic::AtomicUsize; +use core::sync::atomic::Ordering::SeqCst; +use alloc::boxed::Box; +use alloc::sync::Arc; + +/// A type of futures-powered synchronization primitive which is a mutex between +/// two possible owners. +/// +/// This primitive is not as generic as a full-blown mutex but is sufficient for +/// many use cases where there are only two possible owners of a resource. The +/// implementation of `BiLock` can be more optimized for just the two possible +/// owners. +/// +/// Note that it's possible to use this lock through a poll-style interface with +/// the `poll_lock` method but you can also use it as a future with the `lock` +/// method that consumes a `BiLock` and returns a future that will resolve when +/// it's locked. +/// +/// A `BiLock` is typically used for "split" operations where data which serves +/// two purposes wants to be split into two to be worked with separately. For +/// example a TCP stream could be both a reader and a writer or a framing layer +/// could be both a stream and a sink for messages. A `BiLock` enables splitting +/// these two and then using each independently in a futures-powered fashion. +/// +/// This type is only available when the `bilock` feature of this +/// library is activated. +#[derive(Debug)] +pub struct BiLock<T> { + arc: Arc<Inner<T>>, +} + +#[derive(Debug)] +struct Inner<T> { + state: AtomicUsize, + value: Option<UnsafeCell<T>>, +} + +unsafe impl<T: Send> Send for Inner<T> {} +unsafe impl<T: Send> Sync for Inner<T> {} + +impl<T> BiLock<T> { + /// Creates a new `BiLock` protecting the provided data. + /// + /// Two handles to the lock are returned, and these are the only two handles + /// that will ever be available to the lock. These can then be sent to separate + /// tasks to be managed there. + /// + /// The data behind the bilock is considered to be pinned, which allows `Pin` + /// references to locked data. However, this means that the locked value + /// will only be available through `Pin<&mut T>` (not `&mut T`) unless `T` is `Unpin`. + /// Similarly, reuniting the lock and extracting the inner value is only + /// possible when `T` is `Unpin`. + pub fn new(t: T) -> (BiLock<T>, BiLock<T>) { + let arc = Arc::new(Inner { + state: AtomicUsize::new(0), + value: Some(UnsafeCell::new(t)), + }); + + (BiLock { arc: arc.clone() }, BiLock { arc }) + } + + /// Attempt to acquire this lock, returning `Pending` if it can't be + /// acquired. + /// + /// This function will acquire the lock in a nonblocking fashion, returning + /// immediately if the lock is already held. If the lock is successfully + /// acquired then `Poll::Ready` is returned with a value that represents + /// the locked value (and can be used to access the protected data). The + /// lock is unlocked when the returned `BiLockGuard` is dropped. + /// + /// If the lock is already held then this function will return + /// `Poll::Pending`. In this case the current task will also be scheduled + /// to receive a notification when the lock would otherwise become + /// available. + /// + /// # Panics + /// + /// This function will panic if called outside the context of a future's + /// task. + pub fn poll_lock(&self, cx: &mut Context<'_>) -> Poll<BiLockGuard<'_, T>> { + loop { + match self.arc.state.swap(1, SeqCst) { + // Woohoo, we grabbed the lock! + 0 => return Poll::Ready(BiLockGuard { bilock: self }), + + // Oops, someone else has locked the lock + 1 => {} + + // A task was previously blocked on this lock, likely our task, + // so we need to update that task. + n => unsafe { + drop(Box::from_raw(n as *mut Waker)); + } + } + + // type ascription for safety's sake! + let me: Box<Waker> = Box::new(cx.waker().clone()); + let me = Box::into_raw(me) as usize; + + match self.arc.state.compare_exchange(1, me, SeqCst, SeqCst) { + // The lock is still locked, but we've now parked ourselves, so + // just report that we're scheduled to receive a notification. + Ok(_) => return Poll::Pending, + + // Oops, looks like the lock was unlocked after our swap above + // and before the compare_exchange. Deallocate what we just + // allocated and go through the loop again. + Err(0) => unsafe { + drop(Box::from_raw(me as *mut Waker)); + }, + + // The top of this loop set the previous state to 1, so if we + // failed the CAS above then it's because the previous value was + // *not* zero or one. This indicates that a task was blocked, + // but we're trying to acquire the lock and there's only one + // other reference of the lock, so it should be impossible for + // that task to ever block itself. + Err(n) => panic!("invalid state: {}", n), + } + } + } + + /// Perform a "blocking lock" of this lock, consuming this lock handle and + /// returning a future to the acquired lock. + /// + /// This function consumes the `BiLock<T>` and returns a sentinel future, + /// `BiLockAcquire<T>`. The returned future will resolve to + /// `BiLockAcquired<T>` which represents a locked lock similarly to + /// `BiLockGuard<T>`. + /// + /// Note that the returned future will never resolve to an error. + #[cfg(feature = "bilock")] + pub fn lock(&self) -> BiLockAcquire<'_, T> { + BiLockAcquire { + bilock: self, + } + } + + /// Attempts to put the two "halves" of a `BiLock<T>` back together and + /// recover the original value. Succeeds only if the two `BiLock<T>`s + /// originated from the same call to `BiLock::new`. + #[cfg(any(feature = "bilock", feature = "sink"))] + pub fn reunite(self, other: Self) -> Result<T, ReuniteError<T>> + where + T: Unpin, + { + if Arc::ptr_eq(&self.arc, &other.arc) { + drop(other); + let inner = Arc::try_unwrap(self.arc) + .ok() + .expect("futures: try_unwrap failed in BiLock<T>::reunite"); + Ok(unsafe { inner.into_value() }) + } else { + Err(ReuniteError(self, other)) + } + } + + fn unlock(&self) { + match self.arc.state.swap(0, SeqCst) { + // we've locked the lock, shouldn't be possible for us to see an + // unlocked lock. + 0 => panic!("invalid unlocked state"), + + // Ok, no one else tried to get the lock, we're done. + 1 => {} + + // Another task has parked themselves on this lock, let's wake them + // up as its now their turn. + n => unsafe { + Box::from_raw(n as *mut Waker).wake(); + } + } + } +} + +#[cfg(any(feature = "bilock", feature = "sink"))] +impl<T: Unpin> Inner<T> { + unsafe fn into_value(mut self) -> T { + self.value.take().unwrap().into_inner() + } +} + +impl<T> Drop for Inner<T> { + fn drop(&mut self) { + assert_eq!(self.state.load(SeqCst), 0); + } +} + +/// Error indicating two `BiLock<T>`s were not two halves of a whole, and +/// thus could not be `reunite`d. +#[cfg(any(feature = "bilock", feature = "sink"))] +pub struct ReuniteError<T>(pub BiLock<T>, pub BiLock<T>); + +#[cfg(any(feature = "bilock", feature = "sink"))] +impl<T> fmt::Debug for ReuniteError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("ReuniteError") + .field(&"...") + .finish() + } +} + +#[cfg(any(feature = "bilock", feature = "sink"))] +impl<T> fmt::Display for ReuniteError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "tried to reunite two BiLocks that don't form a pair") + } +} + +#[cfg(any(feature = "bilock", feature = "sink"))] +#[cfg(feature = "std")] +impl<T: core::any::Any> std::error::Error for ReuniteError<T> {} + +/// Returned RAII guard from the `poll_lock` method. +/// +/// This structure acts as a sentinel to the data in the `BiLock<T>` itself, +/// implementing `Deref` and `DerefMut` to `T`. When dropped, the lock will be +/// unlocked. +#[derive(Debug)] +pub struct BiLockGuard<'a, T> { + bilock: &'a BiLock<T>, +} + +impl<T> Deref for BiLockGuard<'_, T> { + type Target = T; + fn deref(&self) -> &T { + unsafe { &*self.bilock.arc.value.as_ref().unwrap().get() } + } +} + +impl<T: Unpin> DerefMut for BiLockGuard<'_, T> { + fn deref_mut(&mut self) -> &mut T { + unsafe { &mut *self.bilock.arc.value.as_ref().unwrap().get() } + } +} + +impl<T> BiLockGuard<'_, T> { + /// Get a mutable pinned reference to the locked value. + pub fn as_pin_mut(&mut self) -> Pin<&mut T> { + // Safety: we never allow moving a !Unpin value out of a bilock, nor + // allow mutable access to it + unsafe { Pin::new_unchecked(&mut *self.bilock.arc.value.as_ref().unwrap().get()) } + } +} + +impl<T> Drop for BiLockGuard<'_, T> { + fn drop(&mut self) { + self.bilock.unlock(); + } +} + +/// Future returned by `BiLock::lock` which will resolve when the lock is +/// acquired. +#[cfg(feature = "bilock")] +#[must_use = "futures do nothing unless you `.await` or poll them"] +#[derive(Debug)] +pub struct BiLockAcquire<'a, T> { + bilock: &'a BiLock<T>, +} + +// Pinning is never projected to fields +#[cfg(feature = "bilock")] +impl<T> Unpin for BiLockAcquire<'_, T> {} + +#[cfg(feature = "bilock")] +impl<'a, T> Future for BiLockAcquire<'a, T> { + type Output = BiLockGuard<'a, T>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + self.bilock.poll_lock(cx) + } +} diff --git a/src/lock/mod.rs b/src/lock/mod.rs new file mode 100644 index 0000000..3db5e5b --- /dev/null +++ b/src/lock/mod.rs @@ -0,0 +1,18 @@ +//! Futures-powered synchronization primitives. +//! +//! This module is only available when the `std` or `alloc` feature of this +//! library is activated, and it is activated by default. + +#[cfg(feature = "std")] +mod mutex; +#[cfg(feature = "std")] +pub use self::mutex::{MappedMutexGuard, Mutex, MutexLockFuture, MutexGuard}; + +#[cfg(any(feature = "bilock", feature = "sink", feature = "io"))] +#[cfg_attr(not(feature = "bilock"), allow(unreachable_pub))] +mod bilock; +#[cfg(feature = "bilock")] +pub use self::bilock::{BiLock, BiLockAcquire, BiLockGuard, ReuniteError}; +#[cfg(any(feature = "sink", feature = "io"))] +#[cfg(not(feature = "bilock"))] +pub(crate) use self::bilock::BiLock; diff --git a/src/lock/mutex.rs b/src/lock/mutex.rs new file mode 100644 index 0000000..ccfbf73 --- /dev/null +++ b/src/lock/mutex.rs @@ -0,0 +1,407 @@ +use futures_core::future::{FusedFuture, Future}; +use futures_core::task::{Context, Poll, Waker}; +use slab::Slab; +use std::{fmt, mem}; +use std::cell::UnsafeCell; +use std::ops::{Deref, DerefMut}; +use std::pin::Pin; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::{AtomicUsize, Ordering}; + +/// A futures-aware mutex. +pub struct Mutex<T: ?Sized> { + state: AtomicUsize, + waiters: StdMutex<Slab<Waiter>>, + value: UnsafeCell<T>, +} + +impl<T: ?Sized> fmt::Debug for Mutex<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let state = self.state.load(Ordering::SeqCst); + f.debug_struct("Mutex") + .field("is_locked", &((state & IS_LOCKED) != 0)) + .field("has_waiters", &((state & HAS_WAITERS) != 0)) + .finish() + } +} + +impl<T> From<T> for Mutex<T> { + fn from(t: T) -> Self { + Self::new(t) + } +} + +impl<T: Default> Default for Mutex<T> { + fn default() -> Mutex<T> { + Mutex::new(Default::default()) + } +} + +enum Waiter { + Waiting(Waker), + Woken, +} + +impl Waiter { + fn register(&mut self, waker: &Waker) { + match self { + Waiter::Waiting(w) if waker.will_wake(w) => {}, + _ => *self = Waiter::Waiting(waker.clone()), + } + } + + fn wake(&mut self) { + match mem::replace(self, Waiter::Woken) { + Waiter::Waiting(waker) => waker.wake(), + Waiter::Woken => {}, + } + } +} + +#[allow(clippy::identity_op)] // https://github.com/rust-lang/rust-clippy/issues/3445 +const IS_LOCKED: usize = 1 << 0; +const HAS_WAITERS: usize = 1 << 1; + +impl<T> Mutex<T> { + /// Creates a new futures-aware mutex. + pub fn new(t: T) -> Mutex<T> { + Mutex { + state: AtomicUsize::new(0), + waiters: StdMutex::new(Slab::new()), + value: UnsafeCell::new(t), + } + } + + /// Consumes this mutex, returning the underlying data. + /// + /// # Examples + /// + /// ``` + /// use futures::lock::Mutex; + /// + /// let mutex = Mutex::new(0); + /// assert_eq!(mutex.into_inner(), 0); + /// ``` + pub fn into_inner(self) -> T { + self.value.into_inner() + } +} + +impl<T: ?Sized> Mutex<T> { + /// Attempt to acquire the lock immediately. + /// + /// If the lock is currently held, this will return `None`. + pub fn try_lock(&self) -> Option<MutexGuard<'_, T>> { + let old_state = self.state.fetch_or(IS_LOCKED, Ordering::Acquire); + if (old_state & IS_LOCKED) == 0 { + Some(MutexGuard { mutex: self }) + } else { + None + } + } + + /// Acquire the lock asynchronously. + /// + /// This method returns a future that will resolve once the lock has been + /// successfully acquired. + pub fn lock(&self) -> MutexLockFuture<'_, T> { + MutexLockFuture { + mutex: Some(self), + wait_key: WAIT_KEY_NONE, + } + } + + /// Returns a mutable reference to the underlying data. + /// + /// Since this call borrows the `Mutex` mutably, no actual locking needs to + /// take place -- the mutable borrow statically guarantees no locks exist. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::lock::Mutex; + /// + /// let mut mutex = Mutex::new(0); + /// *mutex.get_mut() = 10; + /// assert_eq!(*mutex.lock().await, 10); + /// # }); + /// ``` + pub fn get_mut(&mut self) -> &mut T { + // We know statically that there are no other references to `self`, so + // there's no need to lock the inner mutex. + unsafe { &mut *self.value.get() } + } + + fn remove_waker(&self, wait_key: usize, wake_another: bool) { + if wait_key != WAIT_KEY_NONE { + let mut waiters = self.waiters.lock().unwrap(); + match waiters.remove(wait_key) { + Waiter::Waiting(_) => {}, + Waiter::Woken => { + // We were awoken, but then dropped before we could + // wake up to acquire the lock. Wake up another + // waiter. + if wake_another { + if let Some((_i, waiter)) = waiters.iter_mut().next() { + waiter.wake(); + } + } + } + } + if waiters.is_empty() { + self.state.fetch_and(!HAS_WAITERS, Ordering::Relaxed); // released by mutex unlock + } + } + } + + // Unlocks the mutex. Called by MutexGuard and MappedMutexGuard when they are + // dropped. + fn unlock(&self) { + let old_state = self.state.fetch_and(!IS_LOCKED, Ordering::AcqRel); + if (old_state & HAS_WAITERS) != 0 { + let mut waiters = self.waiters.lock().unwrap(); + if let Some((_i, waiter)) = waiters.iter_mut().next() { + waiter.wake(); + } + } + } +} + +// Sentinel for when no slot in the `Slab` has been dedicated to this object. +const WAIT_KEY_NONE: usize = usize::max_value(); + +/// A future which resolves when the target mutex has been successfully acquired. +pub struct MutexLockFuture<'a, T: ?Sized> { + // `None` indicates that the mutex was successfully acquired. + mutex: Option<&'a Mutex<T>>, + wait_key: usize, +} + +impl<T: ?Sized> fmt::Debug for MutexLockFuture<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MutexLockFuture") + .field("was_acquired", &self.mutex.is_none()) + .field("mutex", &self.mutex) + .field("wait_key", &( + if self.wait_key == WAIT_KEY_NONE { + None + } else { + Some(self.wait_key) + } + )) + .finish() + } +} + +impl<T: ?Sized> FusedFuture for MutexLockFuture<'_, T> { + fn is_terminated(&self) -> bool { + self.mutex.is_none() + } +} + +impl<'a, T: ?Sized> Future for MutexLockFuture<'a, T> { + type Output = MutexGuard<'a, T>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let mutex = self.mutex.expect("polled MutexLockFuture after completion"); + + if let Some(lock) = mutex.try_lock() { + mutex.remove_waker(self.wait_key, false); + self.mutex = None; + return Poll::Ready(lock); + } + + { + let mut waiters = mutex.waiters.lock().unwrap(); + if self.wait_key == WAIT_KEY_NONE { + self.wait_key = waiters.insert(Waiter::Waiting(cx.waker().clone())); + if waiters.len() == 1 { + mutex.state.fetch_or(HAS_WAITERS, Ordering::Relaxed); // released by mutex unlock + } + } else { + waiters[self.wait_key].register(cx.waker()); + } + } + + // Ensure that we haven't raced `MutexGuard::drop`'s unlock path by + // attempting to acquire the lock again. + if let Some(lock) = mutex.try_lock() { + mutex.remove_waker(self.wait_key, false); + self.mutex = None; + return Poll::Ready(lock); + } + + Poll::Pending + } +} + +impl<T: ?Sized> Drop for MutexLockFuture<'_, T> { + fn drop(&mut self) { + if let Some(mutex) = self.mutex { + // This future was dropped before it acquired the mutex. + // + // Remove ourselves from the map, waking up another waiter if we + // had been awoken to acquire the lock. + mutex.remove_waker(self.wait_key, true); + } + } +} + +/// An RAII guard returned by the `lock` and `try_lock` methods. +/// When this structure is dropped (falls out of scope), the lock will be +/// unlocked. +pub struct MutexGuard<'a, T: ?Sized> { + mutex: &'a Mutex<T>, +} + +impl<'a, T: ?Sized> MutexGuard<'a, T> { + /// Returns a locked view over a portion of the locked data. + /// + /// # Example + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::lock::{Mutex, MutexGuard}; + /// + /// let data = Mutex::new(Some("value".to_string())); + /// { + /// let locked_str = MutexGuard::map(data.lock().await, |opt| opt.as_mut().unwrap()); + /// assert_eq!(&*locked_str, "value"); + /// } + /// # }); + /// ``` + #[inline] + pub fn map<U: ?Sized, F>(this: Self, f: F) -> MappedMutexGuard<'a, T, U> + where + F: FnOnce(&mut T) -> &mut U, + { + let mutex = this.mutex; + let value = f(unsafe { &mut *this.mutex.value.get() }); + // Don't run the `drop` method for MutexGuard. The ownership of the underlying + // locked state is being moved to the returned MappedMutexGuard. + mem::forget(this); + MappedMutexGuard { mutex, value } + } +} + +impl<T: ?Sized + fmt::Debug> fmt::Debug for MutexGuard<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MutexGuard") + .field("value", &&**self) + .field("mutex", &self.mutex) + .finish() + } +} + +impl<T: ?Sized> Drop for MutexGuard<'_, T> { + fn drop(&mut self) { + self.mutex.unlock() + } +} + +impl<T: ?Sized> Deref for MutexGuard<'_, T> { + type Target = T; + fn deref(&self) -> &T { + unsafe { &*self.mutex.value.get() } + } +} + +impl<T: ?Sized> DerefMut for MutexGuard<'_, T> { + fn deref_mut(&mut self) -> &mut T { + unsafe { &mut *self.mutex.value.get() } + } +} + +/// An RAII guard returned by the `MutexGuard::map` and `MappedMutexGuard::map` methods. +/// When this structure is dropped (falls out of scope), the lock will be unlocked. +pub struct MappedMutexGuard<'a, T: ?Sized, U: ?Sized> { + mutex: &'a Mutex<T>, + value: *mut U, +} + +impl<'a, T: ?Sized, U: ?Sized> MappedMutexGuard<'a, T, U> { + /// Returns a locked view over a portion of the locked data. + /// + /// # Example + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::lock::{MappedMutexGuard, Mutex, MutexGuard}; + /// + /// let data = Mutex::new(Some("value".to_string())); + /// { + /// let locked_str = MutexGuard::map(data.lock().await, |opt| opt.as_mut().unwrap()); + /// let locked_char = MappedMutexGuard::map(locked_str, |s| s.get_mut(0..1).unwrap()); + /// assert_eq!(&*locked_char, "v"); + /// } + /// # }); + /// ``` + #[inline] + pub fn map<V: ?Sized, F>(this: Self, f: F) -> MappedMutexGuard<'a, T, V> + where + F: FnOnce(&mut U) -> &mut V, + { + let mutex = this.mutex; + let value = f(unsafe { &mut *this.value }); + // Don't run the `drop` method for MappedMutexGuard. The ownership of the underlying + // locked state is being moved to the returned MappedMutexGuard. + mem::forget(this); + MappedMutexGuard { mutex, value } + } +} + +impl<T: ?Sized, U: ?Sized + fmt::Debug> fmt::Debug for MappedMutexGuard<'_, T, U> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MappedMutexGuard") + .field("value", &&**self) + .field("mutex", &self.mutex) + .finish() + } +} + +impl<T: ?Sized, U: ?Sized> Drop for MappedMutexGuard<'_, T, U> { + fn drop(&mut self) { + self.mutex.unlock() + } +} + +impl<T: ?Sized, U: ?Sized> Deref for MappedMutexGuard<'_, T, U> { + type Target = U; + fn deref(&self) -> &U { + unsafe { &*self.value } + } +} + +impl<T: ?Sized, U: ?Sized> DerefMut for MappedMutexGuard<'_, T, U> { + fn deref_mut(&mut self) -> &mut U { + unsafe { &mut *self.value } + } +} + +// Mutexes can be moved freely between threads and acquired on any thread so long +// as the inner value can be safely sent between threads. +unsafe impl<T: ?Sized + Send> Send for Mutex<T> {} +unsafe impl<T: ?Sized + Send> Sync for Mutex<T> {} + +// It's safe to switch which thread the acquire is being attempted on so long as +// `T` can be accessed on that thread. +unsafe impl<T: ?Sized + Send> Send for MutexLockFuture<'_, T> {} +// doesn't have any interesting `&self` methods (only Debug) +unsafe impl<T: ?Sized> Sync for MutexLockFuture<'_, T> {} + +// Safe to send since we don't track any thread-specific details-- the inner +// lock is essentially spinlock-equivalent (attempt to flip an atomic bool) +unsafe impl<T: ?Sized + Send> Send for MutexGuard<'_, T> {} +unsafe impl<T: ?Sized + Sync> Sync for MutexGuard<'_, T> {} +unsafe impl<T: ?Sized + Send, U: ?Sized> Send for MappedMutexGuard<'_, T, U> {} +unsafe impl<T: ?Sized + Sync, U: ?Sized> Sync for MappedMutexGuard<'_, T, U> {} + +#[test] +fn test_mutex_guard_debug_not_recurse() { + let mutex = Mutex::new(42); + let guard = mutex.try_lock().unwrap(); + let _ = format!("{:?}", guard); + let guard = MutexGuard::map(guard, |n| n); + let _ = format!("{:?}", guard); +} diff --git a/src/never.rs b/src/never.rs new file mode 100644 index 0000000..767c5af --- /dev/null +++ b/src/never.rs @@ -0,0 +1,17 @@ +//! Definition of the `Never` type, +//! a stand-in for the `!` type until it becomes stable. + +/// A type with no possible values. +/// +/// This is used to indicate values which can never be created, such as the +/// error type of infallible futures. +/// +/// This type is a stable equivalent to the `!` type from `std`. +/// +/// This is currently an alias for [`std::convert::Infallible`], but in +/// the future it may be an alias for [`!`][never]. +/// See ["Future compatibility" section of `std::convert::Infallible`][infallible] for more. +/// +/// [never]: https://doc.rust-lang.org/nightly/std/primitive.never.html +/// [infallible]: https://doc.rust-lang.org/nightly/std/convert/enum.Infallible.html#future-compatibility +pub type Never = core::convert::Infallible; diff --git a/src/sink/buffer.rs b/src/sink/buffer.rs new file mode 100644 index 0000000..d2a3f90 --- /dev/null +++ b/src/sink/buffer.rs @@ -0,0 +1,142 @@ +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; +use core::pin::Pin; +use alloc::collections::VecDeque; + +/// Sink for the [`buffer`](super::SinkExt::buffer) method. +#[derive(Debug)] +#[must_use = "sinks do nothing unless polled"] +pub struct Buffer<Si, Item> { + sink: Si, + buf: VecDeque<Item>, + + // Track capacity separately from the `VecDeque`, which may be rounded up + capacity: usize, +} + +impl<Si: Unpin, Item> Unpin for Buffer<Si, Item> {} + +impl<Si: Sink<Item>, Item> Buffer<Si, Item> { + unsafe_pinned!(sink: Si); + unsafe_unpinned!(buf: VecDeque<Item>); + unsafe_unpinned!(capacity: usize); + + pub(super) fn new(sink: Si, capacity: usize) -> Self { + Buffer { + sink, + buf: VecDeque::with_capacity(capacity), + capacity, + } + } + + /// Get a shared reference to the inner sink. + pub fn get_ref(&self) -> &Si { + &self.sink + } + + /// Get a mutable reference to the inner sink. + pub fn get_mut(&mut self) -> &mut Si { + &mut self.sink + } + + /// Get a pinned mutable reference to the inner sink. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut Si> { + self.sink() + } + + /// Consumes this combinator, returning the underlying sink. + /// + /// 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) -> Si { + self.sink + } + + fn try_empty_buffer( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Si::Error>> { + ready!(self.as_mut().sink().poll_ready(cx))?; + while let Some(item) = self.as_mut().buf().pop_front() { + self.as_mut().sink().start_send(item)?; + if !self.buf.is_empty() { + ready!(self.as_mut().sink().poll_ready(cx))?; + } + } + Poll::Ready(Ok(())) + } +} + +// Forwarding impl of Stream from the underlying sink +impl<S, Item> Stream for Buffer<S, Item> where S: Sink<Item> + Stream { + type Item = S::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<S::Item>> { + self.sink().poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.sink.size_hint() + } +} + +impl<S, Item> FusedStream for Buffer<S, Item> where S: Sink<Item> + FusedStream { + fn is_terminated(&self) -> bool { + self.sink.is_terminated() + } +} + +impl<Si: Sink<Item>, Item> Sink<Item> for Buffer<Si, Item> { + type Error = Si::Error; + + fn poll_ready( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + if self.capacity == 0 { + return self.as_mut().sink().poll_ready(cx); + } + + let _ = self.as_mut().try_empty_buffer(cx)?; + + if self.buf.len() >= self.capacity { + Poll::Pending + } else { + Poll::Ready(Ok(())) + } + } + + fn start_send( + mut self: Pin<&mut Self>, + item: Item, + ) -> Result<(), Self::Error> { + if self.capacity == 0 { + self.as_mut().sink().start_send(item) + } else { + self.as_mut().buf().push_back(item); + Ok(()) + } + } + + #[allow(clippy::debug_assert_with_mut_call)] + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + ready!(self.as_mut().try_empty_buffer(cx))?; + debug_assert!(self.as_mut().buf().is_empty()); + self.as_mut().sink().poll_flush(cx) + } + + #[allow(clippy::debug_assert_with_mut_call)] + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + ready!(self.as_mut().try_empty_buffer(cx))?; + debug_assert!(self.as_mut().buf().is_empty()); + self.as_mut().sink().poll_close(cx) + } +} diff --git a/src/sink/close.rs b/src/sink/close.rs new file mode 100644 index 0000000..1514b41 --- /dev/null +++ b/src/sink/close.rs @@ -0,0 +1,38 @@ +use core::marker::PhantomData; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; + +/// Future for the [`close`](super::SinkExt::close) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Close<'a, Si: ?Sized, Item> { + sink: &'a mut Si, + _phantom: PhantomData<fn(Item)>, +} + +impl<Si: Unpin + ?Sized, Item> Unpin for Close<'_, Si, Item> {} + +/// A future that completes when the sink has finished closing. +/// +/// The sink itself is returned after closeing is complete. +impl<'a, Si: Sink<Item> + Unpin + ?Sized, Item> Close<'a, Si, Item> { + pub(super) fn new(sink: &'a mut Si) -> Self { + Close { + sink, + _phantom: PhantomData, + } + } +} + +impl<Si: Sink<Item> + Unpin + ?Sized, Item> Future for Close<'_, Si, Item> { + type Output = Result<(), Si::Error>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + Pin::new(&mut self.sink).poll_close(cx) + } +} diff --git a/src/sink/drain.rs b/src/sink/drain.rs new file mode 100644 index 0000000..46de83b --- /dev/null +++ b/src/sink/drain.rs @@ -0,0 +1,64 @@ +use crate::never::Never; +use core::marker::PhantomData; +use core::pin::Pin; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; + +/// Sink for the [`drain`] function. +#[derive(Debug)] +#[must_use = "sinks do nothing unless polled"] +pub struct Drain<T> { + marker: PhantomData<T>, +} + +/// Create a sink that will just discard all items given to it. +/// +/// Similar to [`io::Sink`](::std::io::Sink). +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::sink::{self, SinkExt}; +/// +/// let mut drain = sink::drain(); +/// drain.send(5).await?; +/// # Ok::<(), futures::never::Never>(()) }).unwrap(); +/// ``` +pub fn drain<T>() -> Drain<T> { + Drain { marker: PhantomData } +} + +impl<T> Unpin for Drain<T> {} + +impl<T> Sink<T> for Drain<T> { + type Error = Never; + + fn poll_ready( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + Poll::Ready(Ok(())) + } + + fn start_send( + self: Pin<&mut Self>, + _item: T, + ) -> Result<(), Self::Error> { + Ok(()) + } + + fn poll_flush( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + Poll::Ready(Ok(())) + } + + fn poll_close( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + Poll::Ready(Ok(())) + } +} diff --git a/src/sink/err_into.rs b/src/sink/err_into.rs new file mode 100644 index 0000000..af198fe --- /dev/null +++ b/src/sink/err_into.rs @@ -0,0 +1,86 @@ +use crate::sink::{SinkExt, SinkMapErr}; +use core::pin::Pin; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +use futures_sink::{Sink}; +use pin_utils::unsafe_pinned; + +/// Sink for the [`sink_err_into`](super::SinkExt::sink_err_into) method. +#[derive(Debug)] +#[must_use = "sinks do nothing unless polled"] +pub struct SinkErrInto<Si: Sink<Item>, Item, E> { + sink: SinkMapErr<Si, fn(Si::Error) -> E>, +} + +impl<Si, E, Item> SinkErrInto<Si, Item, E> + where Si: Sink<Item>, + Si::Error: Into<E>, +{ + unsafe_pinned!(sink: SinkMapErr<Si, fn(Si::Error) -> E>); + + pub(super) fn new(sink: Si) -> Self { + SinkErrInto { + sink: SinkExt::sink_map_err(sink, Into::into), + } + } + + /// Get a shared reference to the inner sink. + pub fn get_ref(&self) -> &Si { + self.sink.get_ref() + } + + /// Get a mutable reference to the inner sink. + pub fn get_mut(&mut self) -> &mut Si { + self.sink.get_mut() + } + + /// Get a pinned mutable reference to the inner sink. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut Si> { + self.sink().get_pin_mut() + } + + /// Consumes this combinator, returning the underlying sink. + /// + /// 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) -> Si { + self.sink.into_inner() + } +} + +impl<Si, Item, E> Sink<Item> for SinkErrInto<Si, Item, E> + where Si: Sink<Item>, + Si::Error: Into<E>, +{ + type Error = E; + + delegate_sink!(sink, Item); +} + +// Forwarding impl of Stream from the underlying sink +impl<S, Item, E> Stream for SinkErrInto<S, Item, E> + where S: Sink<Item> + Stream, + S::Error: Into<E> +{ + type Item = S::Item; + + fn poll_next( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<S::Item>> { + self.sink().poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.sink.size_hint() + } +} + +impl<S, Item, E> FusedStream for SinkErrInto<S, Item, E> + where S: Sink<Item> + FusedStream, + S::Error: Into<E> +{ + fn is_terminated(&self) -> bool { + self.sink.is_terminated() + } +} diff --git a/src/sink/fanout.rs b/src/sink/fanout.rs new file mode 100644 index 0000000..24e4de9 --- /dev/null +++ b/src/sink/fanout.rs @@ -0,0 +1,106 @@ +use core::fmt::{Debug, Formatter, Result as FmtResult}; +use core::pin::Pin; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; +use pin_utils::unsafe_pinned; + +/// Sink that clones incoming items and forwards them to two sinks at the same time. +/// +/// Backpressure from any downstream sink propagates up, which means that this sink +/// can only process items as fast as its _slowest_ downstream sink. +#[must_use = "sinks do nothing unless polled"] +pub struct Fanout<Si1, Si2> { + sink1: Si1, + sink2: Si2 +} + +impl<Si1, Si2> Fanout<Si1, Si2> { + unsafe_pinned!(sink1: Si1); + unsafe_pinned!(sink2: Si2); + + pub(super) fn new(sink1: Si1, sink2: Si2) -> Fanout<Si1, Si2> { + Fanout { sink1, sink2 } + } + + /// Get a shared reference to the inner sinks. + pub fn get_ref(&self) -> (&Si1, &Si2) { + (&self.sink1, &self.sink2) + } + + /// Get a mutable reference to the inner sinks. + pub fn get_mut(&mut self) -> (&mut Si1, &mut Si2) { + (&mut self.sink1, &mut self.sink2) + } + + /// Get a pinned mutable reference to the inner sinks. + pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut Si1>, Pin<&mut Si2>) { + unsafe { + let Self { sink1, sink2 } = self.get_unchecked_mut(); + (Pin::new_unchecked(sink1), Pin::new_unchecked(sink2)) + } + } + + /// Consumes this combinator, returning the underlying sinks. + /// + /// 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) -> (Si1, Si2) { + (self.sink1, self.sink2) + } +} + +impl<Si1: Debug, Si2: Debug> Debug for Fanout<Si1, Si2> { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + f.debug_struct("Fanout") + .field("sink1", &self.sink1) + .field("sink2", &self.sink2) + .finish() + } +} + +impl<Si1, Si2, Item> Sink<Item> for Fanout<Si1, Si2> + where Si1: Sink<Item>, + Item: Clone, + Si2: Sink<Item, Error=Si1::Error> +{ + type Error = Si1::Error; + + fn poll_ready( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + let sink1_ready = self.as_mut().sink1().poll_ready(cx)?.is_ready(); + let sink2_ready = self.as_mut().sink2().poll_ready(cx)?.is_ready(); + let ready = sink1_ready && sink2_ready; + if ready { Poll::Ready(Ok(())) } else { Poll::Pending } + } + + fn start_send( + mut self: Pin<&mut Self>, + item: Item, + ) -> Result<(), Self::Error> { + self.as_mut().sink1().start_send(item.clone())?; + self.as_mut().sink2().start_send(item)?; + Ok(()) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + let sink1_ready = self.as_mut().sink1().poll_flush(cx)?.is_ready(); + let sink2_ready = self.as_mut().sink2().poll_flush(cx)?.is_ready(); + let ready = sink1_ready && sink2_ready; + if ready { Poll::Ready(Ok(())) } else { Poll::Pending } + } + + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + let sink1_ready = self.as_mut().sink1().poll_close(cx)?.is_ready(); + let sink2_ready = self.as_mut().sink2().poll_close(cx)?.is_ready(); + let ready = sink1_ready && sink2_ready; + if ready { Poll::Ready(Ok(())) } else { Poll::Pending } + } +} diff --git a/src/sink/flush.rs b/src/sink/flush.rs new file mode 100644 index 0000000..86fecc4 --- /dev/null +++ b/src/sink/flush.rs @@ -0,0 +1,42 @@ +use core::marker::PhantomData; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; + +/// Future for the [`flush`](super::SinkExt::flush) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Flush<'a, Si: ?Sized, Item> { + sink: &'a mut Si, + _phantom: PhantomData<fn(Item)>, +} + +// Pin is never projected to a field. +impl<Si: Unpin + ?Sized, Item> Unpin for Flush<'_, Si, Item> {} + +/// A future that completes when the sink has finished processing all +/// pending requests. +/// +/// The sink itself is returned after flushing is complete; this adapter is +/// intended to be used when you want to stop sending to the sink until +/// all current requests are processed. +impl<'a, Si: Sink<Item> + Unpin + ?Sized, Item> Flush<'a, Si, Item> { + pub(super) fn new(sink: &'a mut Si) -> Self { + Flush { + sink, + _phantom: PhantomData, + } + } +} + +impl<Si: Sink<Item> + Unpin + ?Sized, Item> Future for Flush<'_, Si, Item> { + type Output = Result<(), Si::Error>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + Pin::new(&mut self.sink).poll_flush(cx) + } +} diff --git a/src/sink/map_err.rs b/src/sink/map_err.rs new file mode 100644 index 0000000..e999bea --- /dev/null +++ b/src/sink/map_err.rs @@ -0,0 +1,108 @@ +use core::pin::Pin; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +use futures_sink::{Sink}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Sink for the [`sink_map_err`](super::SinkExt::sink_map_err) method. +#[derive(Debug, Clone)] +#[must_use = "sinks do nothing unless polled"] +pub struct SinkMapErr<Si, F> { + sink: Si, + f: Option<F>, +} + +impl<Si: Unpin, F> Unpin for SinkMapErr<Si, F> {} + +impl<Si, F> SinkMapErr<Si, F> { + unsafe_pinned!(sink: Si); + unsafe_unpinned!(f: Option<F>); + + pub(super) fn new(sink: Si, f: F) -> SinkMapErr<Si, F> { + SinkMapErr { sink, f: Some(f) } + } + + /// Get a shared reference to the inner sink. + pub fn get_ref(&self) -> &Si { + &self.sink + } + + /// Get a mutable reference to the inner sink. + pub fn get_mut(&mut self) -> &mut Si { + &mut self.sink + } + + /// Get a pinned mutable reference to the inner sink. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut Si> { + self.sink() + } + + /// Consumes this combinator, returning the underlying sink. + /// + /// 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) -> Si { + self.sink + } + + fn take_f(self: Pin<&mut Self>) -> F { + self.f().take().expect("polled MapErr after completion") + } +} + +impl<Si, F, E, Item> Sink<Item> for SinkMapErr<Si, F> + where Si: Sink<Item>, + F: FnOnce(Si::Error) -> E, +{ + type Error = E; + + fn poll_ready( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + self.as_mut().sink().poll_ready(cx).map_err(|e| self.as_mut().take_f()(e)) + } + + fn start_send( + mut self: Pin<&mut Self>, + item: Item, + ) -> Result<(), Self::Error> { + self.as_mut().sink().start_send(item).map_err(|e| self.as_mut().take_f()(e)) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + self.as_mut().sink().poll_flush(cx).map_err(|e| self.as_mut().take_f()(e)) + } + + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + self.as_mut().sink().poll_close(cx).map_err(|e| self.as_mut().take_f()(e)) + } +} + +// Forwarding impl of Stream from the underlying sink +impl<S: Stream, F> Stream for SinkMapErr<S, F> { + type Item = S::Item; + + fn poll_next( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<S::Item>> { + self.sink().poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.sink.size_hint() + } +} + +impl<S: FusedStream, F> FusedStream for SinkMapErr<S, F> { + fn is_terminated(&self) -> bool { + self.sink.is_terminated() + } +} diff --git a/src/sink/mod.rs b/src/sink/mod.rs new file mode 100644 index 0000000..bc59ce5 --- /dev/null +++ b/src/sink/mod.rs @@ -0,0 +1,264 @@ +//! Sinks +//! +//! This module contains a number of functions for working with `Sink`s, +//! including the `SinkExt` trait which adds methods to `Sink` types. +//! +//! This module is only available when the `sink` feature of this +//! library is activated, and it is activated by default. + +use futures_core::future::Future; +use futures_core::stream::{Stream, TryStream}; +use crate::future::Either; + +#[cfg(feature = "compat")] +use crate::compat::CompatSink; + +pub use futures_sink::Sink; + +mod close; +pub use self::close::Close; + +mod drain; +pub use self::drain::{drain, Drain}; + +mod fanout; +pub use self::fanout::Fanout; + +mod flush; +pub use self::flush::Flush; + +mod err_into; +pub use self::err_into::SinkErrInto; + +mod map_err; +pub use self::map_err::SinkMapErr; + +mod send; +pub use self::send::Send; + +mod send_all; +pub use self::send_all::SendAll; + +mod with; +pub use self::with::With; + +mod with_flat_map; +pub use self::with_flat_map::WithFlatMap; + +#[cfg(feature = "alloc")] +mod buffer; +#[cfg(feature = "alloc")] +pub use self::buffer::Buffer; + +impl<T: ?Sized, Item> SinkExt<Item> for T where T: Sink<Item> {} + +/// An extension trait for `Sink`s that provides a variety of convenient +/// combinator functions. +pub trait SinkExt<Item>: Sink<Item> { + /// Composes a function *in front of* the sink. + /// + /// This adapter produces a new sink that passes each value through the + /// given function `f` before sending it to `self`. + /// + /// To process each value, `f` produces a *future*, which is then polled to + /// completion before passing its result down to the underlying sink. If the + /// future produces an error, that error is returned by the new sink. + /// + /// Note that this function consumes the given sink, returning a wrapped + /// version, much like `Iterator::map`. + fn with<U, Fut, F, E>(self, f: F) -> With<Self, Item, U, Fut, F> + where F: FnMut(U) -> Fut, + Fut: Future<Output = Result<Item, E>>, + E: From<Self::Error>, + Self: Sized + { + With::new(self, f) + } + + /// Composes a function *in front of* the sink. + /// + /// This adapter produces a new sink that passes each value through the + /// given function `f` before sending it to `self`. + /// + /// To process each value, `f` produces a *stream*, of which each value + /// is passed to the underlying sink. A new value will not be accepted until + /// the stream has been drained + /// + /// Note that this function consumes the given sink, returning a wrapped + /// version, much like `Iterator::flat_map`. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::sink::SinkExt; + /// use futures::stream::{self, StreamExt}; + /// + /// let (tx, rx) = mpsc::channel(5); + /// + /// let mut tx = tx.with_flat_map(|x| { + /// stream::iter(vec![Ok(42); x]) + /// }); + /// + /// tx.send(5).await.unwrap(); + /// drop(tx); + /// let received: Vec<i32> = rx.collect().await; + /// assert_eq!(received, vec![42, 42, 42, 42, 42]); + /// # }); + /// ``` + fn with_flat_map<U, St, F>(self, f: F) -> WithFlatMap<Self, Item, U, St, F> + where F: FnMut(U) -> St, + St: Stream<Item = Result<Item, Self::Error>>, + Self: Sized + { + WithFlatMap::new(self, f) + } + + /* + fn with_map<U, F>(self, f: F) -> WithMap<Self, U, F> + where F: FnMut(U) -> Self::SinkItem, + Self: Sized; + + fn with_filter<F>(self, f: F) -> WithFilter<Self, F> + where F: FnMut(Self::SinkItem) -> bool, + Self: Sized; + + fn with_filter_map<U, F>(self, f: F) -> WithFilterMap<Self, U, F> + where F: FnMut(U) -> Option<Self::SinkItem>, + Self: Sized; + */ + + /// Transforms the error returned by the sink. + fn sink_map_err<E, F>(self, f: F) -> SinkMapErr<Self, F> + where F: FnOnce(Self::Error) -> E, + Self: Sized, + { + SinkMapErr::new(self, f) + } + + /// Map this sink's error to a different error type using the `Into` trait. + /// + /// If wanting to map errors of a `Sink + Stream`, use `.sink_err_into().err_into()`. + fn sink_err_into<E>(self) -> err_into::SinkErrInto<Self, Item, E> + where Self: Sized, + Self::Error: Into<E>, + { + SinkErrInto::new(self) + } + + + /// Adds a fixed-size buffer to the current sink. + /// + /// The resulting sink will buffer up to `capacity` items when the + /// underlying sink is unwilling to accept additional items. Calling `flush` + /// on the buffered sink will attempt to both empty the buffer and complete + /// processing on the underlying sink. + /// + /// Note that this function consumes the given sink, returning a wrapped + /// version, much like `Iterator::map`. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg(feature = "alloc")] + fn buffer(self, capacity: usize) -> Buffer<Self, Item> + where Self: Sized, + { + Buffer::new(self, capacity) + } + + /// Close the sink. + fn close(&mut self) -> Close<'_, Self, Item> + where Self: Unpin, + { + Close::new(self) + } + + /// Fanout items to multiple sinks. + /// + /// This adapter clones each incoming item and forwards it to both this as well as + /// the other sink at the same time. + fn fanout<Si>(self, other: Si) -> Fanout<Self, Si> + where Self: Sized, + Item: Clone, + Si: Sink<Item, Error=Self::Error> + { + Fanout::new(self, other) + } + + /// Flush the sink, processing all pending items. + /// + /// This adapter is intended to be used when you want to stop sending to the sink + /// until all current requests are processed. + fn flush(&mut self) -> Flush<'_, Self, Item> + where Self: Unpin, + { + Flush::new(self) + } + + /// A future that completes after the given item has been fully processed + /// into the sink, including flushing. + /// + /// Note that, **because of the flushing requirement, it is usually better + /// to batch together items to send via `send_all`, rather than flushing + /// between each item.** + fn send(&mut self, item: Item) -> Send<'_, Self, Item> + where Self: Unpin, + { + Send::new(self, item) + } + + /// A future that completes after the given stream has been fully processed + /// into the sink, including flushing. + /// + /// This future will drive the stream to keep producing items until it is + /// exhausted, sending each item to the sink. It will complete once both the + /// stream is exhausted, the sink has received all items, and the sink has + /// been flushed. Note that the sink is **not** closed. + /// + /// Doing `sink.send_all(stream)` is roughly equivalent to + /// `stream.forward(sink)`. The returned future will exhaust all items from + /// `stream` and send them to `self`. + fn send_all<'a, St>( + &'a mut self, + stream: &'a mut St + ) -> SendAll<'a, Self, St> + where St: TryStream<Ok = Item, Error = Self::Error> + Stream + Unpin + ?Sized, + Self: Unpin, + { + SendAll::new(self, stream) + } + + /// Wrap this sink in an `Either` sink, making it the left-hand variant + /// of that `Either`. + /// + /// This can be used in combination with the `right_sink` method to write `if` + /// statements that evaluate to different streams in different branches. + fn left_sink<Si2>(self) -> Either<Self, Si2> + where Si2: Sink<Item, Error = Self::Error>, + Self: Sized + { + Either::Left(self) + } + + /// Wrap this stream in an `Either` stream, making it the right-hand variant + /// of that `Either`. + /// + /// This can be used in combination with the `left_sink` method to write `if` + /// statements that evaluate to different streams in different branches. + fn right_sink<Si1>(self) -> Either<Si1, Self> + where Si1: Sink<Item, Error = Self::Error>, + Self: Sized + { + Either::Right(self) + } + + /// Wraps a [`Sink`] into a sink compatible with libraries using + /// futures 0.1 `Sink`. Requires the `compat` feature to be enabled. + #[cfg(feature = "compat")] + fn compat(self) -> CompatSink<Self, Item> + where Self: Sized + Unpin, + { + CompatSink::new(self) + } +} diff --git a/src/sink/send.rs b/src/sink/send.rs new file mode 100644 index 0000000..dc7f0be --- /dev/null +++ b/src/sink/send.rs @@ -0,0 +1,51 @@ +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; + +/// Future for the [`send`](super::SinkExt::send) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Send<'a, Si: ?Sized, Item> { + sink: &'a mut Si, + item: Option<Item>, +} + +// Pinning is never projected to children +impl<Si: Unpin + ?Sized, Item> Unpin for Send<'_, Si, Item> {} + +impl<'a, Si: Sink<Item> + Unpin + ?Sized, Item> Send<'a, Si, Item> { + pub(super) fn new(sink: &'a mut Si, item: Item) -> Self { + Send { + sink, + item: Some(item), + } + } +} + +impl<Si: Sink<Item> + Unpin + ?Sized, Item> Future for Send<'_, Si, Item> { + type Output = Result<(), Si::Error>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + let this = &mut *self; + if let Some(item) = this.item.take() { + let mut sink = Pin::new(&mut this.sink); + match sink.as_mut().poll_ready(cx)? { + Poll::Ready(()) => sink.as_mut().start_send(item)?, + Poll::Pending => { + this.item = Some(item); + return Poll::Pending; + } + } + } + + // we're done sending the item, but want to block on flushing the + // sink + ready!(Pin::new(&mut this.sink).poll_flush(cx))?; + + Poll::Ready(Ok(())) + } +} diff --git a/src/sink/send_all.rs b/src/sink/send_all.rs new file mode 100644 index 0000000..255df4d --- /dev/null +++ b/src/sink/send_all.rs @@ -0,0 +1,112 @@ +use crate::stream::{StreamExt, TryStreamExt, Fuse}; +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::stream::{TryStream, Stream}; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; + +/// Future for the [`send_all`](super::SinkExt::send_all) method. +#[allow(explicit_outlives_requirements)] // https://github.com/rust-lang/rust/issues/60993 +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct SendAll<'a, Si, St> +where + Si: ?Sized, + St: ?Sized + TryStream, +{ + sink: &'a mut Si, + stream: Fuse<&'a mut St>, + buffered: Option<St::Ok>, +} + +impl<Si, St> fmt::Debug for SendAll<'_, Si, St> +where + Si: fmt::Debug + ?Sized, + St: fmt::Debug + ?Sized + TryStream, + St::Ok: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SendAll") + .field("sink", &self.sink) + .field("stream", &self.stream) + .field("buffered", &self.buffered) + .finish() + } +} + +// Pinning is never projected to any fields +impl<Si, St> Unpin for SendAll<'_, Si, St> +where + Si: Unpin + ?Sized, + St: TryStream + Unpin + ?Sized, +{} + +impl<'a, Si, St, Ok, Error> SendAll<'a, Si, St> +where + Si: Sink<Ok, Error = Error> + Unpin + ?Sized, + St: TryStream<Ok = Ok, Error = Error> + Stream + Unpin + ?Sized, +{ + pub(super) fn new( + sink: &'a mut Si, + stream: &'a mut St, + ) -> SendAll<'a, Si, St> { + SendAll { + sink, + stream: stream.fuse(), + buffered: None, + } + } + + fn try_start_send( + &mut self, + cx: &mut Context<'_>, + item: St::Ok, + ) -> Poll<Result<(), Si::Error>> { + debug_assert!(self.buffered.is_none()); + match Pin::new(&mut self.sink).poll_ready(cx)? { + Poll::Ready(()) => { + Poll::Ready(Pin::new(&mut self.sink).start_send(item)) + } + Poll::Pending => { + self.buffered = Some(item); + Poll::Pending + } + } + } +} + +impl<Si, St, Ok, Error> Future for SendAll<'_, Si, St> +where + Si: Sink<Ok, Error = Error> + Unpin + ?Sized, + St: Stream<Item = Result<Ok, Error>> + Unpin + ?Sized, +{ + type Output = Result<(), Error>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + let this = &mut *self; + // If we've got an item buffered already, we need to write it to the + // sink before we can do anything else + if let Some(item) = this.buffered.take() { + ready!(this.try_start_send(cx, item))? + } + + loop { + match this.stream.try_poll_next_unpin(cx)? { + Poll::Ready(Some(item)) => { + ready!(this.try_start_send(cx, item))? + } + Poll::Ready(None) => { + ready!(Pin::new(&mut this.sink).poll_flush(cx))?; + return Poll::Ready(Ok(())) + } + Poll::Pending => { + ready!(Pin::new(&mut this.sink).poll_flush(cx))?; + return Poll::Pending + } + } + } + } +} diff --git a/src/sink/with.rs b/src/sink/with.rs new file mode 100644 index 0000000..5acac01 --- /dev/null +++ b/src/sink/with.rs @@ -0,0 +1,168 @@ +use core::fmt; +use core::marker::PhantomData; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Sink for the [`with`](super::SinkExt::with) method. +#[must_use = "sinks do nothing unless polled"] +pub struct With<Si, Item, U, Fut, F> { + sink: Si, + f: F, + state: Option<Fut>, + _phantom: PhantomData<fn(U) -> Item>, +} + +impl<Si, Item, U, Fut, F> Unpin for With<Si, Item, U, Fut, F> +where + Si: Unpin, + Fut: Unpin, +{} + +impl<Si, Item, U, Fut, F> fmt::Debug for With<Si, Item, U, Fut, F> +where + Si: fmt::Debug, + Fut: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("With") + .field("sink", &self.sink) + .field("state", &self.state) + .finish() + } +} + +impl<Si, Item, U, Fut, F> With<Si, Item, U, Fut, F> +where Si: Sink<Item>, + F: FnMut(U) -> Fut, + Fut: Future, +{ + unsafe_pinned!(sink: Si); + unsafe_unpinned!(f: F); + unsafe_pinned!(state: Option<Fut>); + + pub(super) fn new<E>(sink: Si, f: F) -> Self + where + Fut: Future<Output = Result<Item, E>>, + E: From<Si::Error>, + { + With { + state: None, + sink, + f, + _phantom: PhantomData, + } + } +} + +// Forwarding impl of Stream from the underlying sink +impl<S, Item, U, Fut, F> Stream for With<S, Item, U, Fut, F> + where S: Stream + Sink<Item>, + F: FnMut(U) -> Fut, + Fut: Future +{ + type Item = S::Item; + + fn poll_next( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<S::Item>> { + self.sink().poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.sink.size_hint() + } +} + +impl<Si, Item, U, Fut, F, E> With<Si, Item, U, Fut, F> + where Si: Sink<Item>, + F: FnMut(U) -> Fut, + Fut: Future<Output = Result<Item, E>>, + E: From<Si::Error>, +{ + /// Get a shared reference to the inner sink. + pub fn get_ref(&self) -> &Si { + &self.sink + } + + /// Get a mutable reference to the inner sink. + pub fn get_mut(&mut self) -> &mut Si { + &mut self.sink + } + + /// Get a pinned mutable reference to the inner sink. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut Si> { + self.sink() + } + + /// Consumes this combinator, returning the underlying sink. + /// + /// 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) -> Si { + self.sink + } + + /// Completes the processing of previous item if any. + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), E>> { + let item = match self.as_mut().state().as_pin_mut() { + None => return Poll::Ready(Ok(())), + Some(fut) => ready!(fut.poll(cx))?, + }; + self.as_mut().state().set(None); + self.as_mut().sink().start_send(item)?; + Poll::Ready(Ok(())) + } +} + +impl<Si, Item, U, Fut, F, E> Sink<U> for With<Si, Item, U, Fut, F> + where Si: Sink<Item>, + F: FnMut(U) -> Fut, + Fut: Future<Output = Result<Item, E>>, + E: From<Si::Error>, +{ + type Error = E; + + fn poll_ready( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + ready!(self.as_mut().poll(cx))?; + ready!(self.as_mut().sink().poll_ready(cx)?); + Poll::Ready(Ok(())) + } + + fn start_send( + mut self: Pin<&mut Self>, + item: U, + ) -> Result<(), Self::Error> { + let future = (self.as_mut().f())(item); + self.as_mut().state().set(Some(future)); + Ok(()) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + ready!(self.as_mut().poll(cx))?; + ready!(self.as_mut().sink().poll_flush(cx))?; + Poll::Ready(Ok(())) + } + + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + ready!(self.as_mut().poll(cx))?; + ready!(self.as_mut().sink().poll_close(cx))?; + Poll::Ready(Ok(())) + } +} diff --git a/src/sink/with_flat_map.rs b/src/sink/with_flat_map.rs new file mode 100644 index 0000000..05ab77e --- /dev/null +++ b/src/sink/with_flat_map.rs @@ -0,0 +1,185 @@ +use core::fmt; +use core::marker::PhantomData; +use core::pin::Pin; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Sink for the [`with_flat_map`](super::SinkExt::with_flat_map) method. +#[must_use = "sinks do nothing unless polled"] +pub struct WithFlatMap<Si, Item, U, St, F> { + sink: Si, + f: F, + stream: Option<St>, + buffer: Option<Item>, + _marker: PhantomData<fn(U)>, +} + +impl<Si, Item, U, St, F> Unpin for WithFlatMap<Si, Item, U, St, F> +where + Si: Unpin, + St: Unpin, +{} + +impl<Si, Item, U, St, F> fmt::Debug for WithFlatMap<Si, Item, U, St, F> +where + Si: fmt::Debug, + St: fmt::Debug, + Item: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("WithFlatMap") + .field("sink", &self.sink) + .field("stream", &self.stream) + .field("buffer", &self.buffer) + .finish() + } +} + +impl<Si, Item, U, St, F> WithFlatMap<Si, Item, U, St, F> +where + Si: Sink<Item>, + F: FnMut(U) -> St, + St: Stream<Item = Result<Item, Si::Error>>, +{ + unsafe_pinned!(sink: Si); + unsafe_unpinned!(f: F); + unsafe_pinned!(stream: Option<St>); + + pub(super) fn new(sink: Si, f: F) -> Self { + WithFlatMap { + sink, + f, + stream: None, + buffer: None, + _marker: PhantomData, + } + } + + /// Get a shared reference to the inner sink. + pub fn get_ref(&self) -> &Si { + &self.sink + } + + /// Get a mutable reference to the inner sink. + pub fn get_mut(&mut self) -> &mut Si { + &mut self.sink + } + + /// Get a pinned mutable reference to the inner sink. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut Si> { + self.sink() + } + + /// Consumes this combinator, returning the underlying sink. + /// + /// 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) -> Si { + self.sink + } + + fn try_empty_stream( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Si::Error>> { + let WithFlatMap { sink, stream, buffer, .. } = + unsafe { self.get_unchecked_mut() }; + let mut sink = unsafe { Pin::new_unchecked(sink) }; + let mut stream = unsafe { Pin::new_unchecked(stream) }; + + if buffer.is_some() { + ready!(sink.as_mut().poll_ready(cx))?; + let item = buffer.take().unwrap(); + sink.as_mut().start_send(item)?; + } + if let Some(mut some_stream) = stream.as_mut().as_pin_mut() { + while let Some(item) = ready!(some_stream.as_mut().poll_next(cx)?) { + match sink.as_mut().poll_ready(cx)? { + Poll::Ready(()) => sink.as_mut().start_send(item)?, + Poll::Pending => { + *buffer = Some(item); + return Poll::Pending; + } + }; + } + } + stream.set(None); + Poll::Ready(Ok(())) + } +} + +// Forwarding impl of Stream from the underlying sink +impl<S, Item, U, St, F> Stream for WithFlatMap<S, Item, U, St, F> +where + S: Stream + Sink<Item>, + F: FnMut(U) -> St, + St: Stream<Item = Result<Item, S::Error>>, +{ + type Item = S::Item; + + fn poll_next( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<S::Item>> { + self.sink().poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.sink.size_hint() + } +} + +impl<S, Item, U, St, F> FusedStream for WithFlatMap<S, Item, U, St, F> +where + S: FusedStream + Sink<Item>, + F: FnMut(U) -> St, + St: Stream<Item = Result<Item, S::Error>>, +{ + fn is_terminated(&self) -> bool { + self.sink.is_terminated() + } +} + +impl<Si, Item, U, St, F> Sink<U> for WithFlatMap<Si, Item, U, St, F> +where + Si: Sink<Item>, + F: FnMut(U) -> St, + St: Stream<Item = Result<Item, Si::Error>>, +{ + type Error = Si::Error; + + fn poll_ready( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + self.try_empty_stream(cx) + } + + fn start_send( + mut self: Pin<&mut Self>, + item: U, + ) -> Result<(), Self::Error> { + assert!(self.stream.is_none()); + let stream = (self.as_mut().f())(item); + self.stream().set(Some(stream)); + Ok(()) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + ready!(self.as_mut().try_empty_stream(cx)?); + self.as_mut().sink().poll_flush(cx) + } + + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<(), Self::Error>> { + ready!(self.as_mut().try_empty_stream(cx)?); + self.as_mut().sink().poll_close(cx) + } +} diff --git a/src/stream/empty.rs b/src/stream/empty.rs new file mode 100644 index 0000000..903af68 --- /dev/null +++ b/src/stream/empty.rs @@ -0,0 +1,40 @@ +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> { + 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)) + } +} diff --git a/src/stream/futures_ordered.rs b/src/stream/futures_ordered.rs new file mode 100644 index 0000000..a30cbaa --- /dev/null +++ b/src/stream/futures_ordered.rs @@ -0,0 +1,205 @@ +use crate::stream::{FuturesUnordered, StreamExt}; +use futures_core::future::Future; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; +use core::cmp::Ordering; +use core::fmt::{self, Debug}; +use core::iter::FromIterator; +use core::pin::Pin; +use alloc::collections::binary_heap::{BinaryHeap, PeekMut}; + +#[must_use = "futures do nothing unless you `.await` or poll them"] +#[derive(Debug)] +struct OrderWrapper<T> { + data: T, // A future or a future's output + index: usize, +} + +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> OrderWrapper<T> { + unsafe_pinned!(data: T); +} + +impl<T> Future for OrderWrapper<T> + where T: Future +{ + type Output = OrderWrapper<T::Output>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + self.as_mut().data().as_mut().poll(cx) + .map(|output| OrderWrapper { data: output, index: self.index }) + } +} + +/// An unbounded queue of futures. +/// +/// This "combinator" is similar to `FuturesUnordered`, but it imposes an order +/// on top of the set of futures. While futures in the set will race to +/// completion in parallel, results will only be returned in the order their +/// originating futures were added to the queue. +/// +/// Futures are pushed into this queue and their realized values are yielded in +/// order. This structure is optimized to manage a large number of futures. +/// Futures managed by `FuturesOrdered` will only be polled when they generate +/// notifications. This reduces the required amount of work needed to coordinate +/// large numbers of futures. +/// +/// When a `FuturesOrdered` is first created, it does not contain any futures. +/// Calling `poll` in this state will result in `Poll::Ready(None))` to be +/// returned. Futures are submitted to the queue using `push`; however, the +/// future will **not** be polled at this point. `FuturesOrdered` will only +/// poll managed futures when `FuturesOrdered::poll` is called. As such, it +/// is important to call `poll` after pushing new futures. +/// +/// If `FuturesOrdered::poll` returns `Poll::Ready(None)` this means that +/// the queue is currently not managing any futures. A future may be submitted +/// to the queue at a later time. At that point, a call to +/// `FuturesOrdered::poll` will either return the future's resolved value +/// **or** `Poll::Pending` if the future has not yet completed. When +/// multiple futures are submitted to the queue, `FuturesOrdered::poll` will +/// return `Poll::Pending` until the first future completes, even if +/// some of the later futures have already completed. +/// +/// Note that you can create a ready-made `FuturesOrdered` via the +/// [`collect`](Iterator::collect) method, or you can start with an empty queue +/// with the `FuturesOrdered::new` constructor. +/// +/// This type is only available when the `std` or `alloc` feature of this +/// library is activated, and it is activated by default. +#[must_use = "streams do nothing unless polled"] +pub struct FuturesOrdered<T: Future> { + in_progress_queue: FuturesUnordered<OrderWrapper<T>>, + queued_outputs: BinaryHeap<OrderWrapper<T::Output>>, + next_incoming_index: usize, + next_outgoing_index: usize, +} + +impl<T: Future> Unpin for FuturesOrdered<T> {} + +impl<Fut: Future> FuturesOrdered<Fut> { + /// Constructs a new, empty `FuturesOrdered` + /// + /// The returned `FuturesOrdered` does not contain any futures and, in this + /// state, `FuturesOrdered::poll_next` will return `Poll::Ready(None)`. + pub fn new() -> FuturesOrdered<Fut> { + FuturesOrdered { + in_progress_queue: FuturesUnordered::new(), + queued_outputs: BinaryHeap::new(), + next_incoming_index: 0, + next_outgoing_index: 0, + } + } + + /// Returns the number of futures contained in the queue. + /// + /// This represents the total number of in-flight futures, both + /// those currently processing and those that have completed but + /// which are waiting for earlier futures to complete. + pub fn len(&self) -> usize { + self.in_progress_queue.len() + self.queued_outputs.len() + } + + /// Returns `true` if the queue contains no futures + pub fn is_empty(&self) -> bool { + self.in_progress_queue.is_empty() && self.queued_outputs.is_empty() + } + + /// Push a future into the queue. + /// + /// This function submits the given future to the internal set for managing. + /// This function will not call `poll` on the submitted future. The caller + /// must ensure that `FuturesOrdered::poll` is called in order to receive + /// task notifications. + pub fn push(&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); + } +} + +impl<Fut: Future> Default for FuturesOrdered<Fut> { + fn default() -> FuturesOrdered<Fut> { + FuturesOrdered::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 = FuturesOrdered::new(); + iter.into_iter().fold(acc, |mut acc, item| { acc.push(item); acc }) + } +} diff --git a/src/stream/futures_unordered/abort.rs b/src/stream/futures_unordered/abort.rs new file mode 100644 index 0000000..1a42d24 --- /dev/null +++ b/src/stream/futures_unordered/abort.rs @@ -0,0 +1,12 @@ +pub(super) fn abort(s: &str) -> ! { + struct DoublePanic; + + impl Drop for DoublePanic { + fn drop(&mut self) { + panic!("panicking twice to abort the program"); + } + } + + let _bomb = DoublePanic; + panic!("{}", s); +} diff --git a/src/stream/futures_unordered/iter.rs b/src/stream/futures_unordered/iter.rs new file mode 100644 index 0000000..ef7b15a --- /dev/null +++ b/src/stream/futures_unordered/iter.rs @@ -0,0 +1,117 @@ +use super::FuturesUnordered; +use super::task::Task; +use core::marker::PhantomData; +use core::pin::Pin; +use core::sync::atomic::Ordering::Relaxed; + +#[derive(Debug)] +/// Mutable iterator over all futures in the unordered set. +pub struct IterPinMut<'a, Fut> { + pub(super) task: *const Task<Fut>, + pub(super) len: usize, + pub(super) _marker: PhantomData<&'a mut FuturesUnordered<Fut>> +} + +#[derive(Debug)] +/// Mutable iterator over all futures in the unordered set. +pub struct IterMut<'a, Fut: Unpin> (pub(super) IterPinMut<'a, Fut>); + +#[derive(Debug)] +/// Immutable iterator over all futures in the unordered set. +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>> +} + +#[derive(Debug)] +/// Immutable iterator over all the futures in the unordered set. +pub struct Iter<'a, Fut: Unpin> (pub(super) IterPinRef<'a, Fut>); + +impl<'a, Fut> Iterator for IterPinMut<'a, Fut> { + type Item = Pin<&'a mut Fut>; + + fn next(&mut self) -> Option<Pin<&'a mut Fut>> { + 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<&'a mut Fut> { + 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<Pin<&'a Fut>> { + 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<&'a Fut> { + 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> {} diff --git a/src/stream/futures_unordered/mod.rs b/src/stream/futures_unordered/mod.rs new file mode 100644 index 0000000..5d93a64 --- /dev/null +++ b/src/stream/futures_unordered/mod.rs @@ -0,0 +1,624 @@ +//! 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 futures_core::future::Future; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use futures_task::{FutureObj, LocalFutureObj, Spawn, LocalSpawn, SpawnError}; +use crate::task::AtomicWaker; +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::{AtomicPtr, AtomicBool}; +use alloc::sync::{Arc, Weak}; + +mod abort; + +mod iter; +pub use self::iter::{Iter, IterMut, IterPinMut, IterPinRef}; + +mod task; +use self::task::Task; + +mod ready_to_run_queue; +use self::ready_to_run_queue::{ReadyToRunQueue, Dequeue}; + +/// Constant used for a `FuturesUnordered` to determine how many times it is +/// allowed to poll underlying futures without yielding. +/// +/// A single call to `poll_next` may potentially do a lot of work before +/// yielding. This happens in particular if the underlying futures are awoken +/// frequently but continue to return `Pending`. This is problematic if other +/// tasks are waiting on the executor, since they do not get to run. This value +/// caps the number of calls to `poll` on underlying futures a single call to +/// `poll_next` is allowed to make. +/// +/// The value itself is chosen somewhat arbitrarily. It needs to be high enough +/// that amortize wakeup and scheduling costs, but low enough that we do not +/// starve other tasks for long. +/// +/// See also https://github.com/rust-lang/futures-rs/issues/2047. +const YIELD_EVERY: usize = 32; + +/// A set of futures which may complete in any order. +/// +/// This structure is optimized to manage a large number of futures. +/// Futures managed by [`FuturesUnordered`] will only be polled when they +/// generate wake-up notifications. This reduces the required amount of work +/// needed to poll large numbers of futures. +/// +/// [`FuturesUnordered`] can be filled by [`collect`](Iterator::collect)ing an +/// iterator of futures into a [`FuturesUnordered`], or by +/// [`push`](FuturesUnordered::push)ing futures onto an existing +/// [`FuturesUnordered`]. When new futures are added, +/// [`poll_next`](Stream::poll_next) must be called in order to begin receiving +/// wake-ups for new futures. +/// +/// Note that you can create a ready-made [`FuturesUnordered`] via the +/// [`collect`](Iterator::collect) method, or you can start with an empty set +/// with the [`FuturesUnordered::new`] constructor. +/// +/// This type is only available when the `std` or `alloc` feature of this +/// library is activated, and it is activated by default. +#[must_use = "streams do nothing unless polled"] +pub struct FuturesUnordered<Fut> { + ready_to_run_queue: Arc<ReadyToRunQueue<Fut>>, + head_all: AtomicPtr<Task<Fut>>, + is_terminated: AtomicBool, +} + +unsafe impl<Fut: Send> Send for FuturesUnordered<Fut> {} +unsafe impl<Fut: Sync> Sync for FuturesUnordered<Fut> {} +impl<Fut> Unpin for FuturesUnordered<Fut> {} + +impl Spawn for FuturesUnordered<FutureObj<'_, ()>> { + fn spawn_obj(&self, future_obj: FutureObj<'static, ()>) + -> Result<(), SpawnError> + { + self.push(future_obj); + Ok(()) + } +} + +impl LocalSpawn for FuturesUnordered<LocalFutureObj<'_, ()>> { + fn spawn_local_obj(&self, future_obj: LocalFutureObj<'static, ()>) + -> Result<(), SpawnError> + { + self.push(future_obj); + Ok(()) + } +} + +// FuturesUnordered is implemented using two linked lists. One which links all +// futures managed by a `FuturesUnordered` and one that tracks futures that have +// been scheduled for polling. The first linked list allows for thread safe +// insertion of nodes at the head as well as forward iteration, but is otherwise +// not thread safe and is only accessed by the thread that owns the +// `FuturesUnordered` value for any other operations. The second linked list is +// an implementation of the intrusive MPSC queue algorithm described by +// 1024cores.net. +// +// When a future is submitted to the set, a task is allocated and inserted in +// both linked lists. The next call to `poll_next` will (eventually) see this +// task and call `poll` on the future. +// +// Before a managed future is polled, the current context's waker is replaced +// with one that is aware of the specific future being run. This ensures that +// wake-up notifications generated by that specific future are visible to +// `FuturesUnordered`. When a wake-up notification is received, the task is +// inserted into the ready to run queue, so that its future can be polled later. +// +// Each task is wrapped in an `Arc` and thereby atomically reference counted. +// Also, each task contains an `AtomicBool` which acts as a flag that indicates +// whether the task is currently inserted in the atomic queue. When a wake-up +// notifiaction is received, the task will only be inserted into the ready to +// run queue if it isn't inserted already. + +impl<Fut: Future> 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() -> FuturesUnordered<Fut> { + 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(), + }); + let stub_ptr = &*stub as *const Task<Fut>; + let ready_to_run_queue = Arc::new(ReadyToRunQueue { + waker: AtomicWaker::new(), + head: AtomicPtr::new(stub_ptr as *mut _), + tail: UnsafeCell::new(stub_ptr), + stub, + }); + + FuturesUnordered { + head_all: AtomicPtr::new(ptr::null_mut()), + ready_to_run_queue, + is_terminated: AtomicBool::new(false), + } + } +} + +impl<Fut: Future> Default for FuturesUnordered<Fut> { + fn default() -> FuturesUnordered<Fut> { + FuturesUnordered::new() + } +} + +impl<Fut> FuturesUnordered<Fut> { + /// 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), + }); + + // 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. + fn iter_pin_ref(self: Pin<&Self>) -> IterPinRef<'_, Fut> { + let (task, len) = self.atomic_load_head_and_len_all(); + + IterPinRef { + task, + len, + pending_next_all: self.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 destorys the future inside and either drops + /// the `Arc<Task>` or transfers ownership to the ready to run queue. + /// The task this method is called on must have been unlinked before. + fn release_task(&mut self, task: Arc<Task<Fut>>) { + // `release_task` must only be called on unlinked tasks + debug_assert_eq!(task.next_all.load(Relaxed), self.pending_next_all()); + unsafe { + debug_assert!((*task.prev_all.get()).is_null()); + } + + // The future is done, try to reset the queued flag. This will prevent + // `wake` from doing any work in the future + let prev = task.queued.swap(true, SeqCst); + + // Drop the future, even if it hasn't finished yet. This is safe + // because we're dropping the future on the thread that owns + // `FuturesUnordered`, which correctly tracks `Fut`'s lifetimes and + // such. + unsafe { + // Set to `None` rather than `take()`ing to prevent moving the + // future. + *task.future.get() = None; + } + + // If the queued flag was previously set, then it means that this task + // is still in our internal ready to run queue. We then transfer + // ownership of our reference count to the ready to run queue, and it'll + // come along and free it later, noticing that the future is `None`. + // + // If, however, the queued flag was *not* set then we're safe to + // release our reference count on the task. The queued flag was set + // above so all future `enqueue` operations will not actually + // enqueue the task, so our task will never see the ready to run queue + // again. The task itself will be deallocated once all reference counts + // have been dropped elsewhere by the various wakers that contain it. + if prev { + mem::forget(task); + } + } + + /// Insert a new task into the internal linked list. + fn link(&self, task: Arc<Task<Fut>>) -> *const Task<Fut> { + // `next_all` should already be reset to the pending state before this + // function is called. + debug_assert_eq!(task.next_all.load(Relaxed), self.pending_next_all()); + let ptr = Arc::into_raw(task); + + // Atomically swap out the old head node to get the node that should be + // assigned to `next_all`. + let next = self.head_all.swap(ptr as *mut _, AcqRel); + + unsafe { + // Store the new list length in the new node. + let new_len = if next.is_null() { + 1 + } else { + // Make sure `next_all` has been written to signal that it is + // safe to read `len_all`. + (*next).spin_next_all(self.pending_next_all(), Acquire); + *(*next).len_all.get() + 1 + }; + *(*ptr).len_all.get() = new_len; + + // Write the old head as the next node pointer, signaling to other + // threads that `len_all` and `next_all` are ready to read. + (*ptr).next_all.store(next, Release); + + // `prev_all` updates don't need to be synchronized, as the field is + // only ever used after exclusive access has been acquired. + if !next.is_null() { + *(*next).prev_all.get() = ptr; + } + } + + ptr + } + + /// Remove the task from the linked list tracking all tasks currently + /// managed by `FuturesUnordered`. + /// This method is unsafe because it has be guaranteed that `task` is a + /// valid pointer. + unsafe fn unlink(&mut self, task: *const Task<Fut>) -> Arc<Task<Fut>> { + // Compute the new list length now in case we're removing the head node + // and won't be able to retrieve the correct length later. + let head = *self.head_all.get_mut(); + debug_assert!(!head.is_null()); + let new_len = *(*head).len_all.get() - 1; + + let task = Arc::from_raw(task); + let next = task.next_all.load(Relaxed); + let prev = *task.prev_all.get(); + task.next_all.store(self.pending_next_all(), Relaxed); + *task.prev_all.get() = ptr::null_mut(); + + if !next.is_null() { + *(*next).prev_all.get() = prev; + } + + if !prev.is_null() { + (*prev).next_all.store(next, Relaxed); + } else { + *self.head_all.get_mut() = next; + } + + // Store the new list length in the head node. + let head = *self.head_all.get_mut(); + if !head.is_null() { + *(*head).len_all.get() = new_len; + } + + task + } + + /// Returns the reserved value for `Task::next_all` to indicate a pending + /// assignment from the thread that inserted the task. + /// + /// `FuturesUnordered::link` needs to update `Task` pointers in an order + /// that ensures any iterators created on other threads can correctly + /// traverse the entire `Task` list using the chain of `next_all` pointers. + /// This could be solved with a compare-exchange loop that stores the + /// current `head_all` in `next_all` and swaps out `head_all` with the new + /// `Task` pointer if the head hasn't already changed. Under heavy thread + /// contention, this compare-exchange loop could become costly. + /// + /// An alternative is to initialize `next_all` to a reserved pending state + /// first, perform an atomic swap on `head_all`, and finally update + /// `next_all` with the old head node. Iterators will then either see the + /// pending state value or the correct next node pointer, and can reload + /// `next_all` as needed until the correct value is loaded. The number of + /// retries needed (if any) would be small and will always be finite, so + /// this should generally perform better than the compare-exchange loop. + /// + /// A valid `Task` pointer in the `head_all` list is guaranteed to never be + /// this value, so it is safe to use as a reserved value until the correct + /// value can be written. + fn pending_next_all(&self) -> *mut Task<Fut> { + // The `ReadyToRunQueue` stub is never inserted into the `head_all` + // list, and its pointer value will remain valid for the lifetime of + // this `FuturesUnordered`, so we can make use of its value here. + &*self.ready_to_run_queue.stub as *const _ 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>> + { + // Keep track of how many child futures we have polled, + // in case we want to forcibly yield. + let mut polled = 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 waker = Task::waker_ref(bomb.task.as_ref().unwrap()); + 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(); + bomb.queue.link(task); + + if polled == YIELD_EVERY { + // We have polled a large number of futures in a row without yielding. + // To ensure we do not starve other tasks waiting on the executor, + // we yield here, but immediately wake ourselves up to continue. + 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> Drop for FuturesUnordered<Fut> { + fn drop(&mut self) { + // When a `FuturesUnordered` is dropped we want to drop all futures + // associated with it. At the same time though there may be tons of + // wakers flying around which contain `Task<Fut>` references + // inside them. We'll let those naturally get deallocated. + unsafe { + while !self.head_all.get_mut().is_null() { + let head = *self.head_all.get_mut(); + let task = self.unlink(head); + self.release_task(task); + } + } + + // 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<Fut: Future> FromIterator<Fut> for FuturesUnordered<Fut> { + fn from_iter<I>(iter: I) -> Self + where + I: IntoIterator<Item = Fut>, + { + let acc = FuturesUnordered::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) + } +} diff --git a/src/stream/futures_unordered/ready_to_run_queue.rs b/src/stream/futures_unordered/ready_to_run_queue.rs new file mode 100644 index 0000000..2105195 --- /dev/null +++ b/src/stream/futures_unordered/ready_to_run_queue.rs @@ -0,0 +1,109 @@ +use crate::task::AtomicWaker; +use core::cell::UnsafeCell; +use core::ptr; +use core::sync::atomic::AtomicPtr; +use core::sync::atomic::Ordering::{Relaxed, Acquire, Release, AcqRel}; +use alloc::sync::Arc; + +use super::abort::abort; +use super::task::Task; + +pub(super) enum Dequeue<Fut> { + Data(*const Task<Fut>), + Empty, + Inconsistent, +} + +pub(super) struct ReadyToRunQueue<Fut> { + // The waker of the task using `FuturesUnordered`. + pub(super) waker: AtomicWaker, + + // Head/tail of the readiness queue + pub(super) head: AtomicPtr<Task<Fut>>, + pub(super) tail: UnsafeCell<*const Task<Fut>>, + pub(super) stub: Arc<Task<Fut>>, +} + +/// An MPSC queue into which the tasks containing the futures are inserted +/// whenever the future inside is scheduled for polling. +impl<Fut> ReadyToRunQueue<Fut> { + /// The enqueue function from the 1024cores intrusive MPSC queue algorithm. + pub(super) fn enqueue(&self, task: *const Task<Fut>) { + unsafe { + debug_assert!((*task).queued.load(Relaxed)); + + // This action does not require any coordination + (*task).next_ready_to_run.store(ptr::null_mut(), Relaxed); + + // Note that these atomic orderings come from 1024cores + let task = task as *mut _; + let prev = self.head.swap(task, AcqRel); + (*prev).next_ready_to_run.store(task, Release); + } + } + + /// The dequeue function from the 1024cores intrusive MPSC queue algorithm + /// + /// Note that this is unsafe as it required mutual exclusion (only one + /// thread can call this) to be guaranteed elsewhere. + pub(super) unsafe fn dequeue(&self) -> Dequeue<Fut> { + let mut tail = *self.tail.get(); + let mut next = (*tail).next_ready_to_run.load(Acquire); + + if tail == self.stub() { + if next.is_null() { + return Dequeue::Empty; + } + + *self.tail.get() = next; + tail = next; + next = (*next).next_ready_to_run.load(Acquire); + } + + if !next.is_null() { + *self.tail.get() = next; + debug_assert!(tail != self.stub()); + return Dequeue::Data(tail); + } + + if self.head.load(Acquire) as *const _ != tail { + return Dequeue::Inconsistent; + } + + self.enqueue(self.stub()); + + next = (*tail).next_ready_to_run.load(Acquire); + + if !next.is_null() { + *self.tail.get() = next; + return Dequeue::Data(tail); + } + + Dequeue::Inconsistent + } + + pub(super) fn stub(&self) -> *const Task<Fut> { + &*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)), + } + } + } + } +} diff --git a/src/stream/futures_unordered/task.rs b/src/stream/futures_unordered/task.rs new file mode 100644 index 0000000..9277229 --- /dev/null +++ b/src/stream/futures_unordered/task.rs @@ -0,0 +1,118 @@ +use core::cell::UnsafeCell; +use core::sync::atomic::{AtomicPtr, AtomicBool}; +use core::sync::atomic::Ordering::{self, SeqCst}; +use alloc::sync::{Arc, Weak}; + +use crate::task::{ArcWake, WakerRef, waker_ref}; +use super::ReadyToRunQueue; +use super::abort::abort; + +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, +} + +// `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, + }; + + // 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_self); + inner.waker.wake(); + } + } +} + +impl<Fut> Task<Fut> { + /// Returns a waker reference for this task without cloning the Arc. + pub(super) fn waker_ref<'a>(this: &'a Arc<Task<Fut>>) -> WakerRef<'a> { + waker_ref(this) + } + + /// Spins until `next_all` is no longer set to `pending_next_all`. + /// + /// The temporary `pending_next_all` value is typically overwritten fairly + /// quickly after a node is inserted into the list of all futures, so this + /// should rarely spin much. + /// + /// When it returns, the correct `next_all` value is returned. + /// + /// `Relaxed` or `Acquire` ordering can be used. `Acquire` ordering must be + /// used before `len_all` can be safely read. + #[inline] + pub(super) fn spin_next_all( + &self, + pending_next_all: *mut Self, + ordering: Ordering, + ) -> *const Self { + loop { + let next = self.next_all.load(ordering); + if next != pending_next_all { + return next; + } + } + } +} + +impl<Fut> Drop for Task<Fut> { + fn drop(&mut self) { + // Since `Task<Fut>` is sent across all threads for any lifetime, + // regardless of `Fut`, we, to guarantee memory safety, can't actually + // touch `Fut` at any time except when we have a reference to the + // `FuturesUnordered` itself . + // + // Consequently it *should* be the case that we always drop futures from + // the `FuturesUnordered` instance. This is a bomb, just in case there's + // a bug in that logic. + unsafe { + if (*self.future.get()).is_some() { + abort("future still here when dropping"); + } + } + } +} diff --git a/src/stream/iter.rs b/src/stream/iter.rs new file mode 100644 index 0000000..e9df81c --- /dev/null +++ b/src/stream/iter.rs @@ -0,0 +1,48 @@ +use core::pin::Pin; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; + +/// Stream for the [`iter`] function. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Iter<I> { + iter: I, +} + +impl<I> Unpin for Iter<I> {} + +/// Converts an `Iterator` into a `Stream` which is always ready +/// to yield the next value. +/// +/// Iterators in Rust don't express the ability to block, so this adapter +/// simply always calls `iter.next()` and returns that. +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::stream::{self, StreamExt}; +/// +/// let stream = stream::iter(vec![17, 19]); +/// assert_eq!(vec![17, 19], stream.collect::<Vec<i32>>().await); +/// # }); +/// ``` +pub fn iter<I>(i: I) -> Iter<I::IntoIter> + where I: IntoIterator, +{ + Iter { + iter: i.into_iter(), + } +} + +impl<I> Stream for Iter<I> + where I: Iterator, +{ + type Item = I::Item; + + fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<I::Item>> { + Poll::Ready(self.iter.next()) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.iter.size_hint() + } +} diff --git a/src/stream/mod.rs b/src/stream/mod.rs new file mode 100644 index 0000000..ba3575b --- /dev/null +++ b/src/stream/mod.rs @@ -0,0 +1,96 @@ +//! Streams +//! +//! This module contains a number of functions for working with `Stream`s, +//! including the [`StreamExt`] trait and the [`TryStreamExt`] trait which add +//! methods to `Stream` types + +#[cfg(feature = "alloc")] +pub use futures_core::stream::{BoxStream, LocalBoxStream}; +pub use futures_core::stream::{FusedStream, Stream, TryStream}; + +// Extension traits and combinators + +#[allow(clippy::module_inception)] +mod stream; +pub use self::stream::{ + Chain, Collect, Concat, Enumerate, Filter, FilterMap, Flatten, Fold, ForEach, Fuse, Inspect, + Map, Next, Peek, Peekable, Scan, SelectNextSome, Skip, SkipWhile, StreamExt, StreamFuture, Take, + TakeWhile, Then, Zip, +}; + +#[cfg(feature = "std")] +pub use self::stream::CatchUnwind; + +#[cfg(feature = "alloc")] +pub use self::stream::Chunks; + +#[cfg(feature = "sink")] +pub use self::stream::Forward; + +#[cfg_attr(feature = "cfg-target-has-atomic", cfg(target_has_atomic = "ptr"))] +#[cfg(feature = "alloc")] +pub use self::stream::{BufferUnordered, Buffered, ForEachConcurrent}; + +#[cfg_attr(feature = "cfg-target-has-atomic", cfg(target_has_atomic = "ptr"))] +#[cfg(feature = "sink")] +#[cfg(feature = "alloc")] +pub use self::stream::{ReuniteError, SplitSink, SplitStream}; + +mod try_stream; +pub use self::try_stream::{ + try_unfold, AndThen, ErrInto, InspectErr, InspectOk, IntoStream, MapErr, MapOk, OrElse, + TryCollect, TryConcat, TryFilter, TryFilterMap, TryFlatten, TryFold, TryForEach, TryNext, + TrySkipWhile, TryStreamExt, TryUnfold, +}; + +#[cfg(feature = "io")] +#[cfg(feature = "std")] +pub use self::try_stream::IntoAsyncRead; + +#[cfg_attr(feature = "cfg-target-has-atomic", cfg(target_has_atomic = "ptr"))] +#[cfg(feature = "alloc")] +pub use self::try_stream::{TryBufferUnordered, TryForEachConcurrent}; + +// Primitive streams + +mod iter; +pub use self::iter::{iter, Iter}; + +mod repeat; +pub use self::repeat::{repeat, Repeat}; + +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 select; +pub use self::select::{select, Select}; + +mod unfold; +pub use self::unfold::{unfold, Unfold}; + +cfg_target_has_atomic! { + #[cfg(feature = "alloc")] + mod futures_ordered; + #[cfg(feature = "alloc")] + pub use self::futures_ordered::FuturesOrdered; + + #[cfg(feature = "alloc")] + pub mod futures_unordered; + #[cfg(feature = "alloc")] + #[doc(inline)] + pub use self::futures_unordered::FuturesUnordered; + + #[cfg(feature = "alloc")] + mod select_all; + #[cfg(feature = "alloc")] + pub use self::select_all::{select_all, SelectAll}; +} diff --git a/src/stream/once.rs b/src/stream/once.rs new file mode 100644 index 0000000..4f68b0c --- /dev/null +++ b/src/stream/once.rs @@ -0,0 +1,63 @@ +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// 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> { + Once { future: Some(future) } +} + +/// A stream which emits single element and then EOF. +/// +/// This stream will never block and is always ready. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Once<Fut> { + future: Option<Fut> +} + +impl<Fut: Unpin> Unpin for Once<Fut> {} + +impl<Fut> Once<Fut> { + unsafe_pinned!(future: Option<Fut>); +} + +impl<Fut: Future> Stream for Once<Fut> { + type Item = Fut::Output; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let v = match self.as_mut().future().as_pin_mut() { + Some(fut) => ready!(fut.poll(cx)), + None => return Poll::Ready(None), + }; + + self.as_mut().future().set(None); + Poll::Ready(Some(v)) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.future.is_some() { + (1, Some(1)) + } else { + (0, Some(0)) + } + } +} + +impl<Fut: Future> FusedStream for Once<Fut> { + fn is_terminated(&self) -> bool { + self.future.is_none() + } +} diff --git a/src/stream/pending.rs b/src/stream/pending.rs new file mode 100644 index 0000000..bbe0750 --- /dev/null +++ b/src/stream/pending.rs @@ -0,0 +1,38 @@ +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> { + 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)) + } +} diff --git a/src/stream/poll_fn.rs b/src/stream/poll_fn.rs new file mode 100644 index 0000000..e33ca57 --- /dev/null +++ b/src/stream/poll_fn.rs @@ -0,0 +1,56 @@ +//! Definition of the `PollFn` combinator + +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>>, +{ + PollFn { f } +} + +impl<T, F> Stream for PollFn<F> +where + F: FnMut(&mut Context<'_>) -> Poll<Option<T>>, +{ + type Item = T; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> { + (&mut self.f)(cx) + } +} diff --git a/src/stream/repeat.rs b/src/stream/repeat.rs new file mode 100644 index 0000000..21749eb --- /dev/null +++ b/src/stream/repeat.rs @@ -0,0 +1,54 @@ +use core::pin::Pin; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; + +/// Stream for the [`repeat`] function. +#[derive(Debug)] +#[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 +{ + Repeat { item } +} + +impl<T> Unpin for Repeat<T> {} + +impl<T> Stream for Repeat<T> + where T: Clone +{ + type Item = T; + + fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> { + Poll::Ready(Some(self.item.clone())) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (usize::max_value(), None) + } +} + +impl<T> FusedStream for Repeat<T> + where T: Clone, +{ + fn is_terminated(&self) -> bool { + false + } +} diff --git a/src/stream/select.rs b/src/stream/select.rs new file mode 100644 index 0000000..b5fb813 --- /dev/null +++ b/src/stream/select.rs @@ -0,0 +1,132 @@ +use crate::stream::{StreamExt, Fuse}; +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; + +/// Stream for the [`select()`] function. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Select<St1, St2> { + stream1: Fuse<St1>, + stream2: Fuse<St2>, + flag: bool, +} + +impl<St1: Unpin, St2: Unpin> Unpin for Select<St1, St2> {} + +/// 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 stream 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. +pub fn select<St1, St2>(stream1: St1, stream2: St2) -> Select<St1, St2> + where St1: Stream, + St2: Stream<Item = St1::Item> +{ + Select { + stream1: stream1.fuse(), + stream2: stream2.fuse(), + flag: false, + } +} + +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.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>) { + unsafe { + let Self { stream1, stream2, .. } = self.get_unchecked_mut(); + (Pin::new_unchecked(stream1).get_pin_mut(), Pin::new_unchecked(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 Select<St1, St2> + where St1: Stream, + St2: Stream<Item = St1::Item> +{ + fn is_terminated(&self) -> bool { + self.stream1.is_terminated() && self.stream2.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 Select { flag, stream1, stream2 } = + unsafe { self.get_unchecked_mut() }; + let stream1 = unsafe { Pin::new_unchecked(stream1) }; + let stream2 = unsafe { Pin::new_unchecked(stream2) }; + + if !*flag { + poll_inner(flag, stream1, stream2, cx) + } else { + poll_inner(flag, stream2, stream1, cx) + } + } +} + +fn poll_inner<St1, St2>( + flag: &mut bool, + a: Pin<&mut St1>, + b: Pin<&mut St2>, + cx: &mut Context<'_> +) -> Poll<Option<St1::Item>> + where St1: Stream, St2: Stream<Item = St1::Item> +{ + let a_done = match a.poll_next(cx) { + Poll::Ready(Some(item)) => { + // give the other stream a chance to go first next time + *flag = !*flag; + return Poll::Ready(Some(item)) + }, + Poll::Ready(None) => true, + Poll::Pending => false, + }; + + match b.poll_next(cx) { + Poll::Ready(Some(item)) => { + Poll::Ready(Some(item)) + } + Poll::Ready(None) if a_done => Poll::Ready(None), + Poll::Ready(None) | Poll::Pending => Poll::Pending, + } +} diff --git a/src/stream/select_all.rs b/src/stream/select_all.rs new file mode 100644 index 0000000..d257689 --- /dev/null +++ b/src/stream/select_all.rs @@ -0,0 +1,133 @@ +//! An unbounded set of streams + +use core::fmt::{self, Debug}; +use core::iter::FromIterator; +use core::pin::Pin; + +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; + +use crate::stream::{StreamExt, StreamFuture, FuturesUnordered}; + +/// 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() -> SelectAll<St> { + SelectAll { 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()); + } +} + +impl<St: Stream + Unpin> Default for SelectAll<St> { + fn default() -> SelectAll<St> { + SelectAll::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 +/// futures 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); + } + + set +} + +impl<St: Stream + Unpin> FromIterator<St> for SelectAll<St> { + fn from_iter<T: IntoIterator<Item = St>>(iter: T) -> Self { + select_all(iter) + } +} diff --git a/src/stream/stream/buffer_unordered.rs b/src/stream/stream/buffer_unordered.rs new file mode 100644 index 0000000..bea6e5b --- /dev/null +++ b/src/stream/stream/buffer_unordered.rs @@ -0,0 +1,160 @@ +use crate::stream::{Fuse, FuturesUnordered, StreamExt}; +use futures_core::future::Future; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; +use core::fmt; +use core::pin::Pin; + +/// Stream for the [`buffer_unordered`](super::StreamExt::buffer_unordered) +/// method. +#[must_use = "streams do nothing unless polled"] +pub struct BufferUnordered<St> +where + St: Stream, +{ + stream: Fuse<St>, + in_progress_queue: FuturesUnordered<St::Item>, + max: usize, +} + +impl<St> Unpin for BufferUnordered<St> +where + St: Stream + Unpin, +{} + +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, +{ + unsafe_pinned!(stream: Fuse<St>); + unsafe_unpinned!(in_progress_queue: FuturesUnordered<St::Item>); + + pub(super) fn new(stream: St, n: usize) -> BufferUnordered<St> + where + St: Stream, + St::Item: Future, + { + BufferUnordered { + stream: super::Fuse::new(stream), + in_progress_queue: FuturesUnordered::new(), + max: n, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + self.stream.get_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. + pub fn get_mut(&mut self) -> &mut St { + self.stream.get_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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream().get_pin_mut() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream.into_inner() + } +} + +impl<St> Stream for BufferUnordered<St> +where + St: Stream, + St::Item: Future, +{ + type Item = <St::Item as Future>::Output; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + // First up, try to spawn off as many futures as possible by filling up + // our queue of futures. + while self.in_progress_queue.len() < self.max { + match self.as_mut().stream().poll_next(cx) { + Poll::Ready(Some(fut)) => self.as_mut().in_progress_queue().push(fut), + Poll::Ready(None) | Poll::Pending => break, + } + } + + // Attempt to pull the next value from the in_progress_queue + match self.as_mut().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 self.stream.is_done() { + Poll::Ready(None) + } else { + Poll::Pending + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let queue_len = self.in_progress_queue.len(); + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(queue_len); + let upper = match upper { + Some(x) => x.checked_add(queue_len), + None => None, + }; + (lower, upper) + } +} + +impl<St> FusedStream for BufferUnordered<St> +where + St: Stream, + St::Item: Future, +{ + fn is_terminated(&self) -> bool { + self.in_progress_queue.is_terminated() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for BufferUnordered<S> +where + S: Stream + Sink<Item>, + S::Item: Future, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/buffered.rs b/src/stream/stream/buffered.rs new file mode 100644 index 0000000..2445a85 --- /dev/null +++ b/src/stream/stream/buffered.rs @@ -0,0 +1,148 @@ +use crate::stream::{Fuse, FuturesOrdered, StreamExt}; +use futures_core::future::Future; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; +use core::fmt; +use core::pin::Pin; + +/// 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, +{ + stream: Fuse<St>, + in_progress_queue: FuturesOrdered<St::Item>, + max: usize, +} + +impl<St> Unpin for Buffered<St> +where + St: Stream + Unpin, + St::Item: Future, +{} + +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, +{ + unsafe_pinned!(stream: Fuse<St>); + unsafe_unpinned!(in_progress_queue: FuturesOrdered<St::Item>); + + pub(super) fn new(stream: St, n: usize) -> Buffered<St> { + Buffered { + stream: super::Fuse::new(stream), + in_progress_queue: FuturesOrdered::new(), + max: n, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + self.stream.get_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. + pub fn get_mut(&mut self) -> &mut St { + self.stream.get_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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream().get_pin_mut() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream.into_inner() + } +} + +impl<St> Stream for Buffered<St> +where + St: Stream, + St::Item: Future, +{ + type Item = <St::Item as Future>::Output; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + // Try to spawn off as many futures as possible by filling up + // our in_progress_queue of futures. + while self.in_progress_queue.len() < self.max { + match self.as_mut().stream().poll_next(cx) { + Poll::Ready(Some(fut)) => self.as_mut().in_progress_queue().push(fut), + Poll::Ready(None) | Poll::Pending => break, + } + } + + // Attempt to pull the next value from the in_progress_queue + let res = self.as_mut().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 self.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) + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Buffered<S> +where + S: Stream + Sink<Item>, + S::Item: Future, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/catch_unwind.rs b/src/stream/stream/catch_unwind.rs new file mode 100644 index 0000000..8d2dcf7 --- /dev/null +++ b/src/stream/stream/catch_unwind.rs @@ -0,0 +1,62 @@ +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; +use std::any::Any; +use std::pin::Pin; +use std::panic::{catch_unwind, UnwindSafe, AssertUnwindSafe}; + +/// Stream for the [`catch_unwind`](super::StreamExt::catch_unwind) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct CatchUnwind<St> { + stream: St, + caught_unwind: bool, +} + +impl<St: Stream + UnwindSafe> CatchUnwind<St> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(caught_unwind: bool); + + pub(super) fn new(stream: St) -> CatchUnwind<St> { + CatchUnwind { stream, caught_unwind: false } + } +} + +impl<St: Stream + UnwindSafe> Stream for CatchUnwind<St> { + type Item = Result<St::Item, Box<dyn Any + Send>>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.caught_unwind { + Poll::Ready(None) + } else { + let res = catch_unwind(AssertUnwindSafe(|| { + self.as_mut().stream().poll_next(cx) + })); + + match res { + Ok(poll) => poll.map(|opt| opt.map(Ok)), + Err(e) => { + *self.as_mut().caught_unwind() = true; + Poll::Ready(Some(Err(e))) + }, + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.caught_unwind { + (0, Some(0)) + } else { + self.stream.size_hint() + } + } +} + +impl<St: FusedStream + UnwindSafe> FusedStream for CatchUnwind<St> { + fn is_terminated(&self) -> bool { + self.caught_unwind || self.stream.is_terminated() + } +} diff --git a/src/stream/stream/chain.rs b/src/stream/stream/chain.rs new file mode 100644 index 0000000..b2ada69 --- /dev/null +++ b/src/stream/stream/chain.rs @@ -0,0 +1,75 @@ +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// Stream for the [`chain`](super::StreamExt::chain) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Chain<St1, St2> { + first: Option<St1>, + 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>, +{ + unsafe_pinned!(first: Option<St1>); + unsafe_pinned!(second: St2); + + pub(super) fn new(stream1: St1, stream2: St2) -> Chain<St1, St2> { + Chain { + 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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if let Some(first) = self.as_mut().first().as_pin_mut() { + if let Some(item) = ready!(first.poll_next(cx)) { + return Poll::Ready(Some(item)) + } + } + self.as_mut().first().set(None); + self.as_mut().second().poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if let Some(first) = &self.first { + let (first_lower, first_upper) = first.size_hint(); + let (second_lower, second_upper) = self.second.size_hint(); + + let lower = first_lower.saturating_add(second_lower); + + let upper = match (first_upper, second_upper) { + (Some(x), Some(y)) => x.checked_add(y), + _ => None + }; + + (lower, upper) + } else { + self.second.size_hint() + } + } +} diff --git a/src/stream/stream/chunks.rs b/src/stream/stream/chunks.rs new file mode 100644 index 0000000..b42d1d1 --- /dev/null +++ b/src/stream/stream/chunks.rs @@ -0,0 +1,136 @@ +use crate::stream::Fuse; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; +use core::mem; +use core::pin::Pin; +use alloc::vec::Vec; + +/// Stream for the [`chunks`](super::StreamExt::chunks) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Chunks<St: Stream> { + stream: Fuse<St>, + items: Vec<St::Item>, + cap: usize, // https://github.com/rust-lang/futures-rs/issues/1475 +} + +impl<St: Unpin + Stream> Unpin for Chunks<St> {} + +impl<St: Stream> Chunks<St> where St: Stream { + unsafe_unpinned!(items: Vec<St::Item>); + unsafe_pinned!(stream: Fuse<St>); + + pub(super) fn new(stream: St, capacity: usize) -> Chunks<St> { + assert!(capacity > 0); + + Chunks { + stream: super::Fuse::new(stream), + items: Vec::with_capacity(capacity), + cap: capacity, + } + } + + fn take(mut self: Pin<&mut Self>) -> Vec<St::Item> { + let cap = self.cap; + mem::replace(self.as_mut().items(), Vec::with_capacity(cap)) + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + self.stream.get_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. + pub fn get_mut(&mut self) -> &mut St { + self.stream.get_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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream().get_pin_mut() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream.into_inner() + } +} + +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>> { + loop { + match ready!(self.as_mut().stream().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) => { + self.as_mut().items().push(item); + if self.items.len() >= self.cap { + return Poll::Ready(Some(self.as_mut().take())) + } + } + + // Since the underlying stream ran out of values, return what we + // have buffered, if we have anything. + None => { + let last = if self.items.is_empty() { + None + } else { + let full_buf = mem::replace(self.as_mut().items(), Vec::new()); + Some(full_buf) + }; + + return Poll::Ready(last); + } + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let chunk_len = if self.items.is_empty() { 0 } else { 1 }; + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(chunk_len); + let upper = match upper { + Some(x) => x.checked_add(chunk_len), + None => None, + }; + (lower, upper) + } +} + +impl<St: FusedStream> FusedStream for Chunks<St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() && self.items.is_empty() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Chunks<S> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/collect.rs b/src/stream/stream/collect.rs new file mode 100644 index 0000000..127a3f7 --- /dev/null +++ b/src/stream/stream/collect.rs @@ -0,0 +1,57 @@ +use core::mem; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: St, + collection: C, +} + +impl<St: Unpin, C> Unpin for Collect<St, C> {} + +impl<St: Stream, C: Default> Collect<St, C> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(collection: C); + + fn finish(mut self: Pin<&mut Self>) -> C { + mem::replace(self.as_mut().collection(), Default::default()) + } + + pub(super) fn new(stream: St) -> Collect<St, C> { + Collect { + 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> { + loop { + match ready!(self.as_mut().stream().poll_next(cx)) { + Some(e) => self.as_mut().collection().extend(Some(e)), + None => return Poll::Ready(self.as_mut().finish()), + } + } + } +} diff --git a/src/stream/stream/concat.rs b/src/stream/stream/concat.rs new file mode 100644 index 0000000..704efc7 --- /dev/null +++ b/src/stream/stream/concat.rs @@ -0,0 +1,69 @@ +use core::pin::Pin; +use futures_core::future::{Future, FusedFuture}; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: St, + accum: Option<St::Item>, +} + +impl<St: Stream + Unpin> Unpin for Concat<St> {} + +impl<St> Concat<St> +where St: Stream, + St::Item: Extend<<St::Item as IntoIterator>::Item> + + IntoIterator + Default, +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(accum: Option<St::Item>); + + pub(super) fn new(stream: St) -> Concat<St> { + Concat { + 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( + mut self: Pin<&mut Self>, cx: &mut Context<'_> + ) -> Poll<Self::Output> { + loop { + match ready!(self.as_mut().stream().poll_next(cx)) { + None => { + return Poll::Ready(self.as_mut().accum().take().unwrap_or_default()) + } + Some(e) => { + let accum = self.as_mut().accum(); + if let Some(a) = accum { + a.extend(e) + } else { + *accum = Some(e) + } + } + } + } + } +} + +impl<St> FusedFuture for Concat<St> +where St: FusedStream, + St::Item: Extend<<St::Item as IntoIterator>::Item> + + IntoIterator + Default, +{ + fn is_terminated(&self) -> bool { + self.accum.is_none() && self.stream.is_terminated() + } +} diff --git a/src/stream/stream/enumerate.rs b/src/stream/stream/enumerate.rs new file mode 100644 index 0000000..6366c8b --- /dev/null +++ b/src/stream/stream/enumerate.rs @@ -0,0 +1,99 @@ +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`enumerate`](super::StreamExt::enumerate) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Enumerate<St> { + stream: St, + count: usize, +} + +impl<St: Unpin> Unpin for Enumerate<St> {} + +impl<St: Stream> Enumerate<St> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(count: usize); + + pub(super) fn new(stream: St) -> Enumerate<St> { + Enumerate { + stream, + count: 0, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + match ready!(self.as_mut().stream().poll_next(cx)) { + Some(item) => { + let count = self.count; + *self.as_mut().count() += 1; + Poll::Ready(Some((count, item))) + } + None => Poll::Ready(None), + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Enumerate<S> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/filter.rs b/src/stream/stream/filter.rs new file mode 100644 index 0000000..06335f1 --- /dev/null +++ b/src/stream/stream/filter.rs @@ -0,0 +1,157 @@ +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`filter`](super::StreamExt::filter) method. +#[must_use = "streams do nothing unless polled"] +pub struct Filter<St, Fut, F> + where St: Stream, +{ + stream: St, + f: F, + pending_fut: Option<Fut>, + pending_item: Option<St::Item>, +} + +impl<St, Fut, F> Unpin for Filter<St, Fut, F> +where + St: Stream + Unpin, + Fut: Unpin, +{} + +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() + } +} + +impl<St, Fut, F> Filter<St, Fut, F> +where St: Stream, + F: FnMut(&St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + unsafe_pinned!(pending_fut: Option<Fut>); + unsafe_unpinned!(pending_item: Option<St::Item>); + + pub(super) fn new(stream: St, f: F) -> Filter<St, Fut, F> { + Filter { + stream, + f, + pending_fut: None, + pending_item: None, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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() + } +} + +impl<St, Fut, F> Stream for Filter<St, Fut, F> + where St: Stream, + F: FnMut(&St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + type Item = St::Item; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<St::Item>> { + loop { + if self.pending_fut.is_none() { + let item = match ready!(self.as_mut().stream().poll_next(cx)) { + Some(e) => e, + None => return Poll::Ready(None), + }; + let fut = (self.as_mut().f())(&item); + self.as_mut().pending_fut().set(Some(fut)); + *self.as_mut().pending_item() = Some(item); + } + + let yield_item = ready!(self.as_mut().pending_fut().as_pin_mut().unwrap().poll(cx)); + self.as_mut().pending_fut().set(None); + let item = self.as_mut().pending_item().take().unwrap(); + + if yield_item { + return Poll::Ready(Some(item)); + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = if self.pending_item.is_some() { 1 } else { 0 }; + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for Filter<S, Fut, F> + where S: Stream + Sink<Item>, + F: FnMut(&S::Item) -> Fut, + Fut: Future<Output = bool>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/filter_map.rs b/src/stream/stream/filter_map.rs new file mode 100644 index 0000000..532e6ca --- /dev/null +++ b/src/stream/stream/filter_map.rs @@ -0,0 +1,143 @@ +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`filter_map`](super::StreamExt::filter_map) method. +#[must_use = "streams do nothing unless polled"] +pub struct FilterMap<St, Fut, F> { + stream: St, + f: F, + pending: Option<Fut>, +} + +impl<St, Fut, F> Unpin for FilterMap<St, Fut, F> +where + St: Unpin, + Fut: Unpin, +{} + +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, +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + unsafe_pinned!(pending: Option<Fut>); + + pub(super) fn new(stream: St, f: F) -> FilterMap<St, Fut, F> { + FilterMap { stream, f, pending: None } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +impl<St, Fut, F, T> FusedStream for FilterMap<St, Fut, F> + where St: Stream + FusedStream, + F: FnMut(St::Item) -> 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: FnMut(St::Item) -> Fut, + Fut: Future<Output = Option<T>>, +{ + type Item = T; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<T>> { + loop { + if self.pending.is_none() { + let item = match ready!(self.as_mut().stream().poll_next(cx)) { + Some(e) => e, + None => return Poll::Ready(None), + }; + let fut = (self.as_mut().f())(item); + self.as_mut().pending().set(Some(fut)); + } + + let item = ready!(self.as_mut().pending().as_pin_mut().unwrap().poll(cx)); + self.as_mut().pending().set(None); + if item.is_some() { + return Poll::Ready(item); + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = if self.pending.is_some() { 1 } else { 0 }; + 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: FnMut(S::Item) -> Fut, + Fut: Future, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/flatten.rs b/src/stream/stream/flatten.rs new file mode 100644 index 0000000..b19ffc0 --- /dev/null +++ b/src/stream/stream/flatten.rs @@ -0,0 +1,120 @@ +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_utils::unsafe_pinned; + +/// Stream for the [`flatten`](super::StreamExt::flatten) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Flatten<St> +where + St: Stream, +{ + stream: St, + next: Option<St::Item>, +} + +impl<St> Unpin for Flatten<St> +where + St: Stream + Unpin, + St::Item: Unpin, +{ +} + +impl<St> Flatten<St> +where + St: Stream, +{ + unsafe_pinned!(stream: St); + unsafe_pinned!(next: Option<St::Item>); +} + +impl<St> Flatten<St> +where + St: Stream, + St::Item: Stream, +{ + pub(super) fn new(stream: St) -> Self { + Self { stream, next: None } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +impl<St> FusedStream for Flatten<St> +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> +where + St: Stream, + St::Item: Stream, +{ + type Item = <St::Item as Stream>::Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + loop { + if self.next.is_none() { + match ready!(self.as_mut().stream().poll_next(cx)) { + Some(e) => self.as_mut().next().set(Some(e)), + None => return Poll::Ready(None), + } + } + + if let Some(item) = ready!(self.as_mut().next().as_pin_mut().unwrap().poll_next(cx)) { + return Poll::Ready(Some(item)); + } else { + self.as_mut().next().set(None); + } + } + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Flatten<S> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/fold.rs b/src/stream/stream/fold.rs new file mode 100644 index 0000000..e92a72e --- /dev/null +++ b/src/stream/stream/fold.rs @@ -0,0 +1,92 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: St, + f: F, + accum: Option<T>, + future: Option<Fut>, +} + +impl<St: Unpin, Fut: Unpin, T, F> Unpin for Fold<St, Fut, T, F> {} + +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>, +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + unsafe_unpinned!(accum: Option<T>); + unsafe_pinned!(future: Option<Fut>); + + pub(super) fn new(stream: St, f: F, t: T) -> Fold<St, Fut, T, F> { + Fold { + 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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { + loop { + // we're currently processing a future to produce a new accum value + if self.accum.is_none() { + let accum = ready!(self.as_mut().future().as_pin_mut().unwrap().poll(cx)); + *self.as_mut().accum() = Some(accum); + self.as_mut().future().set(None); + } + + let item = ready!(self.as_mut().stream().poll_next(cx)); + let accum = self.as_mut().accum().take() + .expect("Fold polled after completion"); + + if let Some(e) = item { + let future = (self.as_mut().f())(accum, e); + self.as_mut().future().set(Some(future)); + } else { + return Poll::Ready(accum) + } + } + } +} diff --git a/src/stream/stream/for_each.rs b/src/stream/stream/for_each.rs new file mode 100644 index 0000000..f8adcb2 --- /dev/null +++ b/src/stream/stream/for_each.rs @@ -0,0 +1,88 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: St, + f: F, + future: Option<Fut>, +} + +impl<St, Fut, F> Unpin for ForEach<St, Fut, F> +where + St: Unpin, + Fut: Unpin, +{} + +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 = ()>, +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + unsafe_pinned!(future: Option<Fut>); + + pub(super) fn new(stream: St, f: F) -> ForEach<St, Fut, F> { + ForEach { + 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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + loop { + if let Some(future) = self.as_mut().future().as_pin_mut() { + ready!(future.poll(cx)); + } + self.as_mut().future().set(None); + + match ready!(self.as_mut().stream().poll_next(cx)) { + Some(e) => { + let future = (self.as_mut().f())(e); + self.as_mut().future().set(Some(future)); + } + None => { + return Poll::Ready(()); + } + } + } + } +} diff --git a/src/stream/stream/for_each_concurrent.rs b/src/stream/stream/for_each_concurrent.rs new file mode 100644 index 0000000..18ca4bd --- /dev/null +++ b/src/stream/stream/for_each_concurrent.rs @@ -0,0 +1,125 @@ +use crate::stream::{FuturesUnordered, StreamExt}; +use core::fmt; +use core::pin::Pin; +use core::num::NonZeroUsize; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: Option<St>, + f: F, + futures: FuturesUnordered<Fut>, + limit: Option<NonZeroUsize>, +} + +impl<St, Fut, F> Unpin for ForEachConcurrent<St, Fut, F> +where St: Unpin, + Fut: Unpin, +{} + +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 = ()>, +{ + unsafe_pinned!(stream: Option<St>); + unsafe_unpinned!(f: F); + unsafe_unpinned!(futures: FuturesUnordered<Fut>); + unsafe_unpinned!(limit: Option<NonZeroUsize>); + + pub(super) fn new(stream: St, limit: Option<usize>, f: F) -> ForEachConcurrent<St, Fut, F> { + ForEachConcurrent { + 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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + loop { + let mut made_progress_this_iter = false; + + // Try and pull an item from the stream + let current_len = self.futures.len(); + // Check if we've already created a number of futures greater than `limit` + if self.limit.map(|limit| limit.get() > current_len).unwrap_or(true) { + let mut stream_completed = false; + let elem = if let Some(stream) = self.as_mut().stream().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 { + self.as_mut().stream().set(None); + } + if let Some(elem) = elem { + let next_future = (self.as_mut().f())(elem); + self.as_mut().futures().push(next_future); + } + } + + match self.as_mut().futures().poll_next_unpin(cx) { + Poll::Ready(Some(())) => made_progress_this_iter = true, + Poll::Ready(None) => { + if self.stream.is_none() { + return Poll::Ready(()) + } + }, + Poll::Pending => {} + } + + if !made_progress_this_iter { + return Poll::Pending; + } + } + } +} diff --git a/src/stream/stream/forward.rs b/src/stream/stream/forward.rs new file mode 100644 index 0000000..fd89625 --- /dev/null +++ b/src/stream/stream/forward.rs @@ -0,0 +1,99 @@ +use crate::stream::{StreamExt, Fuse}; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::{Stream, TryStream}; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +const INVALID_POLL: &str = "polled `Forward` after completion"; + +/// Future for the [`forward`](super::StreamExt::forward) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Forward<St: TryStream, Si> { + sink: Option<Si>, + stream: Fuse<St>, + buffered_item: Option<St::Ok>, +} + +impl<St: TryStream + Unpin, Si: Unpin> Unpin for Forward<St, Si> {} + +impl<St, Si, E> Forward<St, Si> +where + Si: Sink<St::Ok, Error = E>, + St: TryStream<Error = E> + Stream, +{ + unsafe_pinned!(sink: Option<Si>); + unsafe_pinned!(stream: Fuse<St>); + unsafe_unpinned!(buffered_item: Option<St::Ok>); + + pub(super) fn new(stream: St, sink: Si) -> Self { + Forward { + sink: Some(sink), + stream: stream.fuse(), + buffered_item: None, + } + } + + fn try_start_send( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + item: St::Ok, + ) -> Poll<Result<(), E>> { + debug_assert!(self.buffered_item.is_none()); + { + let mut sink = self.as_mut().sink().as_pin_mut().unwrap(); + if sink.as_mut().poll_ready(cx)?.is_ready() { + return Poll::Ready(sink.start_send(item)); + } + } + *self.as_mut().buffered_item() = Some(item); + Poll::Pending + } +} + +impl<St, Si, Item, E> FusedFuture for Forward<St, Si> +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> +where + Si: Sink<Item, Error = E>, + St: Stream<Item = Result<Item, E>>, +{ + type Output = Result<(), E>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + // If we've got an item buffered already, we need to write it to the + // sink before we can do anything else + if let Some(item) = self.as_mut().buffered_item().take() { + ready!(self.as_mut().try_start_send(cx, item))?; + } + + loop { + match self.as_mut().stream().poll_next(cx)? { + Poll::Ready(Some(item)) => + ready!(self.as_mut().try_start_send(cx, item))?, + Poll::Ready(None) => { + ready!(self.as_mut().sink().as_pin_mut().expect(INVALID_POLL).poll_close(cx))?; + self.as_mut().sink().set(None); + return Poll::Ready(Ok(())) + } + Poll::Pending => { + ready!(self.as_mut().sink().as_pin_mut().expect(INVALID_POLL).poll_flush(cx))?; + return Poll::Pending + } + } + } + } +} diff --git a/src/stream/stream/fuse.rs b/src/stream/stream/fuse.rs new file mode 100644 index 0000000..9085dc5 --- /dev/null +++ b/src/stream/stream/fuse.rs @@ -0,0 +1,107 @@ +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`fuse`](super::StreamExt::fuse) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Fuse<St> { + stream: St, + done: bool, +} + +impl<St: Unpin> Unpin for Fuse<St> {} + +impl<St> Fuse<St> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(done: bool); + + pub(super) fn new(stream: St) -> Fuse<St> { + Fuse { 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 + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<S::Item>> { + if self.done { + return Poll::Ready(None); + } + + let item = ready!(self.as_mut().stream().poll_next(cx)); + if item.is_none() { + *self.as_mut().done() = true; + } + Poll::Ready(item) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.done { + (0, Some(0)) + } else { + self.stream.size_hint() + } + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S: Stream + Sink<Item>, Item> Sink<Item> for Fuse<S> { + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/inspect.rs b/src/stream/stream/inspect.rs new file mode 100644 index 0000000..e34970a --- /dev/null +++ b/src/stream/stream/inspect.rs @@ -0,0 +1,119 @@ +use core::fmt; +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`inspect`](super::StreamExt::inspect) method. +#[must_use = "streams do nothing unless polled"] +pub struct Inspect<St, F> { + stream: St, + f: F, +} + +impl<St: Unpin, F> Unpin for Inspect<St, F> {} + +impl<St, F> fmt::Debug for Inspect<St, F> +where + St: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Inspect") + .field("stream", &self.stream) + .finish() + } +} + +impl<St, F> Inspect<St, F> + where St: Stream, + F: FnMut(&St::Item), +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + + pub(super) fn new(stream: St, f: F) -> Inspect<St, F> { + Inspect { stream, f } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +impl<St, F> FusedStream for Inspect<St, F> + where St: FusedStream, + F: FnMut(&St::Item), +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +// used by `TryStreamExt::{inspect_ok, inspect_err}` +#[inline] +pub(crate) fn inspect<T, F: FnMut(&T)>(x: T, mut f: F) -> T { + f(&x); + x +} + +impl<St, F> Stream for Inspect<St, F> + where St: Stream, + F: FnMut(&St::Item), +{ + type Item = St::Item; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<St::Item>> { + self.as_mut() + .stream() + .poll_next(cx) + .map(|opt| opt.map(|e| inspect(e, self.as_mut().f()))) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, F, Item> Sink<Item> for Inspect<S, F> + where S: Stream + Sink<Item>, + F: FnMut(&S::Item), +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/into_future.rs b/src/stream/stream/into_future.rs new file mode 100644 index 0000000..abae98c --- /dev/null +++ b/src/stream/stream/into_future.rs @@ -0,0 +1,97 @@ +use crate::stream::StreamExt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_utils::unsafe_pinned; + +/// 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: Unpin> Unpin for StreamFuture<St> {} + +impl<St: Stream + Unpin> StreamFuture<St> { + unsafe_pinned!(stream: Option<St>); + + pub(super) fn new(stream: St) -> StreamFuture<St> { + StreamFuture { 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.stream().as_pin_mut() + } + + /// Consumes this combinator, returning the underlying stream. + /// + /// Note that this may discard intermediate state of this combinator, so + /// care should be taken to avoid losing resources when this is called. + /// + /// This method returns an `Option` to account for the fact that `StreamFuture`'s + /// implementation of `Future::poll` consumes the underlying stream during polling + /// in order to return it to the caller of `Future::poll` if the stream yielded + /// an element. + pub fn into_inner(self) -> Option<St> { + self.stream + } +} + +impl<St: Stream + Unpin> FusedFuture for StreamFuture<St> { + fn is_terminated(&self) -> bool { + self.stream.is_none() + } +} + +impl<St: Stream + Unpin> Future for StreamFuture<St> { + type Output = (Option<St::Item>, St); + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + let item = { + let s = self.stream.as_mut().expect("polling StreamFuture twice"); + ready!(s.poll_next_unpin(cx)) + }; + let stream = self.stream.take().unwrap(); + Poll::Ready((item, stream)) + } +} diff --git a/src/stream/stream/map.rs b/src/stream/stream/map.rs new file mode 100644 index 0000000..8119434 --- /dev/null +++ b/src/stream/stream/map.rs @@ -0,0 +1,112 @@ +use core::fmt; +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`map`](super::StreamExt::map) method. +#[must_use = "streams do nothing unless polled"] +pub struct Map<St, F> { + stream: St, + f: F, +} + +impl<St: Unpin, F> Unpin for Map<St, 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, T, F> Map<St, F> + where St: Stream, + F: FnMut(St::Item) -> T, +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + + pub(super) fn new(stream: St, f: F) -> Map<St, F> { + Map { stream, f } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +impl<St, F, T> FusedStream for Map<St, F> + where St: FusedStream, + F: FnMut(St::Item) -> T, +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St, F, T> Stream for Map<St, F> + where St: Stream, + F: FnMut(St::Item) -> T, +{ + type Item = T; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<T>> { + self.as_mut() + .stream() + .poll_next(cx) + .map(|opt| opt.map(|x| self.as_mut().f()(x))) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, F, T, Item> Sink<Item> for Map<S, F> + where S: Stream + Sink<Item>, + F: FnMut(S::Item) -> T, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/mod.rs b/src/stream/stream/mod.rs new file mode 100644 index 0000000..da5ade8 --- /dev/null +++ b/src/stream/stream/mod.rs @@ -0,0 +1,1241 @@ +//! Streams +//! +//! This module contains a number of functions for working with `Stream`s, +//! including the `StreamExt` trait which adds methods to `Stream` types. + +use crate::future::Either; +#[cfg(feature = "alloc")] +use alloc::boxed::Box; +use core::pin::Pin; +#[cfg(feature = "sink")] +use futures_core::stream::TryStream; +#[cfg(feature = "alloc")] +use futures_core::stream::{BoxStream, LocalBoxStream}; +use futures_core::{ + future::Future, + stream::{FusedStream, Stream}, + task::{Context, Poll}, +}; +#[cfg(feature = "sink")] +use futures_sink::Sink; + +mod chain; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::chain::Chain; + +mod collect; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::collect::Collect; + +mod concat; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::concat::Concat; + +mod enumerate; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::enumerate::Enumerate; + +mod filter; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::filter::Filter; + +mod filter_map; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::filter_map::FilterMap; + +mod flatten; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::flatten::Flatten; + +mod fold; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::fold::Fold; + +#[cfg(feature = "sink")] +mod forward; +#[cfg(feature = "sink")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::forward::Forward; + +mod for_each; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::for_each::ForEach; + +mod fuse; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::fuse::Fuse; + +mod into_future; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::into_future::StreamFuture; + +mod inspect; +pub(crate) use self::inspect::inspect; // used by `TryStreamExt::{inspect_ok, inspect_err}` +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::inspect::Inspect; + +mod map; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::map::Map; + +mod next; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::next::Next; + +mod select_next_some; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::select_next_some::SelectNextSome; + +mod peek; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::peek::{Peek, Peekable}; + +mod skip; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::skip::Skip; + +mod skip_while; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::skip_while::SkipWhile; + +mod take; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::take::Take; + +mod take_while; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::take_while::TakeWhile; + +mod then; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::then::Then; + +mod zip; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::zip::Zip; + +#[cfg(feature = "alloc")] +mod chunks; +#[cfg(feature = "alloc")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::chunks::Chunks; + +mod scan; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::scan::Scan; + +cfg_target_has_atomic! { + #[cfg(feature = "alloc")] + mod buffer_unordered; + #[cfg(feature = "alloc")] + #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 + pub use self::buffer_unordered::BufferUnordered; + + #[cfg(feature = "alloc")] + mod buffered; + #[cfg(feature = "alloc")] + #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 + pub use self::buffered::Buffered; + + #[cfg(feature = "alloc")] + mod for_each_concurrent; + #[cfg(feature = "alloc")] + #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 + pub use self::for_each_concurrent::ForEachConcurrent; + + #[cfg(feature = "sink")] + #[cfg(feature = "alloc")] + mod split; + #[cfg(feature = "sink")] + #[cfg(feature = "alloc")] + #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 + pub use self::split::{SplitStream, SplitSink, ReuniteError}; +} + +#[cfg(feature = "std")] +mod catch_unwind; +#[cfg(feature = "std")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::catch_unwind::CatchUnwind; + +impl<T: ?Sized> StreamExt for T where T: Stream {} + +/// An extension trait for `Stream`s that provides a variety of convenient +/// combinator functions. +pub trait StreamExt: Stream { + /// Creates a future that resolves to the next item in the stream. + /// + /// Note that because `next` doesn't take ownership over the stream, + /// the [`Stream`] type must be [`Unpin`]. If you want to use `next` with a + /// [`!Unpin`](Unpin) stream, you'll first have to pin the stream. This can + /// be done by boxing the stream using [`Box::pin`] or + /// pinning it to the stack using the `pin_mut!` macro from the `pin_utils` + /// crate. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let mut stream = stream::iter(1..=3); + /// + /// assert_eq!(stream.next().await, Some(1)); + /// assert_eq!(stream.next().await, Some(2)); + /// assert_eq!(stream.next().await, Some(3)); + /// assert_eq!(stream.next().await, None); + /// # }); + /// ``` + fn next(&mut self) -> Next<'_, Self> + where + Self: Unpin, + { + Next::new(self) + } + + /// Converts this stream into a future of `(next_item, tail_of_stream)`. + /// If the stream terminates, then the next item is [`None`]. + /// + /// The returned future can be used to compose streams and futures together + /// by placing everything into the "world of futures". + /// + /// Note that because `into_future` moves the stream, the [`Stream`] type + /// must be [`Unpin`]. If you want to use `into_future` with a + /// [`!Unpin`](Unpin) stream, you'll first have to pin the stream. This can + /// be done by boxing the stream using [`Box::pin`] or + /// pinning it to the stack using the `pin_mut!` macro from the `pin_utils` + /// crate. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=3); + /// + /// let (item, stream) = stream.into_future().await; + /// assert_eq!(Some(1), item); + /// + /// let (item, stream) = stream.into_future().await; + /// assert_eq!(Some(2), item); + /// # }); + /// ``` + fn into_future(self) -> StreamFuture<Self> + where + Self: Sized + Unpin, + { + StreamFuture::new(self) + } + + /// Maps this stream's items to a different type, returning a new stream of + /// the resulting type. + /// + /// The provided closure is executed over all elements of this stream as + /// they are made available. It is executed inline with calls to + /// [`poll_next`](Stream::poll_next). + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `map` methods in the + /// standard library. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=3); + /// let stream = stream.map(|x| x + 3); + /// + /// assert_eq!(vec![4, 5, 6], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn map<T, F>(self, f: F) -> Map<Self, F> + where + F: FnMut(Self::Item) -> T, + Self: Sized, + { + Map::new(self, f) + } + + /// Creates a stream which gives the current iteration count as well as + /// the next value. + /// + /// The stream returned yields pairs `(i, val)`, where `i` is the + /// current index of iteration and `val` is the value returned by the + /// stream. + /// + /// `enumerate()` keeps its count as a [`usize`]. If you want to count by a + /// different sized integer, the [`zip`](StreamExt::zip) function provides similar + /// functionality. + /// + /// # Overflow Behavior + /// + /// The method does no guarding against overflows, so enumerating more than + /// [`usize::max_value()`] elements either produces the wrong result or panics. If + /// debug assertions are enabled, a panic is guaranteed. + /// + /// # Panics + /// + /// The returned stream might panic if the to-be-returned index would + /// overflow a [`usize`]. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(vec!['a', 'b', 'c']); + /// + /// let mut stream = stream.enumerate(); + /// + /// assert_eq!(stream.next().await, Some((0, 'a'))); + /// assert_eq!(stream.next().await, Some((1, 'b'))); + /// assert_eq!(stream.next().await, Some((2, 'c'))); + /// assert_eq!(stream.next().await, None); + /// # }); + /// ``` + fn enumerate(self) -> Enumerate<Self> + where + Self: Sized, + { + Enumerate::new(self) + } + + /// Filters the values produced by this stream according to the provided + /// asynchronous predicate. + /// + /// As values of this stream are made available, the provided predicate `f` + /// will be run against them. If the predicate returns a `Future` which + /// resolves to `true`, then the stream will yield the value, but if the + /// predicate returns a `Future` which resolves to `false`, then the value + /// will be discarded and the next value will be produced. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `filter` methods in the + /// standard library. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// let evens = stream.filter(|x| future::ready(x % 2 == 0)); + /// + /// assert_eq!(vec![2, 4, 6, 8, 10], evens.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn filter<Fut, F>(self, f: F) -> Filter<Self, Fut, F> + where + F: FnMut(&Self::Item) -> Fut, + Fut: Future<Output = bool>, + Self: Sized, + { + Filter::new(self, f) + } + + /// Filters the values produced by this stream while simultaneously mapping + /// them to a different type according to the provided asynchronous closure. + /// + /// As values of this stream are made available, the provided function will + /// be run on them. If the future returned by the predicate `f` resolves to + /// [`Some(item)`](Some) then the stream will yield the value `item`, but if + /// it resolves to [`None`] then the next value will be produced. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `filter_map` methods in + /// the standard library. + /// + /// # Examples + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// let evens = stream.filter_map(|x| async move { + /// if x % 2 == 0 { Some(x + 1) } else { None } + /// }); + /// + /// assert_eq!(vec![3, 5, 7, 9, 11], evens.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn filter_map<Fut, T, F>(self, f: F) -> FilterMap<Self, Fut, F> + where + F: FnMut(Self::Item) -> Fut, + Fut: Future<Output = Option<T>>, + Self: Sized, + { + FilterMap::new(self, f) + } + + /// Computes from this stream's items new items of a different type using + /// an asynchronous closure. + /// + /// The provided closure `f` will be called with an `Item` once a value is + /// ready, it returns a future which will then be run to completion + /// to produce the next value on this stream. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=3); + /// let stream = stream.then(|x| async move { x + 3 }); + /// + /// assert_eq!(vec![4, 5, 6], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn then<Fut, F>(self, f: F) -> Then<Self, Fut, F> + where + F: FnMut(Self::Item) -> Fut, + Fut: Future, + Self: Sized, + { + Then::new(self, f) + } + + /// Transforms a stream into a collection, returning a + /// future representing the result of that computation. + /// + /// The returned future will be resolved when the stream terminates. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::StreamExt; + /// use std::thread; + /// + /// let (tx, rx) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// for i in 1..=5 { + /// tx.unbounded_send(i).unwrap(); + /// } + /// }); + /// + /// let output = rx.collect::<Vec<i32>>().await; + /// assert_eq!(output, vec![1, 2, 3, 4, 5]); + /// # }); + /// ``` + fn collect<C: Default + Extend<Self::Item>>(self) -> Collect<Self, C> + where + Self: Sized, + { + Collect::new(self) + } + + /// Concatenate all items of a stream into a single extendable + /// destination, returning a future representing the end result. + /// + /// This combinator will extend the first item with the contents + /// of all the subsequent results of the stream. If the stream is + /// empty, the default value will be returned. + /// + /// Works with all collections that implement the + /// [`Extend`](std::iter::Extend) trait. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::StreamExt; + /// use std::thread; + /// + /// let (tx, rx) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// for i in (0..3).rev() { + /// let n = i * 3; + /// tx.unbounded_send(vec![n + 1, n + 2, n + 3]).unwrap(); + /// } + /// }); + /// + /// let result = rx.concat().await; + /// + /// assert_eq!(result, vec![7, 8, 9, 4, 5, 6, 1, 2, 3]); + /// # }); + /// ``` + fn concat(self) -> Concat<Self> + where + Self: Sized, + Self::Item: Extend<<<Self as Stream>::Item as IntoIterator>::Item> + IntoIterator + Default, + { + Concat::new(self) + } + + /// Execute an accumulating asynchronous computation over a stream, + /// collecting all the values into one final result. + /// + /// This combinator will accumulate all values returned by this stream + /// according to the closure provided. The initial state is also provided to + /// this method and then is returned again by each execution of the closure. + /// Once the entire stream has been exhausted the returned future will + /// resolve to this value. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let number_stream = stream::iter(0..6); + /// let sum = number_stream.fold(0, |acc, x| async move { acc + x }); + /// assert_eq!(sum.await, 15); + /// # }); + /// ``` + fn fold<T, Fut, F>(self, init: T, f: F) -> Fold<Self, Fut, T, F> + where + F: FnMut(T, Self::Item) -> Fut, + Fut: Future<Output = T>, + Self: Sized, + { + Fold::new(self, f, init) + } + + /// Flattens a stream of streams into just one continuous stream. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::StreamExt; + /// use std::thread; + /// + /// let (tx1, rx1) = mpsc::unbounded(); + /// let (tx2, rx2) = mpsc::unbounded(); + /// let (tx3, rx3) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// tx1.unbounded_send(1).unwrap(); + /// tx1.unbounded_send(2).unwrap(); + /// }); + /// thread::spawn(move || { + /// tx2.unbounded_send(3).unwrap(); + /// tx2.unbounded_send(4).unwrap(); + /// }); + /// thread::spawn(move || { + /// tx3.unbounded_send(rx1).unwrap(); + /// tx3.unbounded_send(rx2).unwrap(); + /// }); + /// + /// let output = rx3.flatten().collect::<Vec<i32>>().await; + /// assert_eq!(output, vec![1, 2, 3, 4]); + /// # }); + /// ``` + fn flatten(self) -> Flatten<Self> + where + Self::Item: Stream, + Self: Sized, + { + Flatten::new(self) + } + + /// Combinator similar to [`StreamExt::fold`] that holds internal state and produces a new stream. + /// + /// Accepts initial state and closure which will be applied to each element of the stream until provided closure + /// returns `None`. Once `None` is returned, stream will be terminated. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// + /// let stream = stream.scan(0, |state, x| { + /// *state += x; + /// future::ready(if *state < 10 { Some(x) } else { None }) + /// }); + /// + /// assert_eq!(vec![1, 2, 3], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn scan<S, B, Fut, F>(self, initial_state: S, f: F) -> Scan<Self, S, Fut, F> + where + F: FnMut(&mut S, Self::Item) -> Fut, + Fut: Future<Output = Option<B>>, + Self: Sized, + { + Scan::new(self, initial_state, f) + } + + /// Skip elements on this stream while the provided asynchronous predicate + /// resolves to `true`. + /// + /// This function, like `Iterator::skip_while`, will skip elements on the + /// stream until the predicate `f` resolves to `false`. Once one element + /// returns false all future elements will be returned from the underlying + /// stream. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// + /// let stream = stream.skip_while(|x| future::ready(*x <= 5)); + /// + /// assert_eq!(vec![6, 7, 8, 9, 10], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn skip_while<Fut, F>(self, f: F) -> SkipWhile<Self, Fut, F> + where + F: FnMut(&Self::Item) -> Fut, + Fut: Future<Output = bool>, + Self: Sized, + { + SkipWhile::new(self, f) + } + + /// Take elements from this stream while the provided asynchronous predicate + /// resolves to `true`. + /// + /// This function, like `Iterator::take_while`, will take elements from the + /// stream until the predicate `f` resolves to `false`. Once one element + /// returns false it will always return that the stream is done. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10); + /// + /// let stream = stream.take_while(|x| future::ready(*x <= 5)); + /// + /// assert_eq!(vec![1, 2, 3, 4, 5], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn take_while<Fut, F>(self, f: F) -> TakeWhile<Self, Fut, F> + where + F: FnMut(&Self::Item) -> Fut, + Fut: Future<Output = bool>, + Self: Sized, + { + TakeWhile::new(self, f) + } + + /// Runs this stream to completion, executing the provided asynchronous + /// closure for each element on the stream. + /// + /// The closure provided will be called for each item this stream produces, + /// yielding a future. That future will then be executed to completion + /// before moving on to the next item. + /// + /// The returned value is a `Future` where the `Output` type is `()`; it is + /// executed entirely for its side effects. + /// + /// To process each item in the stream and produce another stream instead + /// of a single future, use `then` instead. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt}; + /// + /// let mut x = 0; + /// + /// { + /// let fut = stream::repeat(1).take(3).for_each(|item| { + /// x += item; + /// future::ready(()) + /// }); + /// fut.await; + /// } + /// + /// assert_eq!(x, 3); + /// # }); + /// ``` + fn for_each<Fut, F>(self, f: F) -> ForEach<Self, Fut, F> + where + F: FnMut(Self::Item) -> Fut, + Fut: Future<Output = ()>, + Self: Sized, + { + ForEach::new(self, f) + } + + /// Runs this stream to completion, executing the provided asynchronous + /// closure for each element on the stream concurrently as elements become + /// available. + /// + /// This is similar to [`StreamExt::for_each`], but the futures + /// produced by the closure are run concurrently (but not in parallel-- + /// this combinator does not introduce any threads). + /// + /// The closure provided will be called for each item this stream produces, + /// yielding a future. That future will then be executed to completion + /// concurrently with the other futures produced by the closure. + /// + /// The first argument is an optional limit on the number of concurrent + /// futures. If this limit is not `None`, no more than `limit` futures + /// will be run concurrently. The `limit` argument is of type + /// `Into<Option<usize>>`, and so can be provided as either `None`, + /// `Some(10)`, or just `10`. Note: a limit of zero is interpreted as + /// no limit at all, and will have the same result as passing in `None`. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::oneshot; + /// use futures::stream::{self, StreamExt}; + /// + /// let (tx1, rx1) = oneshot::channel(); + /// let (tx2, rx2) = oneshot::channel(); + /// let (tx3, rx3) = oneshot::channel(); + /// + /// let fut = stream::iter(vec![rx1, rx2, rx3]).for_each_concurrent( + /// /* limit */ 2, + /// |rx| async move { + /// rx.await.unwrap(); + /// } + /// ); + /// tx1.send(()).unwrap(); + /// tx2.send(()).unwrap(); + /// tx3.send(()).unwrap(); + /// fut.await; + /// # }) + /// ``` + #[cfg_attr(feature = "cfg-target-has-atomic", cfg(target_has_atomic = "ptr"))] + #[cfg(feature = "alloc")] + fn for_each_concurrent<Fut, F>( + self, + limit: impl Into<Option<usize>>, + f: F, + ) -> ForEachConcurrent<Self, Fut, F> + where + F: FnMut(Self::Item) -> Fut, + Fut: Future<Output = ()>, + Self: Sized, + { + ForEachConcurrent::new(self, limit.into(), f) + } + + /// Creates a new stream of at most `n` items of the underlying stream. + /// + /// Once `n` items have been yielded from this stream then it will always + /// return that the stream is done. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10).take(3); + /// + /// assert_eq!(vec![1, 2, 3], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn take(self, n: usize) -> Take<Self> + where + Self: Sized, + { + Take::new(self, n) + } + + /// Creates a new stream which skips `n` items of the underlying stream. + /// + /// Once `n` items have been skipped from this stream then it will always + /// return the remaining items on this stream. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=10).skip(5); + /// + /// assert_eq!(vec![6, 7, 8, 9, 10], stream.collect::<Vec<_>>().await); + /// # }); + /// ``` + fn skip(self, n: usize) -> Skip<Self> + where + Self: Sized, + { + Skip::new(self, n) + } + + /// Fuse a stream such that [`poll_next`](Stream::poll_next) will never + /// again be called once it has finished. This method can be used to turn + /// any `Stream` into a `FusedStream`. + /// + /// Normally, once a stream has returned [`None`] from + /// [`poll_next`](Stream::poll_next) any further calls could exhibit bad + /// behavior such as block forever, panic, never return, etc. If it is known + /// that [`poll_next`](Stream::poll_next) may be called after stream + /// has already finished, then this method can be used to ensure that it has + /// defined semantics. + /// + /// The [`poll_next`](Stream::poll_next) method of a `fuse`d stream + /// is guaranteed to return [`None`] after the underlying stream has + /// finished. + /// + /// # Examples + /// + /// ``` + /// use futures::executor::block_on_stream; + /// use futures::stream::{self, StreamExt}; + /// use futures::task::Poll; + /// + /// let mut x = 0; + /// let stream = stream::poll_fn(|_| { + /// x += 1; + /// match x { + /// 0..=2 => Poll::Ready(Some(x)), + /// 3 => Poll::Ready(None), + /// _ => panic!("should not happen") + /// } + /// }).fuse(); + /// + /// let mut iter = block_on_stream(stream); + /// assert_eq!(Some(1), iter.next()); + /// assert_eq!(Some(2), iter.next()); + /// assert_eq!(None, iter.next()); + /// assert_eq!(None, iter.next()); + /// // ... + /// ``` + fn fuse(self) -> Fuse<Self> + where + Self: Sized, + { + Fuse::new(self) + } + + /// Borrows a stream, rather than consuming it. + /// + /// This is useful to allow applying stream adaptors while still retaining + /// ownership of the original stream. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let mut stream = stream::iter(1..5); + /// + /// let sum = stream.by_ref() + /// .take(2) + /// .fold(0, |a, b| async move { a + b }) + /// .await; + /// assert_eq!(sum, 3); + /// + /// // You can use the stream again + /// let sum = stream.take(2) + /// .fold(0, |a, b| async move { a + b }) + /// .await; + /// assert_eq!(sum, 7); + /// # }); + /// ``` + fn by_ref(&mut self) -> &mut Self { + self + } + + /// Catches unwinding panics while polling the stream. + /// + /// Caught panic (if any) will be the last element of the resulting stream. + /// + /// In general, panics within a stream can propagate all the way out to the + /// task level. This combinator makes it possible to halt unwinding within + /// the stream itself. It's most commonly used within task executors. This + /// method should not be used for error handling. + /// + /// Note that this method requires the `UnwindSafe` bound from the standard + /// library. This isn't always applied automatically, and the standard + /// library provides an `AssertUnwindSafe` wrapper type to apply it + /// after-the fact. To assist using this method, the [`Stream`] trait is + /// also implemented for `AssertUnwindSafe<St>` where `St` implements + /// [`Stream`]. + /// + /// This method is only available when the `std` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(vec![Some(10), None, Some(11)]); + /// // Panic on second element + /// let stream_panicking = stream.map(|o| o.unwrap()); + /// // Collect all the results + /// let stream = stream_panicking.catch_unwind(); + /// + /// let results: Vec<Result<i32, _>> = stream.collect().await; + /// match results[0] { + /// Ok(10) => {} + /// _ => panic!("unexpected result!"), + /// } + /// assert!(results[1].is_err()); + /// assert_eq!(results.len(), 2); + /// # }); + /// ``` + #[cfg(feature = "std")] + fn catch_unwind(self) -> CatchUnwind<Self> + where + Self: Sized + std::panic::UnwindSafe, + { + CatchUnwind::new(self) + } + + /// Wrap the stream in a Box, pinning it. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg(feature = "alloc")] + fn boxed<'a>(self) -> BoxStream<'a, Self::Item> + where + Self: Sized + Send + 'a, + { + Box::pin(self) + } + + /// Wrap the stream in a Box, pinning it. + /// + /// Similar to `boxed`, but without the `Send` requirement. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg(feature = "alloc")] + fn boxed_local<'a>(self) -> LocalBoxStream<'a, Self::Item> + where + Self: Sized + 'a, + { + Box::pin(self) + } + + /// An adaptor for creating a buffered list of pending futures. + /// + /// If this stream's item can be converted into a future, then this adaptor + /// will buffer up to at most `n` futures and then return the outputs in the + /// same order as the underlying stream. No more than `n` futures will be + /// buffered at any point in time, and less than `n` may also be buffered + /// depending on the state of each future. + /// + /// The returned stream will be a stream of each future's output. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg_attr(feature = "cfg-target-has-atomic", cfg(target_has_atomic = "ptr"))] + #[cfg(feature = "alloc")] + fn buffered(self, n: usize) -> Buffered<Self> + where + Self::Item: Future, + Self: Sized, + { + Buffered::new(self, n) + } + + /// An adaptor for creating a buffered list of pending futures (unordered). + /// + /// If this stream's item can be converted into a future, then this adaptor + /// will buffer up to `n` futures and then return the outputs in the order + /// in which they complete. No more than `n` futures will be buffered at + /// any point in time, and less than `n` may also be buffered depending on + /// the state of each future. + /// + /// The returned stream will be a stream of each future's output. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::oneshot; + /// use futures::stream::{self, StreamExt}; + /// + /// let (send_one, recv_one) = oneshot::channel(); + /// let (send_two, recv_two) = oneshot::channel(); + /// + /// let stream_of_futures = stream::iter(vec![recv_one, recv_two]); + /// let mut buffered = stream_of_futures.buffer_unordered(10); + /// + /// send_two.send(2i32)?; + /// assert_eq!(buffered.next().await, Some(Ok(2i32))); + /// + /// send_one.send(1i32)?; + /// assert_eq!(buffered.next().await, Some(Ok(1i32))); + /// + /// assert_eq!(buffered.next().await, None); + /// # Ok::<(), i32>(()) }).unwrap(); + /// ``` + #[cfg_attr(feature = "cfg-target-has-atomic", cfg(target_has_atomic = "ptr"))] + #[cfg(feature = "alloc")] + fn buffer_unordered(self, n: usize) -> BufferUnordered<Self> + where + Self::Item: Future, + Self: Sized, + { + BufferUnordered::new(self, n) + } + + /// An adapter for zipping two streams together. + /// + /// The zipped stream waits for both streams to produce an item, and then + /// returns that pair. If either stream ends then the zipped stream will + /// also end. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream1 = stream::iter(1..=3); + /// let stream2 = stream::iter(5..=10); + /// + /// let vec = stream1.zip(stream2) + /// .collect::<Vec<_>>() + /// .await; + /// assert_eq!(vec![(1, 5), (2, 6), (3, 7)], vec); + /// # }); + /// ``` + /// + fn zip<St>(self, other: St) -> Zip<Self, St> + where + St: Stream, + Self: Sized, + { + Zip::new(self, other) + } + + /// Adapter for chaining two streams. + /// + /// The resulting stream emits elements from the first stream, and when + /// first stream reaches the end, emits the elements from the second stream. + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// + /// let stream1 = stream::iter(vec![Ok(10), Err(false)]); + /// let stream2 = stream::iter(vec![Err(true), Ok(20)]); + /// + /// let stream = stream1.chain(stream2); + /// + /// let result: Vec<_> = stream.collect().await; + /// assert_eq!(result, vec![ + /// Ok(10), + /// Err(false), + /// Err(true), + /// Ok(20), + /// ]); + /// # }); + /// ``` + fn chain<St>(self, other: St) -> Chain<Self, St> + where + St: Stream<Item = Self::Item>, + Self: Sized, + { + Chain::new(self, other) + } + + /// Creates a new stream which exposes a `peek` method. + /// + /// Calling `peek` returns a reference to the next item in the stream. + fn peekable(self) -> Peekable<Self> + where + Self: Sized, + { + Peekable::new(self) + } + + /// An adaptor for chunking up items of the stream inside a vector. + /// + /// This combinator will attempt to pull items from this stream and buffer + /// them into a local vector. At most `capacity` items will get buffered + /// before they're yielded from the returned stream. + /// + /// Note that the vectors returned from this iterator may not always have + /// `capacity` elements. If the underlying stream ended and only a partial + /// vector was created, it'll be returned. Additionally if an error happens + /// from the underlying stream then the currently buffered items will be + /// yielded. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Panics + /// + /// This method will panic if `capacity` is zero. + #[cfg(feature = "alloc")] + fn chunks(self, capacity: usize) -> Chunks<Self> + where + Self: Sized, + { + Chunks::new(self, capacity) + } + + /// A future that completes after the given stream has been fully processed + /// into the sink and the sink has been flushed and closed. + /// + /// This future will drive the stream to keep producing items until it is + /// exhausted, sending each item to the sink. It will complete once the + /// stream is exhausted, the sink has received and flushed all items, and + /// the sink is closed. Note that neither the original stream nor provided + /// sink will be output by this future. Pass the sink by `Pin<&mut S>` + /// (for example, via `forward(&mut sink)` inside an `async` fn/block) in + /// order to preserve access to the `Sink`. + #[cfg(feature = "sink")] + fn forward<S>(self, sink: S) -> Forward<Self, S> + where + S: Sink<<Self as TryStream>::Ok>, + Self: TryStream<Error = S::Error> + Sized, + { + Forward::new(self, sink) + } + + /// Splits this `Stream + Sink` object into separate `Stream` and `Sink` + /// objects. + /// + /// This can be useful when you want to split ownership between tasks, or + /// allow direct interaction between the two objects (e.g. via + /// `Sink::send_all`). + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + #[cfg(feature = "sink")] + #[cfg_attr(feature = "cfg-target-has-atomic", cfg(target_has_atomic = "ptr"))] + #[cfg(feature = "alloc")] + fn split<Item>(self) -> (SplitSink<Self, Item>, SplitStream<Self>) + where + Self: Sink<Item> + Sized, + { + split::split(self) + } + + /// Do something with each item of this stream, afterwards passing it on. + /// + /// This is similar to the `Iterator::inspect` method in the standard + /// library where it allows easily inspecting each value as it passes + /// through the stream, for example to debug what's going on. + fn inspect<F>(self, f: F) -> Inspect<Self, F> + where + F: FnMut(&Self::Item), + Self: Sized, + { + Inspect::new(self, f) + } + + /// Wrap this stream in an `Either` stream, making it the left-hand variant + /// of that `Either`. + /// + /// This can be used in combination with the `right_stream` method to write `if` + /// statements that evaluate to different streams in different branches. + fn left_stream<B>(self) -> Either<Self, B> + where + B: Stream<Item = Self::Item>, + Self: Sized, + { + Either::Left(self) + } + + /// Wrap this stream in an `Either` stream, making it the right-hand variant + /// of that `Either`. + /// + /// This can be used in combination with the `left_stream` method to write `if` + /// statements that evaluate to different streams in different branches. + fn right_stream<B>(self) -> Either<B, Self> + where + B: Stream<Item = Self::Item>, + Self: Sized, + { + Either::Right(self) + } + + /// A convenience method for calling [`Stream::poll_next`] on [`Unpin`] + /// stream types. + fn poll_next_unpin(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> + where + Self: Unpin, + { + Pin::new(self).poll_next(cx) + } + + /// Returns a [`Future`] that resolves when the next item in this stream is + /// ready. + /// + /// This is similar to the [`next`][StreamExt::next] method, but it won't + /// resolve to [`None`] if used on an empty [`Stream`]. Instead, the + /// returned future type will return `true` from + /// [`FusedFuture::is_terminated`][] when the [`Stream`] is empty, allowing + /// [`select_next_some`][StreamExt::select_next_some] to be easily used with + /// the [`select!`] macro. + /// + /// If the future is polled after this [`Stream`] is empty it will panic. + /// Using the future with a [`FusedFuture`][]-aware primitive like the + /// [`select!`] macro will prevent this. + /// + /// [`FusedFuture`]: futures_core::future::FusedFuture + /// [`FusedFuture::is_terminated`]: futures_core::future::FusedFuture::is_terminated + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::{future, select}; + /// use futures::stream::{StreamExt, FuturesUnordered}; + /// + /// let mut fut = future::ready(1); + /// let mut async_tasks = FuturesUnordered::new(); + /// let mut total = 0; + /// loop { + /// select! { + /// num = fut => { + /// // First, the `ready` future completes. + /// total += num; + /// // Then we spawn a new task onto `async_tasks`, + /// async_tasks.push(async { 5 }); + /// }, + /// // On the next iteration of the loop, the task we spawned + /// // completes. + /// num = async_tasks.select_next_some() => { + /// total += num; + /// } + /// // Finally, both the `ready` future and `async_tasks` have + /// // finished, so we enter the `complete` branch. + /// complete => break, + /// } + /// } + /// assert_eq!(total, 6); + /// # }); + /// ``` + fn select_next_some(&mut self) -> SelectNextSome<'_, Self> + where + Self: Unpin + FusedStream, + { + SelectNextSome::new(self) + } +} diff --git a/src/stream/stream/next.rs b/src/stream/stream/next.rs new file mode 100644 index 0000000..2d74632 --- /dev/null +++ b/src/stream/stream/next.rs @@ -0,0 +1,37 @@ +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 { + Next { stream } + } +} + +impl<St: ?Sized + FusedStream + Unpin> FusedFuture for Next<'_, St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St: ?Sized + Stream + Unpin> Future for Next<'_, St> { + type Output = Option<St::Item>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + self.stream.poll_next_unpin(cx) + } +} diff --git a/src/stream/stream/peek.rs b/src/stream/stream/peek.rs new file mode 100644 index 0000000..9272baf --- /dev/null +++ b/src/stream/stream/peek.rs @@ -0,0 +1,197 @@ +use crate::future::Either; +use crate::stream::{Fuse, StreamExt}; +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: Fuse<St>, + peeked: Option<St::Item>, +} + +impl<St: Stream + Unpin> Unpin for Peekable<St> {} + +impl<St: Stream> Peekable<St> { + unsafe_pinned!(stream: Fuse<St>); + unsafe_unpinned!(peeked: Option<St::Item>); + + pub(super) fn new(stream: St) -> Peekable<St> { + Peekable { + stream: stream.fuse(), + peeked: None, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + self.stream.get_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. + pub fn get_mut(&mut self) -> &mut St { + self.stream.get_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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream().get_pin_mut() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream.into_inner() + } + + /// Produces a `Peek` 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) } + } + + /// Attempt to poll the underlying stream, and return the mutable borrow + /// in case that is desirable to try for another time. + /// In case a peeking poll is successful, the reference to the next item + /// will be in the `Either::Right` variant; otherwise, the mutable borrow + /// will be in the `Either::Left` variant. + fn do_poll_peek( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Either<Pin<&mut Self>, Option<&St::Item>> { + if self.peeked.is_some() { + let this: &Self = self.into_ref().get_ref(); + return Either::Right(this.peeked.as_ref()); + } + match self.as_mut().stream().poll_next(cx) { + Poll::Ready(None) => Either::Right(None), + Poll::Ready(Some(item)) => { + *self.as_mut().peeked() = Some(item); + let this: &Self = self.into_ref().get_ref(); + Either::Right(this.peeked.as_ref()) + } + _ => Either::Left(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>> { + match self.do_poll_peek(cx) { + Either::Left(_) => Poll::Pending, + Either::Right(poll) => Poll::Ready(poll), + } + } +} + +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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + if let Some(item) = self.as_mut().peeked().take() { + return Poll::Ready(Some(item)); + } + self.as_mut().stream().poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let peek_len = if self.peeked.is_some() { 1 } else { 0 }; + 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); +} + +/// Future for the [`Peekable::peek()`](self::Peekable::peek) function from [`Peekable`] +#[must_use = "futures do nothing unless polled"] +pub struct Peek<'a, St: Stream> { + inner: Option<Pin<&'a mut Peekable<St>>>, +} + +impl<St: Stream> Unpin for Peek<'_, 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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + if let Some(peekable) = self.inner.take() { + match peekable.do_poll_peek(cx) { + Either::Left(peekable) => { + self.inner = Some(peekable); + Poll::Pending + } + Either::Right(peek) => Poll::Ready(peek), + } + } else { + panic!("Peek polled after completion") + } + } +} diff --git a/src/stream/stream/scan.rs b/src/stream/stream/scan.rs new file mode 100644 index 0000000..4f937f4 --- /dev/null +++ b/src/stream/stream/scan.rs @@ -0,0 +1,165 @@ +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_utils::{unsafe_pinned, unsafe_unpinned}; + +struct StateFn<S, F> { + state: S, + f: F, +} + +/// Stream for the [`scan`](super::StreamExt::scan) method. +#[must_use = "streams do nothing unless polled"] +pub struct Scan<St: Stream, S, Fut, F> { + stream: St, + state_f: Option<StateFn<S, F>>, + future: Option<Fut>, +} + +impl<St: Unpin + Stream, S, Fut: Unpin, F> Unpin for Scan<St, S, Fut, F> {} + +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> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(state_f: Option<StateFn<S, F>>); + unsafe_pinned!(future: Option<Fut>); + + /// 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) -> Scan<St, S, Fut, F> { + Scan { + stream, + state_f: Some(StateFn { + state: initial_state, + f, + }), + future: None, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<B>> { + if self.is_done_taking() { + return Poll::Ready(None); + } + + if self.future.is_none() { + let item = match ready!(self.as_mut().stream().poll_next(cx)) { + Some(e) => e, + None => return Poll::Ready(None), + }; + let state_f = self.as_mut().state_f().as_mut().unwrap(); + let fut = (state_f.f)(&mut state_f.state, item); + self.as_mut().future().set(Some(fut)); + } + + let item = ready!(self.as_mut().future().as_pin_mut().unwrap().poll(cx)); + self.as_mut().future().set(None); + + if item.is_none() { + self.as_mut().state_f().take(); + } + + Poll::Ready(item) + } + + 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<S, Fut, F, Item> Sink<Item> for Scan<S, S, Fut, F> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/select_next_some.rs b/src/stream/stream/select_next_some.rs new file mode 100644 index 0000000..884f252 --- /dev/null +++ b/src/stream/stream/select_next_some.rs @@ -0,0 +1,41 @@ +use core::pin::Pin; +use futures_core::stream::FusedStream; +use futures_core::future::{Future, FusedFuture}; +use futures_core::task::{Context, Poll}; +use crate::stream::StreamExt; + +/// 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 { + SelectNextSome { stream } + } +} + +impl<St: ?Sized + FusedStream + Unpin> FusedFuture for SelectNextSome<'_, St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St: ?Sized + FusedStream + Unpin> Future for SelectNextSome<'_, St> { + type Output = St::Item; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + assert!(!self.stream.is_terminated(), "SelectNextSome polled after terminated"); + + if let Some(item) = ready!(self.stream.poll_next_unpin(cx)) { + Poll::Ready(item) + } else { + debug_assert!(self.stream.is_terminated()); + cx.waker().wake_by_ref(); + Poll::Pending + } + } +} diff --git a/src/stream/stream/skip.rs b/src/stream/stream/skip.rs new file mode 100644 index 0000000..0b7c632 --- /dev/null +++ b/src/stream/stream/skip.rs @@ -0,0 +1,107 @@ +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`skip`](super::StreamExt::skip) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Skip<St> { + stream: St, + remaining: usize, +} + +impl<St: Unpin> Unpin for Skip<St> {} + +impl<St: Stream> Skip<St> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(remaining: usize); + + pub(super) fn new(stream: St, n: usize) -> Skip<St> { + Skip { + stream, + remaining: n, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<St::Item>> { + while self.remaining > 0 { + match ready!(self.as_mut().stream().poll_next(cx)) { + Some(_) => *self.as_mut().remaining() -= 1, + None => return Poll::Ready(None), + } + } + + self.as_mut().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 as usize); + let upper = match upper { + Some(x) => Some(x.saturating_sub(self.remaining as usize)), + None => None, + }; + + (lower, upper) + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Skip<S> +where + S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/skip_while.rs b/src/stream/stream/skip_while.rs new file mode 100644 index 0000000..666d9de --- /dev/null +++ b/src/stream/stream/skip_while.rs @@ -0,0 +1,160 @@ +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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 { + stream: St, + f: F, + pending_fut: Option<Fut>, + pending_item: Option<St::Item>, + done_skipping: bool, +} + +impl<St: Unpin + Stream, Fut: Unpin, F> Unpin for SkipWhile<St, Fut, F> {} + +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>, +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + unsafe_pinned!(pending_fut: Option<Fut>); + unsafe_unpinned!(pending_item: Option<St::Item>); + unsafe_unpinned!(done_skipping: bool); + + pub(super) fn new(stream: St, f: F) -> SkipWhile<St, Fut, F> { + SkipWhile { + stream, + f, + pending_fut: None, + pending_item: None, + done_skipping: false, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<St::Item>> { + if self.done_skipping { + return self.as_mut().stream().poll_next(cx); + } + + loop { + if self.pending_item.is_none() { + let item = match ready!(self.as_mut().stream().poll_next(cx)) { + Some(e) => e, + None => return Poll::Ready(None), + }; + let fut = (self.as_mut().f())(&item); + self.as_mut().pending_fut().set(Some(fut)); + *self.as_mut().pending_item() = Some(item); + } + + let skipped = ready!(self.as_mut().pending_fut().as_pin_mut().unwrap().poll(cx)); + let item = self.as_mut().pending_item().take().unwrap(); + self.as_mut().pending_fut().set(None); + + if !skipped { + *self.as_mut().done_skipping() = true; + return Poll::Ready(Some(item)) + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = if self.pending_item.is_some() { 1 } else { 0 }; + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for SkipWhile<S, Fut, F> + where S: Stream + Sink<Item>, + F: FnMut(&S::Item) -> Fut, + Fut: Future<Output = bool>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/split.rs b/src/stream/stream/split.rs new file mode 100644 index 0000000..4118b33 --- /dev/null +++ b/src/stream/stream/split.rs @@ -0,0 +1,139 @@ +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use futures_sink::Sink; +use core::fmt; +use core::pin::Pin; + +use crate::lock::BiLock; + +/// A `Stream` part of the split pair +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct SplitStream<S>(BiLock<S>); + +impl<S> Unpin for SplitStream<S> {} + +impl<S: Unpin> SplitStream<S> { + /// Attempts to put the two "halves" of a split `Stream + Sink` back + /// together. Succeeds only if the `SplitStream<S>` and `SplitSink<S>` are + /// a matching pair originating from the same call to `Stream::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(bad_style)] +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"] +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 `Stream::split`. + pub fn reunite(self, other: SplitStream<S>) -> Result<S, ReuniteError<S, Item>> { + self.lock.reunite(other.0).map_err(|err| { + ReuniteError(SplitSink(err.0), SplitStream(err.1)) + }) + } +} + +impl<S: Sink<Item>, Item> SplitSink<S, Item> { + fn poll_flush_slot(mut inner: Pin<&mut S>, slot: &mut Option<Item>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> { + if slot.is_some() { + ready!(inner.as_mut().poll_ready(cx))?; + Poll::Ready(inner.start_send(slot.take().unwrap())) + } else { + Poll::Ready(Ok(())) + } + } + + fn poll_lock_and_flush_slot(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> { + let this = &mut *self; + let mut inner = ready!(this.lock.poll_lock(cx)); + Self::poll_flush_slot(inner.as_pin_mut(), &mut this.slot, cx) + } +} + +impl<S: Sink<Item>, Item> Sink<Item> for SplitSink<S, Item> { + type Error = S::Error; + + fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> { + loop { + if self.slot.is_none() { + return Poll::Ready(Ok(())); + } + ready!(self.as_mut().poll_lock_and_flush_slot(cx))?; + } + } + + fn start_send(mut self: Pin<&mut Self>, item: Item) -> Result<(), S::Error> { + self.slot = Some(item); + Ok(()) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> { + let this = &mut *self; + let mut inner = ready!(this.lock.poll_lock(cx)); + ready!(Self::poll_flush_slot(inner.as_pin_mut(), &mut this.slot, cx))?; + inner.as_pin_mut().poll_flush(cx) + } + + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> { + let this = &mut *self; + let mut inner = ready!(this.lock.poll_lock(cx)); + ready!(Self::poll_flush_slot(inner.as_pin_mut(), &mut this.slot, cx))?; + inner.as_pin_mut().poll_close(cx) + } +} + +pub(super) fn split<S: Stream + Sink<Item>, Item>(s: S) -> (SplitSink<S, Item>, SplitStream<S>) { + let (a, b) = BiLock::new(s); + let read = SplitStream(a); + let write = SplitSink(b); + (write, read) +} + +/// Error indicating a `SplitSink<S>` and `SplitStream<S>` were not two halves +/// of a `Stream + Split`, and thus could not be `reunite`d. +pub struct ReuniteError<T, Item>(pub SplitSink<T, Item>, pub SplitStream<T>); + +impl<T, Item> fmt::Debug for ReuniteError<T, Item> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("ReuniteError") + .field(&"...") + .finish() + } +} + +impl<T, Item> fmt::Display for ReuniteError<T, Item> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "tried to reunite a SplitStream and SplitSink that don't form a pair") + } +} + +#[cfg(feature = "std")] +impl<T: core::any::Any, Item> std::error::Error for ReuniteError<T, Item> {} diff --git a/src/stream/stream/take.rs b/src/stream/stream/take.rs new file mode 100644 index 0000000..1109a4a --- /dev/null +++ b/src/stream/stream/take.rs @@ -0,0 +1,118 @@ +use core::cmp; +use core::pin::Pin; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`take`](super::StreamExt::take) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Take<St> { + stream: St, + remaining: usize, +} + +impl<St: Unpin> Unpin for Take<St> {} + +impl<St: Stream> Take<St> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(remaining: usize); + + pub(super) fn new(stream: St, n: usize) -> Take<St> { + Take { + stream, + remaining: n, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +impl<St> Stream for Take<St> + where St: Stream, +{ + type Item = St::Item; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<St::Item>> { + if self.remaining == 0 { + Poll::Ready(None) + } else { + let next = ready!(self.as_mut().stream().poll_next(cx)); + match next { + Some(_) => *self.as_mut().remaining() -= 1, + None => *self.as_mut().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 as usize); + + let upper = match upper { + Some(x) if x < self.remaining as usize => Some(x), + _ => Some(self.remaining as usize) + }; + + (lower, upper) + } +} + +impl<St> FusedStream for Take<St> + where St: FusedStream, +{ + fn is_terminated(&self) -> bool { + self.remaining == 0 || self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for Take<S> + where S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/take_while.rs b/src/stream/stream/take_while.rs new file mode 100644 index 0000000..68606ec --- /dev/null +++ b/src/stream/stream/take_while.rs @@ -0,0 +1,164 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::stream::{Stream, FusedStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`take_while`](super::StreamExt::take_while) method. +#[must_use = "streams do nothing unless polled"] +pub struct TakeWhile<St: Stream , Fut, F> { + stream: St, + f: F, + pending_fut: Option<Fut>, + pending_item: Option<St::Item>, + done_taking: bool, +} + +impl<St: Unpin + Stream, Fut: Unpin, F> Unpin for TakeWhile<St, Fut, F> {} + +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: Stream, Fut, F> TakeWhile<St, Fut, F> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + unsafe_pinned!(pending_fut: Option<Fut>); + unsafe_unpinned!(pending_item: Option<St::Item>); + unsafe_unpinned!(done_taking: bool); +} + +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) -> TakeWhile<St, Fut, F> { + TakeWhile { + stream, + f, + pending_fut: None, + pending_item: None, + done_taking: false, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<St::Item>> { + if self.done_taking { + return Poll::Ready(None); + } + + if self.pending_item.is_none() { + let item = match ready!(self.as_mut().stream().poll_next(cx)) { + Some(e) => e, + None => return Poll::Ready(None), + }; + let fut = (self.as_mut().f())(&item); + self.as_mut().pending_fut().set(Some(fut)); + *self.as_mut().pending_item() = Some(item); + } + + let take = ready!(self.as_mut().pending_fut().as_pin_mut().unwrap().poll(cx)); + self.as_mut().pending_fut().set(None); + let item = self.as_mut().pending_item().take().unwrap(); + + if take { + Poll::Ready(Some(item)) + } else { + *self.as_mut().done_taking() = true; + Poll::Ready(None) + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.done_taking { + return (0, Some(0)); + } + + let pending_len = if self.pending_item.is_some() { 1 } else { 0 }; + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +impl<St, Fut, F> FusedStream for TakeWhile<St, Fut, F> + where St: FusedStream, + F: FnMut(&St::Item) -> Fut, + Fut: Future<Output = bool>, +{ + fn is_terminated(&self) -> bool { + self.done_taking || self.pending_item.is_none() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for TakeWhile<S, Fut, F> + where S: Stream + Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/then.rs b/src/stream/stream/then.rs new file mode 100644 index 0000000..39843b2 --- /dev/null +++ b/src/stream/stream/then.rs @@ -0,0 +1,139 @@ +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`then`](super::StreamExt::then) method. +#[must_use = "streams do nothing unless polled"] +pub struct Then<St, Fut, F> { + stream: St, + future: Option<Fut>, + f: F, +} + +impl<St: Unpin, Fut: Unpin, F> Unpin for Then<St, Fut, 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> { + unsafe_pinned!(stream: St); + unsafe_pinned!(future: Option<Fut>); + unsafe_unpinned!(f: F); +} + +impl<St, Fut, F> Then<St, Fut, F> + where St: Stream, + F: FnMut(St::Item) -> Fut, +{ + pub(super) fn new(stream: St, f: F) -> Then<St, Fut, F> { + Then { + stream, + future: None, + f, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Fut::Output>> { + if self.future.is_none() { + let item = match ready!(self.as_mut().stream().poll_next(cx)) { + None => return Poll::Ready(None), + Some(e) => e, + }; + let fut = (self.as_mut().f())(item); + self.as_mut().future().set(Some(fut)); + } + + let e = ready!(self.as_mut().future().as_pin_mut().unwrap().poll(cx)); + self.as_mut().future().set(None); + Poll::Ready(Some(e)) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let future_len = if self.future.is_some() { 1 } else { 0 }; + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(future_len); + let upper = match upper { + Some(x) => x.checked_add(future_len), + None => None, + }; + (lower, upper) + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for Then<S, Fut, F> + where S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/stream/zip.rs b/src/stream/stream/zip.rs new file mode 100644 index 0000000..f97ac17 --- /dev/null +++ b/src/stream/stream/zip.rs @@ -0,0 +1,143 @@ +use crate::stream::{StreamExt, Fuse}; +use core::cmp; +use core::pin::Pin; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`zip`](super::StreamExt::zip) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct Zip<St1: Stream, St2: Stream> { + stream1: Fuse<St1>, + stream2: Fuse<St2>, + queued1: Option<St1::Item>, + queued2: Option<St2::Item>, +} + +#[allow(clippy::type_repetition_in_bounds)] // https://github.com/rust-lang/rust-clippy/issues/4323 +impl<St1, St2> Unpin for Zip<St1, St2> +where + St1: Stream, + Fuse<St1>: Unpin, + St2: Stream, + Fuse<St2>: Unpin, +{} + +impl<St1: Stream, St2: Stream> Zip<St1, St2> { + unsafe_pinned!(stream1: Fuse<St1>); + unsafe_pinned!(stream2: Fuse<St2>); + unsafe_unpinned!(queued1: Option<St1::Item>); + unsafe_unpinned!(queued2: Option<St2::Item>); + + pub(super) fn new(stream1: St1, stream2: St2) -> Zip<St1, St2> { + Zip { + 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>) { + unsafe { + let Self { stream1, stream2, .. } = self.get_unchecked_mut(); + (Pin::new_unchecked(stream1).get_pin_mut(), Pin::new_unchecked(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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.queued1.is_none() { + match self.as_mut().stream1().poll_next(cx) { + Poll::Ready(Some(item1)) => *self.as_mut().queued1() = Some(item1), + Poll::Ready(None) | Poll::Pending => {} + } + } + if self.queued2.is_none() { + match self.as_mut().stream2().poll_next(cx) { + Poll::Ready(Some(item2)) => *self.as_mut().queued2() = Some(item2), + Poll::Ready(None) | Poll::Pending => {} + } + } + + if self.queued1.is_some() && self.queued2.is_some() { + let pair = (self.as_mut().queued1().take().unwrap(), + self.as_mut().queued2().take().unwrap()); + Poll::Ready(Some(pair)) + } else if self.stream1.is_done() || self.stream2.is_done() { + Poll::Ready(None) + } else { + Poll::Pending + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let queued1_len = if self.queued1.is_some() { 1 } else { 0 }; + let queued2_len = if self.queued2.is_some() { 1 } else { 0 }; + let (stream1_lower, stream1_upper) = self.stream1.size_hint(); + let (stream2_lower, stream2_upper) = self.stream2.size_hint(); + + let stream1_lower = stream1_lower.saturating_add(queued1_len); + let stream2_lower = stream2_lower.saturating_add(queued2_len); + + let lower = cmp::min(stream1_lower, stream2_lower); + + let upper = match (stream1_upper, stream2_upper) { + (Some(x), Some(y)) => { + let x = x.saturating_add(queued1_len); + let y = y.saturating_add(queued2_len); + Some(cmp::min(x, y)) + } + (Some(x), None) => x.checked_add(queued1_len), + (None, Some(y)) => y.checked_add(queued2_len), + (None, None) => None + }; + + (lower, upper) + } +} diff --git a/src/stream/try_stream/and_then.rs b/src/stream/try_stream/and_then.rs new file mode 100644 index 0000000..809c32a --- /dev/null +++ b/src/stream/try_stream/and_then.rs @@ -0,0 +1,136 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::stream::{Stream, TryStream, FusedStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`and_then`](super::TryStreamExt::and_then) method. +#[must_use = "streams do nothing unless polled"] +pub struct AndThen<St, Fut, F> { + stream: St, + future: Option<Fut>, + f: F, +} + +impl<St: Unpin, Fut: Unpin, F> Unpin for AndThen<St, Fut, 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> { + unsafe_pinned!(stream: St); + unsafe_pinned!(future: Option<Fut>); + unsafe_unpinned!(f: F); +} + +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 } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.future.is_none() { + let item = match ready!(self.as_mut().stream().try_poll_next(cx)?) { + None => return Poll::Ready(None), + Some(e) => e, + }; + let fut = (self.as_mut().f())(item); + self.as_mut().future().set(Some(fut)); + } + + let e = ready!(self.as_mut().future().as_pin_mut().unwrap().try_poll(cx)); + self.as_mut().future().set(None); + Poll::Ready(Some(e)) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let future_len = if self.future.is_some() { 1 } else { 0 }; + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(future_len); + let upper = match upper { + Some(x) => x.checked_add(future_len), + None => None, + }; + (lower, upper) + } +} + +impl<St, Fut, F> FusedStream for AndThen<St, Fut, F> + where St: TryStream + FusedStream, + F: FnMut(St::Ok) -> Fut, + Fut: TryFuture<Error = St::Error>, +{ + fn is_terminated(&self) -> bool { + self.future.is_none() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for AndThen<S, Fut, F> + where S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/err_into.rs b/src/stream/try_stream/err_into.rs new file mode 100644 index 0000000..f5d9294 --- /dev/null +++ b/src/stream/try_stream/err_into.rs @@ -0,0 +1,98 @@ +use core::marker::PhantomData; +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_utils::unsafe_pinned; + +/// Stream for the [`err_into`](super::TryStreamExt::err_into) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct ErrInto<St, E> { + stream: St, + _marker: PhantomData<E>, +} + +impl<St: Unpin, E> Unpin for ErrInto<St, E> {} + +impl<St, E> ErrInto<St, E> { + unsafe_pinned!(stream: St); + + pub(super) fn new(stream: St) -> Self { + ErrInto { stream, _marker: PhantomData } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +impl<St, E> FusedStream for ErrInto<St, E> +where + St: TryStream + FusedStream, + St::Error: Into<E>, +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St, E> Stream for ErrInto<St, E> +where + St: TryStream, + St::Error: Into<E>, +{ + type Item = Result<St::Ok, E>; + + fn poll_next( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.stream().try_poll_next(cx) + .map(|res| res.map(|some| some.map_err(Into::into))) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, E, Item> Sink<Item> for ErrInto<S, E> +where + S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/inspect_err.rs b/src/stream/try_stream/inspect_err.rs new file mode 100644 index 0000000..3c23ae0 --- /dev/null +++ b/src/stream/try_stream/inspect_err.rs @@ -0,0 +1,118 @@ +use crate::stream::stream::inspect; +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`inspect_err`](super::TryStreamExt::inspect_err) method. +#[must_use = "streams do nothing unless polled"] +pub struct InspectErr<St, F> { + stream: St, + f: F, +} + +impl<St: Unpin, F> Unpin for InspectErr<St, F> {} + +impl<St, F> fmt::Debug for InspectErr<St, F> +where + St: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("InspectErr") + .field("stream", &self.stream) + .finish() + } +} + +impl<St, F> InspectErr<St, F> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); +} + +impl<St, F> InspectErr<St, F> +where + St: TryStream, + F: FnMut(&St::Error), +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +impl<St, F> FusedStream for InspectErr<St, F> +where + St: TryStream + FusedStream, + F: FnMut(&St::Error), +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St, F> Stream for InspectErr<St, F> +where + St: TryStream, + F: FnMut(&St::Error), +{ + type Item = Result<St::Ok, St::Error>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.as_mut() + .stream() + .try_poll_next(cx) + .map(|opt| opt.map(|res| res.map_err(|e| inspect(e, self.as_mut().f())))) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, F, Item> Sink<Item> for InspectErr<S, F> +where + S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/inspect_ok.rs b/src/stream/try_stream/inspect_ok.rs new file mode 100644 index 0000000..89fb459 --- /dev/null +++ b/src/stream/try_stream/inspect_ok.rs @@ -0,0 +1,118 @@ +use crate::stream::stream::inspect; +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`inspect_ok`](super::TryStreamExt::inspect_ok) method. +#[must_use = "streams do nothing unless polled"] +pub struct InspectOk<St, F> { + stream: St, + f: F, +} + +impl<St: Unpin, F> Unpin for InspectOk<St, F> {} + +impl<St, F> fmt::Debug for InspectOk<St, F> +where + St: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("InspectOk") + .field("stream", &self.stream) + .finish() + } +} + +impl<St, F> InspectOk<St, F> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); +} + +impl<St, F> InspectOk<St, F> +where + St: TryStream, + F: FnMut(&St::Ok), +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +impl<St, F> FusedStream for InspectOk<St, F> +where + St: TryStream + FusedStream, + F: FnMut(&St::Ok), +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St, F> Stream for InspectOk<St, F> +where + St: TryStream, + F: FnMut(&St::Ok), +{ + type Item = Result<St::Ok, St::Error>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.as_mut() + .stream() + .try_poll_next(cx) + .map(|opt| opt.map(|res| res.map(|e| inspect(e, self.as_mut().f())))) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, F, Item> Sink<Item> for InspectOk<S, F> +where + S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/into_async_read.rs b/src/stream/try_stream/into_async_read.rs new file mode 100644 index 0000000..71bbfd1 --- /dev/null +++ b/src/stream/try_stream/into_async_read.rs @@ -0,0 +1,186 @@ +use crate::stream::TryStreamExt; +use core::pin::Pin; +use futures_core::stream::TryStream; +use futures_core::task::{Context, Poll}; +use futures_io::{AsyncRead, AsyncWrite, AsyncBufRead}; +use std::cmp; +use std::io::{Error, Result}; + +/// Reader for the [`into_async_read`](super::TryStreamExt::into_async_read) method. +#[derive(Debug)] +#[must_use = "readers do nothing unless polled"] +pub struct IntoAsyncRead<St> +where + St: TryStream<Error = Error> + Unpin, + St::Ok: AsRef<[u8]>, +{ + stream: St, + state: ReadState<St::Ok>, +} + +impl<St> Unpin for IntoAsyncRead<St> +where + St: TryStream<Error = Error> + Unpin, + St::Ok: AsRef<[u8]>, +{ +} + +#[derive(Debug)] +enum ReadState<T: AsRef<[u8]>> { + Ready { chunk: T, chunk_start: usize }, + PendingChunk, + Eof, +} + +impl<St> IntoAsyncRead<St> +where + St: TryStream<Error = Error> + Unpin, + St::Ok: AsRef<[u8]>, +{ + pub(super) fn new(stream: St) -> Self { + IntoAsyncRead { + stream, + state: ReadState::PendingChunk, + } + } +} + +impl<St> AsyncRead for IntoAsyncRead<St> +where + St: TryStream<Error = Error> + Unpin, + St::Ok: AsRef<[u8]>, +{ + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll<Result<usize>> { + loop { + match &mut self.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 { + self.state = ReadState::PendingChunk; + } + + return Poll::Ready(Ok(len)); + } + ReadState::PendingChunk => { + match ready!(self.stream.try_poll_next_unpin(cx)) { + Some(Ok(chunk)) => { + if !chunk.as_ref().is_empty() { + self.state = ReadState::Ready { + chunk, + chunk_start: 0, + }; + } + } + Some(Err(err)) => { + self.state = ReadState::Eof; + return Poll::Ready(Err(err)); + } + None => { + self.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 + Unpin, + St::Ok: AsRef<[u8]>, +{ + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8] + ) -> Poll<Result<usize>> { + Pin::new( &mut self.stream ).poll_write( cx, buf ) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_> + ) -> Poll<Result<()>> { + Pin::new( &mut self.stream ).poll_flush( cx ) + } + + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_> + ) -> Poll<Result<()>> { + Pin::new( &mut self.stream ).poll_close( cx ) + } +} + +impl<St> AsyncBufRead for IntoAsyncRead<St> +where + St: TryStream<Error = Error> + Unpin, + St::Ok: AsRef<[u8]>, +{ + fn poll_fill_buf( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Result<&[u8]>> { + while let ReadState::PendingChunk = self.state { + match ready!(self.stream.try_poll_next_unpin(cx)) { + Some(Ok(chunk)) => { + if !chunk.as_ref().is_empty() { + self.state = ReadState::Ready { + chunk, + chunk_start: 0, + }; + } + } + Some(Err(err)) => { + self.state = ReadState::Eof; + return Poll::Ready(Err(err)); + } + None => { + self.state = ReadState::Eof; + return Poll::Ready(Ok(&[])); + } + } + } + + if let ReadState::Ready { ref chunk, chunk_start } = self.into_ref().get_ref().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( + mut self: Pin<&mut Self>, + amount: usize, + ) { + // https://github.com/rust-lang/futures-rs/pull/1556#discussion_r281644295 + if amount == 0 { return } + if let ReadState::Ready { chunk, chunk_start } = &mut self.state { + *chunk_start += amount; + debug_assert!(*chunk_start <= chunk.as_ref().len()); + if *chunk_start >= chunk.as_ref().len() { + self.state = ReadState::PendingChunk; + } + } else { + debug_assert!(false, "Attempted to consume from IntoAsyncRead without chunk"); + } + } +} diff --git a/src/stream/try_stream/into_stream.rs b/src/stream/try_stream/into_stream.rs new file mode 100644 index 0000000..b0fa07a --- /dev/null +++ b/src/stream/try_stream/into_stream.rs @@ -0,0 +1,84 @@ +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_utils::unsafe_pinned; + +/// Stream for the [`into_stream`](super::TryStreamExt::into_stream) method. +#[derive(Debug)] +#[must_use = "streams do nothing unless polled"] +pub struct IntoStream<St> { + stream: St, +} + +impl<St> IntoStream<St> { + unsafe_pinned!(stream: St); + + #[inline] + pub(super) fn new(stream: St) -> Self { + IntoStream { stream } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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.stream().try_poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S: Sink<Item>, Item> Sink<Item> for IntoStream<S> { + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/map_err.rs b/src/stream/try_stream/map_err.rs new file mode 100644 index 0000000..1b98d6b --- /dev/null +++ b/src/stream/try_stream/map_err.rs @@ -0,0 +1,112 @@ +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`map_err`](super::TryStreamExt::map_err) method. +#[must_use = "streams do nothing unless polled"] +pub struct MapErr<St, F> { + stream: St, + f: F, +} + +impl<St: Unpin, F> Unpin for MapErr<St, F> {} + +impl<St, F> fmt::Debug for MapErr<St, F> +where + St: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MapErr") + .field("stream", &self.stream) + .finish() + } +} + +impl<St, F> MapErr<St, F> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + + /// Creates a new MapErr. + pub(super) fn new(stream: St, f: F) -> Self { + MapErr { stream, f } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +impl<St, F, E> FusedStream for MapErr<St, F> +where + St: TryStream + FusedStream, + F: FnMut(St::Error) -> E, +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St, F, E> Stream for MapErr<St, F> +where + St: TryStream, + F: FnMut(St::Error) -> E, +{ + type Item = Result<St::Ok, E>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.as_mut() + .stream() + .try_poll_next(cx) + .map(|opt| opt.map(|res| res.map_err(|e| self.as_mut().f()(e)))) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, F, Item> Sink<Item> for MapErr<S, F> +where + S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/map_ok.rs b/src/stream/try_stream/map_ok.rs new file mode 100644 index 0000000..19d01be --- /dev/null +++ b/src/stream/try_stream/map_ok.rs @@ -0,0 +1,112 @@ +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_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`map_ok`](super::TryStreamExt::map_ok) method. +#[must_use = "streams do nothing unless polled"] +pub struct MapOk<St, F> { + stream: St, + f: F, +} + +impl<St: Unpin, F> Unpin for MapOk<St, F> {} + +impl<St, F> fmt::Debug for MapOk<St, F> +where + St: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MapOk") + .field("stream", &self.stream) + .finish() + } +} + +impl<St, F> MapOk<St, F> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + + /// Creates a new MapOk. + pub(super) fn new(stream: St, f: F) -> Self { + MapOk { stream, f } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +impl<St, F, T> FusedStream for MapOk<St, F> +where + St: TryStream + FusedStream, + F: FnMut(St::Ok) -> T, +{ + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St, F, T> Stream for MapOk<St, F> +where + St: TryStream, + F: FnMut(St::Ok) -> T, +{ + type Item = Result<T, St::Error>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.as_mut() + .stream() + .try_poll_next(cx) + .map(|opt| opt.map(|res| res.map(|x| self.as_mut().f()(x)))) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.stream.size_hint() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, F, Item> Sink<Item> for MapOk<S, F> +where + S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/mod.rs b/src/stream/try_stream/mod.rs new file mode 100644 index 0000000..6a7ced4 --- /dev/null +++ b/src/stream/try_stream/mod.rs @@ -0,0 +1,855 @@ +//! Streams +//! +//! This module contains a number of functions for working with `Streams`s +//! that return `Result`s, allowing for short-circuiting computations. + +#[cfg(feature = "compat")] +use crate::compat::Compat; +use core::pin::Pin; +use futures_core::{ + future::{Future, TryFuture}, + stream::TryStream, + task::{Context, Poll}, +}; + +mod and_then; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::and_then::AndThen; + +mod err_into; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::err_into::ErrInto; + +mod inspect_ok; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::inspect_ok::InspectOk; + +mod inspect_err; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::inspect_err::InspectErr; + +mod into_stream; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::into_stream::IntoStream; + +mod map_ok; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::map_ok::MapOk; + +mod map_err; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::map_err::MapErr; + +mod or_else; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::or_else::OrElse; + +mod try_next; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_next::TryNext; + +mod try_for_each; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_for_each::TryForEach; + +mod try_filter; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_filter::TryFilter; + +mod try_filter_map; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_filter_map::TryFilterMap; + +mod try_flatten; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_flatten::TryFlatten; + +mod try_collect; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_collect::TryCollect; + +mod try_concat; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_concat::TryConcat; + +mod try_fold; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_fold::TryFold; + +mod try_unfold; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_unfold::{try_unfold, TryUnfold}; + +mod try_skip_while; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::try_skip_while::TrySkipWhile; + +cfg_target_has_atomic! { + #[cfg(feature = "alloc")] + mod try_buffer_unordered; + #[cfg(feature = "alloc")] + #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 + pub use self::try_buffer_unordered::TryBufferUnordered; + + #[cfg(feature = "alloc")] + mod try_for_each_concurrent; + #[cfg(feature = "alloc")] + #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 + pub use self::try_for_each_concurrent::TryForEachConcurrent; +} + +#[cfg(feature = "io")] +#[cfg(feature = "std")] +mod into_async_read; +#[cfg(feature = "io")] +#[cfg(feature = "std")] +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::into_async_read::IntoAsyncRead; + +impl<S: ?Sized + TryStream> TryStreamExt for S {} + +/// Adapters specific to `Result`-returning streams +pub trait TryStreamExt: TryStream { + /// Wraps the current stream in a new stream which converts the error type + /// into the one provided. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// + /// let mut stream = + /// stream::iter(vec![Ok(()), Err(5i32)]) + /// .err_into::<i64>(); + /// + /// assert_eq!(stream.try_next().await, Ok(Some(()))); + /// assert_eq!(stream.try_next().await, Err(5i64)); + /// # }) + /// ``` + fn err_into<E>(self) -> ErrInto<Self, E> + where + Self: Sized, + Self::Error: Into<E>, + { + ErrInto::new(self) + } + + /// Wraps the current stream in a new stream which maps the success value + /// using the provided closure. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// + /// let mut stream = + /// stream::iter(vec![Ok(5), Err(0)]) + /// .map_ok(|x| x + 2); + /// + /// assert_eq!(stream.try_next().await, Ok(Some(7))); + /// assert_eq!(stream.try_next().await, Err(0)); + /// # }) + /// ``` + fn map_ok<T, F>(self, f: F) -> MapOk<Self, F> + where + Self: Sized, + F: FnMut(Self::Ok) -> T, + { + MapOk::new(self, f) + } + + /// Wraps the current stream in a new stream which maps the error value + /// using the provided closure. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// + /// let mut stream = + /// stream::iter(vec![Ok(5), Err(0)]) + /// .map_err(|x| x + 2); + /// + /// assert_eq!(stream.try_next().await, Ok(Some(5))); + /// assert_eq!(stream.try_next().await, Err(2)); + /// # }) + /// ``` + fn map_err<E, F>(self, f: F) -> MapErr<Self, F> + where + Self: Sized, + F: FnMut(Self::Error) -> E, + { + MapErr::new(self, f) + } + + /// Chain on a computation for when a value is ready, passing the successful + /// results to the provided closure `f`. + /// + /// This function can be used to run a unit of work when the next successful + /// value on a stream is ready. The closure provided will be yielded a value + /// when ready, and the returned future will then be run to completion to + /// produce the next value on this stream. + /// + /// Any errors produced by this stream will not be passed to the closure, + /// and will be passed through. + /// + /// The returned value of the closure must implement the `TryFuture` trait + /// and can represent some more work to be done before the composed stream + /// is finished. + /// + /// Note that this function consumes the receiving stream and returns a + /// wrapped version of it. + /// + /// To process the entire stream and return a single future representing + /// success or error, use `try_for_each` instead. + /// + /// # Examples + /// + /// ``` + /// use futures::channel::mpsc; + /// use futures::future; + /// use futures::stream::TryStreamExt; + /// + /// let (_tx, rx) = mpsc::channel::<Result<i32, ()>>(1); + /// + /// let rx = rx.and_then(|result| { + /// future::ok(if result % 2 == 0 { + /// Some(result) + /// } else { + /// None + /// }) + /// }); + /// ``` + fn and_then<Fut, F>(self, f: F) -> AndThen<Self, Fut, F> + where + F: FnMut(Self::Ok) -> Fut, + Fut: TryFuture<Error = Self::Error>, + Self: Sized, + { + AndThen::new(self, f) + } + + /// Chain on a computation for when an error happens, passing the + /// erroneous result to the provided closure `f`. + /// + /// This function can be used to run a unit of work and attempt to recover from + /// an error if one happens. The closure provided will be yielded an error + /// when one appears, and the returned future will then be run to completion + /// to produce the next value on this stream. + /// + /// Any successful values produced by this stream will not be passed to the + /// closure, and will be passed through. + /// + /// The returned value of the closure must implement the [`TryFuture`](futures_core::future::TryFuture) trait + /// and can represent some more work to be done before the composed stream + /// is finished. + /// + /// Note that this function consumes the receiving stream and returns a + /// wrapped version of it. + fn or_else<Fut, F>(self, f: F) -> OrElse<Self, Fut, F> + where + F: FnMut(Self::Error) -> Fut, + Fut: TryFuture<Ok = Self::Ok>, + Self: Sized, + { + OrElse::new(self, f) + } + + /// Do something with the success value of this stream, afterwards passing + /// it on. + /// + /// This is similar to the `StreamExt::inspect` method where it allows + /// easily inspecting the success value as it passes through the stream, for + /// example to debug what's going on. + fn inspect_ok<F>(self, f: F) -> InspectOk<Self, F> + where + F: FnMut(&Self::Ok), + Self: Sized, + { + InspectOk::new(self, f) + } + + /// Do something with the error value of this stream, afterwards passing it on. + /// + /// This is similar to the `StreamExt::inspect` method where it allows + /// easily inspecting the error value as it passes through the stream, for + /// example to debug what's going on. + fn inspect_err<F>(self, f: F) -> InspectErr<Self, F> + where + F: FnMut(&Self::Error), + Self: Sized, + { + InspectErr::new(self, f) + } + + /// Wraps a [`TryStream`] into a type that implements + /// [`Stream`](futures_core::stream::Stream) + /// + /// [`TryStream`]s currently do not implement the + /// [`Stream`](futures_core::stream::Stream) trait because of limitations + /// of the compiler. + /// + /// # Examples + /// + /// ``` + /// use futures::stream::{Stream, TryStream, TryStreamExt}; + /// + /// # type T = i32; + /// # type E = (); + /// fn make_try_stream() -> impl TryStream<Ok = T, Error = E> { // ... } + /// # futures::stream::empty() + /// # } + /// fn take_stream(stream: impl Stream<Item = Result<T, E>>) { /* ... */ } + /// + /// take_stream(make_try_stream().into_stream()); + /// ``` + fn into_stream(self) -> IntoStream<Self> + where + Self: Sized, + { + IntoStream::new(self) + } + + /// Creates a future that attempts to resolve the next item in the stream. + /// If an error is encountered before the next item, the error is returned + /// instead. + /// + /// This is similar to the `Stream::next` combinator, but returns a + /// `Result<Option<T>, E>` rather than an `Option<Result<T, E>>`, making + /// for easy use with the `?` operator. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// + /// let mut stream = stream::iter(vec![Ok(()), Err(())]); + /// + /// assert_eq!(stream.try_next().await, Ok(Some(()))); + /// assert_eq!(stream.try_next().await, Err(())); + /// # }) + /// ``` + fn try_next(&mut self) -> TryNext<'_, Self> + where + Self: Unpin, + { + TryNext::new(self) + } + + /// Attempts to run this stream to completion, executing the provided + /// asynchronous closure for each element on the stream. + /// + /// The provided closure will be called for each item this stream produces, + /// yielding a future. That future will then be executed to completion + /// before moving on to the next item. + /// + /// The returned value is a [`Future`](futures_core::future::Future) where the + /// [`Output`](futures_core::future::Future::Output) type is + /// `Result<(), Self::Error>`. If any of the intermediate + /// futures or the stream returns an error, this future will return + /// immediately with an error. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, TryStreamExt}; + /// + /// let mut x = 0i32; + /// + /// { + /// let fut = stream::repeat(Ok(1)).try_for_each(|item| { + /// x += item; + /// future::ready(if x == 3 { Err(()) } else { Ok(()) }) + /// }); + /// assert_eq!(fut.await, Err(())); + /// } + /// + /// assert_eq!(x, 3); + /// # }) + /// ``` + fn try_for_each<Fut, F>(self, f: F) -> TryForEach<Self, Fut, F> + where + F: FnMut(Self::Ok) -> Fut, + Fut: TryFuture<Ok = (), Error = Self::Error>, + Self: Sized, + { + TryForEach::new(self, f) + } + + /// Skip elements on this stream while the provided asynchronous predicate + /// resolves to `true`. + /// + /// This function is similar to [`StreamExt::skip_while`](crate::stream::StreamExt::skip_while) + /// but exits early if an error occurs. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, TryStreamExt}; + /// + /// let stream = stream::iter(vec![Ok::<i32, i32>(1), Ok(3), Ok(2)]); + /// let stream = stream.try_skip_while(|x| future::ready(Ok(*x < 3))); + /// + /// let output: Result<Vec<i32>, i32> = stream.try_collect().await; + /// assert_eq!(output, Ok(vec![3, 2])); + /// # }) + /// ``` + fn try_skip_while<Fut, F>(self, f: F) -> TrySkipWhile<Self, Fut, F> + where + F: FnMut(&Self::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = Self::Error>, + Self: Sized, + { + TrySkipWhile::new(self, f) + } + + /// Attempts to run this stream to completion, executing the provided asynchronous + /// closure for each element on the stream concurrently as elements become + /// available, exiting as soon as an error occurs. + /// + /// This is similar to + /// [`StreamExt::for_each_concurrent`](crate::stream::StreamExt::for_each_concurrent), + /// but will resolve to an error immediately if the underlying stream or the provided + /// closure return an error. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::oneshot; + /// use futures::stream::{self, StreamExt, TryStreamExt}; + /// + /// let (tx1, rx1) = oneshot::channel(); + /// let (tx2, rx2) = oneshot::channel(); + /// let (_tx3, rx3) = oneshot::channel(); + /// + /// let stream = stream::iter(vec![rx1, rx2, rx3]); + /// let fut = stream.map(Ok).try_for_each_concurrent( + /// /* limit */ 2, + /// |rx| async move { + /// let res: Result<(), oneshot::Canceled> = rx.await; + /// res + /// } + /// ); + /// + /// tx1.send(()).unwrap(); + /// // Drop the second sender so that `rx2` resolves to `Canceled`. + /// drop(tx2); + /// + /// // The final result is an error because the second future + /// // resulted in an error. + /// assert_eq!(Err(oneshot::Canceled), fut.await); + /// # }) + /// ``` + #[cfg_attr(feature = "cfg-target-has-atomic", cfg(target_has_atomic = "ptr"))] + #[cfg(feature = "alloc")] + fn try_for_each_concurrent<Fut, F>( + self, + limit: impl Into<Option<usize>>, + f: F, + ) -> TryForEachConcurrent<Self, Fut, F> + where + F: FnMut(Self::Ok) -> Fut, + Fut: Future<Output = Result<(), Self::Error>>, + Self: Sized, + { + TryForEachConcurrent::new(self, limit.into(), f) + } + + /// Attempt to transform a stream into a collection, + /// returning a future representing the result of that computation. + /// + /// This combinator will collect all successful results of this stream and + /// collect them into the specified collection type. If an error happens then all + /// collected elements will be dropped and the error will be returned. + /// + /// The returned future will be resolved when the stream terminates. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::TryStreamExt; + /// use std::thread; + /// + /// let (tx, rx) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// for i in 1..=5 { + /// tx.unbounded_send(Ok(i)).unwrap(); + /// } + /// tx.unbounded_send(Err(6)).unwrap(); + /// }); + /// + /// let output: Result<Vec<i32>, i32> = rx.try_collect().await; + /// assert_eq!(output, Err(6)); + /// # }) + /// ``` + fn try_collect<C: Default + Extend<Self::Ok>>(self) -> TryCollect<Self, C> + where + Self: Sized, + { + TryCollect::new(self) + } + + /// Attempt to filter the values produced by this stream according to the + /// provided asynchronous closure. + /// + /// As values of this stream are made available, the provided predicate `f` + /// will be run on them. If the predicate returns a `Future` which resolves + /// to `true`, then the stream will yield the value, but if the predicate + /// return a `Future` which resolves to `false`, then the value will be + /// discarded and the next value will be produced. + /// + /// All errors are passed through without filtering in this combinator. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `filter` methods in + /// the standard library. + /// + /// # Examples + /// ``` + /// # futures::executor::block_on(async { + /// use futures::future; + /// use futures::stream::{self, StreamExt, TryStreamExt}; + /// + /// let stream = stream::iter(vec![Ok(1i32), Ok(2i32), Ok(3i32), Err("error")]); + /// let mut evens = stream.try_filter(|x| { + /// future::ready(x % 2 == 0) + /// }); + /// + /// assert_eq!(evens.next().await, Some(Ok(2))); + /// assert_eq!(evens.next().await, Some(Err("error"))); + /// # }) + /// ``` + fn try_filter<Fut, F>(self, f: F) -> TryFilter<Self, Fut, F> + where + Fut: Future<Output = bool>, + F: FnMut(&Self::Ok) -> Fut, + Self: Sized, + { + TryFilter::new(self, f) + } + + /// Attempt to filter the values produced by this stream while + /// simultaneously mapping them to a different type according to the + /// provided asynchronous closure. + /// + /// As values of this stream are made available, the provided function will + /// be run on them. If the future returned by the predicate `f` resolves to + /// [`Some(item)`](Some) then the stream will yield the value `item`, but if + /// it resolves to [`None`] then the next value will be produced. + /// + /// All errors are passed through without filtering in this combinator. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to the existing `filter_map` methods in + /// the standard library. + /// + /// # Examples + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt, TryStreamExt}; + /// use futures::pin_mut; + /// + /// let stream = stream::iter(vec![Ok(1i32), Ok(6i32), Err("error")]); + /// let halves = stream.try_filter_map(|x| async move { + /// let ret = if x % 2 == 0 { Some(x / 2) } else { None }; + /// Ok(ret) + /// }); + /// + /// pin_mut!(halves); + /// assert_eq!(halves.next().await, Some(Ok(3))); + /// assert_eq!(halves.next().await, Some(Err("error"))); + /// # }) + /// ``` + fn try_filter_map<Fut, F, T>(self, f: F) -> TryFilterMap<Self, Fut, F> + where + Fut: TryFuture<Ok = Option<T>, Error = Self::Error>, + F: FnMut(Self::Ok) -> Fut, + Self: Sized, + { + TryFilterMap::new(self, f) + } + + /// Flattens a stream of streams into just one continuous stream. + /// + /// If this stream's elements are themselves streams then this combinator + /// will flatten out the entire stream to one long chain of elements. Any + /// errors are passed through without looking at them, but otherwise each + /// individual stream will get exhausted before moving on to the next. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::{StreamExt, TryStreamExt}; + /// use std::thread; + /// + /// let (tx1, rx1) = mpsc::unbounded(); + /// let (tx2, rx2) = mpsc::unbounded(); + /// let (tx3, rx3) = mpsc::unbounded(); + /// + /// thread::spawn(move || { + /// tx1.unbounded_send(Ok(1)).unwrap(); + /// }); + /// thread::spawn(move || { + /// tx2.unbounded_send(Ok(2)).unwrap(); + /// tx2.unbounded_send(Err(3)).unwrap(); + /// }); + /// thread::spawn(move || { + /// tx3.unbounded_send(Ok(rx1)).unwrap(); + /// tx3.unbounded_send(Ok(rx2)).unwrap(); + /// tx3.unbounded_send(Err(4)).unwrap(); + /// }); + /// + /// let mut stream = rx3.try_flatten(); + /// assert_eq!(stream.next().await, Some(Ok(1))); + /// assert_eq!(stream.next().await, Some(Ok(2))); + /// assert_eq!(stream.next().await, Some(Err(3))); + /// # }); + /// ``` + fn try_flatten(self) -> TryFlatten<Self> + where + Self::Ok: TryStream, + <Self::Ok as TryStream>::Error: From<Self::Error>, + Self: Sized, + { + TryFlatten::new(self) + } + + /// Attempt to execute an accumulating asynchronous computation over a + /// stream, collecting all the values into one final result. + /// + /// This combinator will accumulate all values returned by this stream + /// according to the closure provided. The initial state is also provided to + /// this method and then is returned again by each execution of the closure. + /// Once the entire stream has been exhausted the returned future will + /// resolve to this value. + /// + /// This method is similar to [`fold`](crate::stream::StreamExt::fold), but will + /// exit early if an error is encountered in either the stream or the + /// provided closure. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// + /// let number_stream = stream::iter(vec![Ok::<i32, i32>(1), Ok(2)]); + /// let sum = number_stream.try_fold(0, |acc, x| async move { Ok(acc + x) }); + /// assert_eq!(sum.await, Ok(3)); + /// + /// let number_stream_with_err = stream::iter(vec![Ok::<i32, i32>(1), Err(2), Ok(1)]); + /// let sum = number_stream_with_err.try_fold(0, |acc, x| async move { Ok(acc + x) }); + /// assert_eq!(sum.await, Err(2)); + /// # }) + /// ``` + fn try_fold<T, Fut, F>(self, init: T, f: F) -> TryFold<Self, Fut, T, F> + where + F: FnMut(T, Self::Ok) -> Fut, + Fut: TryFuture<Ok = T, Error = Self::Error>, + Self: Sized, + { + TryFold::new(self, f, init) + } + + /// Attempt to concatenate all items of a stream into a single + /// extendable destination, returning a future representing the end result. + /// + /// This combinator will extend the first item with the contents of all + /// the subsequent successful results of the stream. If the stream is empty, + /// the default value will be returned. + /// + /// Works with all collections that implement the [`Extend`](std::iter::Extend) trait. + /// + /// This method is similar to [`concat`](crate::stream::StreamExt::concat), but will + /// exit early if an error is encountered in the stream. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::TryStreamExt; + /// use std::thread; + /// + /// let (tx, rx) = mpsc::unbounded::<Result<Vec<i32>, ()>>(); + /// + /// thread::spawn(move || { + /// for i in (0..3).rev() { + /// let n = i * 3; + /// tx.unbounded_send(Ok(vec![n + 1, n + 2, n + 3])).unwrap(); + /// } + /// }); + /// + /// let result = rx.try_concat().await; + /// + /// assert_eq!(result, Ok(vec![7, 8, 9, 4, 5, 6, 1, 2, 3])); + /// # }); + /// ``` + fn try_concat(self) -> TryConcat<Self> + where + Self: Sized, + Self::Ok: Extend<<<Self as TryStream>::Ok as IntoIterator>::Item> + IntoIterator + Default, + { + TryConcat::new(self) + } + + /// Attempt to execute several futures from a stream concurrently. + /// + /// This stream's `Ok` type must be a [`TryFuture`](futures_core::future::TryFuture) with an `Error` type + /// that matches the stream's `Error` type. + /// + /// This adaptor will buffer up to `n` futures and then return their + /// outputs in the order in which they complete. If the underlying stream + /// returns an error, it will be immediately propagated. + /// + /// The returned stream will be a stream of results, each containing either + /// an error or a future's output. An error can be produced either by the + /// underlying stream itself or by one of the futures it yielded. + /// + /// This method is only available when the `std` or `alloc` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// Results are returned in the order of completion: + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::oneshot; + /// use futures::stream::{self, StreamExt, TryStreamExt}; + /// + /// let (send_one, recv_one) = oneshot::channel(); + /// let (send_two, recv_two) = oneshot::channel(); + /// + /// let stream_of_futures = stream::iter(vec![Ok(recv_one), Ok(recv_two)]); + /// + /// let mut buffered = stream_of_futures.try_buffer_unordered(10); + /// + /// send_two.send(2i32)?; + /// assert_eq!(buffered.next().await, Some(Ok(2i32))); + /// + /// send_one.send(1i32)?; + /// assert_eq!(buffered.next().await, Some(Ok(1i32))); + /// + /// assert_eq!(buffered.next().await, None); + /// # Ok::<(), i32>(()) }).unwrap(); + /// ``` + /// + /// Errors from the underlying stream itself are propagated: + /// ``` + /// # futures::executor::block_on(async { + /// use futures::channel::mpsc; + /// use futures::stream::{StreamExt, TryStreamExt}; + /// + /// let (sink, stream_of_futures) = mpsc::unbounded(); + /// let mut buffered = stream_of_futures.try_buffer_unordered(10); + /// + /// sink.unbounded_send(Ok(async { Ok(7i32) }))?; + /// assert_eq!(buffered.next().await, Some(Ok(7i32))); + /// + /// sink.unbounded_send(Err("error in the stream"))?; + /// assert_eq!(buffered.next().await, Some(Err("error in the stream"))); + /// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap(); + /// ``` + #[cfg_attr(feature = "cfg-target-has-atomic", cfg(target_has_atomic = "ptr"))] + #[cfg(feature = "alloc")] + fn try_buffer_unordered(self, n: usize) -> TryBufferUnordered<Self> + where + Self::Ok: TryFuture<Error = Self::Error>, + Self: Sized, + { + TryBufferUnordered::new(self, n) + } + + // TODO: false positive warning from rustdoc. Verify once #43466 settles + // + /// A convenience method for calling [`TryStream::try_poll_next`] on [`Unpin`] + /// stream types. + fn try_poll_next_unpin( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<Self::Ok, Self::Error>>> + where + Self: Unpin, + { + Pin::new(self).try_poll_next(cx) + } + + /// Wraps a [`TryStream`] into a stream compatible with libraries using + /// futures 0.1 `Stream`. Requires the `compat` feature to be enabled. + /// ``` + /// use futures::future::{FutureExt, TryFutureExt}; + /// # let (tx, rx) = futures::channel::oneshot::channel(); + /// + /// let future03 = async { + /// println!("Running on the pool"); + /// tx.send(42).unwrap(); + /// }; + /// + /// let future01 = future03 + /// .unit_error() // Make it a TryFuture + /// .boxed() // Make it Unpin + /// .compat(); + /// + /// tokio::run(future01); + /// # assert_eq!(42, futures::executor::block_on(rx).unwrap()); + /// ``` + #[cfg(feature = "compat")] + fn compat(self) -> Compat<Self> + where + Self: Sized + Unpin, + { + Compat::new(self) + } + + /// Adapter that converts this stream into an [`AsyncRead`](crate::io::AsyncRead). + /// + /// Note that because `into_async_read` moves the stream, the [`Stream`](futures_core::stream::Stream) type must be + /// [`Unpin`]. If you want to use `into_async_read` with a [`!Unpin`](Unpin) stream, you'll + /// first have to pin the stream. This can be done by boxing the stream using [`Box::pin`] + /// or pinning it to the stack using the `pin_mut!` macro from the `pin_utils` crate. + /// + /// This method is only available when the `std` feature of this + /// library is activated, and it is activated by default. + /// + /// # Examples + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, TryStreamExt}; + /// use futures::io::AsyncReadExt; + /// + /// let stream = stream::iter(vec![Ok(vec![1, 2, 3, 4, 5])]); + /// let mut reader = stream.into_async_read(); + /// let mut buf = Vec::new(); + /// + /// assert!(reader.read_to_end(&mut buf).await.is_ok()); + /// assert_eq!(buf, &[1, 2, 3, 4, 5]); + /// # }) + /// ``` + #[cfg(feature = "io")] + #[cfg(feature = "std")] + fn into_async_read(self) -> IntoAsyncRead<Self> + where + Self: Sized + TryStreamExt<Error = std::io::Error> + Unpin, + Self::Ok: AsRef<[u8]>, + { + IntoAsyncRead::new(self) + } +} diff --git a/src/stream/try_stream/or_else.rs b/src/stream/try_stream/or_else.rs new file mode 100644 index 0000000..33310d1 --- /dev/null +++ b/src/stream/try_stream/or_else.rs @@ -0,0 +1,137 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::stream::{Stream, TryStream, FusedStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// Stream for the [`or_else`](super::TryStreamExt::or_else) method. +#[must_use = "streams do nothing unless polled"] +pub struct OrElse<St, Fut, F> { + stream: St, + future: Option<Fut>, + f: F, +} + +impl<St: Unpin, Fut: Unpin, F> Unpin for OrElse<St, Fut, 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> { + unsafe_pinned!(stream: St); + unsafe_pinned!(future: Option<Fut>); + unsafe_unpinned!(f: F); +} + +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 } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.future.is_none() { + let item = match ready!(self.as_mut().stream().try_poll_next(cx)) { + None => return Poll::Ready(None), + Some(Ok(e)) => return Poll::Ready(Some(Ok(e))), + Some(Err(e)) => e, + }; + let fut = (self.as_mut().f())(item); + self.as_mut().future().set(Some(fut)); + } + + let e = ready!(self.as_mut().future().as_pin_mut().unwrap().try_poll(cx)); + self.as_mut().future().set(None); + Poll::Ready(Some(e)) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let future_len = if self.future.is_some() { 1 } else { 0 }; + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(future_len); + let upper = match upper { + Some(x) => x.checked_add(future_len), + None => None, + }; + (lower, upper) + } +} + +impl<St, Fut, F> FusedStream for OrElse<St, Fut, F> + where St: TryStream + FusedStream, + F: FnMut(St::Error) -> Fut, + Fut: TryFuture<Ok = St::Ok>, +{ + fn is_terminated(&self) -> bool { + self.future.is_none() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for OrElse<S, Fut, F> + where S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/try_buffer_unordered.rs b/src/stream/try_stream/try_buffer_unordered.rs new file mode 100644 index 0000000..d11e1b4 --- /dev/null +++ b/src/stream/try_stream/try_buffer_unordered.rs @@ -0,0 +1,118 @@ +use crate::stream::{Fuse, FuturesUnordered, StreamExt, IntoStream}; +use crate::future::{IntoFuture, TryFutureExt}; +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_utils::{unsafe_pinned, unsafe_unpinned}; +use core::pin::Pin; + +/// 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 +{ + stream: Fuse<IntoStream<St>>, + in_progress_queue: FuturesUnordered<IntoFuture<St::Ok>>, + max: usize, +} + +impl<St> Unpin for TryBufferUnordered<St> + where St: TryStream + Unpin +{} + +impl<St> TryBufferUnordered<St> + where St: TryStream, + St::Ok: TryFuture, +{ + unsafe_pinned!(stream: Fuse<IntoStream<St>>); + unsafe_unpinned!(in_progress_queue: FuturesUnordered<IntoFuture<St::Ok>>); + + pub(super) fn new(stream: St, n: usize) -> Self { + TryBufferUnordered { + stream: IntoStream::new(stream).fuse(), + in_progress_queue: FuturesUnordered::new(), + max: n, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + self.stream.get_ref().get_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. + pub fn get_mut(&mut self) -> &mut St { + self.stream.get_mut().get_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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream().get_pin_mut().get_pin_mut() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream.into_inner().into_inner() + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + // 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 self.in_progress_queue.len() < self.max { + match self.as_mut().stream().poll_next(cx)? { + Poll::Ready(Some(fut)) => self.as_mut().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 self.as_mut().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 self.stream.is_done() { + Poll::Ready(None) + } else { + Poll::Pending + } + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item, E> Sink<Item> for TryBufferUnordered<S> + where S: TryStream + Sink<Item, Error = E>, + S::Ok: TryFuture<Error = E>, +{ + type Error = E; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/try_collect.rs b/src/stream/try_stream/try_collect.rs new file mode 100644 index 0000000..d22e8e8 --- /dev/null +++ b/src/stream/try_stream/try_collect.rs @@ -0,0 +1,62 @@ +use core::mem; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::{FusedStream, TryStream}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: St, + items: C, +} + +impl<St: TryStream, C: Default> TryCollect<St, C> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(items: C); + + pub(super) fn new(s: St) -> TryCollect<St, C> { + TryCollect { + stream: s, + items: Default::default(), + } + } + + fn finish(self: Pin<&mut Self>) -> C { + mem::replace(self.items(), Default::default()) + } +} + +impl<St: Unpin + TryStream, C> Unpin for TryCollect<St, C> {} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + loop { + match ready!(self.as_mut().stream().try_poll_next(cx)?) { + Some(x) => self.as_mut().items().extend(Some(x)), + None => return Poll::Ready(Ok(self.as_mut().finish())), + } + } + } +} diff --git a/src/stream/try_stream/try_concat.rs b/src/stream/try_stream/try_concat.rs new file mode 100644 index 0000000..395f166 --- /dev/null +++ b/src/stream/try_stream/try_concat.rs @@ -0,0 +1,57 @@ +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::stream::TryStream; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: St, + accum: Option<St::Ok>, +} + +impl<St: TryStream + Unpin> Unpin for TryConcat<St> {} + +impl<St> TryConcat<St> +where + St: TryStream, + St::Ok: Extend<<St::Ok as IntoIterator>::Item> + IntoIterator + Default, +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(accum: Option<St::Ok>); + + pub(super) fn new(stream: St) -> TryConcat<St> { + TryConcat { + 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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + loop { + match ready!(self.as_mut().stream().try_poll_next(cx)?) { + Some(x) => { + let accum = self.as_mut().accum(); + if let Some(a) = accum { + a.extend(x) + } else { + *accum = Some(x) + } + }, + None => { + return Poll::Ready(Ok(self.as_mut().accum().take().unwrap_or_default())) + } + } + } + } +} diff --git a/src/stream/try_stream/try_filter.rs b/src/stream/try_stream/try_filter.rs new file mode 100644 index 0000000..24a9c32 --- /dev/null +++ b/src/stream/try_stream/try_filter.rs @@ -0,0 +1,152 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::stream::{Stream, TryStream, FusedStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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 +{ + stream: St, + f: F, + pending_fut: Option<Fut>, + pending_item: Option<St::Ok>, +} + +impl<St, Fut, F> Unpin for TryFilter<St, Fut, F> + where St: TryStream + Unpin, Fut: Unpin, +{} + +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 +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + unsafe_pinned!(pending_fut: Option<Fut>); + unsafe_unpinned!(pending_item: Option<St::Ok>); + + pub(super) fn new(stream: St, f: F) -> Self { + TryFilter { + stream, + f, + pending_fut: None, + pending_item: None, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<St::Ok, St::Error>>> { + loop { + if self.pending_fut.is_none() { + let item = match ready!(self.as_mut().stream().try_poll_next(cx)?) { + Some(x) => x, + None => return Poll::Ready(None), + }; + let fut = (self.as_mut().f())(&item); + self.as_mut().pending_fut().set(Some(fut)); + *self.as_mut().pending_item() = Some(item); + } + + let yield_item = ready!(self.as_mut().pending_fut().as_pin_mut().unwrap().poll(cx)); + self.as_mut().pending_fut().set(None); + let item = self.as_mut().pending_item().take().unwrap(); + + if yield_item { + return Poll::Ready(Some(Ok(item))); + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = if self.pending_fut.is_some() { 1 } else { 0 }; + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item, E> Sink<Item> for TryFilter<S, Fut, F> + where S: TryStream + Sink<Item, Error = E>, +{ + type Error = E; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/try_filter_map.rs b/src/stream/try_stream/try_filter_map.rs new file mode 100644 index 0000000..ed7eeb2 --- /dev/null +++ b/src/stream/try_stream/try_filter_map.rs @@ -0,0 +1,136 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{TryFuture}; +use futures_core::stream::{Stream, TryStream, FusedStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: St, + f: F, + pending: Option<Fut>, +} + +impl<St, Fut, F> Unpin for TryFilterMap<St, Fut, F> + where St: Unpin, Fut: Unpin, +{} + +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> { + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + unsafe_pinned!(pending: Option<Fut>); + + pub(super) fn new(stream: St, f: F) -> Self { + TryFilterMap { stream, f, pending: None } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<T, St::Error>>> { + loop { + if self.pending.is_none() { + let item = match ready!(self.as_mut().stream().try_poll_next(cx)?) { + Some(x) => x, + None => return Poll::Ready(None), + }; + let fut = (self.as_mut().f())(item); + self.as_mut().pending().set(Some(fut)); + } + + let result = ready!(self.as_mut().pending().as_pin_mut().unwrap().try_poll(cx)); + self.as_mut().pending().set(None); + if let Some(x) = result? { + return Poll::Ready(Some(Ok(x))); + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = if self.pending.is_some() { 1 } else { 0 }; + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item> Sink<Item> for TryFilterMap<S, Fut, F> + where S: Sink<Item>, +{ + type Error = S::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/try_flatten.rs b/src/stream/try_stream/try_flatten.rs new file mode 100644 index 0000000..5f81b22 --- /dev/null +++ b/src/stream/try_stream/try_flatten.rs @@ -0,0 +1,129 @@ +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_utils::unsafe_pinned; + +/// 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, +{ + stream: St, + next: Option<St::Ok>, +} + +impl<St> Unpin for TryFlatten<St> +where + St: TryStream + Unpin, + St::Ok: Unpin, +{ +} + +impl<St> TryFlatten<St> +where + St: TryStream, +{ + unsafe_pinned!(stream: St); + unsafe_pinned!(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 } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + loop { + if self.next.is_none() { + match ready!(self.as_mut().stream().try_poll_next(cx)?) { + Some(e) => self.as_mut().next().set(Some(e)), + None => return Poll::Ready(None), + } + } + + if let Some(item) = ready!(self + .as_mut() + .next() + .as_pin_mut() + .unwrap() + .try_poll_next(cx)?) + { + return Poll::Ready(Some(Ok(item))); + } else { + self.as_mut().next().set(None); + } + } + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Item> Sink<Item> for TryFlatten<S> +where + S: TryStream + Sink<Item>, +{ + type Error = <S as Sink<Item>>::Error; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/try_fold.rs b/src/stream/try_stream/try_fold.rs new file mode 100644 index 0000000..b8b8dc2 --- /dev/null +++ b/src/stream/try_stream/try_fold.rs @@ -0,0 +1,110 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{FusedFuture, Future, TryFuture}; +use futures_core::stream::TryStream; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: St, + f: F, + accum: Option<T>, + future: Option<Fut>, +} + +impl<St: Unpin, Fut: Unpin, T, F> Unpin for TryFold<St, Fut, T, F> {} + +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>, +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + unsafe_unpinned!(accum: Option<T>); + unsafe_pinned!(future: Option<Fut>); + + pub(super) fn new(stream: St, f: F, t: T) -> TryFold<St, Fut, T, F> { + TryFold { + 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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + loop { + // we're currently processing a future to produce a new accum value + if self.accum.is_none() { + let accum = match ready!( + self.as_mut().future().as_pin_mut() + .expect("TryFold polled after completion") + .try_poll(cx) + ) { + Ok(accum) => accum, + Err(e) => { + // Indicate that the future can no longer be polled. + self.as_mut().future().set(None); + return Poll::Ready(Err(e)); + } + }; + *self.as_mut().accum() = Some(accum); + self.as_mut().future().set(None); + } + + let item = match ready!(self.as_mut().stream().try_poll_next(cx)) { + Some(Ok(item)) => Some(item), + Some(Err(e)) => { + // Indicate that the future can no longer be polled. + *self.as_mut().accum() = None; + return Poll::Ready(Err(e)); + } + None => None, + }; + let accum = self.as_mut().accum().take().unwrap(); + + if let Some(e) = item { + let future = (self.as_mut().f())(accum, e); + self.as_mut().future().set(Some(future)); + } else { + return Poll::Ready(Ok(accum)) + } + } + } +} diff --git a/src/stream/try_stream/try_for_each.rs b/src/stream/try_stream/try_for_each.rs new file mode 100644 index 0000000..2c71107 --- /dev/null +++ b/src/stream/try_stream/try_for_each.rs @@ -0,0 +1,72 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::{Future, TryFuture}; +use futures_core::stream::TryStream; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: St, + f: F, + future: Option<Fut>, +} + +impl<St: Unpin, Fut: Unpin, F> Unpin for TryForEach<St, Fut, F> {} + +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>, +{ + unsafe_pinned!(stream: St); + unsafe_unpinned!(f: F); + unsafe_pinned!(future: Option<Fut>); + + pub(super) fn new(stream: St, f: F) -> TryForEach<St, Fut, F> { + TryForEach { + 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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + loop { + if let Some(future) = self.as_mut().future().as_pin_mut() { + ready!(future.try_poll(cx))?; + } + self.as_mut().future().set(None); + + match ready!(self.as_mut().stream().try_poll_next(cx)?) { + Some(e) => { + let future = (self.as_mut().f())(e); + self.as_mut().future().set(Some(future)); + } + None => return Poll::Ready(Ok(())), + } + } + } +} diff --git a/src/stream/try_stream/try_for_each_concurrent.rs b/src/stream/try_stream/try_for_each_concurrent.rs new file mode 100644 index 0000000..19c3e5b --- /dev/null +++ b/src/stream/try_stream/try_for_each_concurrent.rs @@ -0,0 +1,139 @@ +use crate::stream::{FuturesUnordered, StreamExt}; +use core::fmt; +use core::mem; +use core::pin::Pin; +use core::num::NonZeroUsize; +use futures_core::future::{FusedFuture, Future}; +use futures_core::stream::TryStream; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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> { + stream: Option<St>, + f: F, + futures: FuturesUnordered<Fut>, + limit: Option<NonZeroUsize>, +} + +impl<St, Fut, F> Unpin for TryForEachConcurrent<St, Fut, F> +where St: Unpin, + Fut: Unpin, +{} + +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>>, +{ + unsafe_pinned!(stream: Option<St>); + unsafe_unpinned!(f: F); + unsafe_unpinned!(futures: FuturesUnordered<Fut>); + unsafe_unpinned!(limit: Option<NonZeroUsize>); + + pub(super) fn new(stream: St, limit: Option<usize>, f: F) -> TryForEachConcurrent<St, Fut, F> { + TryForEachConcurrent { + 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(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + loop { + let mut made_progress_this_iter = false; + + // Try and pull an item from the stream + let current_len = self.futures.len(); + // Check if we've already created a number of futures greater than `limit` + if self.limit.map(|limit| limit.get() > current_len).unwrap_or(true) { + let poll_res = match self.as_mut().stream().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) => { + self.as_mut().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. + self.as_mut().stream().set(None); + drop(mem::replace(self.as_mut().futures(), FuturesUnordered::new())); + return Poll::Ready(Err(e)); + } + }; + + if let Some(elem) = elem { + let next_future = (self.as_mut().f())(elem); + self.as_mut().futures().push(next_future); + } + } + + match self.as_mut().futures().poll_next_unpin(cx) { + Poll::Ready(Some(Ok(()))) => made_progress_this_iter = true, + Poll::Ready(None) => { + if self.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. + self.as_mut().stream().set(None); + drop(mem::replace(self.as_mut().futures(), FuturesUnordered::new())); + return Poll::Ready(Err(e)); + } + } + + if !made_progress_this_iter { + return Poll::Pending; + } + } + } +} diff --git a/src/stream/try_stream/try_next.rs b/src/stream/try_stream/try_next.rs new file mode 100644 index 0000000..78599ad --- /dev/null +++ b/src/stream/try_stream/try_next.rs @@ -0,0 +1,37 @@ +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 { + TryNext { stream } + } +} + +impl<St: ?Sized + TryStream + Unpin + FusedStream> FusedFuture for TryNext<'_, St> { + fn is_terminated(&self) -> bool { + self.stream.is_terminated() + } +} + +impl<St: ?Sized + TryStream + Unpin> Future for TryNext<'_, St> { + type Output = Result<Option<St::Ok>, St::Error>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + self.stream.try_poll_next_unpin(cx)?.map(Ok) + } +} diff --git a/src/stream/try_stream/try_skip_while.rs b/src/stream/try_stream/try_skip_while.rs new file mode 100644 index 0000000..a3d6803 --- /dev/null +++ b/src/stream/try_stream/try_skip_while.rs @@ -0,0 +1,164 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::stream::{Stream, TryStream, FusedStream}; +use futures_core::task::{Context, Poll}; +#[cfg(feature = "sink")] +use futures_sink::Sink; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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 { + stream: St, + f: F, + pending_fut: Option<Fut>, + pending_item: Option<St::Ok>, + done_skipping: bool, +} + +impl<St: Unpin + TryStream, Fut: Unpin, F> Unpin for TrySkipWhile<St, Fut, F> {} + +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, +{ + unsafe_pinned!(stream: St); +} + +impl<St, Fut, F> TrySkipWhile<St, Fut, F> + where St: TryStream, + F: FnMut(&St::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = St::Error>, +{ + unsafe_unpinned!(f: F); + unsafe_pinned!(pending_fut: Option<Fut>); + unsafe_unpinned!(pending_item: Option<St::Ok>); + unsafe_unpinned!(done_skipping: bool); + + pub(super) fn new(stream: St, f: F) -> TrySkipWhile<St, Fut, F> { + TrySkipWhile { + stream, + f, + pending_fut: None, + pending_item: None, + done_skipping: false, + } + } + + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &St { + &self.stream + } + + /// 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. + pub fn get_mut(&mut self) -> &mut St { + &mut self.stream + } + + /// 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. + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> { + self.stream() + } + + /// 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. + pub fn into_inner(self) -> St { + self.stream + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.done_skipping { + return self.as_mut().stream().try_poll_next(cx); + } + + loop { + if self.pending_item.is_none() { + let item = match ready!(self.as_mut().stream().try_poll_next(cx)?) { + Some(e) => e, + None => return Poll::Ready(None), + }; + let fut = (self.as_mut().f())(&item); + self.as_mut().pending_fut().set(Some(fut)); + *self.as_mut().pending_item() = Some(item); + } + + let skipped = ready!(self.as_mut().pending_fut().as_pin_mut().unwrap().try_poll(cx)?); + let item = self.as_mut().pending_item().take().unwrap(); + self.as_mut().pending_fut().set(None); + + if !skipped { + *self.as_mut().done_skipping() = true; + return Poll::Ready(Some(Ok(item))) + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let pending_len = if self.pending_item.is_some() { 1 } else { 0 }; + let (_, upper) = self.stream.size_hint(); + let upper = match upper { + Some(x) => x.checked_add(pending_len), + None => None, + }; + (0, upper) // can't know a lower bound, due to the predicate + } +} + +impl<St, Fut, F> FusedStream for TrySkipWhile<St, Fut, F> + where St: TryStream + FusedStream, + F: FnMut(&St::Ok) -> Fut, + Fut: TryFuture<Ok = bool, Error = St::Error>, +{ + fn is_terminated(&self) -> bool { + self.pending_item.is_none() && self.stream.is_terminated() + } +} + +// Forwarding impl of Sink from the underlying stream +#[cfg(feature = "sink")] +impl<S, Fut, F, Item, E> Sink<Item> for TrySkipWhile<S, Fut, F> + where S: TryStream + Sink<Item, Error = E>, +{ + type Error = E; + + delegate_sink!(stream, Item); +} diff --git a/src/stream/try_stream/try_unfold.rs b/src/stream/try_stream/try_unfold.rs new file mode 100644 index 0000000..6266274 --- /dev/null +++ b/src/stream/try_stream/try_unfold.rs @@ -0,0 +1,134 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::TryFuture; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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)>>, +{ + TryUnfold { + f, + state: Some(init), + fut: None, + } +} + +/// Stream for the [`try_unfold`] function. +#[must_use = "streams do nothing unless polled"] +pub struct TryUnfold<T, F, Fut> { + f: F, + state: Option<T>, + fut: Option<Fut>, +} + +impl<T, F, Fut: Unpin> Unpin for TryUnfold<T, F, 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> TryUnfold<T, F, Fut> { + unsafe_unpinned!(f: F); + unsafe_unpinned!(state: Option<T>); + unsafe_pinned!(fut: Option<Fut>); +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<Item, Fut::Error>>> { + if let Some(state) = self.as_mut().state().take() { + let fut = (self.as_mut().f())(state); + self.as_mut().fut().set(Some(fut)); + } + + match self.as_mut().fut().as_pin_mut() { + None => { + // The future previously errored + Poll::Ready(None) + } + Some(fut) => { + let step = ready!(fut.try_poll(cx)); + self.as_mut().fut().set(None); + + match step { + Ok(Some((item, next_state))) => { + *self.as_mut().state() = Some(next_state); + Poll::Ready(Some(Ok(item))) + } + Ok(None) => Poll::Ready(None), + Err(e) => Poll::Ready(Some(Err(e))), + } + } + } + } +} diff --git a/src/stream/unfold.rs b/src/stream/unfold.rs new file mode 100644 index 0000000..3153f83 --- /dev/null +++ b/src/stream/unfold.rs @@ -0,0 +1,122 @@ +use core::fmt; +use core::pin::Pin; +use futures_core::future::Future; +use futures_core::stream::{FusedStream, Stream}; +use futures_core::task::{Context, Poll}; +use pin_utils::{unsafe_pinned, unsafe_unpinned}; + +/// 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)>>, +{ + Unfold { + f, + state: Some(init), + fut: None, + } +} + +/// Stream for the [`unfold`] function. +#[must_use = "streams do nothing unless polled"] +pub struct Unfold<T, F, Fut> { + f: F, + state: Option<T>, + fut: Option<Fut>, +} + +impl<T, F, Fut: Unpin> Unpin for Unfold<T, F, 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) + .field("fut", &self.fut) + .finish() + } +} + +impl<T, F, Fut> Unfold<T, F, Fut> { + unsafe_unpinned!(f: F); + unsafe_unpinned!(state: Option<T>); + unsafe_pinned!(fut: Option<Fut>); +} + +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 { + self.state.is_none() && self.fut.is_none() + } +} + +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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if let Some(state) = self.as_mut().state().take() { + let fut = (self.as_mut().f())(state); + self.as_mut().fut().set(Some(fut)); + } + + let step = ready!(self.as_mut().fut().as_pin_mut() + .expect("Unfold must not be polled after it returned `Poll::Ready(None)`").poll(cx)); + self.as_mut().fut().set(None); + + if let Some((item, next_state)) = step { + *self.as_mut().state() = Some(next_state); + Poll::Ready(Some(item)) + } else { + Poll::Ready(None) + } + } +} diff --git a/src/task/mod.rs b/src/task/mod.rs new file mode 100644 index 0000000..fb3b7ad --- /dev/null +++ b/src/task/mod.rs @@ -0,0 +1,28 @@ +//! Task notification + +cfg_target_has_atomic! { + #[cfg(feature = "alloc")] + pub use futures_task::ArcWake; + + #[cfg(feature = "alloc")] + pub use futures_task::waker; + + #[cfg(feature = "alloc")] + pub use futures_task::{waker_ref, WakerRef}; + + pub use futures_core::task::__internal::AtomicWaker; +} + +mod spawn; +pub use self::spawn::{SpawnExt, LocalSpawnExt}; + +pub use futures_core::task::{Context, Poll, Waker, RawWaker, RawWakerVTable}; + +pub use futures_task::{ + Spawn, LocalSpawn, SpawnError, + FutureObj, LocalFutureObj, UnsafeFutureObj, +}; + +pub use futures_task::noop_waker; +#[cfg(feature = "std")] +pub use futures_task::noop_waker_ref; diff --git a/src/task/spawn.rs b/src/task/spawn.rs new file mode 100644 index 0000000..f34445a --- /dev/null +++ b/src/task/spawn.rs @@ -0,0 +1,160 @@ +use futures_task::{LocalSpawn, Spawn}; + +#[cfg(feature = "compat")] +use crate::compat::Compat; + +#[cfg(feature = "channel")] +#[cfg(feature = "std")] +use crate::future::{FutureExt, RemoteHandle}; +#[cfg(feature = "alloc")] +use alloc::boxed::Box; +#[cfg(feature = "alloc")] +use futures_core::future::Future; +#[cfg(feature = "alloc")] +use futures_task::{FutureObj, LocalFutureObj, SpawnError}; + +impl<Sp: ?Sized> SpawnExt for Sp where Sp: Spawn {} +impl<Sp: ?Sized> LocalSpawnExt for Sp where Sp: LocalSpawn {} + +/// Extension trait for `Spawn`. +pub trait SpawnExt: Spawn { + /// Spawns a task that polls the given future with output `()` to + /// completion. + /// + /// This method returns a [`Result`] that contains a [`SpawnError`] if + /// spawning fails. + /// + /// You can use [`spawn_with_handle`](SpawnExt::spawn_with_handle) if + /// you want to spawn a future with output other than `()` or if you want + /// to be able to await its completion. + /// + /// Note this method will eventually be replaced with the upcoming + /// `Spawn::spawn` method which will take a `dyn Future` as input. + /// Technical limitations prevent `Spawn::spawn` from being implemented + /// today. Feel free to use this method in the meantime. + /// + /// ``` + /// use futures::executor::ThreadPool; + /// use futures::task::SpawnExt; + /// + /// let executor = ThreadPool::new().unwrap(); + /// + /// let future = async { /* ... */ }; + /// executor.spawn(future).unwrap(); + /// ``` + #[cfg(feature = "alloc")] + fn spawn<Fut>(&self, future: Fut) -> Result<(), SpawnError> + where + Fut: Future<Output = ()> + Send + 'static, + { + self.spawn_obj(FutureObj::new(Box::new(future))) + } + + /// Spawns a task that polls the given future to completion and returns a + /// future that resolves to the spawned future's output. + /// + /// This method returns a [`Result`] that contains a [`RemoteHandle`](crate::future::RemoteHandle), or, if + /// spawning fails, a [`SpawnError`]. [`RemoteHandle`](crate::future::RemoteHandle) is a future that + /// resolves to the output of the spawned future. + /// + /// ``` + /// use futures::executor::{block_on, ThreadPool}; + /// use futures::future; + /// use futures::task::SpawnExt; + /// + /// let executor = ThreadPool::new().unwrap(); + /// + /// let future = future::ready(1); + /// let join_handle_fut = executor.spawn_with_handle(future).unwrap(); + /// assert_eq!(block_on(join_handle_fut), 1); + /// ``` + #[cfg(feature = "channel")] + #[cfg(feature = "std")] + fn spawn_with_handle<Fut>(&self, future: Fut) -> Result<RemoteHandle<Fut::Output>, SpawnError> + where + Fut: Future + Send + 'static, + Fut::Output: Send, + { + let (future, handle) = future.remote_handle(); + self.spawn(future)?; + Ok(handle) + } + + /// Wraps a [`Spawn`] and makes it usable as a futures 0.1 `Executor`. + /// Requires the `compat` feature to enable. + #[cfg(feature = "compat")] + fn compat(self) -> Compat<Self> + where + Self: Sized, + { + Compat::new(self) + } +} + +/// Extension trait for `LocalSpawn`. +pub trait LocalSpawnExt: LocalSpawn { + /// Spawns a task that polls the given future with output `()` to + /// completion. + /// + /// This method returns a [`Result`] that contains a [`SpawnError`] if + /// spawning fails. + /// + /// You can use [`spawn_with_handle`](SpawnExt::spawn_with_handle) if + /// you want to spawn a future with output other than `()` or if you want + /// to be able to await its completion. + /// + /// Note this method will eventually be replaced with the upcoming + /// `Spawn::spawn` method which will take a `dyn Future` as input. + /// Technical limitations prevent `Spawn::spawn` from being implemented + /// today. Feel free to use this method in the meantime. + /// + /// ``` + /// use futures::executor::LocalPool; + /// use futures::task::LocalSpawnExt; + /// + /// let executor = LocalPool::new(); + /// let spawner = executor.spawner(); + /// + /// let future = async { /* ... */ }; + /// spawner.spawn_local(future).unwrap(); + /// ``` + #[cfg(feature = "alloc")] + fn spawn_local<Fut>(&self, future: Fut) -> Result<(), SpawnError> + where + Fut: Future<Output = ()> + 'static, + { + self.spawn_local_obj(LocalFutureObj::new(Box::new(future))) + } + + /// Spawns a task that polls the given future to completion and returns a + /// future that resolves to the spawned future's output. + /// + /// This method returns a [`Result`] that contains a [`RemoteHandle`](crate::future::RemoteHandle), or, if + /// spawning fails, a [`SpawnError`]. [`RemoteHandle`](crate::future::RemoteHandle) is a future that + /// resolves to the output of the spawned future. + /// + /// ``` + /// use futures::executor::LocalPool; + /// use futures::task::LocalSpawnExt; + /// + /// let mut executor = LocalPool::new(); + /// let spawner = executor.spawner(); + /// + /// let future = async { 1 }; + /// let join_handle_fut = spawner.spawn_local_with_handle(future).unwrap(); + /// assert_eq!(executor.run_until(join_handle_fut), 1); + /// ``` + #[cfg(feature = "channel")] + #[cfg(feature = "std")] + fn spawn_local_with_handle<Fut>( + &self, + future: Fut, + ) -> Result<RemoteHandle<Fut::Output>, SpawnError> + where + Fut: Future + 'static, + { + let (future, handle) = future.remote_handle(); + self.spawn_local(future)?; + Ok(handle) + } +} |