binius_core/transparent/
powers.rs1use 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#[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}