binius_hash/
parallel_compression.rs

1// Copyright 2025 Irreducible Inc.
2// Copyright 2026 The Binius Developers
3
4use std::{array, fmt::Debug, mem::MaybeUninit};
5
6use binius_utils::rayon::prelude::*;
7
8use crate::CompressionFunction;
9
10/// A trait for parallel application of N-to-1 compression functions.
11///
12/// This trait enables efficient batch compression operations where multiple N-element
13/// chunks are compressed in parallel. It's particularly useful for constructing hash trees
14/// and Merkle trees where many compression operations need to be performed simultaneously.
15///
16/// The trait is parameterized by:
17/// - `T`: The type of values being compressed (typically hash digests)
18/// - `N`: The arity of the compression function (number of inputs per compression)
19pub trait ParallelPseudoCompression<T, const N: usize> {
20	/// The underlying compression function that performs N-to-1 compression.
21	type Compression: CompressionFunction<T, N>;
22
23	/// Returns a reference to the underlying compression function.
24	fn compression(&self) -> &Self::Compression;
25
26	/// Compresses multiple N-element chunks in parallel.
27	///
28	/// # Arguments
29	/// * `inputs` - A slice containing the values to compress. Must have length `N * out.len()`.
30	/// * `out` - Output buffer where compressed values will be written.
31	///
32	/// # Behavior
33	/// For each index `i` in `0..out.len()`, this method computes:
34	/// ```text
35	/// out[i] = Compression::compress([inputs[i*N], inputs[i*N+1], ..., inputs[i*N+N-1]])
36	/// ```
37	///
38	/// All compressions are performed in parallel for efficiency.
39	///
40	/// # Post-conditions
41	/// After this method returns, all elements in `out` will be initialized with the
42	/// compressed values from their corresponding N-element chunks in `inputs`.
43	///
44	/// # Panics
45	/// Panics if `inputs.len() != N * out.len()`.
46	fn parallel_compress(&self, inputs: &[T], out: &mut [MaybeUninit<T>]);
47}
48
49/// A simple adapter that wraps any `CompressionFunction` to implement `ParallelCompression`.
50///
51/// This adapter provides a straightforward way to use existing compression functions
52/// in parallel contexts by applying them sequentially to each N-element chunk.
53#[derive(Debug, Clone, Default)]
54pub struct ParallelCompressionAdaptor<C> {
55	compression: C,
56}
57
58impl<C> ParallelCompressionAdaptor<C> {
59	/// Creates a new adapter wrapping the given compression function.
60	pub const fn new(compression: C) -> Self {
61		Self { compression }
62	}
63}
64
65impl<T, C, const ARITY: usize> ParallelPseudoCompression<T, ARITY> for ParallelCompressionAdaptor<C>
66where
67	T: Clone + Send + Sync,
68	C: CompressionFunction<T, ARITY> + Sync,
69{
70	type Compression = C;
71
72	fn compression(&self) -> &Self::Compression {
73		&self.compression
74	}
75
76	fn parallel_compress(&self, inputs: &[T], out: &mut [MaybeUninit<T>]) {
77		assert_eq!(inputs.len(), ARITY * out.len(), "Input length must be N * output length");
78
79		inputs
80			.par_chunks_exact(ARITY)
81			.zip(out.par_iter_mut())
82			.for_each(|(chunk, output)| {
83				// Convert slice to array for compression function
84				let chunk_array: [T; ARITY] = array::from_fn(|j| chunk[j].clone());
85				let compressed = self.compression.compress(chunk_array);
86				output.write(compressed);
87			});
88	}
89}
90
91#[cfg(test)]
92mod tests {
93	use std::mem::MaybeUninit;
94
95	use rand::prelude::*;
96
97	use super::*;
98
99	// Simple test compression function that XORs all inputs
100	#[derive(Clone, Debug)]
101	struct XorCompression;
102
103	impl CompressionFunction<u64, 3> for XorCompression {
104		fn compress(&self, input: [u64; 3]) -> u64 {
105			input[0] ^ input[1] ^ input[2]
106		}
107	}
108
109	#[test]
110	fn test_parallel_compression_adaptor() {
111		let mut rng = StdRng::seed_from_u64(0);
112		let compression = XorCompression;
113		let adaptor = ParallelCompressionAdaptor::new(compression.clone());
114
115		// Test with 4 chunks of 3 elements each
116		const N: usize = 3;
117		const NUM_CHUNKS: usize = 4;
118		let inputs: Vec<u64> = (0..N * NUM_CHUNKS).map(|_| rng.random()).collect();
119
120		// Use the adaptor
121		let mut adaptor_output = [MaybeUninit::<u64>::uninit(); NUM_CHUNKS];
122		adaptor.parallel_compress(&inputs, &mut adaptor_output);
123		let adaptor_results: Vec<u64> = adaptor_output
124			.into_iter()
125			.map(|x| unsafe { x.assume_init() })
126			.collect();
127
128		// Manually compress each chunk
129		let mut manual_results = Vec::new();
130		for chunk_idx in 0..NUM_CHUNKS {
131			let start = chunk_idx * N;
132			let chunk = [inputs[start], inputs[start + 1], inputs[start + 2]];
133			manual_results.push(compression.compress(chunk));
134		}
135
136		// Results should be identical
137		assert_eq!(adaptor_results, manual_results);
138	}
139
140	#[test]
141	#[should_panic(expected = "Input length must be N * output length")]
142	fn test_mismatched_input_length() {
143		let compression = XorCompression;
144		let adaptor = ParallelCompressionAdaptor::new(compression);
145
146		let inputs = vec![1u64, 2, 3, 4]; // 4 elements
147		let mut output = [MaybeUninit::<u64>::uninit(); 2]; // Expecting 6 elements (2 * 3)
148
149		adaptor.parallel_compress(&inputs, &mut output);
150	}
151}
binius_hash/parallel_compression.rs

binius_hash/
parallel_compression.rs
