binius_field/

util.rs

// Copyright 2024-2025 Irreducible Inc.

use std::iter;

use crate::{Field, PackedField, UnderlierType, field::FieldOps};

/// An arithmetic function over field elements, generic in the field it evaluates in.
///
/// A closure `FnOnce(&[F]) -> F` is monomorphic: it runs in one fixed field.
/// Carrying the genericity on the method instead lets one value run in many fields:
/// - natively, in the verifier's own base field `F`;
/// - over any larger field `E` whose scalar is `F`.
pub trait FieldFn<F: Field> {
	/// Evaluates the function on `inputs` in the field `E`, returning one element.
	///
	/// The scalar of `E` is the base field `F`.
	/// The `From<F>` bound lets the function embed base-field constants into `E`.
	fn call<E: FieldOps<Scalar = F> + From<F>>(&self, inputs: &[E]) -> E;

	/// Evaluates the function on `inputs` natively in the base field `F`.
	///
	/// The default is `self.call::<F>(inputs)`; implementors may override with a base-field
	/// specialized fast path (e.g. deferred `WideMul` reduction) that the generic
	/// [`call`](Self::call) — which cannot assume `E: WideMul` — can't express. Callers evaluating
	/// in `F` should prefer this.
	fn call_native(&self, inputs: &[F]) -> F {
		self.call::<F>(inputs)
	}
}

/// Expands an array of field elements into all possible subset sums.
///
/// For an input array `[a, b, c]`, this computes all possible sums of subsets:
/// `[0, a, b, a+b, c, a+c, b+c, a+b+c]`
///
/// This is used to create lookup tables for the Method of Four Russians optimization,
/// where we precompute all possible combinations of a small set of values to avoid
/// doing the additions at runtime.
///
/// ## Type Parameters
///
/// * `F` - The field element type
/// * `N` - Size of the input array
/// * `N_EXP2` - Size of the output array, must be 2^N
///
/// ## Arguments
///
/// * `elems` - Input array of N field elements
///
/// ## Returns
///
/// An array of size N_EXP2 containing all possible subset sums of the input elements
///
/// ## Preconditions
///
/// * N_EXP2 must equal 2^N
///
/// ## Example
///
/// ```ignore
/// let input = [F::ONE, F::from(2)];
/// let sums = expand_subset_sums_array(input);
/// // sums = [F::ZERO, F::ONE, F::from(2), F::from(3)]
/// ```
pub fn expand_subset_sums_array<P: PackedField, const N: usize, const N_EXP2: usize>(
	elems: [P; N],
) -> [P; N_EXP2] {
	const {
		assert!(N_EXP2 == 1 << N, "N_EXP2 must equal 2^N");
	}

	let mut expanded = [P::zero(); N_EXP2];
	for (i, elem_i) in elems.into_iter().enumerate() {
		let span = &mut expanded[..1 << (i + 1)];
		let (lo_half, hi_half) = span.split_at_mut(1 << i);
		for (lo_half_i, hi_half_i) in iter::zip(lo_half, hi_half) {
			*hi_half_i = *lo_half_i + elem_i;
		}
	}
	expanded
}

/// Expands `elems` into all `2^elems.len()` subset sums, indexed by subset bitmask.
///
/// The dynamically sized counterpart of [`expand_subset_sums_array`], for callers whose element
/// count is only known at run time. Entry `mask` holds the sum of `elems[i]` over every bit `i` set
/// in `mask`, so entry `0` is zero and entry `2^i` is `elems[i]`.
///
/// Each entry costs one addition, where summing a subset directly would cost one per set bit.
///
/// ## Preconditions
///
/// * `elems.len()` must be less than `usize::BITS`
pub fn expand_subset_sums<P: PackedField>(elems: &[P]) -> Vec<P> {
	assert!(elems.len() < usize::BITS as usize); // precondition

	let mut expanded = vec![P::zero(); 1 << elems.len()];
	for (i, &elem_i) in elems.iter().enumerate() {
		let (lo_half, hi_half) = expanded[..1 << (i + 1)].split_at_mut(1 << i);
		for (lo_half_i, hi_half_i) in iter::zip(lo_half, hi_half) {
			*hi_half_i = *lo_half_i + elem_i;
		}
	}
	expanded
}

