binius_hash/

binary_merkle_tree.rs

// Copyright 2024-2025 Irreducible Inc.
// Copyright 2026 The Binius Developers

use std::{fmt::Debug, mem::MaybeUninit};

use binius_field::Field;
use binius_utils::{
	checked_arithmetics::log2_strict_usize,
	rand::par_rand,
	rayon::{prelude::*, slice::ParallelSlice},
};
use digest::{Digest, FixedOutputReset, Output, block_api::BlockSizeUser};
use rand::{CryptoRng, Rng, rngs::StdRng};

use super::{
	compress::CompressionFunction, parallel_compression::ParallelPseudoCompression,
	parallel_digest::ParallelDigest,
};

/// A bundle of hash and compression types used to build and verify a binary Merkle tree.
///
/// Most callers want to vary the underlying hash family (SHA-256, etc.) as a single unit
/// rather than independently picking a leaf hash, a compression function, and their parallel
/// counterparts. `HashSuite` bundles the four related types so that user-facing prover and
/// verifier APIs can take a single `H: HashSuite` parameter instead of two or three loose hash
/// trait parameters.
pub trait HashSuite {
	/// Sequential hash used to compute leaf digests during verification.
	type LeafHash: Digest + BlockSizeUser + FixedOutputReset + Send;
	/// Sequential 2-to-1 compression used to fold inner Merkle nodes during verification.
	type Compression: CompressionFunction<Output<Self::LeafHash>, 2> + Default;
	/// Parallel counterpart of [`Self::LeafHash`] used during proving.
	type ParLeafHash: ParallelDigest<Digest = Self::LeafHash> + Default;
	/// Parallel counterpart of [`Self::Compression`] used during proving.
	type ParCompression: ParallelPseudoCompression<Output<Self::LeafHash>, 2, Compression = Self::Compression>
		+ Default;
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
	#[error("Index exceeds Merkle tree base size: {max}")]
	IndexOutOfRange { max: usize },
	#[error("values length must be a multiple of the batch size")]
	IncorrectBatchSize,
	#[error("The argument length must be a power of two.")]
	PowerOfTwoLengthRequired,
	#[error("The layer does not exist in the Merkle tree")]
	IncorrectLayerDepth,
}

/// A binary Merkle tree that commits batches of vectors.
///
/// The vector entries at each index in a batch are hashed together into leaf digests. Then a
/// Merkle tree is constructed over the leaf digests. The implementation requires that the vector
/// lengths are all equal to each other and a power of two.
#[derive(Debug, Clone)]
pub struct BinaryMerkleTree<D, F> {
	/// Base-2 logarithm of the number of leaves
	pub log_len: usize,
	/// The inner nodes, arranged as a flattened array of layers with the root at the end
	pub inner_nodes: Vec<D>,
	/// Salt values for each leaf (if using hiding commitments)
	pub salts: Vec<F>,
}

pub fn build<F, H, R>(
	elements: &[F],
	batch_size: usize,
	salt_len: usize,
	rng: R,
) -> Result<BinaryMerkleTree<Output<H::LeafHash>, F>, Error>
where
	F: Field,
	H: HashSuite,
	R: Rng + CryptoRng,
{
	if !elements.len().is_multiple_of(batch_size) {
		return Err(Error::IncorrectBatchSize);
	}

	let len = elements.len() / batch_size;

	if !len.is_power_of_two() {
		return Err(Error::PowerOfTwoLengthRequired);
	}

	build_from_iterator::<_, H, _, _>(
		elements
			.par_chunks(batch_size)
			.map(|chunk| chunk.iter().copied()),
		batch_size,
		salt_len,
		rng,
	)
}

