feat(mpz-common): scoped! macro (#143)

privacy-scaling-explorations · May 31, 2024 · 2f35271 · 2f35271
1 parent 9b51bd4
commit 2f35271
Show file tree

Hide file tree

Showing 4 changed files with 103 additions and 57 deletions.
diff --git 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
@@ -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::<u8>().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::<u8>().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
@@ -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
@@ -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: futures::Future<Output = u16>>(_: 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, Fut>(_: T)
+        where
+            T: Fn(u8, u16) -> Fut,
+            Fut: Future<Output = u8>,
+        {
+        }
+
+        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
+        });
+    }
+}