binius_ip_prover/prodcheck/

mod.rs

// Copyright 2025-2026 The Binius Developers

use std::iter;

use binius_compute::{Allocator, VecLike};
use binius_field::{Field, PackedField};
use binius_ip::{mlecheck, prodcheck::MultilinearEvalClaim};
use binius_math::{
	FieldBuffer, FieldVec,
	line::extrapolate_line,
	multilinear::eq::{eq_ind_partial_eval, eq_one_var},
};
use binius_utils::rayon::{
	prelude::*,
	task_size::{IndexedParallelIteratorExt, WorkPerItem},
};
use itertools::izip;

use crate::{
	channel::IPProverChannel,
	sumcheck::{
		ProveSingleOutput, bivariate_product_mle, common::MleCheckProver, prove_single_mlecheck,
	},
};

pub mod one_pad_mle;

use one_pad_mle::OnePadMleCheckProver;

/// Witness-based prover for the product check protocol.
///
/// This prover reduces the claim that a multilinear polynomial evaluates to a product over a
/// Boolean hypercube to a single multilinear evaluation claim.
pub struct ProdcheckProver<'a, A: Allocator, P: PackedField> {
	/// Product layers from largest (original witness) to second-smallest.
	/// `layers[0]` is the original witness. The final products layer is returned
	/// separately from the constructor.
	layers: Vec<FieldVec<P, A>>,
	/// Allocator the product layers are drawn from.
	alloc: &'a A,
}

// A manual `Clone` impl (rather than `#[derive(Clone)]`) so the bound lands on the layer buffer
// `A::Vec<P>` rather than on `A` and `P`: a derive would require `A: Clone` yet still fail to clone
// the layers unless `A::Vec<P>: Clone`. Holds for the `Vec`-backed `GlobalAllocator`.
impl<A: Allocator, P: PackedField> Clone for ProdcheckProver<'_, A, P>
where
	A::Vec<P>: Clone,
{
	fn clone(&self) -> Self {
		Self {
			layers: self.layers.clone(),
			alloc: self.alloc,
		}
	}
}

impl<'a, A, F, P> ProdcheckProver<'a, A, P>
where
	A: Allocator,
	F: Field,
	P: PackedField<Scalar = F>,
{
	/// Creates a new [`ProdcheckProver`].
	///
	/// Returns `(prover, products)` where `products` is the final layer containing the
	/// products over all `k` variables.
	///
	/// # Arguments
	/// * `k` - The number of variables over which the product is taken. Each reduction step reduces
	///   one variable by computing pairwise products.
	/// * `witness` - The witness polynomial
	///
	/// # Preconditions
	/// * `witness.log_len() >= k`
	pub fn new(k: usize, alloc: &'a A, witness: FieldVec<P, A>) -> (Self, FieldVec<P, A>) {
		assert!(witness.log_len() >= k); // precondition

		let mut layers = Vec::with_capacity(k + 1);
		layers.push(witness);

		for _ in 0..k {
			let prev_layer = layers.last().expect("layers is non-empty");
			let next_log_len = prev_layer.log_len() - 1;
			let (half_0, half_1) = prev_layer.split_half();

			// A next-layer word is the product of its two sibling words, written into the pool.
			// Each layer is half the width of the one below it, down to a single word.
			// One word is one multiply, so a run must be long enough to pay back handing it off.
			let out_len = half_0.as_ref().len();
			let mut next_data = alloc.alloc::<P>(out_len);
			(next_data.spare_capacity_mut(), half_0.as_ref(), half_1.as_ref())
				.into_par_iter()
				.with_min_task(WorkPerItem::FieldMuls)
				.for_each(|(out, &v0, &v1)| {
					out.write(v0 * v1);
				});
			// Invariant: every zip input holds at least `out_len` words.
			//
			// A parallel zip yields as many items as its shortest input holds.
			// A shorter input would leave trailing slots uninitialized.
			//
			//     spare capacity:  >= out_len   allocated for at least that many
			//     sibling halves:  == out_len   halves of one buffer
			assert!(
				next_data.capacity() - next_data.len() >= out_len,
				"the allocated buffer must hold every claimed slot"
			);
			assert_eq!(
				half_1.as_ref().len(),
				out_len,
				"the two sibling halves must hold exactly one word per claimed slot"
			);
			// Safety: the length claim covers only initialized slots.
			// - The assertions above bound every zip input below by `out_len`.
			// - So the loop ran `out_len` items.
			// - Each item wrote one slot.
			unsafe {
				next_data.set_len(out_len);
			}
			let next_layer = FieldBuffer::new(next_log_len, next_data);

			layers.push(next_layer);
		}

		let products = layers.pop().expect("layers has k+1 elements");
		(Self { layers, alloc }, products)
	}

	/// Returns the number of remaining layers to prove.
	pub const fn n_layers(&self) -> usize {
		self.layers.len()
	}

	/// Pops the widest remaining layer and returns it.
	///
	/// # Preconditions
	/// * `self.n_layers() >= 1`
	pub fn pop_layer(&mut self) -> FieldVec<P, A> {
		self.layers
			.pop()
			.expect("precondition: layers is non-empty")
	}

	/// Pops the last layer and returns an MLE-check prover for it.
	///
	/// Returns `(layer_prover, remaining)` where:
	/// - `layer_prover` is an MLE-check prover for the popped layer
	/// - `remaining` is `Some(self)` if there are more layers, `None` otherwise
	pub fn layer_prover(
		mut self,
		claim: MultilinearEvalClaim<F>,
	) -> (impl MleCheckProver<F> + 'a, Option<Self>) {
		let alloc = self.alloc;
		let layer = self.layers.pop().expect("layers is non-empty");

		let remaining = if self.layers.is_empty() {
			None
		} else {
			Some(self)
		};

		// The layer has one more variable than the claim point.
		// Its low and high halves are the two multilinears whose product this layer reduces.
		// Sharing the buffer between the halves avoids copying this (largest) layer.
		let prover = bivariate_product_mle::new_split_half(alloc, layer, claim.point, claim.eval);

		(prover, remaining)
	}

	/// Runs the product check protocol and returns the final evaluation claim.
	///
	/// This consumes the prover and runs sumcheck reductions from the smallest layer back to
	/// the largest.
	///
	/// # Arguments
	/// * `claim` - The initial multilinear evaluation claim
	/// * `channel` - The channel for sending prover messages and sampling challenges
	///
	/// # Preconditions
	/// * `claim.point.len() == witness.log_len() - k` (where k is the number of reduction layers)
	pub fn prove(
		self,
		claim: MultilinearEvalClaim<F>,
		channel: &mut impl IPProverChannel<F>,
	) -> MultilinearEvalClaim<F> {
		let mut prover_opt = Some(self);
		let mut claim = claim;

		while let Some(prover) = prover_opt {
			let (mle_prover, remaining) = prover.layer_prover(claim.clone());
			prover_opt = remaining;

			let ProveSingleOutput {
				multilinear_evals,
				challenges,
			} = prove_single_mlecheck(mle_prover, channel);

			let [eval_0, eval_1] = multilinear_evals
				.try_into()
				.expect("prover has two multilinears");

			channel.send_many(&[eval_0, eval_1]);

			let r = channel.sample();
			let next_eval = extrapolate_line(eval_0, eval_1, r);

			let mut next_point = challenges;
			next_point.reverse();
			next_point.push(r);

			claim = MultilinearEvalClaim {
				eval: next_eval,
				point: next_point,
			};
		}

		claim
	}
}