pub fn build_from_iterator<F, H, R, ParIter>(
	iterated_chunks: ParIter,
	n_items_per_input: usize,
	salt_len: usize,
	mut rng: R,
) -> Result<BinaryMerkleTree<Output<H::LeafHash>, F>, Error>
where
	F: Field,
	H: HashSuite,
	R: Rng + CryptoRng,
	ParIter: IndexedParallelIterator<Item: IntoIterator<Item = F, IntoIter: Send>>,
{
	let log_len = log2_strict_usize(iterated_chunks.len()); // precondition

	// Generate salts if needed
	let salts =
		par_rand::<StdRng, _, _>(salt_len << log_len, &mut rng, F::random).collect::<Vec<_>>();

	let total_length = (1 << (log_len + 1)) - 1;
	let mut inner_nodes = Vec::with_capacity(total_length);
	hash_leaves::<F, H, _>(
		iterated_chunks,
		n_items_per_input,
		&mut inner_nodes.spare_capacity_mut()[..(1 << log_len)],
		&salts,
	);

	let (prev_layer, mut remaining) = inner_nodes.spare_capacity_mut().split_at_mut(1 << log_len);

	let mut prev_layer = unsafe {
		// SAFETY: prev-layer was initialized by hash_leaves
		prev_layer.assume_init_mut()
	};
	let parallel_compression = H::ParCompression::default();
	for i in 1..(log_len + 1) {
		let (next_layer, next_remaining) = remaining.split_at_mut(1 << (log_len - i));
		remaining = next_remaining;

		parallel_compression.parallel_compress(prev_layer, next_layer);

		prev_layer = unsafe {
			// SAFETY: next_layer was just initialized by compress_layer
			next_layer.assume_init_mut()
		};
	}

	unsafe {
		// SAFETY: inner_nodes should be entirely initialized by now
		// Note that we don't incrementally update inner_nodes.len() since
		// that doesn't play well with using split_at_mut on spare capacity.
		inner_nodes.set_len(total_length);
	}
	Ok(BinaryMerkleTree {
		log_len,
		inner_nodes,
		salts,
	})
}

impl<D: Clone, F> BinaryMerkleTree<D, F> {
	pub fn root(&self) -> D {
		self.inner_nodes
			.last()
			.expect("MerkleTree inner nodes can't be empty")
			.clone()
	}

	/// Returns the salt values associated with a specific leaf index in the Merkle tree.
	///
	/// # Arguments
	/// * `index` - The index of the leaf. Must be less than 2^log_len (the total number of leaves).
	pub fn get_salt(&self, index: usize) -> &[F] {
		assert!(index < (1 << self.log_len));
		let salt_len = self.salts.len() >> self.log_len;
		&self.salts[index * salt_len..(index + 1) * salt_len]
	}

	pub fn layer(&self, layer_depth: usize) -> Result<&[D], Error> {
		if layer_depth > self.log_len {
			return Err(Error::IncorrectLayerDepth);
		}
		let range_start = self.inner_nodes.len() + 1 - (1 << (layer_depth + 1));

		Ok(&self.inner_nodes[range_start..range_start + (1 << layer_depth)])
	}

	/// Get a Merkle branch for the given index
	///
	/// Throws if the index is out of range
	pub fn branch(&self, index: usize, layer_depth: usize) -> Result<Vec<D>, Error> {
		if index >= 1 << self.log_len || layer_depth > self.log_len {
			return Err(Error::IndexOutOfRange {
				max: (1 << self.log_len) - 1,
			});
		}

		let branch = (0..self.log_len - layer_depth)
			.map(|j| {
				let node_index = (((1 << j) - 1) << (self.log_len + 1 - j)) | (index >> j) ^ 1;
				self.inner_nodes[node_index].clone()
			})
			.collect();

		Ok(branch)
	}
}

/// Hashes the elements in chunks of a vector into digests.
///
/// Given a vector of elements and an output buffer of N hash digests, this splits the elements
/// into N equal-sized chunks and hashes each chunks into the corresponding output digest.
///
/// Each leaf is built from exactly `n_items_per_input` data elements (plus the per-leaf salt, when
/// salts are present), so the leaf byte length is constant. This is passed to
/// [`ParallelDigest::digest_with_const_len`] so the hasher can specialize for short leaves.
///
/// # Preconditions
/// - Each iterator in `iterated_chunks` yields exactly `n_items_per_input` elements.
#[tracing::instrument("hash_leaves", skip_all, level = "debug")]
fn hash_leaves<F, H, ParIter>(
	iterated_chunks: ParIter,
	n_items_per_input: usize,
	digests: &mut [MaybeUninit<Output<H::LeafHash>>],
	salts: &[F],
) where
	F: Field,
	H: HashSuite,
	ParIter: IndexedParallelIterator<Item: IntoIterator<Item = F, IntoIter: Send>>,
{
	if salts.is_empty() {
		// Need special-case handling when salts is empty, otherwise salt_len is 0 and par_chunks
		// cannot handle chunk size of 0.
		let hasher = H::ParLeafHash::default();
		hasher.digest_with_const_len(n_items_per_input, iterated_chunks, digests);
	} else {
		assert!(salts.len().is_multiple_of(digests.len()));

		let salt_len = salts.len() / digests.len();

		// Create an iterator that chains each chunk with its salt
		let salted_iter = iterated_chunks
			.zip(salts.par_chunks(salt_len))
			.map(|(chunk, salt)| chunk.into_iter().chain(salt.iter().copied()));

		// Each salted leaf yields the data elements followed by the salt elements.
		let hasher = H::ParLeafHash::default();
		hasher.digest_with_const_len(n_items_per_input + salt_len, salted_iter, digests);
	}
}