use spawn-compute in correlated OT

Author: robinhundt

cryprot-ot/src/extension.rs

@@ -634,33 +634,42 @@ impl<S: Security> CotSender for OtExtensionSender<S> {
     ) -> Result<(), Self::Error>
     where
         B: Buf<Block>,
-        F: FnMut(usize, Block) -> Block + Send,
+        F: FnMut(usize, Block) -> Block + Send + 'static,
     {
         let mut r_ots: B::BufKind<[Block; 2]> = B::BufKind::zeroed(ots.len());
         self.send_into(&mut r_ots).await?;
-        let mut send_buf: Vec<Block> = Vec::zeroed(COR_CHUNK_SIZE);
         let (mut tx, _) = self.conn.byte_stream().await?;
         // TODO benchmark if this should use spawn_compute, I suspect not, since we're
         // not doing that much work per byte sent
-        for (chunk_idx, (ot_chunk, rot_chunk)) in ots
-            .chunks_mut(send_buf.len())
-            .zip(r_ots.chunks(send_buf.len()))
-            .enumerate()
-        {
-            for (idx, ((ot, r_ot), correction)) in ot_chunk
-                .iter_mut()
-                .zip(rot_chunk)
-                .zip(&mut send_buf)
+        let mut ots_owned = mem::take(ots);
+        let (ch_s, mut ch_r) = tokio::sync::mpsc::channel(5);
+        let jh = spawn_compute(move || {
+            for (chunk_idx, (ot_chunk, rot_chunk)) in ots_owned
+                .chunks_mut(COR_CHUNK_SIZE)
+                .zip(r_ots.chunks(COR_CHUNK_SIZE))
                 .enumerate()
             {
-                *ot = r_ot[0];
-                *correction = r_ot[1] ^ correlation(chunk_idx * COR_CHUNK_SIZE + idx, r_ot[0]);
+                let mut send_buf: Vec<Block> = Vec::zeroed(ot_chunk.len());
+                for (idx, ((ot, r_ot), correction)) in ot_chunk
+                    .iter_mut()
+                    .zip(rot_chunk)
+                    .zip(&mut send_buf)
+                    .enumerate()
+                {
+                    *ot = r_ot[0];
+                    *correction = r_ot[1] ^ correlation(chunk_idx * COR_CHUNK_SIZE + idx, r_ot[0]);
+                }
+                ch_s.blocking_send(send_buf).unwrap();
             }
-            tx.write_all(bytemuck::must_cast_slice_mut(
-                &mut send_buf[..ot_chunk.len()],
-            ))
-            .await?;
+            ots_owned
+        });
+
+        while let Some(buf) = ch_r.recv().await {
+            tx.write_all(bytemuck::must_cast_slice(&buf)).await.unwrap();
         }
+
+        *ots = jh.await.unwrap();
+
         Ok(())
     }
 }
@@ -677,21 +686,41 @@ impl<S: Security> CotReceiver for OtExtensionReceiver<S> {
         B: Buf<Block>,
     {
         self.receive_into(ots, choices).await?;
-        let mut recv_buf: Vec<Block> = Vec::zeroed(COR_CHUNK_SIZE);
         let (_, mut rx) = self.conn.byte_stream().await?;
-        for (ot_chunk, choice_chunk) in ots
-            .chunks_mut(COR_CHUNK_SIZE)
-            .zip(choices.chunks(COR_CHUNK_SIZE))
-        {
-            rx.read_exact(bytemuck::must_cast_slice_mut(
-                &mut recv_buf[..ot_chunk.len()],
-            ))
-            .await?;
-            for ((ot, correction), choice) in ot_chunk.iter_mut().zip(&recv_buf).zip(choice_chunk) {
-                let use_correction = Block::conditional_select(&Block::ZERO, &Block::ONES, *choice);
-                *ot ^= use_correction & *correction;
+        let (ch_s, mut ch_r) = tokio::sync::mpsc::channel(5);
+        let batch_block_sizes = iter::repeat_n(COR_CHUNK_SIZE, ots.len() / COR_CHUNK_SIZE);
+        let last_batch_size = ots.len() % COR_CHUNK_SIZE;
+        let batch_block_sizes = batch_block_sizes.chain(if last_batch_size != 0 {
+            Some(last_batch_size)
+        } else {
+            None
+        });
+        let mut ots_owned = mem::take(ots);
+        let choices = choices.to_owned();
+        let jh = spawn_compute(move || {
+            for (ot_chunk, choice_chunk) in ots_owned
+                .chunks_mut(COR_CHUNK_SIZE)
+                .zip(choices.chunks(COR_CHUNK_SIZE))
+            {
+                let recv_buf: Vec<Block> = ch_r.blocking_recv().unwrap();
+                for ((ot, correction), choice) in
+                    ot_chunk.iter_mut().zip(&recv_buf).zip(choice_chunk)
+                {
+                    let use_correction =
+                        Block::conditional_select(&Block::ZERO, &Block::ONES, *choice);
+                    *ot ^= use_correction & *correction;
+                }
             }
+            ots_owned
+        });
+
+        for batch_size in batch_block_sizes {
+            let mut recv_buf: Vec<Block> = Vec::zeroed(batch_size);
+            rx.read_exact(bytemuck::must_cast_slice_mut(&mut recv_buf))
+                .await?;
+            ch_s.send(recv_buf).await.unwrap();
         }
+        *ots = jh.await.unwrap();
         Ok(())
     }
 }
@@ -810,7 +839,7 @@ mod tests {
     #[tokio::test]
     async fn test_correlated_extension() {
         let _g = init_tracing();
-        const COUNT: usize = 128;
+        const COUNT: usize = DEFAULT_OT_BATCH_SIZE;
         let (c1, c2) = local_conn().await.unwrap();
         let rng1 = StdRng::seed_from_u64(42);
         let mut rng2 = StdRng::seed_from_u64(24);

cryprot-ot/src/lib.rs

@@ -187,7 +187,7 @@ pub trait CotSender: Connected + Send {
         correlation: F,
     ) -> impl Future<Output = Result<Vec<Block>, Self::Error>> + Send
     where
-        F: FnMut(usize, Block) -> Block + Send,
+        F: FnMut(usize, Block) -> Block + Send + 'static,
     {
         async move {
             let mut ots = Vec::zeroed(count);
@@ -203,7 +203,7 @@ pub trait CotSender: Connected + Send {
     ) -> impl Future<Output = Result<(), Self::Error>> + Send
     where
         B: Buf<Block>,
-        F: FnMut(usize, Block) -> Block + Send;
+        F: FnMut(usize, Block) -> Block + Send + 'static;
 }
 
 pub trait CotReceiver: Connected + Send {