/// Output of [`batch_prove`].
///
/// After the full `n_layers` reduction, `evals` holds each input prover's reduced evaluation at
/// `eval_point`. The batched claim the verifier checks is the eq(selector)-weighted combination
/// of these evaluations.
pub struct BatchProveOutput<F> {
	/// The reduced evaluation point (`selector ++ content ++ node`) shared by all input provers.
	pub eval_point: Vec<F>,
	/// Each input prover's reduced evaluation at `eval_point`, in input order.
	pub evals: Vec<F>,
}

/// Runs a batched product check protocol for multiple independent prodcheck provers.
///
/// This combines n provers, each for an $m$-variate multilinear, using multilinear interpolation
/// over k selector variables (where $n \le 2^k$). The combined claim is the multilinear
/// extrapolation of the individual claimed products (padded with zeros to $2^k$) evaluated at the
/// given point.
///
/// The claimed products may themselves be evaluations of the $m$-variate product multilinears at a
/// shared `content_point` (of length equal to each prover's reduced product dimension). When the
/// products are scalars (each prover reduces over all of its variables), `content_point` is empty.
///
/// # Arguments
/// * `provers` - Vec of n prodcheck provers. All must have the same `n_layers()`, which is $m$.
/// * `claimed_products` - Vec of n claimed product values, one per prover. Each is the
///   corresponding prover's product multilinear evaluated at `content_point`.
/// * `selector_point` - Evaluation point for the selector variables. Length is $k$.
/// * `content_point` - Shared evaluation point at which the claimed products are taken. Length is
///   the product-multilinear dimension (i.e. `witness.log_len() - n_layers`). Empty for scalar
///   products.
/// * `channel` - The channel for sending prover messages and sampling challenges.
///
/// # Preconditions
/// * `provers` must be non-empty.
/// * All provers must have the same `n_layers()` value.
/// * `2^selector_point.len() >= provers.len()`.
/// * `claimed_products.len() == provers.len()`.
/// * `content_point.len() == witness.log_len() - n_layers` for each prover.
///
/// Returns the reduced per-input-prover evaluations at the reduced evaluation point. The
/// batched claim is checked by the ordinary `binius_ip::prodcheck::verify` recursion over
/// `n_layers` layers (the eq(selector)-weighted combination of the returned evaluations), with
/// the selector coordinates forming the first `k` coordinates of the claim point.
///
/// # Returns
/// A [`BatchProveOutput`] with the reduced `eval_point` and each input prover's reduced eval.
///
/// # Mathematical Description
///
/// Let $f_i \in K[X_0, \ldots, X_{m-1}]$ be multilinear for all $i \in \{0, \ldots, n - 1\}$. The
/// $i$'th prover is a prodcheck prover for $f_i$. Let $p_i \in K$ be the claimed hypercube product
/// of $f_i$.
///
/// Let $y \in K^k$ be the evaluation point. The prover is proving a claim that
///
/// $$
/// \sum_{i \in B_k} \textsf{eq}(i; y) \prod_{j \in B_m} f_i(j) = \sum_{i \in B_k} \textsf{eq}(i; y)
/// p_i, $$
///
/// reducing to an evaluation of the interpolated multilinear
///
/// $$
/// \hat{f}(Y_0, \ldots, Y_{k-1}, X_0, \ldots, X_{m-1}) = \sum_{i \in B_k} \textsf{eq}(i; Y) f_i(X).
/// $$
pub fn batch_prove<'a, A: Allocator, F: Field, P: PackedField<Scalar = F>>(
	provers: Vec<ProdcheckProver<'a, A, P>>,
	claimed_products: Vec<F>,
	selector_point: Vec<F>,
	content_point: Vec<F>,
	channel: &mut impl IPProverChannel<F>,
) -> BatchProveOutput<F> {
	assert!(!provers.is_empty()); // precondition
	assert_eq!(claimed_products.len(), provers.len()); // precondition

	let n = provers.len();
	let k = selector_point.len();
	assert!(provers.len() <= (1 << k)); // precondition

	let n_layers = provers[0].n_layers();
	assert!(n_layers >= 1); // precondition
	assert!(provers.iter().all(|p| p.n_layers() == n_layers)); // precondition

	// Thread the content point as the initial inner (content) coordinates of the evaluation point.
	// `batch_prove_layer` splits `eval_point.split_at(k)` into (selector, content); on the first
	// layer this seeds each layer prover with `claim{point: content_point, eval: claimed_product}`.
	let eval_point = [selector_point, content_point].concat();

	let (provers, evals, eval_point) = (0..n_layers).fold(
		(provers, claimed_products, eval_point),
		|(provers, claimed_products, eval_point), _| {
			batch_prove_layer(provers, claimed_products, &eval_point, k, channel)
		},
	);
	debug_assert!(provers.is_empty(), "the final layer leaves no provers");
	debug_assert_eq!(evals.len(), n);

	BatchProveOutput { eval_point, evals }
}

