Skip to content

Commit

Permalink
feat: lazy ot
Browse files Browse the repository at this point in the history
  • Loading branch information
sinui0 committed Oct 23, 2024
1 parent 65994b0 commit 7cfe967
Show file tree
Hide file tree
Showing 97 changed files with 4,958 additions and 7,238 deletions.
6 changes: 4 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,8 @@ mpz-ole-core = { path = "crates/mpz-ole-core" }
clmul = { path = "crates/clmul" }
matrix-transpose = { path = "crates/matrix-transpose" }

tlsn-utils = { git = "https://github.com/tlsnotary/tlsn-utils", rev = "6e0be94" }
tlsn-utils-aio = { git = "https://github.com/tlsnotary/tlsn-utils", rev = "6e0be94" }
tlsn-utils = { git = "https://github.com/tlsnotary/tlsn-utils", rev = "5899190" }
tlsn-utils-aio = { git = "https://github.com/tlsnotary/tlsn-utils", rev = "5899190" }

# rand
rand_chacha = "0.3"
Expand Down Expand Up @@ -124,3 +124,5 @@ typenum = "1"
generic-array = "0.14"
itybity = "0.2"
enum-try-as-inner = "0.1.0"
bitvec = "1.0"
hashbrown = "0.14.5"
2 changes: 1 addition & 1 deletion crates/mpz-cointoss/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ edition = "2021"

[dependencies]
mpz-core.workspace = true
mpz-common.workspace = true
mpz-common = { workspace = true, features = ["ctx"] }
mpz-cointoss-core.workspace = true

futures.workspace = true
Expand Down
14 changes: 9 additions & 5 deletions crates/mpz-common/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,14 @@ version = "0.1.0"
edition = "2021"

[features]
default = ["sync"]
sync = ["tokio/sync"]
test-utils = ["uid-mux/test-utils"]
ideal = []
default = []
ctx = []
cpu = ["rayon"]
executor = ["ctx", "cpu"]
sync = ["tokio/sync", "ctx"]
future = []
test-utils = ["dep:uid-mux", "uid-mux/test-utils"]
ideal = ["ctx"]
rayon = ["dep:rayon"]
force-st = []

Expand All @@ -20,7 +24,7 @@ pin-project-lite.workspace = true
scoped-futures.workspace = true
thiserror.workspace = true
serio.workspace = true
uid-mux.workspace = true
uid-mux = { workspace = true, optional = true }
serde = { workspace = true, features = ["derive"] }
pollster.workspace = true
rayon = { workspace = true, optional = true }
Expand Down
193 changes: 193 additions & 0 deletions crates/mpz-common/src/future.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,193 @@
//! Future types.
use std::{
future::Future,
mem,
pin::Pin,
task::{ready, Context, Poll},
};

use futures::{channel::oneshot, FutureExt};
use pin_project_lite::pin_project;

/// Creates a new output future.
pub fn new_output<T>() -> (Sender<T>, MaybeDone<T>) {
let (send, recv) = oneshot::channel();
(Sender { send }, MaybeDone { recv })
}

/// A future output value.
///
/// This trait extends [`std::future::Future`] for values which can be received
/// outside of a task context.
pub trait Output: Future<Output = Result<Self::Ok, Canceled>> {
/// Success type.
type Ok;

/// Attempts to receive the output outside of a task context, returning
/// `None` if it is not ready.
fn try_recv(&mut self) -> Result<Option<Self::Ok>, Canceled>;
}

/// An extension trait for [`Output`].
pub trait OutputExt: Output {
/// Maps the output value to a different type.
fn map<F, O>(self, f: F) -> Map<Self, F>
where
Self: Sized,
F: FnOnce(Self::Ok) -> O,
{
Map::new(self, f)
}
}

impl<T> OutputExt for T where T: Output {}

/// Output canceled error.
#[derive(Debug, thiserror::Error)]
#[error("output canceled")]
pub struct Canceled {
_private: (),
}

/// Sender of an output value.
#[derive(Debug)]
pub struct Sender<T> {
send: oneshot::Sender<T>,
}

impl<T> Sender<T> {
/// Sends an output value.
pub fn send(self, value: T) {
let _ = self.send.send(value);
}
}

/// An output value that may be ready.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct MaybeDone<T> {
recv: oneshot::Receiver<T>,
}

impl<T> Output for MaybeDone<T> {
type Ok = T;

fn try_recv(&mut self) -> Result<Option<Self::Ok>, Canceled> {
match self.recv.try_recv() {
Ok(Some(value)) => Ok(Some(value)),
Ok(None) => Ok(None),
Err(oneshot::Canceled) => Err(Canceled { _private: () }),
}
}
}

impl<T> Future for MaybeDone<T> {
type Output = Result<T, Canceled>;

fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.recv
.poll_unpin(cx)
.map_err(|_| Canceled { _private: () })
}
}

pin_project! {
/// Maps an output value to a different type.
///
/// Returned by [`OutputExt::map`].
#[derive(Debug)]
pub struct Map<I, F> {
#[pin]
inner: MapInner<I, F>,
}
}

impl<I, F> Map<I, F> {
fn new(inner: I, f: F) -> Self {
Self {
inner: MapInner::Incomplete { inner, f },
}
}
}

impl<I, F, O> Output for Map<I, F>
where
I: Output,
F: FnOnce(I::Ok) -> O,
{
type Ok = O;

fn try_recv(&mut self) -> Result<Option<Self::Ok>, Canceled> {
self.inner.try_recv()
}
}