/// Expands `elems` into all `2^elems.len()` subset products, indexed by subset bitmask.
///
/// The multiplicative counterpart of [`expand_subset_sums`].
/// Entry `mask` holds the product of `elems[i]` over every bit `i` set in `mask`.
/// So entry `0` is one and entry `2^i` is `elems[i]`.
///
/// This is the tensor expansion `(1, elems[0]) x ... x (1, elems[k-1])`.
/// A caller holding `k` factors of a product basis recovers all `2^k` basis elements from them.
///
/// Each entry costs one multiplication.
/// Multiplying a subset directly would cost one per set bit.
///
/// ## Preconditions
///
/// * `elems.len()` must be less than `usize::BITS`
pub fn expand_subset_products<P: PackedField>(elems: &[P]) -> Vec<P> {
	assert!(elems.len() < usize::BITS as usize); // precondition

	let mut expanded = vec![P::one(); 1 << elems.len()];
	for (i, &elem_i) in elems.iter().enumerate() {
		let (lo_half, hi_half) = expanded[..1 << (i + 1)].split_at_mut(1 << i);
		for (lo_half_i, hi_half_i) in iter::zip(lo_half, hi_half) {
			*hi_half_i = *lo_half_i * elem_i;
		}
	}
	expanded
}

/// Expands `elems` into all `2^N` subset XOR combinations, indexed by subset bitmask.
///
/// Entry `mask` holds the XOR of `elems[i]` over every bit `i` set in `mask`. This is the
/// bitwise-XOR analogue of [`expand_subset_sums_array`] over raw underliers, used to build Method
/// of Four Russians lookup tables.
///
/// ## Preconditions
///
/// * `N_EXP2` must equal `2^N`
pub fn expand_subset_xors<U: UnderlierType, const N: usize, const N_EXP2: usize>(
	elems: [U; N],
) -> [U; N_EXP2] {
	const {
		assert!(N_EXP2 == 1 << N, "N_EXP2 must equal 2^N");
	}

	let mut expanded = [U::ZERO; N_EXP2];
	for (i, elem_i) in elems.into_iter().enumerate() {
		let span = &mut expanded[..1 << (i + 1)];
		let (lo_half, hi_half) = span.split_at_mut(1 << i);
		for (lo_half_i, hi_half_i) in iter::zip(lo_half, hi_half) {
			*hi_half_i = *lo_half_i ^ elem_i;
		}
	}
	expanded
}

#[cfg(test)]
mod tests {
	use std::array;

	use proptest::prelude::*;
	use rand::{SeedableRng, rngs::StdRng};

	use super::*;
	use crate::{Ghash128b, Random};

	type F = Ghash128b;

	/// Expands `N` random elements and asserts that entry `index` of the resulting `2^N`-sized
	/// lookup table equals the subset sum selected by the set bits of `index`.
	fn check_subset_sums<const N: usize, const N_EXP2: usize>(seed: u64, index: usize) {
		let mut rng = StdRng::seed_from_u64(seed);
		let elems: [F; N] = array::from_fn(|_| F::random(&mut rng));

		let result = expand_subset_sums_array::<_, N, N_EXP2>(elems);
		assert_eq!(result.len(), N_EXP2);

		// Compute expected sum based on the binary representation of index.
		let index = index % N_EXP2;
		let mut expected = F::ZERO;
		for (bit_pos, &elem) in elems.iter().enumerate() {
			if (index >> bit_pos) & 1 == 1 {
				expected += elem;
			}
		}

		assert_eq!(
			result[index], expected,
			"index {index} should hold the subset sum for its binary representation"
		);
	}

	proptest! {
		#[test]
		fn test_expand_subset_sums_array_correctness(
			n in 0usize..=8,  // Input length (small to avoid exponential blowup)
			index in 0usize..256,  // Index to check
		) {
			// Dispatch to the const-generic helper: `expand_subset_sums_array` needs the output
			// length `2^n` at compile time.
			match n {
				0 => check_subset_sums::<0, 1>(n as u64, index),
				1 => check_subset_sums::<1, 2>(n as u64, index),
				2 => check_subset_sums::<2, 4>(n as u64, index),
				3 => check_subset_sums::<3, 8>(n as u64, index),
				4 => check_subset_sums::<4, 16>(n as u64, index),
				5 => check_subset_sums::<5, 32>(n as u64, index),
				6 => check_subset_sums::<6, 64>(n as u64, index),
				7 => check_subset_sums::<7, 128>(n as u64, index),
				8 => check_subset_sums::<8, 256>(n as u64, index),
				_ => unreachable!("n is constrained to 0..=8"),
			}
		}
	}
	proptest! {
		#[test]
		fn expand_subset_products_selects_the_product_over_set_bits(seed: u64, n in 0usize..=8) {
			let mut rng = StdRng::seed_from_u64(seed);
			let elems = (0..n)
				.map(|_| F::random(&mut rng))
				.collect::<Vec<_>>();

			let expanded = expand_subset_products(&elems);
			prop_assert_eq!(expanded.len(), 1 << n);

			// Entry `mask` multiplies exactly the elements whose bit is set in `mask`.
			for (mask, &entry) in expanded.iter().enumerate() {
				let expected = (0..n)
					.filter(|i| (mask >> i) & 1 == 1)
					.fold(F::ONE, |acc, i| acc * elems[i]);
				prop_assert_eq!(entry, expected);
			}
		}
	}
}