1use core::ops::Index;
  2
  3use alloc::sync::Arc;
  4use winter_math::{fields::f64::BaseElement, log2, FieldElement, StarkField};
  5
  6mod eq_poly;
  7pub use eq_poly::EqPolynomial;
  8
  9#[derive(Clone, Debug)]
 10pub struct MultiLinear<E: FieldElement> {
 11    pub num_variables: usize,
 12    pub evaluations: Vec<E>,
 13}
 14
 15impl<E: FieldElement> MultiLinear<E> {
 16    pub fn new(values: Vec<E>) -> Self {
 17        Self {
 18            num_variables: log2(values.len()) as usize,
 19            evaluations: values,
 20        }
 21    }
 22
 23    pub fn from_values(values: &[E]) -> Self {
 24        Self {
 25            num_variables: log2(values.len()) as usize,
 26            evaluations: values.to_owned(),
 27        }
 28    }
 29
 30    pub fn num_variables(&self) -> usize {
 31        self.num_variables
 32    }
 33
 34    pub fn evaluations(&self) -> &[E] {
 35        &self.evaluations
 36    }
 37
 38    pub fn len(&self) -> usize {
 39        self.evaluations.len()
 40    }
 41
 42    pub fn evaluate(&self, query: &[E]) -> E {
 43        let tensored_query = tensorize(query);
 44        inner_product(&self.evaluations, &tensored_query)
 45    }
 46
 47    pub fn bind(&self, round_challenge: E) -> Self {
 48        let mut result = vec![E::ZERO; 1 << (self.num_variables() - 1)];
 49        for i in 0..(1 << (self.num_variables() - 1)) {
 50            result[i] = self.evaluations[i << 1]
 51                + round_challenge * (self.evaluations[(i << 1) + 1] - self.evaluations[i << 1]);
 52        }
 53        Self::from_values(&result)
 54    }
 55
 56    pub fn bind_assign(&mut self, round_challenge: E) {
 57        let mut result = vec![E::ZERO; 1 << (self.num_variables() - 1)];
 58        for i in 0..(1 << (self.num_variables() - 1)) {
 59            result[i] = self.evaluations[i << 1]
 60                + round_challenge * (self.evaluations[(i << 1) + 1] - self.evaluations[i << 1]);
 61        }
 62        *self = Self::from_values(&result);
 63    }
 64
 65    pub fn split(&self, at: usize) -> (Self, Self) {
 66        assert!(at < self.len());
 67        (
 68            Self::new(self.evaluations[..at].to_vec()),
 69            Self::new(self.evaluations[at..2 * at].to_vec()),
 70        )
 71    }
 72
 73    pub fn extend(&mut self, other: &MultiLinear<E>) {
 74        let other_vec = other.evaluations.to_vec();
 75        assert_eq!(other_vec.len(), self.len());
 76        self.evaluations.extend(other_vec);
 77        self.num_variables += 1;
 78    }
 79}
 80
 81impl<E: FieldElement> Index<usize> for MultiLinear<E> {
 82    type Output = E;
 83
 84    fn index(&self, index: usize) -> &E {
 85        &(self.evaluations[index])
 86    }
 87}
 88
 89/// A multi-variate polynomial for composing individual multi-linear polynomials
 90pub trait CompositionPolynomial<E: FieldElement>: Sync + Send {
 91    /// The number of variables when interpreted as a multi-variate polynomial.
 92    fn num_variables(&self) -> usize;
 93
 94    /// Maximum degree in all variables.
 95    fn max_degree(&self) -> usize;
 96
 97    /// Given a query, of length equal the number of variables, evaluate [Self] at this query.
 98    fn evaluate(&self, query: &[E]) -> E;
 99}