impl<I, F, O> Future for Map<I, F>
where
I: Output,
F: FnOnce(I::Ok) -> O,
{
type Output = Result<O, Canceled>;

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.project().inner.poll(cx)
}
}

pin_project! {
/// Maps an output value to a different type.
///
/// Returned by [`OutputExt::map`].
#[derive(Debug)]
#[project = MapProj]
#[project_replace = MapProjReplace]
#[must_use = "futures do nothing unless you `.await` or poll them"]
enum MapInner<I, F> {
Incomplete {
#[pin]
inner: I,
f: F,
},
Done,
}
}

impl<I, F, O> Output for MapInner<I, F>
where
I: Output,
F: FnOnce(I::Ok) -> O,
{
type Ok = O;

fn try_recv(&mut self) -> Result<Option<Self::Ok>, Canceled> {
let this = mem::replace(self, MapInner::Done);
match this {
MapInner::Incomplete { mut inner, f } => inner.try_recv().map(|res| res.map(f)),
MapInner::Done => Err(Canceled { _private: () }),
}
}
}

impl<I, F, O> Future for MapInner<I, F>
where
I: Output,
F: FnOnce(I::Ok) -> O,
{
type Output = Result<O, Canceled>;

fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.as_mut().project() {
MapProj::Incomplete { inner, .. } => {
let output = ready!(inner.poll(cx));
match self.project_replace(Self::Done) {
MapProjReplace::Incomplete { f, .. } => Poll::Ready(output.map(f)),
MapProjReplace::Done => unreachable!(),
}
}
MapProj::Done => {
panic!("Map must not be polled after it returned `Poll::Ready`")
}
}
}
}
26 changes: 23 additions & 3 deletions crates/mpz-common/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
//! Common functionality for `mpz`.
//!
//! This crate provides various common functionalities needed for modeling protocol execution, I/O,
//! and multi-threading.
//! This crate provides various common functionalities needed for modeling
//! protocol execution, I/O, and multi-threading.
//!
//! This crate does not provide any cryptographic primitives, see `mpz-core` for that.
//! This crate does not provide any cryptographic primitives, see `mpz-core` for
//! that.
#![deny(
unsafe_code,
Expand All @@ -14,16 +15,22 @@
clippy::all
)]

#[cfg(any(test, feature = "ctx"))]
mod context;
#[cfg(any(test, feature = "cpu"))]
pub mod cpu;
#[cfg(any(test, feature = "executor"))]
pub mod executor;
#[cfg(any(test, feature = "future"))]
pub mod future;
mod id;
#[cfg(any(test, feature = "ideal"))]
pub mod ideal;
#[cfg(feature = "sync")]
pub mod sync;

use async_trait::async_trait;
#[cfg(any(test, feature = "ctx"))]
pub use context::{Context, ContextError};
pub use id::{Counter, ThreadId};

Expand All @@ -46,6 +53,19 @@ pub trait Preprocess<Ctx>: Allocate {
async fn preprocess(&mut self, ctx: &mut Ctx) -> Result<(), Self::Error>;
}

/// A functionality that can be flushed.
#[async_trait]
pub trait Flush<Ctx> {
/// Error type.
type Error: std::error::Error + Send + Sync + 'static;

/// Returns `true` if the functionality wants to be flushed.
fn wants_flush(&self) -> bool;

/// Flushes the functionality.
async fn flush(&mut self, ctx: &mut Ctx) -> Result<(), Self::Error>;
}

/// A convenience macro for creating a closure which returns a scoped future.
///
/// # Example
Expand Down
1 change: 1 addition & 0 deletions crates/mpz-core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ bytemuck = { workspace = true, features = ["derive"] }
generic-array.workspace = true
rayon = { workspace = true, optional = true }
cfg-if.workspace = true
bitvec = { workspace = true, features = ["serde"] }

[dev-dependencies]
rstest.workspace = true
Expand Down
6 changes: 6 additions & 0 deletions crates/mpz-core/src/bitvec.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
//! Bit vectors.
/// Bit vector.
pub type BitVec<T = u32> = bitvec::vec::BitVec<T, bitvec::order::Lsb0>;
/// Bit slice.
pub type BitSlice<T = u32> = bitvec::slice::BitSlice<T, bitvec::order::Lsb0>;
1 change: 1 addition & 0 deletions crates/mpz-core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#![deny(clippy::all)]

pub mod aes;
pub mod bitvec;
pub mod block;
pub mod commit;
pub mod ggm_tree;
Expand Down
5 changes: 5 additions & 0 deletions crates/mpz-core/src/prg.rs
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,11 @@ impl Prg {
Prg::from_seed(rand::random::<Block>())
}

/// Create a new PRG from a seed.
pub fn new_with_seed(seed: [u8; 16]) -> Self {
Prg::from_seed(Block::from(seed))
}

/// Returns the current counter.
pub fn counter(&self) -> u64 {
self.0.core.counter
Expand Down
2 changes: 2 additions & 0 deletions crates/mpz-ot-core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ test-utils = []

[dependencies]
mpz-core.workspace = true
mpz-common = { workspace = true, features = ["future"] }
clmul.workspace = true
matrix-transpose.workspace = true

Expand All @@ -38,6 +39,7 @@ opaque-debug.workspace = true
cfg-if.workspace = true
bytemuck = { workspace = true, features = ["derive"] }
enum-try-as-inner.workspace = true
futures = { workspace = true }

[dev-dependencies]
rstest.workspace = true
Expand Down
Loading

0 comments on commit 7cfe967

Please sign in to comment.