p3_interpolation/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
//! Tools for Lagrange interpolation.

#![no_std]

extern crate alloc;

use alloc::vec::Vec;

use p3_field::{
    batch_multiplicative_inverse, scale_vec, two_adic_coset_zerofier, ExtensionField, TwoAdicField,
};
use p3_matrix::Matrix;
use p3_maybe_rayon::prelude::*;
use p3_util::log2_strict_usize;

/// Given evaluations of a batch of polynomials over the canonical power-of-two subgroup, evaluate
/// the polynomials at `point`.
///
/// This assumes the point is not in the subgroup, otherwise the behavior is undefined.
pub fn interpolate_subgroup<F, EF, Mat>(subgroup_evals: &Mat, point: EF) -> Vec<EF>
where
    F: TwoAdicField,
    EF: ExtensionField<F> + TwoAdicField,
    Mat: Matrix<F>,
{
    interpolate_coset(subgroup_evals, F::ONE, point, None)
}

/// Given evaluations of a batch of polynomials over the given coset of the canonical power-of-two
/// subgroup, evaluate the polynomials at `point`.
///
/// This assumes the point is not in the coset, otherwise the behavior is undefined.
/// If available, reuse denominator diffs that is `1 / (x_i-z)` to avoid batch inversion.
pub fn interpolate_coset<F, EF, Mat>(
    coset_evals: &Mat,
    shift: F,
    point: EF,
    diff_invs: Option<&[EF]>,
) -> Vec<EF>
where
    F: TwoAdicField,
    EF: ExtensionField<F> + TwoAdicField,
    Mat: Matrix<F>,
{
    // Slight variation of this approach: https://hackmd.io/@vbuterin/barycentric_evaluation

    let height = coset_evals.height();
    let log_height = log2_strict_usize(height);

    let g = F::two_adic_generator(log_height).powers().take(height);
    let col_scale: Vec<_> = if let Some(diff_invs) = diff_invs {
        g.zip(diff_invs)
            .map(|(sg, &diff_inv)| diff_inv * sg)
            .collect()
    } else {
        let subgroup = g.collect::<Vec<_>>();
        let diffs: Vec<EF> = subgroup
            .par_iter()
            .map(|&subgroup_i| point - subgroup_i * shift)
            .collect();
        let diff_invs = batch_multiplicative_inverse(&diffs);
        subgroup
            .par_iter()
            .zip(diff_invs)
            .map(|(&sg, diff_inv)| diff_inv * sg)
            .collect()
    };
    let sum = coset_evals.columnwise_dot_product(&col_scale);

    let zerofier = two_adic_coset_zerofier::<EF>(log_height, EF::from_base(shift), point);
    let denominator = F::from_canonical_usize(height) * shift.exp_u64(height as u64 - 1);
    scale_vec(zerofier * denominator.inverse(), sum)
}

#[cfg(test)]
mod tests {
    use alloc::vec;
    use alloc::vec::Vec;

    use p3_baby_bear::BabyBear;
    use p3_field::{batch_multiplicative_inverse, Field, FieldAlgebra};
    use p3_matrix::dense::RowMajorMatrix;
    use p3_util::log2_strict_usize;

    use crate::{interpolate_coset, interpolate_subgroup};

    #[test]
    fn test_interpolate_subgroup() {
        // x^2 + 2 x + 3
        type F = BabyBear;
        let evals = [
            6, 886605102, 1443543107, 708307799, 2, 556938009, 569722818, 1874680944,
        ]
        .map(F::from_canonical_u32);
        let evals_mat = RowMajorMatrix::new(evals.to_vec(), 1);
        let point = F::from_canonical_u32(100);
        let result = interpolate_subgroup(&evals_mat, point);
        assert_eq!(result, vec![F::from_canonical_u32(10203)]);
    }

    #[test]
    fn test_interpolate_coset() {
        // x^2 + 2 x + 3
        type F = BabyBear;
        let shift = F::GENERATOR;
        let evals = [
            1026, 129027310, 457985035, 994890337, 902, 1988942953, 1555278970, 913671254,
        ]
        .map(F::from_canonical_u32);
        let evals_mat = RowMajorMatrix::new(evals.to_vec(), 1);
        let point = F::from_canonical_u32(100);
        let result = interpolate_coset(&evals_mat, shift, point, None);
        assert_eq!(result, vec![F::from_canonical_u32(10203)]);

        use p3_field::TwoAdicField;
        let n = evals.len();
        let k = log2_strict_usize(n);

        let denom: Vec<_> = F::two_adic_generator(k)
            .shifted_powers(shift)
            .take(n)
            .map(|w| point - w)
            .collect();

        let denom = batch_multiplicative_inverse(&denom);
        let result = interpolate_coset(&evals_mat, shift, point, Some(&denom));
        assert_eq!(result, vec![F::from_canonical_u32(10203)]);
    }
}