Add compute roots of unity method

This commit is contained in:
Daniel Sanchez Quiros 2024-06-20 14:42:18 +02:00
parent 8ed03880d4
commit 3208e366eb
3 changed files with 160 additions and 23 deletions

View File

@ -4,10 +4,11 @@ use std::fmt;
// crates
use crate::{FieldElement, BYTES_PER_FIELD_ELEMENT};
use ark_bls12_381::fr::Fr;
use ark_ff::Zero;
use ark_ff::{BigInt, FftField, Field, Zero};
use ark_poly::domain::general::GeneralEvaluationDomain;
use ark_poly::evaluations::univariate::Evaluations;
use ark_poly::univariate::DensePolynomial;
use ark_poly::{EvaluationDomain, Radix2EvaluationDomain};
use blst::BLST_ERROR;
use num_bigint::BigUint;
use thiserror::Error;
@ -122,6 +123,11 @@ pub fn field_element_from_bytes_le(b: &[u8]) -> FieldElement {
FieldElement::from(BigUint::from_bytes_le(b))
}
pub fn compute_roots_of_unity(size: usize) -> Vec<Fr> {
let domain = GeneralEvaluationDomain::new(size).unwrap();
domain.elements().take(size).collect()
}
#[cfg(test)]
mod test {
use super::{bytes_to_evaluations, bytes_to_polynomial, KzgRsError};

View File

@ -106,8 +106,90 @@ pub fn ifft_g1(vals: &[G1Affine], roots_of_unity: &[Fr]) -> Vec<G1Affine> {
.collect()
}
pub fn fft_fr(vals: &[Fr], roots_of_unity: &[Fr]) -> Vec<Fr> {
debug_assert_eq!(vals.len(), roots_of_unity.len());
let original_len = vals.len();
if original_len == 1 {
return vals.to_vec();
}
let half_roots: Vec<_> = roots_of_unity.iter().step_by(2).copied().collect();
let l = || {
crate::fft::fft_fr(
vals.iter()
.step_by(2)
.copied()
.collect::<Vec<_>>()
.as_slice(),
half_roots.as_slice(),
)
};
let r = || {
crate::fft::fft_fr(
vals.iter()
.skip(1)
.step_by(2)
.copied()
.collect::<Vec<_>>()
.as_slice(),
half_roots.as_slice(),
)
};
let [l, r]: [Vec<Fr>; 2] = {
#[cfg(feature = "parallel")]
{
let (l, r) = rayon::join(l, r);
[l, r]
}
#[cfg(not(feature = "parallel"))]
{
[l(), r()]
}
};
// Double sized so we can use iterator later on
let l: Vec<_> = l.into_iter().cycle().take(original_len).collect();
let r: Vec<_> = r.into_iter().cycle().take(original_len).collect();
let y_times_root = {
#[cfg(feature = "parallel")]
{
r.into_par_iter()
}
#[cfg(not(feature = "parallel"))]
{
r.into_iter()
}
}
.enumerate()
.map(|(i, y)| y * roots_of_unity[i % vals.len()]);
{
#[cfg(feature = "parallel")]
{
l.into_par_iter()
}
#[cfg(not(feature = "parallel"))]
{
l.into_iter()
}
}
.zip(y_times_root)
.enumerate()
.map(|(i, (x, y_times_root))| {
if i < vals.len() / 2 {
x + y_times_root
} else {
x - y_times_root
}
})
.collect()
}
#[cfg(test)]
mod test {
use crate::common::compute_roots_of_unity;
use crate::fft::{fft_g1, ifft_g1};
use ark_bls12_381::{Fr, G1Affine};
use ark_ec::{AffineRepr, CurveGroup};
@ -116,10 +198,8 @@ mod test {
#[test]
fn test_fft_ifft_g1() {
for size in [16usize, 32, 64, 128, 256, 512, 1024, 2048, 4096] {
let primitive_root = <Fr as FftField>::get_root_of_unity(size as u64).unwrap();
let roots_of_unity: Vec<_> = (1..=size)
.map(|i| primitive_root.pow::<ark_ff::BigInt<4>>(BigInt::from(i as u64)))
.collect();
let roots_of_unity = compute_roots_of_unity(size);
let r: Vec<_> = roots_of_unity.iter().map(|a| a.to_string()).collect();
let buff: Vec<G1Affine> = (0..size)
.map(|i| {
G1Affine::identity()

View File

@ -1,19 +1,21 @@
use ark_bls12_381::{Fr, G1Affine, G1Projective};
use ark_ec::{AffineRepr, CurveGroup};
use ark_ff::FftField;
use ark_poly::{EvaluationDomain, GeneralEvaluationDomain};
use ark_ff::{FftField, Field};
use ark_poly::{EvaluationDomain, GeneralEvaluationDomain, Polynomial as _};
use num_traits::Zero;
use crate::fft::{fft_g1, ifft_g1};
use crate::common::compute_roots_of_unity;
use crate::fft::{fft_fr, fft_g1, ifft_g1};
use crate::{GlobalParameters, Polynomial, Proof};
fn toeplitz1(global_parameters: &[G1Affine], roots_of_unity: &[Fr]) -> Vec<G1Affine> {
debug_assert_eq!(global_parameters.len(), roots_of_unity.len());
debug_assert!(roots_of_unity.len().is_power_of_two());
fn toeplitz1(global_parameters: &[G1Affine], polynomial_degree: usize) -> Vec<G1Affine> {
debug_assert_eq!(global_parameters.len(), polynomial_degree);
debug_assert!(polynomial_degree.is_power_of_two());
let roots_of_unity = compute_roots_of_unity(polynomial_degree * 2);
let vector_extended: Vec<G1Affine> = global_parameters
.iter()
.copied()
.chain(std::iter::repeat_with(G1Affine::zero).take(global_parameters.len()))
.chain(std::iter::repeat_with(G1Affine::zero).take(polynomial_degree))
.collect();
fft_g1(&vector_extended, &roots_of_unity)
}
@ -23,6 +25,8 @@ fn toeplitz2(coefficients: &[Fr], extended_vector: &[G1Affine]) -> Vec<G1Affine>
let domain: GeneralEvaluationDomain<Fr> =
GeneralEvaluationDomain::new(coefficients.len()).expect("Domain should be able to build");
let toeplitz_coefficients_fft = domain.fft(coefficients);
// let roots_of_unity = compute_roots_of_unity(coefficients.len());
// let toeplitz_coefficients_fft = fft_fr(coefficients, &roots_of_unity);
extended_vector
.iter()
.copied()
@ -32,9 +36,7 @@ fn toeplitz2(coefficients: &[Fr], extended_vector: &[G1Affine]) -> Vec<G1Affine>
}
fn toeplitz3(h_extended_fft: &[G1Affine], polynomial_degree: usize) -> Vec<G1Affine> {
let roots_of_unity: Vec<Fr> = (0..h_extended_fft.len())
.map(|i| Fr::get_root_of_unity(i as u64).expect("Root should be present"))
.collect();
let roots_of_unity: Vec<Fr> = compute_roots_of_unity(h_extended_fft.len());
ifft_g1(h_extended_fft, &roots_of_unity)
.into_iter()
.take(polynomial_degree)
@ -48,19 +50,17 @@ pub fn fk20_batch_generate_elements_proofs(
let polynomial_degree = polynomial.len();
debug_assert!(polynomial_degree <= global_parameters.powers_of_g.len());
debug_assert!(polynomial_degree.is_power_of_two());
let roots_of_unity: Vec<Fr> = (0..polynomial_degree)
.map(|i| Fr::get_root_of_unity(i as u64).expect("Root should be present"))
.collect();
let roots_of_unity: Vec<Fr> = compute_roots_of_unity(polynomial_degree);
let global_parameters: Vec<G1Affine> = global_parameters
.powers_of_g
.iter()
.copied()
.take(polynomial_degree)
.take(polynomial_degree - 1)
.rev()
.chain(std::iter::once(G1Affine::identity() * Fr::from(0)).map(G1Projective::into_affine))
.chain(std::iter::once(G1Affine::identity() * Fr::ZERO).map(G1Projective::into_affine))
.collect();
let extended_vector = toeplitz1(&global_parameters, &roots_of_unity);
let extended_vector = toeplitz1(&global_parameters, polynomial_degree);
let toeplitz_coefficients: Vec<Fr> = std::iter::once(
polynomial
.coeffs
@ -68,8 +68,15 @@ pub fn fk20_batch_generate_elements_proofs(
.copied()
.expect("Polynomial should not be empty"),
)
.chain(std::iter::repeat_with(Fr::zero).take(polynomial_degree + 1))
.chain(polynomial.coeffs.iter().skip(1).rev().copied())
.chain(std::iter::repeat(Fr::ZERO).take(polynomial_degree + 1))
.chain(
polynomial
.coeffs
.iter()
.skip(1)
.take(polynomial_degree - 2)
.copied(),
)
.collect();
let h_extended_vector = toeplitz2(&toeplitz_coefficients, &extended_vector);
let h_vector = toeplitz3(&h_extended_vector, polynomial_degree);
@ -81,3 +88,47 @@ pub fn fk20_batch_generate_elements_proofs(
})
.collect()
}
#[cfg(test)]
mod test {
use crate::common::compute_roots_of_unity;
use crate::fk20::fk20_batch_generate_elements_proofs;
use crate::{
common::bytes_to_polynomial, kzg::generate_element_proof, GlobalParameters, Proof,
BYTES_PER_FIELD_ELEMENT,
};
use ark_bls12_381::{Bls12_381, Fr};
use ark_ff::FftField;
use ark_poly::univariate::DensePolynomial;
use ark_poly::{EvaluationDomain, GeneralEvaluationDomain};
use ark_poly_commit::kzg10::KZG10;
use once_cell::sync::Lazy;
use rand::SeedableRng;
static GLOBAL_PARAMETERS: Lazy<GlobalParameters> = Lazy::new(|| {
let mut rng = rand::rngs::StdRng::seed_from_u64(1987);
KZG10::<Bls12_381, DensePolynomial<Fr>>::setup(4096, true, &mut rng).unwrap()
});
#[test]
fn test_generate_proofs() {
for size in [16usize, 32, 64, 128, 256] {
let mut buff: Vec<_> = (0..BYTES_PER_FIELD_ELEMENT * size)
.map(|i| (i % 255) as u8)
.rev()
.collect();
let domain = GeneralEvaluationDomain::new(size).unwrap();
let (evals, poly) =
bytes_to_polynomial::<BYTES_PER_FIELD_ELEMENT>(&buff, domain).unwrap();
let polynomial_degree = poly.len();
let roots_of_unity: Vec<Fr> = compute_roots_of_unity(size);
let slow_proofs: Vec<Proof> = (0..polynomial_degree)
.map(|i| {
generate_element_proof(i, &poly, &evals, &GLOBAL_PARAMETERS, domain).unwrap()
})
.collect();
let fk20_proofs = fk20_batch_generate_elements_proofs(&poly, &GLOBAL_PARAMETERS);
assert_eq!(slow_proofs, fk20_proofs);
}
}
}