#[allow(clippy::type_complexity)]
fn batch_prove_layer<'a, A: Allocator, F: Field, P: PackedField<Scalar = F>>(
	provers: Vec<ProdcheckProver<'a, A, P>>,
	claimed_products: Vec<F>,
	eval_point: &[F],
	k: usize,
	channel: &mut impl IPProverChannel<F>,
) -> (Vec<ProdcheckProver<'a, A, P>>, Vec<F>, Vec<F>) {
	let alloc = provers[0].alloc;
	let inner_coords = &eval_point[k..];

	let (layer_provers, next_provers): (Vec<_>, Vec<_>) = iter::zip(provers, claimed_products)
		.map(|(prover, prod)| {
			prover.layer_prover(MultilinearEvalClaim {
				eval: prod,
				point: inner_coords.to_vec(),
			})
		})
		.unzip();

	let (next_claimed_products, next_point) =
		prove_layer_rounds::<A, F, P>(layer_provers, eval_point, k, alloc, channel);
	let next_provers = next_provers.into_iter().flatten().collect();

	(next_provers, next_claimed_products, next_point)
}

/// Runs one batched layer reduction over already-constructed layer provers.
///
/// The per-instance reductions run in lockstep, their round polynomials combined with the
/// eq(selector) weights, followed by the shared selector rounds over the per-instance child
/// evaluation pairs. This is the layer loop of both [`batch_prove`], whose instances share a layer
/// count, and [`batch_prove_unequal_depths`], whose don't.
///
/// # Returns
///
/// Each instance's claim on the next layer, in input order, and the reduced evaluation point.
fn prove_layer_rounds<A: Allocator, F: Field, P: PackedField<Scalar = F>>(
	mut layer_provers: Vec<impl MleCheckProver<F>>,
	eval_point: &[F],
	k: usize,
	alloc: &A,
	channel: &mut impl IPProverChannel<F>,
) -> (Vec<F>, Vec<F>) {
	let n = layer_provers.len();
	// Split eval_point into outer (selector) and inner (content) coordinates.
	let (outer_coords, inner_coords) = eval_point.split_at(k);

	// Compute eq weights for batching: eq(i, outer_coords) for all i in B_k.
	let eq_weights = eq_ind_partial_eval::<F>(outer_coords);

	// Content rounds: individual provers operate independently.
	let mut challenges = Vec::with_capacity(eval_point.len());

	for _round in 0..inner_coords.len() {
		// Execute each prover and compute weighted sum of round coefficients.
		let coeffss = layer_provers
			.iter_mut()
			.map(|prover| {
				let mut round_coeffs_vec = prover.execute();
				round_coeffs_vec
					.pop()
					.expect("prodcheck layer provers have round_coeffs_vec.len() == 1")
			})
			.collect::<Vec<_>>();

		let coeffs = iter::zip(coeffss, eq_weights.iter_scalars())
			.map(|(coeffs, weight)| coeffs * weight)
			.sum();

		// Send truncated round proof to channel.
		channel.send_many(mlecheck::RoundProof::truncate(coeffs).coeffs());

		// Sample challenge and fold all provers.
		let challenge = channel.sample();
		challenges.push(challenge);

		for prover in layer_provers.iter_mut() {
			prover.fold(challenge);
		}
	}

	// Finish inner provers to get [eval_0, eval_1] pairs.
	let (mut vals_0, mut vals_1): (Vec<F>, Vec<F>) = layer_provers
		.into_iter()
		.map(|prover| {
			let evals = prover.finish();
			let [e0, e1]: [F; 2] = evals
				.try_into()
				.expect("bivariate product prover has two multilinears");
			(e0, e1)
		})
		.unzip();

	// Pad vals_0 and vals_1 to 2^k with zeros for FieldBuffer::from_values.
	vals_0.resize(1 << k, F::ZERO);
	vals_1.resize(1 << k, F::ZERO);

	// Compute eval from buffers: sum_v eq(v, outer_coords) * vals_0[v] * vals_1[v].
	let eval = izip!(&vals_0, &vals_1, eq_weights.as_ref())
		.map(|(&v0, &v1, &eq_i)| v0 * v1 * eq_i)
		.sum();

	// Selector rounds: pack eval pairs straight into allocator buffers and use a single prover.
	let outer_prover = bivariate_product_mle::new(
		alloc,
		[
			FieldBuffer::<P, _>::from_values_in(alloc, &vals_0),
			FieldBuffer::<P, _>::from_values_in(alloc, &vals_1),
		],
		outer_coords.to_vec(),
		eval,
	);

	let ProveSingleOutput {
		multilinear_evals: outer_evals,
		challenges: outer_challenges,
	} = prove_single_mlecheck(outer_prover, channel);

	challenges.extend(outer_challenges);

	let [merged_eval_0, merged_eval_1]: [F; 2] =
		outer_evals.try_into().expect("prover has two multilinears");

	// Finalize layer: send evals, sample r, compute next claim.
	channel.send_many(&[merged_eval_0, merged_eval_1]);

	let r = channel.sample();

	let mut next_point = challenges;
	next_point.reverse();
	next_point.push(r);

	// Update claimed products for next iteration, dropping the padded (2^k) selector slots.
	let next_claimed_products = iter::zip(&vals_0[..n], &vals_1[..n])
		.map(|(e0, e1)| extrapolate_line(*e0, *e1, r))
		.collect();

	(next_claimed_products, next_point)
}

