diff --git a/crates/core/benches/ggm.rs b/crates/core/benches/ggm.rs index 8c0e4fb5..d12d59a2 100644 --- a/crates/core/benches/ggm.rs +++ b/crates/core/benches/ggm.rs @@ -8,8 +8,9 @@ fn criterion_benchmark(c: &mut Criterion) { let depth = 10; let seed = rand::random::(); let mut leaves = vec![Block::ZERO; 1 << depth]; + let mut buf = vec![Block::ZERO; (1 << depth) - 1]; bench.iter(|| { - GgmTree::new_from_seed(depth, seed, &mut leaves); + GgmTree::new_from_seed(depth, seed, &mut leaves, &mut buf); black_box(&leaves); }); }); @@ -18,8 +19,9 @@ fn criterion_benchmark(c: &mut Criterion) { let depth = 10; let sums = vec![Block::ZERO; depth]; let mut leaves = vec![Block::ZERO; 1 << depth]; + let mut buf = vec![Block::ZERO; (1 << depth) - 1]; bench.iter(|| { - GgmTree::new_partial(depth, &sums, 420, &mut leaves); + GgmTree::new_partial(depth, &sums, 420, &mut leaves, &mut buf); black_box(&leaves); }); }); diff --git a/crates/core/src/block.rs b/crates/core/src/block.rs index d371818b..82014d68 100644 --- a/crates/core/src/block.rs +++ b/crates/core/src/block.rs @@ -100,21 +100,77 @@ impl Block { /// Compute the inner product of two block vectors, without reducing the /// polynomial. + /// + /// Uses 8 independent accumulators to break the carry-less multiply + /// latency-dependency chain (`clmul` is high-latency, fully-pipelined), so + /// the loop runs throughput-bound rather than latency-bound. #[inline] pub fn inn_prdt_no_red(a: &[Block], b: &[Block]) -> (Block, Block) { assert_eq!(a.len(), b.len()); - a.iter() - .zip(b.iter()) - .fold((Block::ZERO, Block::ZERO), |acc, (x, y)| { - let t = x.clmul(*y); - (t.0 ^ acc.0, t.1 ^ acc.1) - }) + + const LANES: usize = 8; + let mut hi = [Block::ZERO; LANES]; + let mut lo = [Block::ZERO; LANES]; + + let mut a_chunks = a.chunks_exact(LANES); + let mut b_chunks = b.chunks_exact(LANES); + for (ac, bc) in a_chunks.by_ref().zip(b_chunks.by_ref()) { + for j in 0..LANES { + let (h, l) = ac[j].clmul(bc[j]); + hi[j] ^= h; + lo[j] ^= l; + } + } + + let mut acc_hi = Block::ZERO; + let mut acc_lo = Block::ZERO; + for j in 0..LANES { + acc_hi ^= hi[j]; + acc_lo ^= lo[j]; + } + + for (x, y) in a_chunks.remainder().iter().zip(b_chunks.remainder()) { + let (h, l) = x.clmul(*y); + acc_hi ^= h; + acc_lo ^= l; + } + + (acc_hi, acc_lo) } /// Compute the inner product of two block vectors. + /// + /// With the `rayon` feature enabled, the (unreduced) inner product is + /// computed in parallel over chunks and combined before a single final + /// reduction. #[inline] pub fn inn_prdt_red(a: &[Block], b: &[Block]) -> Block { - let (x, y) = Block::inn_prdt_no_red(a, b); + assert_eq!(a.len(), b.len()); + + cfg_if::cfg_if! { + if #[cfg(feature = "rayon")] { + use rayon::prelude::*; + + // Large enough that per-chunk overhead is negligible, small + // enough to keep all cores busy on production-sized vectors. + const CHUNK: usize = 1 << 16; + + let (x, y) = if a.len() <= CHUNK { + Block::inn_prdt_no_red(a, b) + } else { + a.par_chunks(CHUNK) + .zip(b.par_chunks(CHUNK)) + .map(|(ac, bc)| Block::inn_prdt_no_red(ac, bc)) + .reduce( + || (Block::ZERO, Block::ZERO), + |p, q| (p.0 ^ q.0, p.1 ^ q.1), + ) + }; + } else { + let (x, y) = Block::inn_prdt_no_red(a, b); + } + } + Block::reduce_gcm(x, y) } diff --git a/crates/core/src/ggm.rs b/crates/core/src/ggm.rs index 35646732..8ae11db6 100644 --- a/crates/core/src/ggm.rs +++ b/crates/core/src/ggm.rs @@ -24,7 +24,10 @@ fn width(n: usize) -> usize { /// GGM tree. pub struct GgmTree<'a> { depth: usize, - buf: Vec, + /// Internal (non-leaf) tree nodes. Caller-provided scratch of length + /// `(1 << depth) - 1`, so it can be reused across trees without + /// reallocating or re-zeroing. + buf: &'a mut [Block], leaves: &'a mut [Block], } @@ -36,10 +39,17 @@ impl<'a> GgmTree<'a> { /// * `depth` - The depth of the tree. /// * `seed` - The seed of the tree. /// * `leaves` - The leaves of the tree. - pub fn new_from_seed(depth: usize, seed: Block, leaves: &'a mut [Block]) -> Self { + /// * `buf` - Scratch for the internal nodes, of length `(1 << depth) - 1`. + /// Its contents are fully overwritten; reuse it across trees to avoid + /// per-tree allocation. + pub fn new_from_seed( + depth: usize, + seed: Block, + leaves: &'a mut [Block], + buf: &'a mut [Block], + ) -> Self { assert_eq!(leaves.len(), 1 << depth, "invalid length of leaves"); - - let mut buf = vec![Block::ZERO; (1 << depth) - 1]; + assert_eq!(buf.len(), (1 << depth) - 1, "invalid length of buf"); let tkprp = TwoKeyPrp::new([Block::ZERO, Block::ONE]); @@ -76,12 +86,19 @@ impl<'a> GgmTree<'a> { /// layer. /// * `idx` - Index of the missing leaf. /// * `leaves` - Leaves of the tree. - pub fn new_partial(depth: usize, sums: &[Block], idx: usize, leaves: &'a mut [Block]) -> Self { + /// * `buf` - Scratch for the internal nodes, of length `(1 << depth) - 1`. + /// Reuse it across trees to avoid per-tree allocation. + pub fn new_partial( + depth: usize, + sums: &[Block], + idx: usize, + leaves: &'a mut [Block], + buf: &'a mut [Block], + ) -> Self { assert_eq!(leaves.len(), 1 << depth, "invalid length of leaves"); assert!(idx < leaves.len(), "index out of bounds"); assert_eq!(sums.len(), depth, "invalid length of sums"); - - let mut buf = vec![Block::ZERO; (1 << depth) - 1]; + assert_eq!(buf.len(), (1 << depth) - 1, "invalid length of buf"); let tkprp = TwoKeyPrp::new([Block::ZERO, Block::ONE]); @@ -189,8 +206,9 @@ mod tests { let depth = 4; let mut leaves = vec![Block::ZERO; 1 << depth]; + let mut buf = vec![Block::ZERO; (1 << depth) - 1]; - GgmTree::new_from_seed(depth, seed, &mut leaves); + GgmTree::new_from_seed(depth, seed, &mut leaves, &mut buf); assert!(leaves.iter().all(|leaf| *leaf != Block::ZERO)); } @@ -201,8 +219,9 @@ mod tests { let depth = 4; let mut leaves = vec![Block::ZERO; 1 << depth]; + let mut buf = vec![Block::ZERO; (1 << depth) - 1]; - let ggm = GgmTree::new_from_seed(depth, seed, &mut leaves); + let ggm = GgmTree::new_from_seed(depth, seed, &mut leaves, &mut buf); for i in 0..depth { let layer = ggm.layer(i).unwrap(); @@ -217,7 +236,8 @@ mod tests { let depth = 4; let mut full_leaves = vec![Block::ZERO; 1 << depth]; - let ggm = GgmTree::new_from_seed(depth, seed, &mut full_leaves); + let mut buf = vec![Block::ZERO; (1 << depth) - 1]; + let ggm = GgmTree::new_from_seed(depth, seed, &mut full_leaves, &mut buf); for i in 0..1 << depth { let path = i as u32; @@ -228,7 +248,8 @@ mod tests { .collect::>(); let mut leaves = vec![Block::ZERO; 1 << depth]; - let ggm_partial = GgmTree::new_partial(depth, &sums, i, &mut leaves); + let mut partial_buf = vec![Block::ZERO; (1 << depth) - 1]; + let ggm_partial = GgmTree::new_partial(depth, &sums, i, &mut leaves, &mut partial_buf); let mut full_leaves = ggm.leaves().to_vec(); full_leaves[i] = Block::ZERO; diff --git a/crates/core/src/prg.rs b/crates/core/src/prg.rs index 1200f571..eaaabe2d 100644 --- a/crates/core/src/prg.rs +++ b/crates/core/src/prg.rs @@ -174,6 +174,147 @@ impl Prg { let bytes: &mut [u8] = bytemuck::cast_slice_mut(buf); self.fill_bytes(bytes); } + + /// Fills `buf` with the same pseudo-random block stream that + /// `Prg::from_seed(seed).random_blocks(buf)` would produce, but computed in + /// parallel. + /// + /// The PRG is AES in counter mode, so block `i` is + /// `AES_seed(i ‖ stream_id=0)` and can be generated independently. With the + /// `rayon` feature disabled this falls back to the sequential path and + /// produces byte-identical output. + pub fn random_blocks_par(seed: Block, buf: &mut [Block]) { + #[inline] + fn fill_chunk(aes: &AesEncryptor, start: u64, chunk: &mut [Block]) { + for (j, blk) in chunk.iter_mut().enumerate() { + let mut b = [0u8; 16]; + b[..8].copy_from_slice(&(start + j as u64).to_le_bytes()); + // stream_id == 0, so the high 8 bytes stay zero. + *blk = Block::from(b); + } + aes.encrypt_blocks(chunk); + } + + cfg_if::cfg_if! { + if #[cfg(feature = "rayon")] { + use rayon::prelude::*; + + const CHUNK: usize = 1 << 14; + let aes = AesEncryptor::new(seed); + buf.par_chunks_mut(CHUNK).enumerate().for_each(|(ci, chunk)| { + fill_chunk(&aes, (ci * CHUNK) as u64, chunk); + }); + } else { + let aes = AesEncryptor::new(seed); + fill_chunk(&aes, 0, buf); + } + } + } + + /// Computes the GF(2^128) inner product `Σ_i chi_i · b_i`, where `chi` is + /// the block stream of `Prg::from_seed(seed)` (matching `random_blocks` + /// and `random_blocks_par`). + /// + /// `chi` is regenerated on the fly in counter mode and never materialized, + /// so only `b` is streamed from memory — this avoids both the allocation + /// and the write+read round-trip of an explicit `chi` vector. The + /// result is identical to `Block::inn_prdt_red(&chi, b)`. With the + /// `rayon` feature the product is computed over parallel chunks and + /// reduced once at the end. + pub fn chi_inner_product(seed: Block, b: &[Block]) -> Block { + let aes = AesEncryptor::new(seed); + + cfg_if::cfg_if! { + if #[cfg(feature = "rayon")] { + use rayon::prelude::*; + + const CHUNK: usize = 1 << 16; + let (hi, lo) = if b.len() <= CHUNK { + chi_clmul_acc(&aes, 0, b) + } else { + b.par_chunks(CHUNK) + .enumerate() + .map(|(ci, chunk)| chi_clmul_acc(&aes, (ci * CHUNK) as u64, chunk)) + .reduce( + || (Block::ZERO, Block::ZERO), + |p, q| (p.0 ^ q.0, p.1 ^ q.1), + ) + }; + } else { + let (hi, lo) = chi_clmul_acc(&aes, 0, b); + } + } + + Block::reduce_gcm(hi, lo) + } + + /// Fills `out[j]` with the block at counter index `positions[j]` of the + /// `Prg::from_seed(seed)` stream (stream id 0), i.e. the same value as + /// `random_blocks(buf)[positions[j]]`. + /// + /// # Panics + /// + /// Panics if `out.len() != positions.len()`. + pub fn blocks_at(seed: Block, positions: &[usize], out: &mut [Block]) { + assert_eq!(positions.len(), out.len()); + let aes = AesEncryptor::new(seed); + for (o, &p) in out.iter_mut().zip(positions) { + let mut blk = [0u8; 16]; + blk[..8].copy_from_slice(&(p as u64).to_le_bytes()); + *o = Block::from(blk); + } + aes.encrypt_blocks(out); + } +} + +/// Accumulates the unreduced inner product `Σ chi_{start+j} · b[j]`, generating +/// `chi` in counter mode in small cache-resident batches and using 8 +/// independent accumulators to break the `clmul` latency-dependency chain. +#[inline] +fn chi_clmul_acc(aes: &AesEncryptor, start: u64, b: &[Block]) -> (Block, Block) { + const BATCH: usize = 256; + let mut hi = [Block::ZERO; 8]; + let mut lo = [Block::ZERO; 8]; + let mut chi = [Block::ZERO; BATCH]; + + let mut off = 0usize; + while off < b.len() { + let len = BATCH.min(b.len() - off); + + // Generate this batch of chi blocks (counter mode, stream id 0). + for (j, c) in chi[..len].iter_mut().enumerate() { + let mut blk = [0u8; 16]; + blk[..8].copy_from_slice(&(start + (off + j) as u64).to_le_bytes()); + *c = Block::from(blk); + } + aes.encrypt_blocks(&mut chi[..len]); + + let bc = &b[off..off + len]; + let mut k = 0; + while k + 8 <= len { + for j in 0..8 { + let (h, l) = chi[k + j].clmul(bc[k + j]); + hi[j] ^= h; + lo[j] ^= l; + } + k += 8; + } + for j in k..len { + let (h, l) = chi[j].clmul(bc[j]); + hi[0] ^= h; + lo[0] ^= l; + } + + off += len; + } + + let mut acc_hi = Block::ZERO; + let mut acc_lo = Block::ZERO; + for j in 0..8 { + acc_hi ^= hi[j]; + acc_lo ^= lo[j]; + } + (acc_hi, acc_lo) } impl Default for Prg { @@ -208,6 +349,53 @@ mod tests { assert_ne!(x[0], y[0]); } + #[test] + fn test_random_blocks_par_matches_sequential() { + let seed = Block::from(*b"0123456789abcdef"); + // Cover non-multiples of the 8-block AES batch and the rayon chunk. + for len in [0usize, 1, 7, 8, 9, 100, 1023, 16384, 16385, 40000] { + let mut seq = vec![Block::ZERO; len]; + Prg::from_seed(seed).random_blocks(&mut seq); + + let mut par = vec![Block::ZERO; len]; + Prg::random_blocks_par(seed, &mut par); + + assert_eq!(seq, par, "mismatch at len {len}"); + } + } + + #[test] + fn test_chi_inner_product_matches_explicit() { + let seed = Block::from(*b"chi_seed_0123456"); + let mut src = Prg::from_seed(Block::from(*b"vs_seed_abcdefgh")); + for len in [0usize, 1, 7, 8, 9, 255, 256, 257, 1000, 70_000] { + let mut b = vec![Block::ZERO; len]; + src.random_blocks(&mut b); + + let mut chis = vec![Block::ZERO; len]; + Prg::random_blocks_par(seed, &mut chis); + let expected = Block::inn_prdt_red(&chis, &b); + + assert_eq!(expected, Prg::chi_inner_product(seed, &b), "len {len}"); + } + } + + #[test] + fn test_blocks_at_matches_stream() { + let seed = Block::from(*b"some_seed_012345"); + let n = 2048; + let mut full = vec![Block::ZERO; n]; + Prg::from_seed(seed).random_blocks(&mut full); + + let positions = [0usize, 1, 5, 8, 100, 255, 256, 1023, 2047]; + let mut out = vec![Block::ZERO; positions.len()]; + Prg::blocks_at(seed, &positions, &mut out); + + for (o, &p) in out.iter().zip(&positions) { + assert_eq!(*o, full[p], "position {p}"); + } + } + #[test] fn test_prg_state_persisted() { let mut prg = Prg::from_seed(Block::ZERO); diff --git a/crates/ot-core/src/ferret.rs b/crates/ot-core/src/ferret.rs index 1734f1e6..3454a135 100644 --- a/crates/ot-core/src/ferret.rs +++ b/crates/ot-core/src/ferret.rs @@ -3,6 +3,8 @@ mod config; pub(crate) mod cuckoo; pub(crate) mod mpcot; +#[cfg(test)] +mod profile; mod receiver; mod sender; pub(crate) mod spcot; diff --git a/crates/ot-core/src/ferret/profile.rs b/crates/ot-core/src/ferret/profile.rs new file mode 100644 index 00000000..69a74930 --- /dev/null +++ b/crates/ot-core/src/ferret/profile.rs @@ -0,0 +1,139 @@ +//! Temporary per-phase profiling harness for the Ferret extension. +//! +//! Run with: +//! cargo test -p mpz-ot-core --release --features rayon \ +//! ferret::profile -- --nocapture --ignored +//! +//! Drives the *real* phase functions (cuckoo bucket construction, SPCOT GGM +//! generation, MPCOT combine, consistency check, LPN encode) at production LPN +//! parameters and reports the wall-clock spent in each, so we can attribute +//! cost without protocol/network/serialization overhead. + +use std::time::Instant; + +use rand::{RngExt, SeedableRng, rngs::StdRng}; + +use mpz_core::{ + Block, + bitvec::BitVec, + lpn::{LpnEncoder, LpnParameters, LpnType, sample_error_indices}, +}; + +use crate::ferret::{ + config::{CSP, REGULAR_PARAMS, UNIFORM_PARAMS}, + mpcot::{MPCOTReceiver, MPCOTSender}, + spcot::SPCOTSender, +}; + +fn ms(d: std::time::Duration) -> f64 { + d.as_secs_f64() * 1e3 +} + +#[test] +#[ignore = "profiling harness, run explicitly with --ignored --nocapture"] +fn profile_ferret_phases() { + let param_indices = [0usize, 2, 4]; + + for (lpn_type, params) in [ + (LpnType::Uniform, UNIFORM_PARAMS), + (LpnType::Regular, REGULAR_PARAMS), + ] { + println!("\n=== LpnType::{lpn_type:?} ==="); + println!( + "{:>9} {:>9} {:>11} | {:>10} {:>10} {:>8} {:>8} {:>8} {:>8} | {:>10}", + "n", + "k", + "ggm_leaves", + "cuckoo_snd", + "cuckoo_rcv", + "spcot", + "combine", + "check", + "lpn_enc", + "SIDE_TOTAL", + ); + + for &pi in ¶m_indices { + profile_one(lpn_type, params[pi]); + } + } + println!("\n(all times in ms; rayon = {})", cfg!(feature = "rayon")); +} + +fn profile_one(lpn_type: LpnType, params: LpnParameters) { + let LpnParameters { n, k, t } = params; + { + let mut rng = StdRng::seed_from_u64(0); + let delta: Block = rng.random(); + let cuckoo_seed: Block = rng.random(); + + // ---- sender cuckoo bucket construction ---- + let t0 = Instant::now(); + let (mpcot_send, log2_lengths) = MPCOTSender::new(cuckoo_seed, lpn_type) + .start_extend(t, n) + .unwrap(); + let t_cuckoo_send = t0.elapsed(); + + // ---- receiver cuckoo (CuckooHash insert + Buckets) ---- + let idxs = sample_error_indices(&mut rng, lpn_type, n, t); + let t0 = Instant::now(); + let _ = MPCOTReceiver::new(cuckoo_seed, lpn_type) + .start_extend(&idxs, n) + .unwrap(); + let t_cuckoo_recv = t0.elapsed(); + + // SPCOT inputs sized exactly as the protocol would. + let sum_log2: usize = log2_lengths.iter().sum(); + let keys: Vec = (0..sum_log2).map(|_| rng.random()).collect(); + let masks: BitVec = (0..sum_log2).map(|_| rng.random::()).collect(); + let ggm_leaves: usize = log2_lengths.iter().map(|l| 1usize << l).sum(); + + // ---- SPCOT sender extend (GGM gen + fixed-key AES) ---- + let mut spcot = SPCOTSender::new(delta); + let t0 = Instant::now(); + let (vs, _ms_out, _sums) = spcot + .extend(&mut rng, &log2_lengths, &keys, &masks) + .unwrap(); + let t_spcot = t0.elapsed(); + + // ---- MPCOT combine (random-access XOR gather) ---- + let t0 = Instant::now(); + let res = mpcot_send.extend(vs).unwrap(); + let t_combine = t0.elapsed(); + debug_assert_eq!(res.len(), n); + + // ---- consistency check (chi gen + O(leaves) inner product) ---- + let check_keys: Vec = (0..CSP).map(|_| rng.random()).collect(); + let check_masks: BitVec = (0..CSP).map(|_| rng.random::()).collect(); + let t0 = Instant::now(); + let _hashed = spcot.check(&check_keys, &check_masks).unwrap(); + let t_check = t0.elapsed(); + + // ---- LPN encode y = A*v + s ---- + let x: Vec = (0..k).map(|_| rng.random()).collect(); + let mut y = res; + let enc = LpnEncoder::<10>::new(k as u32); + let lpn_seed: Block = rng.random(); + let t0 = Instant::now(); + enc.compute(lpn_seed, &mut y, &x); + let t_lpn = t0.elapsed(); + std::hint::black_box(&y); + + // One party (sender) pays: cuckoo + spcot + combine + check + lpn. + let side_total = t_cuckoo_send + t_spcot + t_combine + t_check + t_lpn; + + println!( + "{:>9} {:>9} {:>11} | {:>10.1} {:>10.1} {:>8.1} {:>8.1} {:>8.1} {:>8.1} | {:>10.1}", + n, + k, + ggm_leaves, + ms(t_cuckoo_send), + ms(t_cuckoo_recv), + ms(t_spcot), + ms(t_combine), + ms(t_check), + ms(t_lpn), + ms(side_total), + ); + } +} diff --git a/crates/ot-core/src/ferret/spcot/receiver.rs b/crates/ot-core/src/ferret/spcot/receiver.rs index 2576939a..e948ad52 100644 --- a/crates/ot-core/src/ferret/spcot/receiver.rs +++ b/crates/ot-core/src/ferret/spcot/receiver.rs @@ -1,7 +1,6 @@ use blake3::{Hash, Hasher, hash}; use cfg_if::cfg_if; use itybity::ToBits; -use rand::SeedableRng; #[cfg(feature = "rayon")] use rayon::prelude::*; @@ -23,7 +22,7 @@ type Result = core::result::Result; #[derive(Debug)] struct Check { z: Block, - chis: Vec, + chi_seed: Block, } #[derive(Debug)] @@ -173,25 +172,30 @@ impl SPCOTReceiver { let ggm_sums = slices_from_lengths(&ggm_sums, log2_lengths); let ws = slices_from_lengths_mut(&mut self.ws[start..], &spcot_lengths); - let iter = { - cfg_if! { - if #[cfg(feature = "rayon")] { - ws.into_par_iter() - } else { - ws.into_iter() - } + // `recover_tree` reuses a per-thread scratch buffer for the GGM internal + // nodes, so we don't allocate one tree's worth of buffer per bucket. + cfg_if! { + if #[cfg(feature = "rayon")] { + ws.into_par_iter() + .zip(ggm_sums) + .zip(sums) + .zip(log2_lengths) + .zip(idxs) + .for_each_init(Vec::new, |scratch, ((((w, gsums), sum), &length), &idx)| { + recover_tree(scratch, length, gsums, idx, w, *sum); + }); + } else { + let mut scratch = Vec::new(); + ws.into_iter() + .zip(ggm_sums) + .zip(sums) + .zip(log2_lengths) + .zip(idxs) + .for_each(|((((w, gsums), sum), &length), &idx)| { + recover_tree(&mut scratch, length, gsums, idx, w, *sum); + }); } - }; - - iter.zip(ggm_sums) - .zip(sums) - .zip(log2_lengths) - .zip(idxs) - .for_each(|((((w, sums), sum), &length), &idx)| { - GgmTree::new_partial(length, sums, idx, w); - - w[idx] = w.iter().fold(*sum, |acc, &x| acc ^ x); - }); + } self.transcript.update(Block::array_as_flattened_bytes(ms)); self.transcript.update(Block::as_flattened_bytes(sums)); @@ -221,19 +225,23 @@ impl SPCOTReceiver { } let seed = *self.transcript.finalize().as_bytes(); - let mut prg = Prg::from_seed(Block::try_from(&seed[0..16]).unwrap()); - - // The sum of all the chi[alpha]. - let mut sum_chi_alpha = Block::ZERO; - - let mut chis = vec![Block::ZERO; self.ws.len()]; - prg.random_blocks(&mut chis); - - let mut i = 0; - for (length, idx) in self.lengths.iter().zip(&self.indices) { - sum_chi_alpha ^= chis[i + idx]; - i += 1 << length; - } + let chi_seed = Block::try_from(&seed[0..16]).unwrap(); + + // The sum of all the chi[alpha]. Only the `t` chi values at the chosen + // indices are needed here, so generate just those instead of the whole + // chi vector (the rest is regenerated on the fly in `check`). + let alpha_positions: Vec = { + let mut positions = Vec::with_capacity(self.indices.len()); + let mut i = 0; + for (length, idx) in self.lengths.iter().zip(&self.indices) { + positions.push(i + idx); + i += 1 << length; + } + positions + }; + let mut chi_alpha = vec![Block::ZERO; alpha_positions.len()]; + Prg::blocks_at(chi_seed, &alpha_positions, &mut chi_alpha); + let sum_chi_alpha = chi_alpha.iter().fold(Block::ZERO, |acc, &x| acc ^ x); let x_prime = BitVec::from_iter( sum_chi_alpha @@ -244,18 +252,19 @@ impl SPCOTReceiver { let z = Block::inn_prdt_red(macs, &Block::MONOMIAL); - self.check = Some(Check { z, chis }); + self.check = Some(Check { z, chi_seed }); Ok(Derandomize { flip: x_prime }) } pub(crate) fn check(&mut self, hashed_v: Hash) -> Result<()> { - let Some(Check { z, chis }) = self.check.take() else { + let Some(Check { z, chi_seed }) = self.check.take() else { return Err(ErrorRepr::State("check not started".to_string()).into()); }; - // Computes W. - let w = z ^ Block::inn_prdt_red(&chis, &self.ws); + // Computes W. The chi vector is regenerated on the fly inside the inner + // product (same seed as `start_check`), so it is never materialized. + let w = z ^ Prg::chi_inner_product(chi_seed, &self.ws); // Computes H'(W) let hashed_w = hash(&w.to_bytes()); @@ -273,6 +282,28 @@ impl SPCOTReceiver { } } +/// Recovers one partial GGM tree into `w` and fixes up the missing leaf at +/// `idx`. `scratch` is grown as needed and reused across calls to avoid +/// per-tree allocation. +#[inline] +fn recover_tree( + scratch: &mut Vec, + length: usize, + gsums: &[Block], + idx: usize, + w: &mut [Block], + sum: Block, +) { + let buf_len = (1usize << length) - 1; + if scratch.len() < buf_len { + scratch.resize(buf_len, Block::ZERO); + } + + GgmTree::new_partial(length, gsums, idx, w, &mut scratch[..buf_len]); + + w[idx] = w.iter().fold(sum, |acc, &x| acc ^ x); +} + #[derive(Debug, thiserror::Error)] #[error(transparent)] pub(crate) struct SPCOTReceiverError(#[from] ErrorRepr); diff --git a/crates/ot-core/src/ferret/spcot/sender.rs b/crates/ot-core/src/ferret/spcot/sender.rs index 9e66f7dc..2f32ee19 100644 --- a/crates/ot-core/src/ferret/spcot/sender.rs +++ b/crates/ot-core/src/ferret/spcot/sender.rs @@ -1,6 +1,6 @@ use blake3::{Hash, Hasher, hash}; use cfg_if::cfg_if; -use rand::{Rng, RngExt, SeedableRng}; +use rand::{Rng, RngExt}; #[cfg(feature = "rayon")] use rayon::prelude::*; @@ -104,36 +104,35 @@ impl SPCOTSender { let vs = slices_from_lengths_mut(&mut self.vs[start..], &spcot_lengths); let ks = slices_from_lengths_mut(&mut ms, log2_lengths); - let iter = { + let delta = self.delta; + + // `gen_tree` reuses a per-thread scratch buffer for the GGM internal + // nodes, so we don't allocate one tree's worth of buffer per bucket. + let sums: Vec = { cfg_if! { if #[cfg(feature = "rayon")] { vs.into_par_iter() + .zip(ks) + .zip(log2_lengths) + .zip(seeds) + .map_init(Vec::new, |scratch, (((v, ks), &depth), seed)| { + gen_tree(scratch, depth, seed, v, ks, delta) + }) + .collect() } else { + let mut scratch = Vec::new(); vs.into_iter() + .zip(ks) + .zip(log2_lengths) + .zip(seeds) + .map(|(((v, ks), &depth), seed)| { + gen_tree(&mut scratch, depth, seed, v, ks, delta) + }) + .collect() } } }; - let sums: Vec<_> = iter - .zip(ks) - .zip(log2_lengths) - .zip(seeds) - .map(|(((v, ks), &depth), seed)| { - // Generate the SPCOT vector from GGM leaves. - let tree = GgmTree::new_from_seed(depth, seed, v); - - // Encrypt the OT messages. - tree.layer_sums().zip(ks).for_each(|(sums, ks)| { - // `sums` is K_0 and K_1 in Fig. 6 Step 3. - ks[0] ^= sums[0]; - ks[1] ^= sums[1]; - }); - - // Compute the sum of the leaves. - tree.leaves().iter().fold(self.delta, |acc, x| acc ^ x) - }) - .collect(); - let masks_len = masks.len(); self.transcript .update(&masks.as_raw_slice().as_bytes()[..masks_len.div_ceil(8)]); @@ -177,14 +176,12 @@ impl SPCOTSender { // Computes Y let mut v = Block::inn_prdt_red(&y, &Block::MONOMIAL); - // Computes V + // Computes V. The chi vector is regenerated on the fly inside the inner + // product, so it is never materialized. let seed = *self.transcript.finalize().as_bytes(); - let mut prg = Prg::from_seed(Block::try_from(&seed[0..16]).unwrap()); - - let mut chis = vec![Block::ZERO; self.vs.len()]; - prg.random_blocks(&mut chis); + let chi_seed = Block::try_from(&seed[0..16]).unwrap(); - v ^= Block::inn_prdt_red(&chis, &self.vs); + v ^= Prg::chi_inner_product(chi_seed, &self.vs); // Computes H'(V) let hashed_v = hash(&v.to_bytes()); @@ -196,6 +193,34 @@ impl SPCOTSender { } } +/// Generates one SPCOT vector from a GGM tree, folds the layer sums into the OT +/// messages `ks`, and returns the (delta-seeded) sum of the leaves. `scratch` +/// is grown as needed and reused across calls to avoid per-tree allocation. +#[inline] +fn gen_tree( + scratch: &mut Vec, + depth: usize, + seed: Block, + v: &mut [Block], + ks: &mut [[Block; 2]], + delta: Block, +) -> Block { + let buf_len = (1usize << depth) - 1; + if scratch.len() < buf_len { + scratch.resize(buf_len, Block::ZERO); + } + + let tree = GgmTree::new_from_seed(depth, seed, v, &mut scratch[..buf_len]); + + // `sums` is K_0 and K_1 in Fig. 6 Step 3. + tree.layer_sums().zip(ks.iter_mut()).for_each(|(sums, ks)| { + ks[0] ^= sums[0]; + ks[1] ^= sums[1]; + }); + + tree.leaves().iter().fold(delta, |acc, &x| acc ^ x) +} + #[derive(Debug, thiserror::Error)] #[error(transparent)] pub(crate) struct SPCOTSenderError(#[from] ErrorRepr);