100
101pub struct ComposedMultiLinears<E: FieldElement> {
102    pub composer: Arc<dyn CompositionPolynomial<E>>,
103    pub multi_linears: Vec<MultiLinear<E>>,
104}
105
106impl<E: FieldElement> ComposedMultiLinears<E> {
107    pub fn new(
108        composer: Arc<dyn CompositionPolynomial<E>>,
109        multi_linears: Vec<MultiLinear<E>>,
110    ) -> Self {
111        Self { composer, multi_linears }
112    }
113
114    pub fn num_ml(&self) -> usize {
115        self.multi_linears.len()
116    }
117
118    pub fn num_variables(&self) -> usize {
119        self.composer.num_variables()
120    }
121
122    pub fn num_variables_ml(&self) -> usize {
123        self.multi_linears[0].num_variables
124    }
125
126    pub fn degree(&self) -> usize {
127        self.composer.max_degree()
128    }
129
130    pub fn bind(&self, round_challenge: E) -> ComposedMultiLinears<E> {
131        let result: Vec<MultiLinear<E>> =
132            self.multi_linears.iter().map(|f| f.bind(round_challenge)).collect();
133
134        Self {
135            composer: self.composer.clone(),
136            multi_linears: result,
137        }
138    }
139}
140
141#[derive(Clone)]
142pub struct ComposedMultiLinearsOracle<E: FieldElement> {
143    pub composer: Arc<dyn CompositionPolynomial<E>>,
144    pub multi_linears: Vec<MultiLinearOracle>,
145}
146
147#[derive(Debug, Clone)]
148pub struct MultiLinearOracle {
149    pub id: usize,
150}
151
152// Composition polynomials
153
154pub struct IdentityComposition {
155    num_variables: usize,
156}
157
158impl IdentityComposition {
159    pub fn new() -> Self {
160        Self { num_variables: 1 }
161    }
162}
163
164impl<E> CompositionPolynomial<E> for IdentityComposition
165where
166    E: FieldElement,
167{
168    fn num_variables(&self) -> usize {
169        self.num_variables
170    }
171
172    fn max_degree(&self) -> usize {
173        self.num_variables
174    }
175
176    fn evaluate(&self, query: &[E]) -> E {
177        assert_eq!(query.len(), 1);
178        query[0]
179    }
180}
181
182pub struct ProjectionComposition {
183    coordinate: usize,
184}
185
186impl ProjectionComposition {
187    pub fn new(coordinate: usize) -> Self {
188        Self { coordinate }
189    }
190}
191
192impl<E> CompositionPolynomial<E> for ProjectionComposition
193where
194    E: FieldElement,
195{
196    fn num_variables(&self) -> usize {
197        1
198    }
199
200    fn max_degree(&self) -> usize {
201        1
202    }
203
204    fn evaluate(&self, query: &[E]) -> E {
205        query[self.coordinate]
206    }
207}
208
209pub struct LogUpDenominatorTableComposition<E>
210where
211    E: FieldElement,
212{
213    projection_coordinate: usize,
214    alpha: E,
215}
216
217impl<E> LogUpDenominatorTableComposition<E>
218where
219    E: FieldElement,
220{
221    pub fn new(projection_coordinate: usize, alpha: E) -> Self {
222        Self { projection_coordinate, alpha }
223    }
224}
225
226impl<E> CompositionPolynomial<E> for LogUpDenominatorTableComposition<E>
227where
228    E: FieldElement,
229{
230    fn num_variables(&self) -> usize {
231        1
232    }
233
234    fn max_degree(&self) -> usize {
235        1
236    }
237
238    fn evaluate(&self, query: &[E]) -> E {
239        query[self.projection_coordinate] + self.alpha
240    }
241}
242
243pub struct LogUpDenominatorWitnessComposition<E>
244where
245    E: FieldElement,
246{
247    projection_coordinate: usize,
248    alpha: E,
249}
250
251impl<E> LogUpDenominatorWitnessComposition<E>
252where
253    E: FieldElement,
254{
255    pub fn new(projection_coordinate: usize, alpha: E) -> Self {
256        Self { projection_coordinate, alpha }
257    }
258}
259
260impl<E> CompositionPolynomial<E> for LogUpDenominatorWitnessComposition<E>
261where
262    E: FieldElement,
263{
264    fn num_variables(&self) -> usize {
265        1
266    }
267
268    fn max_degree(&self) -> usize {
269        1
270    }
271
272    fn evaluate(&self, query: &[E]) -> E {
273        -(query[self.projection_coordinate] + self.alpha)
274    }
275}
276
277pub struct ProductComposition {
278    num_variables: usize,
279}
280
281impl ProductComposition {
282    pub fn new(num_variables: usize) -> Self {
283        Self { num_variables }
284    }
285}
286
287impl<E> CompositionPolynomial<E> for ProductComposition
288where
289    E: FieldElement,
290{
291    fn num_variables(&self) -> usize {
292        self.num_variables
293    }
294
295    fn max_degree(&self) -> usize {
296        self.num_variables
297    }
298
299    fn evaluate(&self, query: &[E]) -> E {
300        query.iter().fold(E::ONE, |acc, x| acc * *x)
301    }
302}
303
304pub struct SumComposition {
305    num_variables: usize,
306}
307
308impl SumComposition {
309    pub fn new(num_variables: usize) -> Self {
310        Self { num_variables }
311    }
312}
313
314impl<E> CompositionPolynomial<E> for SumComposition
315where
316    E: FieldElement,
317{
318    fn num_variables(&self) -> usize {
319        self.num_variables
320    }
321
322    fn max_degree(&self) -> usize {
323        self.num_variables
324    }
325
326    fn evaluate(&self, query: &[E]) -> E {
327        query.iter().fold(E::ZERO, |acc, x| acc + *x)
328    }
329}
330
331pub struct GkrCompositionVanilla<E: 'static>
332where
333    E: FieldElement,
334{
335    num_variables_ml: usize,
336    num_variables_merge: usize,
337    combining_randomness: E,
338    gkr_randomness: Vec<E>,
339}
340
341impl<E> GkrCompositionVanilla<E>
342where
343    E: FieldElement,
344{
345    pub fn new(
346        num_variables_ml: usize,
347        num_variables_merge: usize,
348        combining_randomness: E,
349        gkr_randomness: Vec<E>,
350    ) -> Self {
351        Self {
352            num_variables_ml,
353            num_variables_merge,
354            combining_randomness,
355            gkr_randomness,
356        }
357    }
358}
359
360impl<E> CompositionPolynomial<E> for GkrCompositionVanilla<E>
361where
362    E: FieldElement,
363{
364    fn num_variables(&self) -> usize {
365        self.num_variables_ml // + TODO
366    }
367
368    fn max_degree(&self) -> usize {
369        self.num_variables_ml //TODO
370    }
371
372    fn evaluate(&self, query: &[E]) -> E {
373        let eval_left_numerator = query[0];
374        let eval_right_numerator = query[1];
375        let eval_left_denominator = query[2];
376        let eval_right_denominator = query[3];
377        let eq_eval = query[4];
378
379        eq_eval
380            * ((eval_left_numerator * eval_right_denominator
381                + eval_right_numerator * eval_left_denominator)
382                + eval_left_denominator * eval_right_denominator * self.combining_randomness)
383    }
384}
385
386#[derive(Clone)]
387pub struct GkrComposition<E>
388where
389    E: FieldElement<BaseField = BaseElement>,
390{
391    pub num_variables_ml: usize,
392    pub combining_randomness: E,
393
394    eq_composer: Arc<dyn CompositionPolynomial<E>>,
395    right_numerator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
396    left_numerator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
397    right_denominator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
398    left_denominator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
399}
400
401impl<E> GkrComposition<E>
402where
403    E: FieldElement<BaseField = BaseElement>,
404{
405    pub fn new(
406        num_variables_ml: usize,
407        combining_randomness: E,
408        eq_composer: Arc<dyn CompositionPolynomial<E>>,
409        right_numerator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
410        left_numerator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
411        right_denominator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
412        left_denominator_composer: Vec<Arc<dyn CompositionPolynomial<E>>>,
413    ) -> Self {
414        Self {
415            num_variables_ml,
416            combining_randomness,
417            eq_composer,
418            right_numerator_composer,
419            left_numerator_composer,
420            right_denominator_composer,
421            left_denominator_composer,
422        }
423    }
424}
425
426impl<E> CompositionPolynomial<E> for GkrComposition<E>
427where
428    E: FieldElement<BaseField = BaseElement>,
429{
430    fn num_variables(&self) -> usize {
431        self.num_variables_ml // + TODO
432    }
433
434    fn max_degree(&self) -> usize {
435        3 // TODO
436    }
437
438    fn evaluate(&self, query: &[E]) -> E {
439        let eval_right_numerator = self.right_numerator_composer[0].evaluate(query);
440        let eval_left_numerator = self.left_numerator_composer[0].evaluate(query);
441        let eval_right_denominator = self.right_denominator_composer[0].evaluate(query);
442        let eval_left_denominator = self.left_denominator_composer[0].evaluate(query);
443        let eq_eval = self.eq_composer.evaluate(query);
444
445        let res = eq_eval
446            * ((eval_left_numerator * eval_right_denominator
447                + eval_right_numerator * eval_left_denominator)
448                + eval_left_denominator * eval_right_denominator * self.combining_randomness);
449        res
450    }
451}
452
453/// Generates a composed ML polynomial for the initial GKR layer from a vector of composition
454/// polynomials.
455/// The composition polynomials are divided into LeftNumerator, RightNumerator, LeftDenominator
456/// and RightDenominator.
457/// TODO: Generalize this to the case where each numerator/denominator contains more than one
458/// composition polynomial i.e., a merged composed ML polynomial.
459pub fn gkr_composition_from_composition_polys<
460    E: FieldElement<BaseField = BaseElement> + 'static,
461>(
462    composition_polys: &Vec<Vec<Arc<dyn CompositionPolynomial<E>>>>,
463    combining_randomness: E,
464    num_variables: usize,
465) -> GkrComposition<E> {
466    let eq_composer = Arc::new(ProjectionComposition::new(4));
467    let left_numerator = composition_polys[0].to_owned();
468    let right_numerator = composition_polys[1].to_owned();
469    let left_denominator = composition_polys[2].to_owned();
470    let right_denominator = composition_polys[3].to_owned();
471    GkrComposition::new(
472        num_variables,
473        combining_randomness,
474        eq_composer,
475        right_numerator,
476        left_numerator,
477        right_denominator,
478        left_denominator,
479    )
480}
481
482/// Generates a plain oracle for the sum-check protocol except the final one.
483pub fn gen_plain_gkr_oracle<E: FieldElement<BaseField = BaseElement> + 'static>(
484    num_rounds: usize,
485    r_sum_check: E,
486) -> ComposedMultiLinearsOracle<E> {
487    let gkr_composer = Arc::new(GkrCompositionVanilla::new(num_rounds, 0, r_sum_check, vec![]));
488
489    let ml_oracles = vec![
490        MultiLinearOracle { id: 0 },
491        MultiLinearOracle { id: 1 },
492        MultiLinearOracle { id: 2 },
493        MultiLinearOracle { id: 3 },
494        MultiLinearOracle { id: 4 },
495    ];
496
497    let oracle = ComposedMultiLinearsOracle {
498        composer: gkr_composer,
499        multi_linears: ml_oracles,
500    };
501    oracle
502}
503
504fn to_index<E: FieldElement<BaseField = BaseElement>>(index: &[E]) -> usize {
505    let res = index.iter().fold(E::ZERO, |acc, term| acc * E::ONE.double() + (*term));
506    let res = res.base_element(0);
507    res.as_int() as usize
508}
509
510fn inner_product<E: FieldElement>(evaluations: &[E], tensored_query: &[E]) -> E {
511    assert_eq!(evaluations.len(), tensored_query.len());
512    evaluations
513        .iter()
514        .zip(tensored_query.iter())
515        .fold(E::ZERO, |acc, (x_i, y_i)| acc + *x_i * *y_i)
516}
517
518pub fn tensorize<E: FieldElement>(query: &[E]) -> Vec<E> {
519    let nu = query.len();
520    let n = 1 << nu;
521
522    (0..n).map(|i| lagrange_basis_eval(query, i)).collect()
523}
524
525fn lagrange_basis_eval<E: FieldElement>(query: &[E], i: usize) -> E {
526    query
527        .iter()
528        .enumerate()
529        .map(|(j, x_j)| if i & (1 << j) == 0 { E::ONE - *x_j } else { *x_j })
530        .fold(E::ONE, |acc, v| acc * v)
531}
532
533pub fn compute_claim<E: FieldElement>(poly: &ComposedMultiLinears<E>) -> E {
534    let cube_size = 1 << poly.num_variables_ml();
535    let mut res = E::ZERO;
536
537    for i in 0..cube_size {
538        let eval_point: Vec<E> =
539            poly.multi_linears.iter().map(|poly| poly.evaluations[i]).collect();
540        res += poly.composer.evaluate(&eval_point);
541    }
542    res
543}