/// Output of [`batch_prove_unequal_depths`].
///
/// After `n_layers - 1` reduction layers, each tree retains its final (widest) layer, wrapped in
/// the one-padding prover that corrects it to the batch's depth.
pub struct BatchProveUnequalDepthsOutput<F, Prover> {
	/// The reduced evaluation point shared by all remaining provers.
	pub eval_point: Vec<F>,
	/// Each input prover's reduced (padded) claim paired with the not-yet-run prover for its final
	/// layer, in input order.
	pub provers: Vec<(F, Prover)>,
}

/// Runs a batched product check for trees of *unequal* depths.
///
/// This is [`batch_prove`] without the requirement that every prover have the same layer count.
/// Each tree shallower than the deepest is proved as a product check over the one-padding of its
/// witness — the same witness with constant-1 leaves filling the extra depth, which leaves its
/// product unchanged. The transcript is then exactly that of an equal-depth batch of the maximum
/// depth: the verifier runs the ordinary [`binius_ip::prodcheck::verify`] over `n_layers` layers
/// and never learns the individual depths.
///
/// Unlike [`batch_prove`], every prover must reduce over *all* of its witness variables, so each
/// product is a scalar and there is no content point. The one-padding is only worth its
/// bookkeeping on full trees, and dropping the content dimension keeps that bookkeeping to two
/// scalars per layer.
///
/// The prover does not materialize the padded witnesses. Each layer's per-tree reduction runs
/// through [`one_pad_mle`], which corrects the unpadded layer's messages in $O(1)$ per round.
///
/// The protocol and the round polynomials are specified in the *Batched Product Checks of Unequal
/// Depths* appendix of the Binius64 whitepaper.
///
/// # Arguments
///
/// As [`batch_prove`], except that the provers' layer counts may differ and there is no
/// `content_point`.
///
/// # Preconditions
/// * `provers` must be non-empty.
/// * Every prover's witness must have exactly `prover.n_layers()` variables, which must be at least
///   1.
/// * `2^selector_point.len() >= provers.len()`.
/// * `claimed_products.len() == provers.len()`.
///
/// # Returns
///
/// This stops one layer short, returning each tree's reduced claim beside the prover for its
/// final (widest) layer — the layer whose reduction
/// dominates the cost, which a caller can therefore batch with other sumchecks. Running those
/// provers and the selector rounds that follow finishes the check.
///
/// The returned claims and, after the final layer, the per-tree leaf evaluations are claims on the
/// *padded* witnesses. [`unpad_leaf_claim`] reduces one to the claim on the
/// tree's own witness.
pub fn batch_prove_unequal_depths<'a, A, F, P, Channel>(
	mut provers: Vec<ProdcheckProver<'a, A, P>>,
	claimed_products: Vec<F>,
	selector_point: Vec<F>,
	channel: &mut Channel,
) -> BatchProveUnequalDepthsOutput<F, impl MleCheckProver<F> + use<'a, A, F, P, Channel>>
where
	A: Allocator,
	F: Field,
	P: PackedField<Scalar = F>,
	Channel: IPProverChannel<F>,
{
	assert!(!provers.is_empty()); // precondition
	assert_eq!(claimed_products.len(), provers.len()); // precondition

	let k = selector_point.len();
	assert!(provers.len() <= (1 << k)); // precondition
	assert!(provers.iter().all(|prover| prover.n_layers() >= 1)); // precondition

	let alloc = provers[0].alloc;
	let n_layers = provers
		.iter()
		.map(ProdcheckProver::n_layers)
		.max()
		.expect("provers is non-empty");
	// How much depth each tree is padded by.
	let pad_lens = provers
		.iter()
		.map(|prover| n_layers - prover.n_layers())
		.collect::<Vec<_>>();

	// The tree products stay in hand for the padding layers, which are one-paddings of them; the
	// input vector itself becomes the running per-tree claims.
	let products = claimed_products.clone();
	let mut claims = claimed_products;
	let mut eval_point = selector_point;

	// Reduce down to the final layer. Iteration `node_len` reduces the layer holding `node_len`
	// node variables, which is the point's suffix past the selector coordinates.
	for _ in 0..n_layers - 1 {
		let layer_provers =
			layer_provers(&mut provers, &pad_lens, &products, &claims, &eval_point[k..]);
		let (next_claims, next_point) =
			prove_layer_rounds::<A, F, P>(layer_provers, &eval_point, k, alloc, channel);
		claims = next_claims;
		eval_point = next_point;
	}

	let provers = layer_provers(&mut provers, &pad_lens, &products, &claims, &eval_point[k..]);

	BatchProveUnequalDepthsOutput {
		eval_point,
		provers: iter::zip(claims, provers).collect(),
	}
}

/// Builds one padded layer prover per tree, for the layer claimed at `node_point`.
fn layer_provers<'a, A: Allocator, F: Field, P: PackedField<Scalar = F>>(
	provers: &mut [ProdcheckProver<'a, A, P>],
	pad_lens: &[usize],
	products: &[F],
	claims: &[F],
	node_point: &[F],
) -> Vec<OnePadMleCheckProver<F, impl MleCheckProver<F> + use<'a, A, F, P>>> {
	let node_len = node_point.len();

	izip!(provers, pad_lens, products, claims)
		.map(|(prover, &pad_len, &product, &claim)| {
			let alloc = prover.alloc;
			// While the batch is still above this tree, the layer is a one-padding of the tree's
			// product, whose two children are that product and the constant one. Only once the
			// batch reaches the tree does it start consuming its layers.
			let layer = if node_len < pad_len {
				FieldBuffer::from_values_in(alloc, &[product, F::ONE])
			} else {
				let layer = prover.pop_layer();
				assert_eq!(
					layer.log_len(),
					node_len - pad_len + 1,
					"precondition: the witness has exactly n_layers variables"
				);
				layer
			};
			one_pad_mle::new(alloc, layer, pad_len.min(node_len), node_point.to_vec(), claim)
		})
		.collect()
}

