From 2f352710b5b5021041055e671e813c20dce6312b Mon Sep 17 00:00:00 2001 From: "sinu.eth" <65924192+sinui0@users.noreply.github.com> Date: Fri, 31 May 2024 06:36:31 -0700 Subject: [PATCH] feat(mpz-common): scoped! macro (#143) --- crates/mpz-common/src/executor/dummy.rs | 17 ++---- crates/mpz-common/src/executor/mt.rs | 40 ++++-------- crates/mpz-common/src/executor/st.rs | 22 ++----- crates/mpz-common/src/lib.rs | 81 +++++++++++++++++++++++++ 4 files changed, 103 insertions(+), 57 deletions(-) diff --git a/crates/mpz-common/src/executor/dummy.rs b/crates/mpz-common/src/executor/dummy.rs index 26a7ad5f..a95a458a 100644 --- a/crates/mpz-common/src/executor/dummy.rs +++ b/crates/mpz-common/src/executor/dummy.rs @@ -125,7 +125,8 @@ impl Context for DummyExecutor { #[cfg(test)] mod tests { use futures::executor::block_on; - use scoped_futures::ScopedFutureExt; + + use crate::scoped; use super::*; @@ -142,18 +143,8 @@ mod tests { let a = &mut self.a; let b = &mut self.b; ctx.join( - |ctx| { - async move { - *a = ctx.id().clone(); - } - .scope_boxed() - }, - |ctx| { - async move { - *b = ctx.id().clone(); - } - .scope_boxed() - }, + scoped!(|ctx| *a = ctx.id().clone()), + scoped!(|ctx| *b = ctx.id().clone()), ) .await .unwrap(); diff --git a/crates/mpz-common/src/executor/mt.rs b/crates/mpz-common/src/executor/mt.rs index 7fe0d5b3..54f96781 100644 --- a/crates/mpz-common/src/executor/mt.rs +++ b/crates/mpz-common/src/executor/mt.rs @@ -298,10 +298,9 @@ where #[cfg(test)] mod tests { - use scoped_futures::ScopedFutureExt; use serio::{stream::IoStreamExt, SinkExt}; - use crate::{executor::test_mt_executor, join}; + use crate::{executor::test_mt_executor, scoped}; use super::*; @@ -318,15 +317,11 @@ mod tests { let a = &mut self.a; let b = &mut self.b; - join! { - ctx, - async { - *a = ctx.id().clone(); - }, - async { - *b = ctx.id().clone(); - } - } + ctx.join( + scoped!(|ctx| *a = ctx.id().clone()), + scoped!(|ctx| *b = ctx.id().clone()), + ) + .await .unwrap(); // Make sure we can mutate the fields after borrowing them in the async closures. @@ -355,22 +350,13 @@ mod tests { let (mut ctx_a, mut ctx_b) = futures::try_join!(exec_a.new_thread(), exec_b.new_thread()).unwrap(); - let (_, received) = futures::join!( - async { - ctx_a - .blocking(|ctx| async { ctx.io_mut().send(1u8).await.unwrap() }.scope_boxed()) - .await - .unwrap() - }, - async { - ctx_b - .blocking(|ctx| { - async { ctx.io_mut().expect_next::().await.unwrap() }.scope_boxed() - }) - .await - .unwrap() - } - ); + let (_, received) = futures::try_join!( + ctx_a.blocking(scoped!(|ctx| ctx.io_mut().send(1u8).await.unwrap())), + ctx_b.blocking(scoped!(|ctx| async move { + ctx.io_mut().expect_next::().await.unwrap() + })) + ) + .unwrap(); assert_eq!(received, 1u8); assert!(ctx_a.inner.is_some()); diff --git a/crates/mpz-common/src/executor/st.rs b/crates/mpz-common/src/executor/st.rs index 004aa86b..ee655533 100644 --- a/crates/mpz-common/src/executor/st.rs +++ b/crates/mpz-common/src/executor/st.rs @@ -125,9 +125,10 @@ where #[cfg(test)] mod tests { use futures::executor::block_on; - use scoped_futures::ScopedFutureExt; use serio::channel::duplex; + use crate::scoped; + use super::*; #[derive(Debug, Default)] @@ -143,18 +144,8 @@ mod tests { let a = &mut self.a; let b = &mut self.b; ctx.join( - |ctx| { - async move { - *a = ctx.id().clone(); - } - .scope_boxed() - }, - |ctx| { - async move { - *b = ctx.id().clone(); - } - .scope_boxed() - }, + scoped!(|ctx| *a = ctx.id().clone()), + scoped!(|ctx| *b = ctx.id().clone()), ) .await .unwrap(); @@ -180,10 +171,7 @@ mod tests { let mut ctx = STExecutor::new(io); block_on(async { - let id = ctx - .blocking(|ctx| async move { ctx.id().clone() }.scope_boxed()) - .await - .unwrap(); + let id = ctx.blocking(scoped!(|ctx| ctx.id().clone())).await.unwrap(); assert_eq!(&id, ctx.id()); assert!(ctx.inner.is_some()); diff --git a/crates/mpz-common/src/lib.rs b/crates/mpz-common/src/lib.rs index 6d100ecb..f22aca74 100644 --- a/crates/mpz-common/src/lib.rs +++ b/crates/mpz-common/src/lib.rs @@ -28,3 +28,84 @@ pub use id::{Counter, ThreadId}; // Re-export scoped-futures for use with the callback-like API in `Context`. pub use scoped_futures; + +/// A convenience macro for creating a closure which returns a scoped future. +/// +/// # Example +/// +/// ``` +/// # use mpz_common::scoped; +/// +/// let closure = scoped!(|a: u8, b: u16| a as u16 + b); +/// let fut = closure(1, 2); +/// +/// fn is_future>(_: T) {} +/// +/// is_future(fut); +/// ``` +#[macro_export] +macro_rules! scoped { + // Async move block. + (| $($arg:ident $(: $typ:ty)?),* | async move $body:block) => {{ + #[allow(unused_imports)] + use $crate::scoped_futures::ScopedFutureExt; + | $($arg $( : $typ )?),* | async move $body.scope_boxed() + }}; + // Async move block, move. + (move | $($arg:ident $(: $typ:ty)?),* | async move $body:block) => {{ + #[allow(unused_imports)] + use $crate::scoped_futures::ScopedFutureExt; + move | $($arg $( : $typ )?),* | async move $body.scope_boxed() + }}; + // No async block. + (| $($arg:ident $(: $typ:ty)?),* | $body:expr) => {{ + #[allow(unused_imports)] + use $crate::scoped_futures::ScopedFutureExt; + | $($arg $( : $typ )?),* | async move { $body }.scope_boxed() + }}; + // No async block, move. + (move | $($arg:ident $(: $typ:ty)?),* | $body:expr) => {{ + #[allow(unused_imports)] + use $crate::scoped_futures::ScopedFutureExt; + move | $($arg $( : $typ )?),* | async move { $body }.scope_boxed() + }}; +} + +#[cfg(test)] +mod tests { + use futures::Future; + + #[test] + fn test_scoped_macro() { + fn assert_signature(_: T) + where + T: Fn(u8, u16) -> Fut, + Fut: Future, + { + } + + assert_signature(scoped! { + |a: u8, _b: u16| async move { a } + }); + + assert_signature(scoped! { + move |a, _b| async move { a } + }); + + assert_signature(scoped! { + |a, _b| a + }); + + assert_signature(scoped! { + move |a: u8, _b| a + }); + + assert_signature(scoped! { + |a: u8, _b| a + }); + + assert_signature(scoped! { + |a, _b: u16| a + }); + } +}