33
44//! Functions for polynomial interpolation and evaluation
55
6- use crate :: field:: NttFriendlyFieldElement ;
76#[ cfg( all( feature = "crypto-dependencies" , feature = "experimental" ) ) ]
87use crate :: ntt:: { ntt, ntt_inv_finish} ;
8+ use crate :: {
9+ field:: { FieldElement , NttFriendlyFieldElement } ,
10+ fp:: log2,
11+ } ;
912
1013use std:: convert:: TryFrom ;
1114
@@ -60,6 +63,90 @@ pub fn poly_interpret_eval<F: NttFriendlyFieldElement>(
6063 poly_eval ( & tmp_coeffs[ ..points. len ( ) ] , eval_at)
6164}
6265
66+ /// Returns the element `1/n` on `F`, where `n` must be a power of two.
67+ #[ inline]
68+ fn inv_pow2 < F : FieldElement > ( n : usize ) -> F {
69+ let log2_n = usize:: try_from ( log2 ( n as u128 ) ) . unwrap ( ) ;
70+ assert_eq ! ( n, 1 << log2_n) ;
71+
72+ let half = F :: half ( ) ;
73+ let mut x = F :: one ( ) ;
74+ for _ in 0 ..log2_n {
75+ x *= half
76+ }
77+ x
78+ }
79+
80+ /// Evaluates multiple polynomials given in the Lagrange basis.
81+ ///
82+ /// This is Algorithm 7 of rhizomes paper.
83+ /// <https://eprint.iacr.org/2025/1727>.
84+ pub ( crate ) fn poly_eval_batched < F : FieldElement > (
85+ polynomials : & [ Vec < F > ] ,
86+ roots : & [ F ] ,
87+ x : F ,
88+ ) -> Vec < F > {
89+ let mut l = F :: one ( ) ;
90+ let mut u = Vec :: with_capacity ( polynomials. len ( ) ) ;
91+ u. extend ( polynomials. iter ( ) . map ( |poly| poly[ 0 ] ) ) ;
92+ let mut d = roots[ 0 ] - x;
93+ for ( i, wn_i) in ( 1 ..) . zip ( & roots[ 1 ..] ) {
94+ l *= d;
95+ d = * wn_i - x;
96+ let t = l * * wn_i;
97+ for ( u_j, poly) in u. iter_mut ( ) . zip ( polynomials) {
98+ * u_j *= d;
99+ if let Some ( yi) = poly. get ( i) {
100+ * u_j += t * * yi;
101+ }
102+ }
103+ }
104+
105+ if roots. len ( ) > 1 {
106+ let num_roots_inv = -inv_pow2 :: < F > ( roots. len ( ) ) ;
107+ u. iter_mut ( ) . for_each ( |u_j| * u_j *= num_roots_inv) ;
108+ }
109+
110+ u
111+ }
112+
113+ /// Generates the powers of the primitive n-th root of unity.
114+ ///
115+ /// Returns
116+ /// roots\[i\] = w_n^i for 0 ≤ i < n,
117+ /// where
118+ /// w_n is the primitive n-th root of unity in `F`, and
119+ /// `n` must be a power of two.
120+ pub ( crate ) fn nth_root_powers < F : NttFriendlyFieldElement > ( n : usize ) -> Vec < F > {
121+ let log2_n = usize:: try_from ( log2 ( n as u128 ) ) . unwrap ( ) ;
122+ assert_eq ! ( n, 1 << log2_n) ;
123+
124+ let mut roots = vec ! [ F :: zero( ) ; n] ;
125+ roots[ 0 ] = F :: one ( ) ;
126+ if n > 1 {
127+ roots[ 1 ] = -F :: one ( ) ;
128+ for i in 2 ..=log2_n {
129+ let mid = 1 << ( i - 1 ) ;
130+ // Due to w_{2n}^{2j} = w_{n}^j
131+ for j in ( 1 ..mid) . rev ( ) {
132+ roots[ j << 1 ] = roots[ j]
133+ }
134+
135+ let wn = F :: root ( i) . unwrap ( ) ;
136+ roots[ 1 ] = wn;
137+ roots[ 1 + mid] = -wn;
138+
139+ // Due to w_{n}^{j} = -w_{n}^{j+n/2}
140+ for j in ( 3 ..mid) . step_by ( 2 ) {
141+ roots[ j] = wn * roots[ j - 1 ] ;
142+ roots[ j + mid] = -roots[ j]
143+ }
144+ }
145+ }
146+
147+ roots
148+ }
149+
63150/// Returns a polynomial that evaluates to `0` if the input is in range `[start, end)`. Otherwise,
64151/// the output is not `0`.
65152pub ( crate ) fn poly_range_check < F : NttFriendlyFieldElement > ( start : usize , end : usize ) -> Vec < F > {
@@ -76,7 +163,11 @@ pub(crate) fn poly_range_check<F: NttFriendlyFieldElement>(start: usize, end: us
76163mod tests {
77164 use crate :: {
78165 field:: { Field64 , FieldElement , FieldPrio2 , NttFriendlyFieldElement } ,
79- polynomial:: { poly_deg, poly_eval, poly_mul, poly_range_check} ,
166+ fp:: log2,
167+ polynomial:: {
168+ nth_root_powers, poly_deg, poly_eval, poly_eval_batched, poly_interpret_eval, poly_mul,
169+ poly_range_check,
170+ } ,
80171 } ;
81172 use std:: convert:: TryFrom ;
82173
@@ -154,4 +245,77 @@ mod tests {
154245 let y = poly_eval ( & p, x) ;
155246 assert_ne ! ( y, Field64 :: zero( ) ) ;
156247 }
248+
249+ /// Generates the powers of the primitive n-th root of unity.
250+ ///
251+ /// Returns
252+ /// roots\[i\] = w_n^i for 0 ≤ i < n,
253+ /// where
254+ /// w_n is the primitive n-th root of unity in `F`, and
255+ /// `n` must be a power of two.
256+ ///
257+ /// This is the iterative method.
258+ fn nth_root_powers_slow < F : NttFriendlyFieldElement > ( n : usize ) -> Vec < F > {
259+ let log2_n = usize:: try_from ( log2 ( n as u128 ) ) . unwrap ( ) ;
260+ let wn = F :: root ( log2_n) . unwrap ( ) ;
261+ core:: iter:: successors ( Some ( F :: one ( ) ) , |& x| Some ( x * wn) )
262+ . take ( n)
263+ . collect ( )
264+ }
265+
266+ #[ test]
267+ fn test_nth_root_powers ( ) {
268+ for i in 0 ..8 {
269+ assert_eq ! (
270+ nth_root_powers:: <Field64 >( 1 << i) ,
271+ nth_root_powers_slow:: <Field64 >( 1 << i)
272+ ) ;
273+ }
274+ }
275+
276+ #[ test]
277+ fn test_poly_eval_batched_ones ( ) {
278+ test_poly_eval_batched ( & [ 1 ] ) ;
279+ test_poly_eval_batched ( & [ 1 , 1 ] ) ;
280+ }
281+
282+ #[ test]
283+ fn test_poly_eval_batched_powers ( ) {
284+ test_poly_eval_batched ( & [ 1 , 2 , 4 , 16 , 64 ] ) ;
285+ }
286+
287+ #[ test]
288+ fn test_poly_eval_batched_arbitrary ( ) {
289+ test_poly_eval_batched ( & [ 1 , 6 , 3 , 9 ] ) ;
290+ }
291+
292+ fn test_poly_eval_batched ( lengths : & [ usize ] ) {
293+ let sizes = lengths
294+ . iter ( )
295+ . map ( |s| s. next_power_of_two ( ) )
296+ . collect :: < Vec < _ > > ( ) ;
297+
298+ let polynomials = sizes
299+ . iter ( )
300+ . map ( |& size| Field64 :: random_vector ( size) )
301+ . collect :: < Vec < _ > > ( ) ;
302+ let x = Field64 :: random_vector ( 1 ) [ 0 ] ;
303+
304+ let & n = sizes. iter ( ) . max ( ) . unwrap ( ) ;
305+ let mut ntt_mem = vec ! [ Field64 :: zero( ) ; n] ;
306+ let roots = nth_root_powers ( n) ;
307+
308+ // Evaluates several polynomials converting them to the monomial basis (iteratively).
309+ let want = polynomials
310+ . iter ( )
311+ . map ( |poly| {
312+ let extended_poly = [ poly. clone ( ) , vec ! [ Field64 :: zero( ) ; n - poly. len( ) ] ] . concat ( ) ;
313+ poly_interpret_eval ( & extended_poly, x, & mut ntt_mem)
314+ } )
315+ . collect :: < Vec < _ > > ( ) ;
316+
317+ // Simultaneouly evaluates several polynomials directly in the Lagrange basis (batched).
318+ let got = poly_eval_batched ( & polynomials, & roots, x) ;
319+ assert_eq ! ( got, want, "sizes: {sizes:?} x: {x} P: {polynomials:?}" ) ;
320+ }
157321}
0 commit comments