/// Reduces a leaf claim on a one-padded witness to the claim on the witness itself.
///
/// A batched product check over trees of unequal depths lifts each shallow tree to the batch's
/// depth by filling `n_pad_vars` extra leaf positions with ones, which leaves its product
/// unchanged. [`binius_ip::prodcheck::verify`] is oblivious to that, so the claim it outputs for
/// such a tree is a claim on the padded witness
///
/// $$
/// M'(X_\text{pad}, X_\text{real}) = 1 + \bigl( M(X_\text{real}) - 1 \bigr) \cdot
/// \text{eq}(0^\nu; X_\text{pad}),
/// $$
///
/// whose padding variables are the lowest ones. This divides out their equality weight and drops
/// them from the point, leaving the claim on $M$.
///
/// # Arguments
///
/// * `eval` - The claimed evaluation of the padded witness.
/// * `point` - The reduced evaluation point, with the batch's selector coordinates already
///   stripped.
/// * `n_pad_vars` - How much depth this tree was padded by: the batch's layer count less the tree's
///   own.
///
/// # Preconditions
/// * `point.len() >= n_pad_vars`
///
/// # Panics
///
/// Panics if the padding coordinates' equality weight is zero, which requires one of them to equal
/// one. They are the verifier's own challenges, so no prover can induce this; it happens with
/// probability at most $\nu / |K|$.
pub fn unpad_leaf_claim<F: Field>(
	eval: F,
	point: &[F],
	n_pad_vars: usize,
) -> MultilinearEvalClaim<F> {
	assert!(point.len() >= n_pad_vars); // precondition

	let pad_eq = point[..n_pad_vars]
		.iter()
		.map(|&coord| eq_one_var(F::ZERO, coord))
		.product::<F>();
	assert!(pad_eq != F::ZERO, "a padding coordinate equals one");

	MultilinearEvalClaim {
		eval: F::ONE + (eval - F::ONE) * pad_eq.invert_or_zero(),
		point: point[n_pad_vars..].to_vec(),
	}
}

#[cfg(test)]
mod tests {
	use binius_field::{PackedField, field::FieldOps};
	use binius_ip::prodcheck;
	use binius_math::{
		inner_product::inner_product,
		multilinear::{eq::eq_ind_partial_eval, evaluate::evaluate},
		test_utils::{Packed128b, random_field_buffer, random_scalars},
	};
	use binius_transcript::{ProverTranscript, fiat_shamir::HasherChallenger};
	use binius_utils::checked_arithmetics::log2_ceil_usize;

	type StdChallenger = HasherChallenger<sha2::Sha256>;
	use binius_compute::GlobalAllocator;
	use rand::prelude::*;

	use super::*;

	/// Combines the per-input-prover evals returned by [`batch_prove`] into the single
	/// [`MultilinearEvalClaim`] the verifier produces: the eq(selector)-weighted sum over the
	/// first `k` (selector) coordinates of the reduced evaluation point.
	fn combine_batch_prove<F: Field, P: PackedField<Scalar = F>>(
		output: BatchProveOutput<F>,
		k: usize,
	) -> MultilinearEvalClaim<F> {
		let BatchProveOutput { eval_point, evals } = output;
		let eq_weights = eq_ind_partial_eval::<P>(&eval_point[..k]);
		let final_eval =
			inner_product(evals.iter().copied(), (0..evals.len()).map(|i| eq_weights.get(i)));

		MultilinearEvalClaim {
			eval: final_eval,
			point: eval_point,
		}
	}

	fn test_prodcheck_prove_verify_helper<P: PackedField>(n: usize, k: usize) {
		let mut rng = StdRng::seed_from_u64(0);
		let alloc = GlobalAllocator;

		// 1. Create random witness with log_len = n + k
		let witness = random_field_buffer::<P>(&mut rng, n + k);

		// 2. Create prover (computes product layers)
		let (prover, products) = ProdcheckProver::new(k, &alloc, witness.clone());

		// 3. Generate random n-dimensional challenge point
		let eval_point = random_scalars::<P::Scalar>(&mut rng, n);

		// 4. Evaluate products layer at challenge point to create claim
		let products_eval = evaluate(&products, &eval_point);
		let claim = MultilinearEvalClaim {
			eval: products_eval,
			point: eval_point,
		};

		// 5. Run prover
		let mut prover_transcript = ProverTranscript::new(StdChallenger::default());
		let prover_output = prover.prove(claim.clone(), &mut prover_transcript);

		// 6. Run verifier
		let mut verifier_transcript = prover_transcript.into_verifier();
		let verifier_output = prodcheck::verify(k, claim, &mut verifier_transcript).unwrap();

		// 7. Check outputs match
		assert_eq!(prover_output, verifier_output);

		// 8. Verify multilinear evaluation of original witness
		let expected_eval = evaluate(&witness, &verifier_output.point);
		assert_eq!(verifier_output.eval, expected_eval);
	}

	#[test]
	fn test_prodcheck_prove_verify() {
		test_prodcheck_prove_verify_helper::<Packed128b>(4, 3);
	}

	#[test]
	fn test_prodcheck_full_prove_verify() {
		test_prodcheck_prove_verify_helper::<Packed128b>(0, 4);
	}

	fn test_prodcheck_layer_computation_helper<P: PackedField>(n: usize, k: usize) {
		let mut rng = StdRng::seed_from_u64(0);
		let alloc = GlobalAllocator;

		// Create random witness with log_len = n + k
		let witness = random_field_buffer::<P>(&mut rng, n + k);

		// Create prover (computes product layers)
		let (_prover, products) = ProdcheckProver::new(k, &alloc, witness.clone());

		// For each index i in the products layer, verify it equals the product of witness values
		// at indices i + z * 2^n for z in 0..2^k (strided access, not contiguous)
		let stride = 1 << n;
		let num_terms = 1 << k;
		for i in 0..(1 << n) {
			let mut expected_product = P::Scalar::ONE;
			for z in 0..num_terms {
				expected_product *= witness.get(i + z * stride);
			}
			let actual = products.get(i);
			assert_eq!(actual, expected_product, "Product mismatch at index {i}");
		}
	}

