Skip to main content

binius_math/multilinear/
hypercube.rs

1// Copyright 2025 Irreducible Inc.
2// Copyright 2026 The Binius Developers
3
4//! Multilinear tensor expansions, generic over the hypercube the coefficients are indexed by.
5//!
6//! An `n`-variate multilinear is stored as `2^n` coefficients.
7//! The basis those coefficients are taken against factors as a tensor product over the variables.
8//! Every variable contributes the same two-element basis `(b_0, b_1)` of linear polynomials.
9//! That single choice fixes the cube, and with it what each coefficient means:
10//!
11//! ```text
12//! basis (1 - X, X)    vertices {0, 1}      coefficients are evaluations
13//! basis (1, X)        vertices {0, inf}    coefficients are monomial coefficients
14//! ```
15//!
16//! The object built over a cube again and again is the equality indicator.
17//! Written `eq(X, Y)`, it extends the predicate `X == Y` multilinearly over the cube.
18//! Fixing one operand to a point leaves `2^n` coefficients, called the expansion of that point.
19//!
20//! Every routine here is generic over the cube.
21//! The Boolean-cube specializations live beside them, in the sibling equality indicator module.
22
23use std::{iter, slice};
24
25use binius_compute::{Allocator, BufferData, VecLike};
26use binius_field::{Field, PackedField, field::FieldOps};
27use binius_utils::rayon::{
28	prelude::*,
29	task_size::{IndexedParallelIteratorExt, WorkPerItem, min_len_for_bytes},
30};
31
32use crate::{FieldBuffer, FieldVec};
33
34/// A hypercube of coefficients for multilinear polynomials.
35///
36/// A cube is fixed by the two-element basis `(b_0, b_1)` that each of its variables contributes.
37/// That basis is a pair of linear polynomials, so a cube is a choice between two of them.
38/// Everything else is derived from that choice, and shared by every implementor.
39pub trait Hypercube {
40	/// Evaluates the basis of one variable at a coordinate.
41	///
42	/// Returns `(b_0(r), b_1(r))` for the coordinate `r`.
43	fn basis<F: FieldOps>(coord: &F) -> [F; 2];
44
45	/// Scales the basis of one variable by a value.
46	///
47	/// Returns `(v * b_0(r), v * b_1(r))` for the value `v` and the coordinate `r`.
48	/// This is the inner loop of every expansion.
49	/// So an implementor beats the two multiplications that scaling the basis directly costs.
50	fn expand_var<F: FieldOps>(value: &F, coord: &F) -> [F; 2];
51
52	/// Strips one variable's basis factor from the two halves of an expansion.
53	///
54	/// The halves hold `v * b_0(r)` and `v * b_1(r)` for the stripped variable's coordinate `r`.
55	/// The low half is overwritten with `v`.
56	///
57	/// Recovering `v` is one fixed linear combination of the two halves:
58	///
59	/// ```text
60	/// sum_i w_i * v * b_i(r) = v    where    sum_i w_i * b_i(X) = 1
61	/// ```
62	///
63	/// Those weights are unique and free of `r`, so the same combination works at any coordinate.
64	fn contract_var<F: FieldOps>(lo: &mut F, hi: &F);
65
66	/// Evaluates the equality indicator of one variable.
67	///
68	/// ```text
69	/// eq(X, Y) = sum_i b_i(X) * b_i(Y)
70	/// ```
71	///
72	/// An implementor overrides this with a cheaper closed form.
73	fn eq_one_var<F: FieldOps>(x: F, y: F) -> F {
74		// The generic definition: pair the two bases at the two coordinates and sum.
75		let [x_0, x_1] = Self::basis(&x);
76		let [y_0, y_1] = Self::basis(&y);
77		x_0 * y_0 + x_1 * y_1
78	}
79}
80
81/// The Boolean cube `{0, 1}^n`, whose per-variable basis is `(1 - X, X)`.
82///
83/// That basis is the pair of Lagrange polynomials on the two vertices `0` and `1`.
84/// So the coefficient indexed by a vertex is the multilinear's evaluation at that vertex.
85#[derive(Debug)]
86pub struct OneCube;
87
88impl Hypercube for OneCube {
89	#[inline(always)]
90	fn basis<F: FieldOps>(coord: &F) -> [F; 2] {
91		[F::one() - coord, coord.clone()]
92	}
93
94	#[inline(always)]
95	fn expand_var<F: FieldOps>(value: &F, coord: &F) -> [F; 2] {
96		// Both halves share the product `value * coord`, so one multiplication covers both.
97		let prod = value.clone() * coord;
98		[value.clone() - &prod, prod]
99	}
100
101	#[inline(always)]
102	fn contract_var<F: FieldOps>(lo: &mut F, hi: &F) {
103		// The two basis polynomials sum to one, so both recovery weights are one.
104		*lo += hi;
105	}
106
107	#[inline(always)]
108	fn eq_one_var<F: FieldOps>(x: F, y: F) -> F {
109		// Over characteristic two the `2 * X * Y` term vanishes, so
110		//
111		//     X * Y + (1 - X) * (1 - Y)  =  X + Y + 1
112		//
113		// The condition is a compile-time constant, so only one arm is ever generated.
114		if F::Scalar::CHARACTERISTIC == 2 {
115			x + y + F::one()
116		} else {
117			let one = F::one();
118			x.clone() * y.clone() + (one.clone() - x) * (one - y)
119		}
120	}
121}
122
123/// The infinity cube `{0, inf}^n`, whose per-variable basis is `(1, X)`.
124///
125/// The vertex `inf` selects a multilinear's leading coefficient in that variable.
126/// So the coefficient indexed by a vertex `v` belongs to the monomial
127///
128/// ```text
129/// prod_{i : v_i = inf} X_i
130/// ```
131#[derive(Debug)]
132pub struct InfCube;
133
134impl Hypercube for InfCube {
135	#[inline(always)]
136	fn basis<F: FieldOps>(coord: &F) -> [F; 2] {
137		[F::one(), coord.clone()]
138	}
139
140	#[inline(always)]
141	fn expand_var<F: FieldOps>(value: &F, coord: &F) -> [F; 2] {
142		// The constant basis polynomial is one, so the low half is the value untouched.
143		[value.clone(), value.clone() * coord]
144	}
145
146	#[inline(always)]
147	fn contract_var<F: FieldOps>(_lo: &mut F, _hi: &F) {
148		// The low half already holds the value, so the weights are one and zero.
149		// Contracting a variable is therefore free for this cube.
150	}
151
152	#[inline(always)]
153	fn eq_one_var<F: FieldOps>(x: F, y: F) -> F {
154		// The constant basis polynomial contributes one, and the linear one the product.
155		F::one() + x * y
156	}
157}
158
159/// Tensor of values with the equality indicator evaluated at extra coordinates.
160///
161/// Take `n` values and the `k` coordinates `r = (r_0, ..., r_{k-1})`.
162/// The result is the tensor product of those values with the basis at every coordinate:
163///
164/// ```text
165/// v (x) b(r_0) (x) ... (x) b(r_{k-1})
166/// ```
167///
168/// It holds `2^(n + k)` coefficients, one variable added per coordinate.
169///
170/// Read as polynomials, the input holds an `n`-variate multilinear `f`.
171/// The output then holds the `(n + k)`-variate multilinear
172///
173/// ```text
174/// g(X_0, ..., X_{n+k-1}) = f(X_0, ..., X_{n-1}) * eq(X_n, ..., X_{n+k-1}, r)
175/// ```
176///
177/// Appending the coordinate `r` doubles the length, turning every value `v` into a pair:
178///
179/// ```text
180/// before   [ v_0            v_1            ]
181/// after    [ v_0 * b_0(r)   v_1 * b_0(r)   |   v_0 * b_1(r)   v_1 * b_1(r) ]
182/// ```
183///
184/// The two halves sit one after the other, so the appended variable is the highest indexed one.
185pub fn tensor_prod_eq_ind<Cube: Hypercube, P: PackedField>(
186	values: FieldBuffer<P, Vec<P>>,
187	extra_query_coordinates: &[P::Scalar],
188) -> FieldBuffer<P, Vec<P>> {
189	let start_log_len = values.log_len();
190	let final_log_len = start_log_len + extra_query_coordinates.len();
191	let mut data = values.into_inner();
192
193	// Reserve the whole final capacity once, so no round reallocates.
194	// Each round then writes its new coefficients straight into the reserved spare capacity,
195	// instead of zero-initializing a region the expansion immediately overwrites.
196	let final_packed_len = packed_words::<P>(final_log_len);
197	data.reserve_exact(final_packed_len.saturating_sub(data.len()));
198
199	tensor_prod_eq_ind_reserved::<Cube, P, _>(
200		FieldBuffer::new(start_log_len, data),
201		extra_query_coordinates,
202	)
203}
204
205/// The number of packed words an expansion of that many variables occupies.
206///
207/// Below one packed word the count is one, since a single word backs any shorter length.
208const fn packed_words<P: PackedField>(log_len: usize) -> usize {
209	1usize << log_len.saturating_sub(P::LOG_WIDTH)
210}
211
212/// Appends one variable per coordinate to a store that already has room for the result.
213///
214/// The public entry point reserves the final capacity before calling here.
215/// The allocator-backed callers pass a buffer drawn at the final size, which cannot grow.
216/// So for them the reservation is a precondition rather than a step of the expansion.
217///
218/// # Preconditions
219///
220/// * The store's capacity must cover the packed length of the final expansion.
221fn tensor_prod_eq_ind_reserved<Cube: Hypercube, P: PackedField, Data: VecLike<P>>(
222	values: FieldBuffer<P, Data>,
223	extra_query_coordinates: &[P::Scalar],
224) -> FieldBuffer<P, Data> {
225	let start_log_len = values.log_len();
226	let final_log_len = start_log_len + extra_query_coordinates.len();
227	let mut data = values.into_inner();
228
229	// precondition
230	debug_assert!(data.capacity() >= packed_words::<P>(final_log_len));
231
232	// The coordinates split cleanly in two at the packing width:
233	//
234	//     narrower than one word   the whole expansion lives in data[0]
235	//     one word or wider        every round doubles the packed length
236	let sub_width_count = extra_query_coordinates
237		.len()
238		.min(P::LOG_WIDTH.saturating_sub(start_log_len));
239	let (sub_width_coords, packed_coords) = extra_query_coordinates.split_at(sub_width_count);
240
241	// Sub-packing-width rounds: both halves of the result share the single word data[0].
242	// Split that word into its two halves, expand them, and interleave them back together.
243	// The backing store stays one element long throughout.
244	for (i, &r_i) in sub_width_coords.iter().enumerate() {
245		let log_len = start_log_len + i;
246		let packed_r_i = P::broadcast(r_i);
247		let (lo, _) = data[0].interleave(P::zero(), log_len);
248		let [lo, hi] = Cube::expand_var(&lo, &packed_r_i);
249		data[0] = lo.interleave(hi, log_len).0;
250	}
251
252	// Packed rounds: the initialized words are exactly the low half of the result.
253	//
254	//     low half    the initialized prefix, expanded in place
255	//     high half   reserved spare capacity, written once
256	for &r_i in packed_coords {
257		let packed_r_i = P::broadcast(r_i);
258		let old_packed = data.len();
259
260		// The safe two-slice split of a Vec into its initialized prefix and its spare capacity
261		// is still unstable (rust-lang/rust#81944).
262		// So the spare half comes from the safe accessor and the initialized half from a raw part.
263		let low_ptr = data.as_mut_ptr();
264		let high = &mut data.spare_capacity_mut()[..old_packed];
265		// SAFETY: `[0, old_packed)` is the initialized low half, disjoint from the spare `high`
266		// half `[old_packed, 2 * old_packed)`; the two slices never overlap.
267		let low = unsafe { slice::from_raw_parts_mut(low_ptr, old_packed) };
268		// Each round doubles the expansion, starting from a single word.
269		// So the first rounds are far too small to be worth splitting across threads.
270		(low, high)
271			.into_par_iter()
272			.with_min_task(WorkPerItem::FieldMuls)
273			.for_each(|(low_i, high_i)| {
274				let [new_low, new_high] = Cube::expand_var(low_i, &packed_r_i);
275				*low_i = new_low;
276				high_i.write(new_high);
277			});
278		// SAFETY: the loop above initialized every one of the `old_packed` spare words.
279		unsafe { data.set_len(2 * old_packed) };
280	}
281
282	FieldBuffer::new(final_log_len, data)
283}
284
285/// Computes the partial evaluation of the equality indicator polynomial.
286///
287/// For the point `r = (r_0, ..., r_{n-1})` the result holds the `2^n` coefficients
288///
289/// ```text
290/// b(r_0) (x) ... (x) b(r_{n-1})
291/// ```
292///
293/// which are the coefficients of the equality indicator `eq(X_0, ..., X_{n-1}, r)` over the cube.
294pub fn eq_ind_partial_eval<Cube: Hypercube, P: PackedField>(point: &[P::Scalar]) -> FieldBuffer<P> {
295	// The unscaled indicator is the scaled indicator with a scale of one.
296	scaled_eq_ind_partial_eval::<Cube, P>(point, P::Scalar::ONE)
297}
298
299/// Builds the equality indicator expansion of a point into a buffer drawn from an allocator.
300///
301/// Backed by a pool, the result is a recyclable buffer rather than a fresh allocation.
302pub fn eq_ind_partial_eval_in<Cube: Hypercube, A: Allocator, P: PackedField>(
303	alloc: &A,
304	point: &[P::Scalar],
305) -> FieldVec<P, A> {
306	// The allocator hands out the final packed length, which the expansion never outgrows.
307	let packed_len = packed_words::<P>(point.len());
308	scaled_eq_ind_partial_eval_into::<Cube, P, _>(
309		point,
310		P::Scalar::ONE,
311		alloc.alloc::<P>(packed_len),
312	)
313}
314
315/// Computes the partial evaluation of the equality indicator polynomial, scaled by a constant.
316///
317/// Every coefficient of the equality indicator is multiplied by the scale.
318/// A scale of one is the identity, since the expansion is linear in it.
319///
320/// # Arguments
321///
322/// * `point` - The evaluation point whose length is the number of variables.
323/// * `scale` - The constant every returned value is multiplied by.
324pub fn scaled_eq_ind_partial_eval<Cube: Hypercube, P: PackedField>(
325	point: &[P::Scalar],
326	scale: P::Scalar,
327) -> FieldBuffer<P> {
328	// Reserving the final packed length keeps the per-variable growth reallocation free.
329	let packed_len = packed_words::<P>(point.len());
330	scaled_eq_ind_partial_eval_into::<Cube, P, _>(point, scale, Vec::with_capacity(packed_len))
331}
332
333/// Builds the scaled equality indicator expansion of a point in a caller-supplied store.
334///
335/// This is the allocation-hoisting form.
336/// The caller owns the store, so it can be drawn from a pool.
337/// It can equally be reserved on a different thread than the one that fills it.
338///
339/// Each of the `2^n` coefficients is written once, so the expansion costs one multiplication per
340/// coefficient.
341///
342/// # Preconditions
343///
344/// * The store's capacity must cover the packed length of the expansion.
345pub fn scaled_eq_ind_partial_eval_into<Cube: Hypercube, P: PackedField, Data: VecLike<P>>(
346	point: &[P::Scalar],
347	scale: P::Scalar,
348	mut buffer: Data,
349) -> FieldBuffer<P, Data> {
350	assert!(
351		buffer.capacity() >= packed_words::<P>(point.len()),
352		"precondition: buffer capacity must cover the packed expansion length"
353	);
354
355	// Seed a one-coefficient expansion with the scale.
356	// Appending the coordinates multiplies it through, so every coefficient ends up scaled.
357	buffer.clear();
358	buffer.push(P::from_scalars(iter::once(scale)));
359	let seed = FieldBuffer::new(0, buffer);
360
361	// A coefficient is a product over the coordinates, so cutting the point cuts the product:
362	//
363	//     index = high * 2^low_len + low
364	//     coeff = high_expansion[high] * low_expansion[low]
365	//
366	// The cut sits at the midpoint, so neither expansion is much larger than the square root of the
367	// result and the two cost the same to build.
368	// A block spans whole packed words, so the cut never falls below the packing width, and it
369	// never passes the end of the point, which leaves the low expansion holding the whole result.
370	let low_len = (point.len() / 2).max(P::LOG_WIDTH).min(point.len());
371	let (low_coords, high_coords) = point.split_at(low_len);
372
373	// The low expansion is built straight into the store's first block, which is where the result
374	// wants it anyway. Nothing else is expanded over the result's own length.
375	let low = tensor_prod_eq_ind_reserved::<Cube, P, Data>(seed, low_coords);
376
377	// A cut that takes the whole point leaves one block, which is the result already.
378	if high_coords.is_empty() {
379		return low;
380	}
381
382	// A block is `2^low_len` scalars, a whole number of packed words exactly because a cut that
383	// leaves a tail sits at or above the packing width.
384	// That is what makes a block's position in the store equal its index into the high expansion.
385	let block = low.as_ref().len();
386	let mut data = low.into_inner();
387
388	// One scalar per block of the result, so the high expansion stays far smaller than it.
389	let high = eq_ind_partial_eval_scalars::<Cube, P::Scalar>(high_coords);
390	let total = block * high.len();
391	debug_assert_eq!(total, packed_words::<P>(point.len()));
392
393	// The safe two-slice split of a Vec into its initialized prefix and its spare capacity is
394	// still unstable (rust-lang/rust#81944).
395	// So the spare blocks come from the safe accessor and the first block from a raw part.
396	let first_ptr = data.as_mut_ptr();
397	let spare = &mut data.spare_capacity_mut()[..total - block];
398	// SAFETY: `[0, block)` is the initialized first block, disjoint from the spare blocks past
399	// it; the two slices never overlap.
400	let first = unsafe { slice::from_raw_parts(first_ptr, block) };
401
402	// Every block past the first is the first scaled by one coefficient of the high expansion.
403	// One item here is a whole block, so the byte floor is divided down by what a block holds.
404	let min_len = (min_len_for_bytes::<P>() / block).max(1);
405	spare
406		.par_chunks_mut(block)
407		.zip(high[1..].par_iter())
408		.with_min_len(min_len)
409		.for_each(|(dst, &coeff)| {
410			let coeff = P::broadcast(coeff);
411			for (dst_i, &src_i) in iter::zip(dst, first) {
412				dst_i.write(src_i * coeff);
413			}
414		});
415	// SAFETY: the loop above initialized every spare word up to the total.
416	unsafe { data.set_len(total) };
417
418	// The first block still holds the low expansion unscaled, since every other block read it.
419	// Its own coefficient therefore lands last.
420	let coeff = P::broadcast(high[0]);
421	for word in &mut data[..block] {
422		*word *= coeff;
423	}
424
425	FieldBuffer::new(point.len(), data)
426}
427
428/// Truncates a built equality indicator expansion to its low indexed variables.
429///
430/// Each step contracts the two halves of the buffer, stripping the highest variable.
431/// That removes the highest variable's basis factor, whatever its coordinate was.
432/// Truncating to `n'` variables therefore leaves the indicator over `r_0, ..., r_{n'-1}`.
433///
434/// The expansion occupies a prefix of the buffer.
435/// Scalars after the truncated length are dropped.
436///
437/// # Preconditions
438///
439/// * the truncated length must be at most the buffer's current length
440pub fn eq_ind_truncate_low_inplace<Cube: Hypercube, P: PackedField, Data: BufferData<P>>(
441	values: &mut FieldBuffer<P, Data>,
442	truncated_log_len: usize,
443) {
444	assert!(
445		truncated_log_len <= values.log_len(),
446		"precondition: truncated_log_len must be at most values.log_len()"
447	);
448
449	// One round per variable stripped, highest first, so the survivors stay in a prefix.
450	for log_len in (truncated_log_len..values.log_len()).rev() {
451		{
452			let mut split = values.split_half_mut();
453			let (mut lo, hi) = split.halves();
454			// Contracting a variable costs additions only.
455			// So the cost of one step is the two words it reads, not its arithmetic.
456			(lo.as_mut(), hi.as_ref())
457				.into_par_iter()
458				.with_min_task_bytes::<[P; 2]>()
459				.for_each(|(zero, one)| {
460					Cube::contract_var(zero, one);
461				});
462		}
463
464		values.truncate(log_len);
465	}
466}
467
468/// Evaluates the equality indicator multilinear at a pair of points.
469///
470/// This is the `2n`-variate multilinear
471///
472/// ```text
473/// eq(X_0, ..., X_{n-1}, Y_0, ..., Y_{n-1}) = prod_i sum_j b_j(X_i) * b_j(Y_i)
474/// ```
475pub fn eq_ind<Cube: Hypercube, F: FieldOps>(x: &[F], y: &[F]) -> F {
476	assert_eq!(x.len(), y.len(), "pre-condition: x and y must be the same length");
477	// The indicator factors over the variables, so one per-variable product suffices.
478	iter::zip(x, y)
479		.map(|(x, y)| Cube::eq_one_var(x.clone(), y.clone()))
480		.product()
481}
482
483/// Evaluates the equality indicator multilinear with one operand fixed to all zeros.
484///
485/// Only the constant basis polynomial survives at a zero coordinate:
486///
487/// ```text
488/// eq(0^n, Y_0, ..., Y_{n-1}) = prod_i b_0(Y_i)
489/// ```
490pub fn eq_ind_zero<Cube: Hypercube, F: FieldOps>(point: &[F]) -> F {
491	// The linear basis polynomial is multiplied by a zero coordinate, so it drops out.
492	point
493		.iter()
494		.map(|y| {
495			let [y_0, _] = Cube::basis(y);
496			y_0
497		})
498		.product()
499}
500
501/// Computes the partial evaluation of the equality indicator polynomial, returning scalars.
502///
503/// This is the scalar-only engine, which never touches a packed store.
504pub fn eq_ind_partial_eval_scalars<Cube: Hypercube, F: FieldOps>(point: &[F]) -> Vec<F> {
505	// The unscaled indicator is the scaled indicator with a scale of one.
506	scaled_eq_ind_partial_eval_scalars::<Cube, F>(point, F::one())
507}
508
509/// Computes the scaled partial evaluation of the equality indicator, returning scalars.
510///
511/// This is the scalar-only engine, which never touches a packed store.
512/// A scale of one is the identity, since the expansion is linear in it.
513pub fn scaled_eq_ind_partial_eval_scalars<Cube: Hypercube, F: FieldOps>(
514	point: &[F],
515	scale: F,
516) -> Vec<F> {
517	// One coefficient per cube vertex, allocated once.
518	let mut result = Vec::with_capacity(1 << point.len());
519	// Seed with the scale, which every later multiplication carries through.
520	result.push(scale);
521
522	for r_i in point {
523		// Each coordinate doubles the length.
524		// The low half takes the constant basis factor, the appended high half the linear one.
525		//
526		//     read index j  ->  overwrite result[j], push its partner past the end
527		//
528		// Walking the low half front to back is safe, since pushing only appends past it.
529		let len = result.len();
530		for j in 0..len {
531			let [lo, hi] = Cube::expand_var(&result[j], r_i);
532			result[j] = lo;
533			result.push(hi);
534		}
535	}
536	result
537}
538
539#[cfg(test)]
540mod tests {
541	use binius_utils::rayon::task_size::{min_len_for_bytes, min_len_for_work};
542	use proptest::prelude::*;
543	use rand::prelude::*;
544
545	use super::*;
546	use crate::test_utils::{B128, Packed128b, random_scalars};
547
548	type P = Packed128b;
549	type F = B128;
550
551	#[test]
552	fn expand_var_matches_scaled_basis() {
553		let mut rng = StdRng::seed_from_u64(0);
554
555		// Each implementor saves a multiplication over scaling the basis the plain way.
556		// So both must land on the same pair.
557		let [value, coord] = [(); 2].map(|_| random_scalars::<F>(&mut rng, 1)[0]);
558		assert_eq!(
559			OneCube::expand_var(&value, &coord),
560			OneCube::basis(&coord).map(|b_i| b_i * value)
561		);
562		assert_eq!(
563			InfCube::expand_var(&value, &coord),
564			InfCube::basis(&coord).map(|b_i| b_i * value)
565		);
566	}
567
568	#[test]
569	fn contract_var_inverts_expand_var() {
570		let mut rng = StdRng::seed_from_u64(0);
571
572		// Expanding a value by a coordinate and contracting it back must be the identity.
573		let [value, coord] = [(); 2].map(|_| random_scalars::<F>(&mut rng, 1)[0]);
574
575		let [mut lo, hi] = OneCube::expand_var(&value, &coord);
576		OneCube::contract_var(&mut lo, &hi);
577		assert_eq!(lo, value);
578
579		let [mut lo, hi] = InfCube::expand_var(&value, &coord);
580		InfCube::contract_var(&mut lo, &hi);
581		assert_eq!(lo, value);
582	}
583
584	#[test]
585	fn eq_one_var_matches_basis_definition() {
586		let mut rng = StdRng::seed_from_u64(0);
587
588		// Each implementor overrides this with a closed form that skips the basis evaluations.
589		// So pin both against the generic pairing of the bases.
590		let [x, y] = [(); 2].map(|_| random_scalars::<F>(&mut rng, 1)[0]);
591		let eq_from_basis = |[x_0, x_1]: [F; 2], [y_0, y_1]: [F; 2]| x_0 * y_0 + x_1 * y_1;
592		assert_eq!(
593			OneCube::eq_one_var(x, y),
594			eq_from_basis(OneCube::basis(&x), OneCube::basis(&y))
595		);
596		assert_eq!(
597			InfCube::eq_one_var(x, y),
598			eq_from_basis(InfCube::basis(&x), InfCube::basis(&y))
599		);
600	}
601
602	#[test]
603	fn inf_cube_eq_ind_zero_is_one() {
604		let mut rng = StdRng::seed_from_u64(0);
605
606		// Every monomial of positive degree vanishes at zero, leaving the constant one.
607		for n_vars in [0, 1, 5] {
608			let point = random_scalars::<F>(&mut rng, n_vars);
609			assert_eq!(eq_ind_zero::<InfCube, F>(&point), F::ONE);
610
611			// The same value as evaluating the full indicator against an all-zero operand.
612			assert_eq!(
613				eq_ind_zero::<InfCube, F>(&point),
614				eq_ind::<InfCube, F>(&vec![F::ZERO; n_vars], &point)
615			);
616		}
617	}
618
619	/// The expansion of a point, straight from the definition of the tensor of bases `(1, r_i)`.
620	fn inf_cube_reference(point: &[F]) -> Vec<F> {
621		// The coefficient at an index is the product of the coordinates its set bits select.
622		(0..1 << point.len())
623			.map(|index| {
624				point
625					.iter()
626					.enumerate()
627					.filter(|(i, _)| index >> i & 1 == 1)
628					.map(|(_, r_i)| *r_i)
629					.product()
630			})
631			.collect()
632	}
633
634	/// Evaluates the multilinear whose monomial coefficients are given, at a point.
635	fn eval_monomial_basis(coeffs: &[F], point: &[F]) -> F {
636		// The coefficient at an index belongs to the monomial its set bits select.
637		coeffs
638			.iter()
639			.enumerate()
640			.map(|(index, coeff)| {
641				*coeff
642					* point
643						.iter()
644						.enumerate()
645						.filter(|(i, _)| index >> i & 1 == 1)
646						.map(|(_, x_i)| *x_i)
647						.product::<F>()
648			})
649			.sum()
650	}
651
652	#[test]
653	fn inf_cube_expansion_matches_the_tensor_of_bases() {
654		let mut rng = StdRng::seed_from_u64(0);
655
656		// Sizes span the empty point up to a 256-coefficient cube.
657		for n_vars in [0, 1, 2, 5, 8] {
658			let point = random_scalars::<F>(&mut rng, n_vars);
659			let expansion = eq_ind_partial_eval::<InfCube, P>(&point);
660			let expansion_scalars = expansion.iter_scalars().collect::<Vec<_>>();
661			assert_eq!(expansion_scalars, inf_cube_reference(&point), "mismatch at {n_vars} vars");
662		}
663	}
664
665	#[test]
666	fn inf_cube_expansion_holds_the_monomial_coefficients_of_the_indicator() {
667		let mut rng = StdRng::seed_from_u64(0);
668
669		// The defining property of this cube: coefficients are monomial coefficients.
670		//
671		//     expansion of r  ->  the monomial coefficients of eq(X, r)
672		//
673		// So reading the expansion in the monomial basis at any x must give the indicator there.
674		for n_vars in [0, 1, 2, 5] {
675			let point = random_scalars::<F>(&mut rng, n_vars);
676			let coeffs = eq_ind_partial_eval_scalars::<InfCube, F>(&point);
677
678			let x = random_scalars::<F>(&mut rng, n_vars);
679			assert_eq!(eval_monomial_basis(&coeffs, &x), eq_ind::<InfCube, F>(&x, &point));
680		}
681	}
682
683	#[test]
684	fn inf_cube_expansion_is_the_evaluation_functional() {
685		let mut rng = StdRng::seed_from_u64(0);
686
687		// Read the other way round, the expansion of a point is the functional that evaluates
688		// any multilinear at that point, given the multilinear's monomial coefficients.
689		for n_vars in [0, 1, 2, 5] {
690			let point = random_scalars::<F>(&mut rng, n_vars);
691			let coeffs = random_scalars::<F>(&mut rng, 1 << n_vars);
692
693			let expansion = eq_ind_partial_eval_scalars::<InfCube, F>(&point);
694			let inner_product = iter::zip(&coeffs, &expansion)
695				.map(|(c, e)| *c * e)
696				.sum::<F>();
697			assert_eq!(inner_product, eval_monomial_basis(&coeffs, &point));
698		}
699	}
700
701	#[test]
702	fn growth_above_the_split_threshold_matches_the_inline_path() {
703		let mut rng = StdRng::seed_from_u64(9);
704
705		// Invariant: a round splits across threads only once it exceeds the minimum task size.
706		// Below that it runs inline, leaving the parallel path unexercised.
707		//
708		//     words in the widest round = 2^(n_vars - 1) / scalars per word
709		//     pick the smallest n_vars whose widest round reaches the minimum
710		//
711		// Every other test here is smaller, so this one covers the split.
712		let min_len = min_len_for_work(WorkPerItem::FieldMuls);
713		let n_vars = (2 * min_len * P::WIDTH).next_power_of_two().ilog2() as usize;
714		let point = random_scalars::<F>(&mut rng, n_vars);
715
716		let packed = eq_ind_partial_eval::<OneCube, P>(&point);
717		let reference = eq_ind_partial_eval_scalars::<OneCube, F>(&point);
718		assert!(packed.iter_scalars().eq(reference.iter().copied()));
719	}
720
721	#[test]
722	fn truncation_above_the_split_threshold_matches_the_inline_path() {
723		let mut rng = StdRng::seed_from_u64(0);
724
725		// Invariant: a round splits across threads only once it exceeds the minimum task size.
726		// Below that it runs inline, leaving the parallel path unexercised.
727		//
728		//     words read in the first round = 2^(n_vars - 1) / scalars per word
729		//     pick the smallest n_vars whose first round reaches the minimum
730		//
731		// Every other truncation test here is smaller, so this one covers the split.
732		let min_len = min_len_for_bytes::<[P; 2]>();
733		let n_vars = (2 * min_len * P::WIDTH).next_power_of_two().ilog2() as usize;
734		let point = random_scalars::<F>(&mut rng, n_vars);
735
736		// Strip the top variable and compare against a direct expansion of the prefix.
737		let mut truncated = eq_ind_partial_eval::<OneCube, P>(&point);
738		eq_ind_truncate_low_inplace::<OneCube, _, _>(&mut truncated, n_vars - 1);
739		assert_eq!(truncated, eq_ind_partial_eval::<OneCube, P>(&point[..n_vars - 1]));
740	}
741
742	proptest! {
743		#![proptest_config(ProptestConfig::with_cases(16))]
744
745		#[test]
746		fn the_two_engines_agree(
747			seed in any::<u64>(),
748			log_n in 0usize..=8,
749		) {
750			let mut rng = StdRng::seed_from_u64(seed);
751			let point = random_scalars::<F>(&mut rng, log_n);
752			let scale = random_scalars::<F>(&mut rng, 1)[0];
753
754			// Both bases, both the plain and the scaled form: the packed engine and the scalar
755			// engine must land on the same coefficients.
756			prop_assert_eq!(
757				eq_ind_partial_eval::<OneCube, P>(&point).iter_scalars().collect::<Vec<_>>(),
758				eq_ind_partial_eval_scalars::<OneCube, F>(&point)
759			);
760			prop_assert_eq!(
761				eq_ind_partial_eval::<InfCube, P>(&point).iter_scalars().collect::<Vec<_>>(),
762				eq_ind_partial_eval_scalars::<InfCube, F>(&point)
763			);
764			prop_assert_eq!(
765				scaled_eq_ind_partial_eval::<OneCube, P>(&point, scale)
766					.iter_scalars()
767					.collect::<Vec<_>>(),
768				scaled_eq_ind_partial_eval_scalars::<OneCube, F>(&point, scale)
769			);
770			prop_assert_eq!(
771				scaled_eq_ind_partial_eval::<InfCube, P>(&point, scale)
772					.iter_scalars()
773					.collect::<Vec<_>>(),
774				scaled_eq_ind_partial_eval_scalars::<InfCube, F>(&point, scale)
775			);
776
777			// A scalar field is a packed field of one lane, so the packed engine also runs at
778			// a packing width of one. That width is its own path through the growth loop:
779			//
780			//     4 lanes per word    two rounds live inside one word, then rounds double
781			//     1 lane per word     no round fits inside a word, so every round doubles
782			//
783			// The one-lane store is exactly the scalars, so the two must agree there too.
784			prop_assert_eq!(
785				eq_ind_partial_eval::<OneCube, F>(&point).into_inner(),
786				eq_ind_partial_eval_scalars::<OneCube, F>(&point)
787			);
788			prop_assert_eq!(
789				eq_ind_partial_eval::<InfCube, F>(&point).into_inner(),
790				eq_ind_partial_eval_scalars::<InfCube, F>(&point)
791			);
792		}
793
794		#[test]
795		fn scaling_commutes_with_the_expansion(
796			seed in any::<u64>(),
797			log_n in 0usize..=8,
798		) {
799			let mut rng = StdRng::seed_from_u64(seed);
800			let point = random_scalars::<F>(&mut rng, log_n);
801			let scale = random_scalars::<F>(&mut rng, 1)[0];
802
803			// Scaling the seed scales every coefficient, for either basis.
804			let one_scaled = scaled_eq_ind_partial_eval::<OneCube, P>(&point, scale);
805			let one_plain = eq_ind_partial_eval::<OneCube, P>(&point);
806			for (got, base) in one_scaled.iter_scalars().zip(one_plain.iter_scalars()) {
807				prop_assert_eq!(got, scale * base);
808			}
809
810			let inf_scaled = scaled_eq_ind_partial_eval::<InfCube, P>(&point, scale);
811			let inf_plain = eq_ind_partial_eval::<InfCube, P>(&point);
812			for (got, base) in inf_scaled.iter_scalars().zip(inf_plain.iter_scalars()) {
813				prop_assert_eq!(got, scale * base);
814			}
815		}
816
817		#[test]
818		fn truncation_strips_trailing_variables(
819			seed in any::<u64>(),
820			log_n in 0usize..=8,
821		) {
822			let mut rng = StdRng::seed_from_u64(seed);
823			let point = random_scalars::<F>(&mut rng, log_n);
824
825			// Truncating to any length must equal expanding that prefix of the point directly.
826			for truncated_log_len in 0..=log_n {
827				let mut one_cube = eq_ind_partial_eval::<OneCube, P>(&point);
828				eq_ind_truncate_low_inplace::<OneCube, _, _>(&mut one_cube, truncated_log_len);
829				prop_assert_eq!(
830					one_cube,
831					eq_ind_partial_eval::<OneCube, P>(&point[..truncated_log_len])
832				);
833
834				let mut inf_cube = eq_ind_partial_eval::<InfCube, P>(&point);
835				eq_ind_truncate_low_inplace::<InfCube, _, _>(&mut inf_cube, truncated_log_len);
836				prop_assert_eq!(
837					inf_cube,
838					eq_ind_partial_eval::<InfCube, P>(&point[..truncated_log_len])
839				);
840			}
841		}
842
843		#[test]
844		fn the_split_agrees_with_the_doubling_rounds(
845			seed in any::<u64>(),
846			n_vars in 0usize..=10,
847		) {
848			// Property: cutting the point and multiplying the two expansions together is the same
849			// map as appending its coordinates one at a time.
850			//
851			// The range covers a cut that takes the whole point, one held at the packing width,
852			// and one at the midpoint.
853			let mut rng = StdRng::seed_from_u64(seed);
854			let point = random_scalars::<F>(&mut rng, n_vars);
855			let scale = random_scalars::<F>(&mut rng, 1)[0];
856
857			// A seed is one coefficient carrying the scale, over a store sized for the result.
858			let seeded = || {
859				let mut buffer = Vec::with_capacity(1 << n_vars.saturating_sub(P::LOG_WIDTH));
860				buffer.push(P::from_scalars(iter::once(scale)));
861				FieldBuffer::new(0, buffer)
862			};
863
864			prop_assert_eq!(
865				scaled_eq_ind_partial_eval::<OneCube, P>(&point, scale),
866				tensor_prod_eq_ind_reserved::<OneCube, P, _>(seeded(), &point)
867			);
868			prop_assert_eq!(
869				scaled_eq_ind_partial_eval::<InfCube, P>(&point, scale),
870				tensor_prod_eq_ind_reserved::<InfCube, P, _>(seeded(), &point)
871			);
872		}
873	}
874
875	#[test]
876	fn the_split_path_agrees_with_the_scalar_engine() {
877		let mut rng = StdRng::seed_from_u64(0);
878
879		// The public entry point picks the cut itself, so this runs against the scalar engine,
880		// which shares no code with either packed path.
881		//
882		// 0 up to the packing width is where the cut takes the whole point and leaves one block.
883		// Past it the packing-width floor holds the cut until the midpoint overtakes it, and
884		// 13, 14 and 19 split many ways, odd and even.
885		for n_vars in (0..=8).chain([13, 14, 19]) {
886			let point = random_scalars::<F>(&mut rng, n_vars);
887			let scale = random_scalars::<F>(&mut rng, 1)[0];
888
889			let packed = scaled_eq_ind_partial_eval::<OneCube, P>(&point, scale);
890			let scalars = scaled_eq_ind_partial_eval_scalars::<OneCube, F>(&point, scale);
891			assert!(packed.iter_scalars().eq(scalars), "one cube at {n_vars} vars");
892
893			let packed = scaled_eq_ind_partial_eval::<InfCube, P>(&point, scale);
894			let scalars = scaled_eq_ind_partial_eval_scalars::<InfCube, F>(&point, scale);
895			assert!(packed.iter_scalars().eq(scalars), "inf cube at {n_vars} vars");
896		}
897	}
898}