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 > {
@@ -74,9 +161,14 @@ pub(crate) fn poly_range_check<F: NttFriendlyFieldElement>(start: usize, end: us
74161
75162#[ cfg( test) ]
76163mod tests {
164+ #[ cfg( all( feature = "crypto-dependencies" , feature = "experimental" ) ) ]
165+ use crate :: polynomial:: poly_interpret_eval;
77166 use crate :: {
78167 field:: { Field64 , FieldElement , FieldPrio2 , NttFriendlyFieldElement } ,
79- polynomial:: { poly_deg, poly_eval, poly_mul, poly_range_check} ,
168+ fp:: log2,
169+ polynomial:: {
170+ nth_root_powers, poly_deg, poly_eval, poly_eval_batched, poly_mul, poly_range_check,
171+ } ,
80172 } ;
81173 use std:: convert:: TryFrom ;
82174
@@ -154,4 +246,75 @@ mod tests {
154246 let y = poly_eval ( & p, x) ;
155247 assert_ne ! ( y, Field64 :: zero( ) ) ;
156248 }
249+
250+ /// Generates the powers of the primitive n-th root of unity.
251+ ///
252+ /// Returns
253+ /// roots\[i\] = w_n^i for 0 ≤ i < n,
254+ /// where
255+ /// w_n is the primitive n-th root of unity in `F`, and
256+ /// `n` must be a power of two.
257+ ///
258+ /// This is the iterative method.
259+ fn nth_root_powers_slow < F : NttFriendlyFieldElement > ( n : usize ) -> Vec < F > {
260+ let log2_n = usize:: try_from ( log2 ( n as u128 ) ) . unwrap ( ) ;
261+ let wn = F :: root ( log2_n) . unwrap ( ) ;
262+ core:: iter:: successors ( Some ( F :: one ( ) ) , |& x| Some ( x * wn) )
263+ . take ( n)
264+ . collect ( )
265+ }
266+
267+ #[ test]
268+ fn test_nth_root_powers ( ) {
269+ for i in 0 ..8 {
270+ assert_eq ! (
271+ nth_root_powers:: <Field64 >( 1 << i) ,
272+ nth_root_powers_slow:: <Field64 >( 1 << i)
273+ ) ;
274+ }
275+ }
276+
277+ #[ cfg( all( feature = "crypto-dependencies" , feature = "experimental" ) ) ]
278+ #[ test]
279+ fn test_poly_eval_batched_arbitrary ( ) {
280+ // Single polynomial with constant terms
281+ test_poly_eval_batched ( & [ 1 ] ) ;
282+ // Constant terms
283+ test_poly_eval_batched ( & [ 1 , 1 ] ) ;
284+ // Powers of two
285+ test_poly_eval_batched ( & [ 1 , 2 , 4 , 16 , 64 ] ) ;
286+ // arbitrary
287+ test_poly_eval_batched ( & [ 1 , 6 , 3 , 9 ] ) ;
288+ }
289+
290+ #[ cfg( all( feature = "crypto-dependencies" , feature = "experimental" ) ) ]
291+ fn test_poly_eval_batched ( lengths : & [ usize ] ) {
292+ let sizes = lengths
293+ . iter ( )
294+ . map ( |s| s. next_power_of_two ( ) )
295+ . collect :: < Vec < _ > > ( ) ;
296+
297+ let polynomials = sizes
298+ . iter ( )
299+ . map ( |& size| Field64 :: random_vector ( size) )
300+ . collect :: < Vec < _ > > ( ) ;
301+ let x = Field64 :: random_vector ( 1 ) [ 0 ] ;
302+
303+ let & n = sizes. iter ( ) . max ( ) . unwrap ( ) ;
304+ let mut ntt_mem = vec ! [ Field64 :: zero( ) ; n] ;
305+ let roots = nth_root_powers ( n) ;
306+
307+ // Evaluates several polynomials converting them to the monomial basis (iteratively).
308+ let want = polynomials
309+ . iter ( )
310+ . map ( |poly| {
311+ let extended_poly = [ poly. clone ( ) , vec ! [ Field64 :: zero( ) ; n - poly. len( ) ] ] . concat ( ) ;
312+ poly_interpret_eval ( & extended_poly, x, & mut ntt_mem)
313+ } )
314+ . collect :: < Vec < _ > > ( ) ;
315+
316+ // Simultaneouly evaluates several polynomials directly in the Lagrange basis (batched).
317+ let got = poly_eval_batched ( & polynomials, & roots, x) ;
318+ assert_eq ! ( got, want, "sizes: {sizes:?} x: {x} P: {polynomials:?}" ) ;
319+ }
157320}
0 commit comments