binius_core/transparent/
powers.rs

1// Copyright 2024-2025 Irreducible Inc.
2
3use std::iter::successors;
4
5use binius_field::{BinaryField128b, PackedField, TowerField};
6use binius_macros::{erased_serialize_bytes, DeserializeBytes, SerializeBytes};
7use binius_math::MultilinearExtension;
8use binius_maybe_rayon::prelude::*;
9use binius_utils::{bail, DeserializeBytes};
10use bytemuck::zeroed_vec;
11use itertools::{izip, Itertools};
12
13use crate::polynomial::{Error, MultivariatePoly};
14
15/// A transparent multilinear polynomial whose evaluation at index $i$ is $g^i$ for
16/// some field element $g$.
17#[derive(Debug, SerializeBytes, DeserializeBytes)]
18pub struct Powers<F: TowerField> {
19	n_vars: usize,
20	base: F,
21}
22
23inventory::submit! {
24	<dyn MultivariatePoly<BinaryField128b>>::register_deserializer(
25		"Powers",
26		|buf, mode| Ok(Box::new(Powers::<BinaryField128b>::deserialize(&mut *buf, mode)?))
27	)
28}
29
30impl<F: TowerField> Powers<F> {
31	pub const fn new(n_vars: usize, base: F) -> Self {
32		Self { n_vars, base }
33	}
34
35	pub fn multilinear_extension<P: PackedField<Scalar = F>>(
36		&self,
37	) -> Result<MultilinearExtension<P>, Error> {
38		let mut values = zeroed_vec(1 << self.n_vars.saturating_sub(P::LOG_WIDTH));
39
40		const CHUNK_LOG_SIZE: usize = 10;
41		values
42			.par_chunks_mut(1 << CHUNK_LOG_SIZE)
43			.enumerate()
44			.for_each(|(chunk_index, chunk)| {
45				let start_power = (chunk_index as u64) << CHUNK_LOG_SIZE;
46				let powers = successors(Some(self.base.pow_vartime([start_power])), |&power| {
47					Some(power * self.base)
48				})
49				.take(1 << self.n_vars);
50
51				for (dest, values) in izip!(chunk, &powers.chunks(P::WIDTH)) {
52					*dest = P::from_scalars(values);
53				}
54			});
55
56		Ok(MultilinearExtension::new(self.n_vars, values)?)
57	}
58}
59
60#[erased_serialize_bytes]
61impl<F: TowerField, P: PackedField<Scalar = F>> MultivariatePoly<P> for Powers<F> {
62	fn n_vars(&self) -> usize {
63		self.n_vars
64	}
65
66	fn degree(&self) -> usize {
67		self.n_vars
68	}
69
70	fn evaluate(&self, query: &[P]) -> Result<P, Error> {
71		let n_vars = self.n_vars;
72		if query.len() != self.n_vars {
73			bail!(Error::IncorrectQuerySize { expected: n_vars });
74		}
75
76		let mut result = P::one();
77		let mut base_power = self.base;
78		for &q_i in query {
79			result *= P::one() - q_i + q_i * base_power;
80			base_power = base_power.square();
81		}
82
83		Ok(result)
84	}
85
86	fn binary_tower_level(&self) -> usize {
87		F::TOWER_LEVEL
88	}
89}
90
91#[cfg(test)]
92mod tests {
93	use binius_field::{BinaryField32b, Field, PackedBinaryField4x32b, PackedField, TowerField};
94	use rand::{prelude::StdRng, SeedableRng};
95
96	use super::Powers;
97	use crate::polynomial::MultivariatePoly;
98
99	fn test_consistency_helper<P: PackedField<Scalar: TowerField>>(n_vars: usize, base: P::Scalar) {
100		let powers = Powers::new(n_vars, base);
101		let mle = powers.multilinear_extension::<P>().unwrap();
102		let mut power = P::Scalar::ONE;
103		for i in 0..1 << n_vars {
104			let query = (0..n_vars)
105				.map(|j| {
106					if (i >> j) & 1 == 0 {
107						P::zero()
108					} else {
109						P::one()
110					}
111				})
112				.collect::<Vec<_>>();
113			assert_eq!(powers.evaluate(&query).unwrap(), P::broadcast(power));
114			assert_eq!(mle.evaluate_on_hypercube(i).unwrap(), power);
115			power *= base;
116		}
117	}
118
119	#[test]
120	fn test_consistency() {
121		type F = BinaryField32b;
122		type P = PackedBinaryField4x32b;
123		let mut rng = StdRng::seed_from_u64(0);
124		for n_vars in 0..12 {
125			test_consistency_helper::<P>(n_vars, <F as Field>::random(&mut rng));
126		}
127	}
128}