	#[test]
	fn test_prodcheck_layer_computation() {
		test_prodcheck_layer_computation_helper::<Packed128b>(4, 3);
	}

	/// Builds the same product layers serially, through a temporary vector.
	///
	/// The layer recurrence is written twice, so it can drift once.
	fn reference_layers<P: PackedField>(
		k: usize,
		alloc: &GlobalAllocator,
		witness: FieldBuffer<P>,
	) -> Vec<FieldBuffer<P>> {
		let mut layers = Vec::with_capacity(k + 1);
		layers.push(witness);

		for _ in 0..k {
			let prev_layer = layers.last().expect("layers is non-empty");
			let next_log_len = prev_layer.log_len() - 1;
			let (half_0, half_1) = prev_layer.split_half();

			// Pair each word of the low half with the word of the high half above it.
			let evals = half_0
				.as_ref()
				.iter()
				.zip(half_1.as_ref())
				.map(|(v0, v1)| *v0 * *v1)
				.collect::<Vec<P>>();

			let mut next_data = alloc.alloc::<P>(evals.len());
			next_data.extend_from_slice(&evals);
			layers.push(FieldBuffer::new(next_log_len, next_data));
		}

		layers
	}

	#[test]
	fn layers_match_the_reference_word_for_word() {
		let mut rng = StdRng::seed_from_u64(7);
		let alloc = GlobalAllocator;

		// Invariant: the layers are compared as packed words, not as scalars.
		// So the lanes past a short layer's length are compared too.
		//
		// Fixture state: four scalars per packed word.
		//
		//     log_len:   6   5   4   3   2   1   0
		//     words:    16   8   4   2   1   1   1
		//
		// From log_len 2 down a layer is one word whose trailing lanes are not elements.
		// Sweeping every depth at every witness size covers both sides of that line.
		for log_len in 1..=6 {
			for k in 1..=log_len {
				let witness = random_field_buffer::<Packed128b>(&mut rng, log_len);

				let (prover, products) = ProdcheckProver::new(k, &alloc, witness.clone());
				let reference = reference_layers(k, &alloc, witness);

				// The constructor keeps the first k layers and hands back the last one.
				let built = prover.layers.iter().chain(iter::once(&products));

				for (depth, (built, reference)) in built.zip(&reference).enumerate() {
					assert_eq!(built.log_len(), reference.log_len(), "log_len at depth {depth}");
					assert_eq!(
						built.as_ref(),
						reference.as_ref(),
						"layer words at depth {depth}, witness log_len {log_len}, k {k}"
					);
				}
			}
		}
	}

	// ==================== batch_prove tests ====================

	/// Helper function for testing batch_prove with ProdcheckProvers.
	///
	/// Each witness has exactly `n_layers` variables so that the products are scalars (0-variate).
	///
	/// # Arguments
	/// * `n_layers` - Number of product reduction layers (= variables per witness)
	/// * `n_provers` - Number of provers to batch
	fn test_batch_prove_verify_helper<P: PackedField>(n_layers: usize, n_provers: usize) {
		let mut rng = StdRng::seed_from_u64(42);
		let alloc = GlobalAllocator;

		let log_n_provers = log2_ceil_usize(n_provers);

		// Each witness has exactly n_layers variables; products are scalars
		let witnesses: Vec<FieldBuffer<P>> = (0..n_provers)
			.map(|_| random_field_buffer::<P>(&mut rng, n_layers))
			.collect();

		// One prover per witness, each paired with the products of its own layers.
		let (provers, individual_products): (Vec<_>, Vec<_>) = witnesses
			.iter()
			.map(|witness| ProdcheckProver::new(n_layers, &alloc, witness.clone()))
			.unzip();

		// Products are 0-variate (scalars): just get the single value
		let claimed_products: Vec<P::Scalar> = individual_products
			.iter()
			.map(|products| {
				assert_eq!(products.log_len(), 0);
				products.get(0)
			})
			.collect();

		// Generate random selector challenge point (length = log_n_provers)
		let selector_challenge = random_scalars::<P::Scalar>(&mut rng, log_n_provers);

		// Compute combined claim using eq weights
		let eq_weights = eq_ind_partial_eval::<P>(&selector_challenge);
		let combined_eval = inner_product(
			claimed_products.iter().copied(),
			(0..n_provers).map(|i| eq_weights.get(i)),
		);

		// The verifier claim has point = selector_challenge (length log_n_provers)
		let claim = MultilinearEvalClaim {
			eval: combined_eval,
			point: selector_challenge.clone(),
		};

		// Run batch_prove (scalar products: empty content point), then finish the final layer.
		let mut prover_transcript = ProverTranscript::new(StdChallenger::default());
		let batch_output = batch_prove(
			provers,
			claimed_products,
			selector_challenge,
			Vec::new(),
			&mut prover_transcript,
		);
		// Each remaining prover has exactly one layer.
		assert_eq!(batch_output.evals.len(), n_provers);
		let prover_output = combine_batch_prove::<_, P>(batch_output, log_n_provers);

		// Run verifier with n_layers layers
		let mut verifier_transcript = prover_transcript.into_verifier();
		let verifier_output = prodcheck::verify(n_layers, claim, &mut verifier_transcript).unwrap();

		// Check prover and verifier outputs match
		assert_eq!(prover_output, verifier_output);

		// Verify final evaluation against multilinear extrapolation of input witnesses
		let final_point = &verifier_output.point;
		assert_eq!(final_point.len(), log_n_provers + n_layers);

		let selector_challenges = &final_point[..log_n_provers];
		let content_challenges = &final_point[log_n_provers..];

		let selector_weights = eq_ind_partial_eval::<P>(selector_challenges);

		let expected_eval: P::Scalar = inner_product(
			(0..n_provers).map(|i| evaluate(&witnesses[i], content_challenges)),
			(0..n_provers).map(|i| selector_weights.get(i)),
		);

		assert_eq!(
			verifier_output.eval, expected_eval,
			"Final evaluation should match batch witness interpolation"
		);
	}

