WIP: Astra Setup for ABY2

crates/seec-channel/src/multi.rs

@@ -1,4 +1,6 @@
-use crate::{multi, sub_channel, tcp, CommunicationError, Receiver, ReceiverT, Sender, SenderT};
+use crate::{
+    channel, multi, sub_channel, tcp, CommunicationError, Receiver, ReceiverT, Sender, SenderT,
+};
 use async_trait::async_trait;
 use futures::future::join;
 use futures::stream::FuturesUnordered;
@@ -54,14 +56,14 @@ pub struct MultiReceiver<T> {
 }
 
 impl<T: RemoteSend + Clone> MultiSender<T> {
-    pub async fn send_to(&self, to: impl IntoIterator<Item = &u32>, msg: T) -> Result<(), Error> {
+    pub async fn send_to(&self, to: impl IntoIterator<Item = u32>, msg: T) -> Result<(), Error> {
         let mut fu = FuturesUnordered::new();
         for to in to {
             debug!(to, "Sending");
             let sender = self
                 .senders
-                .get(to)
-                .ok_or_else(|| Error::UnknownParty(*to))?;
+                .get(&to)
+                .ok_or_else(|| Error::UnknownParty(to))?;
             fu.push(sender.send(msg.clone()));
         }
         let mut errors = vec![];
@@ -81,7 +83,7 @@ impl<T: RemoteSend + Clone> MultiSender<T> {
 
     #[instrument(level = "debug", skip(self, msg), ret)]
     pub async fn send_all(&self, msg: T) -> Result<(), Error> {
-        self.send_to(self.senders.keys(), msg).await
+        self.send_to(self.senders.keys().copied(), msg).await
     }
 
     pub fn sender(&self, to: u32) -> Option<&Sender<T>> {
@@ -100,6 +102,14 @@ pub struct MsgFrom<T> {
 }
 
 impl<T: RemoteSend> MultiReceiver<T> {
+    pub async fn recv_from_single(&mut self, from: u32) -> Result<T, Error> {
+        let receiver = self
+            .receivers
+            .get_mut(&from)
+            .ok_or(Error::UnknownParty(from))?;
+        Ok(map_recv_fut((&from, receiver)).await?.into_msg())
+    }
+
     pub fn recv_from(
         &mut self,
         from: &HashSet<u32>,
@@ -337,6 +347,44 @@ where
     Ok(MultiReceiver { receivers })
 }
 
+impl<T> MsgFrom<T> {
+    pub fn into_msg(self) -> T {
+        self.msg
+    }
+}
+
+pub fn new_local<T: RemoteSend>(parties: usize) -> Vec<(MultiSender<T>, MultiReceiver<T>)> {
+    let mut res: Vec<(MultiSender<T>, MultiReceiver<T>)> =
+        (0..parties).map(|_| Default::default()).collect();
+    for party in 0..parties {
+        for other in 0..parties {
+            if party == other {
+                continue;
+            }
+            let (sender, receiver) = channel(128);
+            res[party].0.senders.insert(other as u32, sender);
+            res[other].1.receivers.insert(party as u32, receiver);
+        }
+    }
+    res
+}
+
+impl<T> Default for MultiSender<T> {
+    fn default() -> Self {
+        Self {
+            senders: Default::default(),
+        }
+    }
+}
+
+impl<T> Default for MultiReceiver<T> {
+    fn default() -> Self {
+        Self {
+            receivers: Default::default(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/seec/src/executor.rs

@@ -42,6 +42,7 @@ pub type DynFDSetup<'c, P, Idx> = Box<
         + 'c,
 >;
 
+#[derive(Debug, Clone)]
 pub struct GateOutputs<Shares> {
     data: Vec<Input<Shares>>,
     // Used as a sanity check in debug builds. Stores for which gates we have set the output,
@@ -502,6 +503,10 @@ impl<Shares> GateOutputs<Shares> {
     pub fn iter(&self) -> impl Iterator<Item = &Input<Shares>> {
         self.data.iter()
     }
+
+    pub fn into_iter(self) -> impl Iterator<Item = Input<Shares>> {
+        self.data.into_iter()
+    }
 }
 
 impl<Shares> Input<Shares> {
@@ -625,6 +630,24 @@ impl<Shares: Clone> GateOutputs<Shares> {
     }
 }
 
+impl<S> Default for GateOutputs<S> {
+    fn default() -> Self {
+        Self {
+            data: vec![],
+            output_set: Default::default(),
+        }
+    }
+}
+
+impl<Shares> FromIterator<Input<Shares>> for GateOutputs<Shares> {
+    fn from_iter<T: IntoIterator<Item = Input<Shares>>>(iter: T) -> Self {
+        Self {
+            data: iter.into_iter().collect(),
+            output_set: Default::default(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::circuit::base_circuit::BaseGate;