Rework simulation

Signed-off-by: Denis Varlakov <denis@dfns.co>

examples/random-generation-protocol/Cargo.toml

@@ -22,7 +22,6 @@ generic-array = { version = "0.14", features = ["serde"] }
 [dev-dependencies]
 round-based = { path = "../../round-based", features = ["derive", "dev", "state-machine"]  }
 tokio = { version = "1.15", features = ["macros", "rt"] }
-futures = "0.3"
 hex = "0.4"
 rand_dev = "0.1"
 rand = "0.8"

examples/random-generation-protocol/src/lib.rs

@@ -173,56 +173,39 @@ pub struct Blame {
 
 #[cfg(test)]
 mod tests {
-    use alloc::{vec, vec::Vec};
-
     use rand::Rng;
-    use round_based::simulation::Simulation;
     use sha2::{Digest, Sha256};
 
-    use super::{protocol_of_random_generation, Msg};
+    use super::protocol_of_random_generation;
 
     #[tokio::test]
     async fn simulation_async() {
         let mut rng = rand_dev::DevRng::new();
 
         let n: u16 = 5;
 
-        let mut simulation = Simulation::<Msg>::new();
-        let mut party_output = vec![];
-
-        for i in 0..n {
-            let party = simulation.add_party();
-            let output = protocol_of_random_generation(party, i, n, rng.fork());
-            party_output.push(output);
-        }
-
-        let output = futures::future::try_join_all(party_output).await.unwrap();
-
-        // Assert that all parties outputed the same randomness
-        for i in 1..n {
-            assert_eq!(output[0], output[usize::from(i)]);
-        }
+        let randomness = round_based::simulation::run_with_setup(
+            core::iter::repeat_with(|| rng.fork()).take(n.into()),
+            |i, party, rng| protocol_of_random_generation(party, i, n, rng),
+        )
+        .await
+        .expect_success()
+        .expect_same();
 
-        std::println!("Output randomness: {}", hex::encode(output[0]));
+        std::println!("Output randomness: {}", hex::encode(randomness));
     }
 
     #[test]
     fn simulation_sync() {
         let mut rng = rand_dev::DevRng::new();
 
-        let simulation = round_based::simulation::SimulationSync::from_async_fn(5, |i, party| {
+        round_based::simulation::SimulationSync::from_async_fn(5, |i, party| {
             protocol_of_random_generation(party, i, 5, rng.fork())
-        });
-
-        let outputs = simulation
-            .run()
-            .unwrap()
-            .into_iter()
-            .collect::<Result<Vec<_>, _>>()
-            .unwrap();
-        for output_i in &outputs {
-            assert_eq!(*output_i, outputs[0]);
-        }
+        })
+        .run()
+        .unwrap()
+        .expect_success()
+        .expect_same();
     }
 
     // Emulate the protocol using the state machine interface

round-based/Cargo.toml

@@ -32,11 +32,15 @@ trybuild = "1"
 matches = "0.1"
 futures = { version = "0.3", default-features = false }
 tokio = { version = "1", features = ["macros"] }
+hex = "0.4"
+
+rand = "0.8"
+rand_dev = "0.1"
 
 [features]
 default = ["std"]
 state-machine = []
-dev = ["std", "tokio/sync", "tokio-stream"]
+dev = ["std", "tokio/sync", "tokio-stream", "futures-util/alloc"]
 derive = ["round-based-derive"]
 runtime-tokio = ["tokio"]
 std = ["thiserror"]

round-based/src/simulation/mod.rs

@@ -1,43 +1,54 @@
 //! Multiparty protocol simulation
 //!
-//! [`Simulation`] is an essential developer tool for testing the multiparty protocol locally.
+//! Simulator is an essential developer tool for testing the multiparty protocol locally.
 //! It covers most of the boilerplate by mocking networking.
 //!
-//! ## Example
+//! The entry point is either [`run`] or [`run_with_setup`] functions. They take a protocol
+//! defined as an async function, provide simulated networking, carry out the simulation,
+//! and return the result.
+//!
+//! If you need more control over execution, you can use [`Network`] to simulate the networking
+//! and carry out the protocol manually.
+//!
+//! When `state-machine` feature is enabled, [`SimulationSync`] is available which can carry out
+//! protocols defined as a state machine.
 //!
-//! ```rust
+//! ## Example
+//! ```rust,no_run
+//! # #[tokio::main(flavor = "current_thread")]
+//! # async fn main() {
 //! use round_based::{Mpc, PartyIndex};
-//! use round_based::simulation::Simulation;
-//! use futures::future::try_join_all;
 //!
 //! # type Result<T, E = ()> = std::result::Result<T, E>;
 //! # type Randomness = [u8; 32];
 //! # type Msg = ();
 //! // Any MPC protocol you want to test
-//! pub async fn protocol_of_random_generation<M>(party: M, i: PartyIndex, n: u16) -> Result<Randomness>
-//! where M: Mpc<ProtocolMessage = Msg>
+//! pub async fn protocol_of_random_generation<M>(
+//!     party: M,
+//!     i: PartyIndex,
+//!     n: u16
+//! ) -> Result<Randomness>
+//! where
+//!     M: Mpc<ProtocolMessage = Msg>
 //! {
 //!     // ...
 //! # todo!()
 //! }
 //!
-//! async fn test_randomness_generation() {
-//!     let n = 3;
-//!
-//!     let mut simulation = Simulation::<Msg>::new();
-//!     let mut outputs = vec![];
-//!     for i in 0..n {
-//!         let party = simulation.add_party();
-//!         outputs.push(protocol_of_random_generation(party, i, n));
-//!     }   
-//!
-//!     // Waits each party to complete the protocol
-//!     let outputs = try_join_all(outputs).await.expect("protocol wasn't completed successfully");
-//!     // Asserts that all parties output the same randomness
-//!     for output in outputs.iter().skip(1) {
-//!         assert_eq!(&outputs[0], output);
-//!     }  
-//! }
+//! let n = 3;
+//!
+//! let output = round_based::simulation::run(
+//!     n,
+//!     |i, party| protocol_of_random_generation(party, i, n),
+//! )
+//! .await
+//! // unwrap `Result`s
+//! .expect_success()
+//! // check that all parties produced the same response
+//! .expect_same();
+//!
+//! println!("Output randomness: {}", hex::encode(output));
+//! # }  
 //! ```
 
 mod sim_async;
@@ -47,3 +58,102 @@ mod sim_sync;
 pub use sim_async::*;
 #[cfg(feature = "state-machine")]
 pub use sim_sync::*;
+
+/// Result of the simulation
+pub struct SimResult<T>(pub alloc::vec::Vec<T>);
+
+impl<T, E> SimResult<Result<T, E>>
+where
+    E: core::fmt::Debug,
+{
+    /// Unwraps `Result<T, E>` produced by each party
+    ///
+    /// Panics if at least one of the parties returned `Err(_)`. In this case,
+    /// a verbose error message will shown specifying which of the parties returned
+    /// an error.
+    pub fn expect_success(self) -> SimResult<T> {
+        let mut oks = alloc::vec::Vec::with_capacity(self.0.len());
+        let mut errs = alloc::vec::Vec::with_capacity(self.0.len());
+
+        for (res, i) in self.0.into_iter().zip(0u16..) {
+            match res {
+                Ok(res) => oks.push(res),
+                Err(res) => errs.push((i, res)),
+            }
+        }
+
+        if !errs.is_empty() {
+            let mut msg = alloc::format!(
+                "Simulation output didn't match expectations.\n\
+                Expected: all parties succeed\n\
+                Actual  : {success} parties succeeded, {failed} parties returned an error\n\
+                Failures:\n",
+                success = oks.len(),
+                failed = errs.len(),
+            );
+
+            for (i, err) in errs {
+                msg += &alloc::format!("- Party {i}: {err:?}\n");
+            }
+
+            panic!("{msg}");
+        }
+
+        SimResult(oks)
+    }
+}
+
+impl<T> SimResult<T>
+where
+    T: PartialEq + core::fmt::Debug,
+{
+    /// Checks that outputs of all parties are equally the same
+    ///
+    /// Returns the output on success (all the outputs are checked to be the same), otherwise
+    /// panics with a verbose error message.
+    ///
+    /// Panics if simulation contained zero parties.
+    pub fn expect_same(mut self) -> T {
+        let Some(first) = self.0.get(0) else {
+            panic!("simulation contained zero parties");
+        };
+
+        if !self.0[1..].iter().all(|i| i == first) {
+            let mut msg = alloc::format!(
+                "Simulation output didn't match expectations.\n\
+                Expected: all parties return the same output\n\
+                Actual  : some of the parties returned a different output\n\
+                Outputs :\n"
+            );
+
+            for (i, res) in self.0.iter().enumerate() {
+                msg += &alloc::format!("- Party {i}: {res:?}");
+            }
+
+            panic!("{msg}")
+        }
+
+        self.0
+            .pop()
+            .expect("we checked that the list contains at least one element")
+    }
+}
+
+impl<T> SimResult<T> {
+    /// Deconstructs the simulation result returning inner list of results
+    pub fn into_vec(self) -> alloc::vec::Vec<T> {
+        self.0
+    }
+}
+
+impl<T> From<alloc::vec::Vec<T>> for SimResult<T> {
+    fn from(list: alloc::vec::Vec<T>) -> Self {
+        Self(list)
+    }
+}
+
+impl<T> From<SimResult<T>> for alloc::vec::Vec<T> {
+    fn from(res: SimResult<T>) -> Self {
+        res.0
+    }
+}