	#[test]
	fn test_batch_prove_power_of_two_provers() {
		// 4 provers, 3 layers
		test_batch_prove_verify_helper::<Packed128b>(3, 4);
	}

	#[test]
	fn test_batch_prove_non_power_of_two_provers() {
		// 3 provers (non-power of 2, requires padding), 4 layers
		test_batch_prove_verify_helper::<Packed128b>(4, 3);
	}

	#[test]
	fn test_batch_prove_single_prover() {
		// 1 prover (edge case), 5 layers
		test_batch_prove_verify_helper::<Packed128b>(5, 1);
	}

	#[test]
	fn test_batch_prove_single_layer() {
		// n_layers=1 edge case (the minimum): batch_prove runs 0 reductions and retains the single
		// (final) layer, which the test then finishes.
		test_batch_prove_verify_helper::<Packed128b>(1, 4);
	}

	/// Helper for testing batch_prove where the claimed products are non-scalar: each prover's
	/// product multilinear is `content_len`-variate, claimed at a shared random content point.
	///
	/// # Arguments
	/// * `n_layers` - Number of product reduction layers (= selector-reduced variables per witness)
	/// * `n_provers` - Number of provers to batch
	/// * `content_len` - Number of variables of the product multilinear (witness has log_len =
	///   content_len + n_layers)
	fn test_batch_prove_with_content_helper<P: PackedField>(
		n_layers: usize,
		n_provers: usize,
		content_len: usize,
	) {
		let mut rng = StdRng::seed_from_u64(7);
		let alloc = GlobalAllocator;

		let log_n_provers = log2_ceil_usize(n_provers);

		// Each witness has log_len = content_len + n_layers; products are content_len-variate.
		let witnesses: Vec<FieldBuffer<P>> = (0..n_provers)
			.map(|_| random_field_buffer::<P>(&mut rng, content_len + n_layers))
			.collect();

		let (provers, individual_products): (Vec<_>, Vec<_>) = witnesses
			.iter()
			.map(|witness| ProdcheckProver::new(n_layers, &alloc, witness.clone()))
			.unzip();

		// Shared content point; each claimed product is its multilinear evaluated there.
		let content_point = random_scalars::<P::Scalar>(&mut rng, content_len);
		let claimed_products: Vec<P::Scalar> = individual_products
			.iter()
			.map(|products| {
				assert_eq!(products.log_len(), content_len);
				evaluate(products, &content_point)
			})
			.collect();

		// Combined verifier claim: eq(selector)-weighted sum of the claimed products, at point
		// selector ++ content.
		let selector_challenge = random_scalars::<P::Scalar>(&mut rng, log_n_provers);
		let eq_weights = eq_ind_partial_eval::<P>(&selector_challenge);
		let combined_eval = inner_product(
			claimed_products.iter().copied(),
			(0..n_provers).map(|i| eq_weights.get(i)),
		);

		let claim = MultilinearEvalClaim {
			eval: combined_eval,
			point: [selector_challenge.clone(), content_point.clone()].concat(),
		};

		// Run batch_prove with non-empty content point, then finish the final layer.
		let mut prover_transcript = ProverTranscript::new(StdChallenger::default());
		let batch_output = batch_prove(
			provers,
			claimed_products,
			selector_challenge,
			content_point,
			&mut prover_transcript,
		);
		assert_eq!(batch_output.evals.len(), n_provers);
		let prover_output = combine_batch_prove::<_, P>(batch_output, log_n_provers);

		// Run verifier with n_layers layers.
		let mut verifier_transcript = prover_transcript.into_verifier();
		let verifier_output = prodcheck::verify(n_layers, claim, &mut verifier_transcript).unwrap();

		assert_eq!(prover_output, verifier_output);

		// The reduced point is [selector (log_n_provers), witness vars (content_len + n_layers)],
		// where the witness-var coordinates are in each witness's own variable order — so
		// `witness_i.evaluate(witness_challenges)` is the per-prover content evaluation that the
		// selector-weighted sum recombines.
		let final_point = &verifier_output.point;
		assert_eq!(final_point.len(), log_n_provers + n_layers + content_len);

		let selector_challenges = &final_point[..log_n_provers];
		let witness_challenges = &final_point[log_n_provers..];

		let selector_weights = eq_ind_partial_eval::<P>(selector_challenges);

		let expected_eval: P::Scalar = inner_product(
			(0..n_provers).map(|i| evaluate(&witnesses[i], witness_challenges)),
			(0..n_provers).map(|i| selector_weights.get(i)),
		);

		assert_eq!(
			verifier_output.eval, expected_eval,
			"Final evaluation should match batch witness interpolation"
		);
	}

	#[test]
	fn test_batch_prove_with_content() {
		// 3 provers (non power of 2), 4 layers, content_len = 2.
		test_batch_prove_with_content_helper::<Packed128b>(4, 3, 2);
	}

	// ==================== batch_prove_unequal_depths tests ====================

	/// One prover per entry of `depths`, each reducing over all of its witness variables.
	#[allow(clippy::type_complexity)]
	fn unequal_depth_provers<'a, P: PackedField>(
		rng: &mut impl Rng,
		alloc: &'a GlobalAllocator,
		depths: &[usize],
	) -> (Vec<FieldBuffer<P>>, Vec<ProdcheckProver<'a, GlobalAllocator, P>>, Vec<P::Scalar>) {
		itertools::multiunzip(depths.iter().map(|&depth| {
			let witness = random_field_buffer::<P>(&mut *rng, depth);
			let (prover, products) = ProdcheckProver::new(depth, alloc, witness.clone());
			assert_eq!(products.log_len(), 0);
			(witness, prover, products.get(0))
		}))
	}

	/// The eq(selector)-weighted combination of per-tree claims, as the verifier forms it.
	fn combine_claims<P: PackedField>(
		claims: &[P::Scalar],
		selector_point: &[P::Scalar],
	) -> P::Scalar {
		let eq_weights = eq_ind_partial_eval::<P>(selector_point);
		inner_product(claims.iter().copied(), (0..claims.len()).map(|i| eq_weights.get(i)))
	}

	/// Proves a batch of unequal-depth trees against the depth-oblivious verifier, then unpads each
	/// tree's leaf claim and checks it against that tree's own witness.
	fn test_unequal_depths_helper<P: PackedField>(depths: &[usize]) {
		let mut rng = StdRng::seed_from_u64(11);
		let alloc = GlobalAllocator;

		let k = log2_ceil_usize(depths.len());
		let n_layers = *depths.iter().max().expect("depths is non-empty");

		let (witnesses, provers, claimed_products) =
			unequal_depth_provers::<P>(&mut rng, &alloc, depths);

		// The verifier's input claim is the eq(selector)-weighted combination of the products.
		let selector_point = random_scalars::<P::Scalar>(&mut rng, k);
		let claim = MultilinearEvalClaim {
			eval: combine_claims::<P>(&claimed_products, &selector_point),
			point: selector_point.clone(),
		};

		let mut prover_transcript = ProverTranscript::new(StdChallenger::default());
		let BatchProveUnequalDepthsOutput {
			eval_point,
			provers,
		} = batch_prove_unequal_depths(
			provers,
			claimed_products,
			selector_point,
			&mut prover_transcript,
		);

		// Finish the retained final layer: run it exactly as an interior reduction layer does.
		let (_claims, provers): (Vec<_>, Vec<_>) = provers.into_iter().unzip();
		let (evals, eval_point) =
			prove_layer_rounds::<_, _, P>(provers, &eval_point, k, &alloc, &mut prover_transcript);
		assert_eq!(evals.len(), depths.len());

		// The verifier's control flow depends only on the maximum depth.
		let mut verifier_transcript = prover_transcript.into_verifier();
		let verifier_output = prodcheck::verify(n_layers, claim, &mut verifier_transcript).unwrap();

		assert_eq!(verifier_output.point, eval_point);
		assert_eq!(verifier_output.eval, combine_claims::<P>(&evals, &eval_point[..k]));

		// Each tree's reduced claim is on its *padded* witness; unpadding it yields a claim on the
		// witness itself, at a suffix of the shared node point.
		for (i, (&depth, witness)) in iter::zip(depths, &witnesses).enumerate() {
			let leaf = unpad_leaf_claim(evals[i], &eval_point[k..], n_layers - depth);
			assert_eq!(leaf.point.len(), depth);
			assert_eq!(leaf.eval, evaluate(witness, &leaf.point), "tree {i}");
		}
	}

	#[test]
	fn test_unequal_depths_mixed() {
		test_unequal_depths_helper::<Packed128b>(&[2, 4, 5]);
	}

	#[test]
	fn test_unequal_depths_single_prover() {
		test_unequal_depths_helper::<Packed128b>(&[3]);
	}

	#[test]
	fn test_unequal_depths_power_of_two_provers() {
		// The shallowest tree is padded by more than one layer, the deepest not at all.
		test_unequal_depths_helper::<Packed128b>(&[1, 2, 5, 5]);
	}

	#[test]
	fn test_unequal_depths_all_minimal() {
		// Depth 1 throughout: every tree retains its final layer immediately.
		test_unequal_depths_helper::<Packed128b>(&[1, 1, 1]);
	}

	#[test]
	fn test_unequal_depths_maximal_padding() {
		// A single-layer tree beside a deep one: all but its last reduction is padding.
		test_unequal_depths_helper::<Packed128b>(&[1, 6]);
	}

	/// At equal depths every tree is padded by nothing, so the unequal-depth driver must emit
	/// byte-for-byte the transcript that [`batch_prove`] does.
	#[test]
	fn test_unequal_depths_matches_batch_prove_at_equal_depths() {
		type P = Packed128b;
		type F = <P as FieldOps>::Scalar;

		let depths = [4; 3];
		let k = log2_ceil_usize(depths.len());
		let alloc = GlobalAllocator;

		let mut rng = StdRng::seed_from_u64(23);
		let selector_point = random_scalars::<F>(&mut rng, k);
		// Both drivers see the same trees, so both rebuild them from the same seed.
		let prover_seed = 24;

		let unequal_proof = {
			let mut rng = StdRng::seed_from_u64(prover_seed);
			let (_, provers, claimed_products) =
				unequal_depth_provers::<P>(&mut rng, &alloc, &depths);

			let mut transcript = ProverTranscript::new(StdChallenger::default());
			let BatchProveUnequalDepthsOutput {
				eval_point,
				provers,
			} = batch_prove_unequal_depths(
				provers,
				claimed_products,
				selector_point.clone(),
				&mut transcript,
			);
			let (_claims, provers): (Vec<_>, Vec<_>) = provers.into_iter().unzip();
			prove_layer_rounds::<_, _, P>(provers, &eval_point, k, &alloc, &mut transcript);
			transcript.finalize()
		};

		let equal_proof = {
			let mut rng = StdRng::seed_from_u64(prover_seed);
			let (_, provers, claimed_products) =
				unequal_depth_provers::<P>(&mut rng, &alloc, &depths);

			let mut transcript = ProverTranscript::new(StdChallenger::default());
			batch_prove(provers, claimed_products, selector_point, Vec::new(), &mut transcript);
			transcript.finalize()
		};

		assert_eq!(unequal_proof, equal_proof);
	}
}