Cleanup/roots of unities setup functions (#37)

* adding bare changes for batch verification

* adding some comments

* adding more comments

* moving back to sha2

* removing a test which is no longer needed. Removing methods no longer needed

* updates to method visibility, updating tests

* fmt fixes

* clean up

* cleanup, optimization, inline docs

* removing unwanted const

* more docs and cleanup

* formatting

* removing unwanted comments

* cargo fmt and clippy

* adding test for point at infinity

* cleaner errors, cleanup

* adding another test case

* removing unwanted errors

* adding fixes per comments

* adding 4844 spec references

* comment fixes

* formatting, adding index out of bound check, removing print statement

* removing unwanted test, adding test for evaluate_polynomial_in_evaluation_form

* moving test to bottom section

* Update src/polynomial.rs

Co-authored-by: Samuel Laferriere <[email protected]>

* Update src/kzg.rs

Co-authored-by: Samuel Laferriere <[email protected]>

* Update src/kzg.rs

Co-authored-by: Samuel Laferriere <[email protected]>

* Update src/kzg.rs

Co-authored-by: Samuel Laferriere <[email protected]>

* Update src/helpers.rs

Co-authored-by: Samuel Laferriere <[email protected]>

* updating deps, and toolchain to 1.84

* removing errors test, no longer useful

* adding to_byte_array arg explanation

* fmt fixes

* fmt and clippy fixes

* fixing function names and fmt

* clippy fixes

* removing unwanted setup functions

* removing vars from struct

* fixing function name and comments

* fixing naming for tests

* fixing naming in benches

* formatting

* removing is_zero

---------

Co-authored-by: anupsv <[email protected]>
Co-authored-by: Samuel Laferriere

benches/bench_kzg_commit.rs

@@ -16,7 +16,7 @@ fn bench_kzg_commit(c: &mut Criterion) {
         let random_blob: Vec<u8> = (0..10000).map(|_| rng.gen_range(32..=126) as u8).collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_coeff_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         b.iter(|| kzg.commit_coeff_form(&input_poly).unwrap());
     });
@@ -25,7 +25,7 @@ fn bench_kzg_commit(c: &mut Criterion) {
         let random_blob: Vec<u8> = (0..30000).map(|_| rng.gen_range(32..=126) as u8).collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_coeff_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         b.iter(|| kzg.commit_coeff_form(&input_poly).unwrap());
     });
@@ -34,7 +34,7 @@ fn bench_kzg_commit(c: &mut Criterion) {
         let random_blob: Vec<u8> = (0..50000).map(|_| rng.gen_range(32..=126) as u8).collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_coeff_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         b.iter(|| kzg.commit_coeff_form(&input_poly).unwrap());
     });

benches/bench_kzg_commit_large_blobs.rs

@@ -18,7 +18,7 @@ fn bench_kzg_commit(c: &mut Criterion) {
             .collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_coeff_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         b.iter(|| kzg.commit_coeff_form(&input_poly).unwrap());
     });
@@ -29,7 +29,7 @@ fn bench_kzg_commit(c: &mut Criterion) {
             .collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_coeff_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         b.iter(|| kzg.commit_coeff_form(&input_poly).unwrap());
     });

benches/bench_kzg_proof.rs

@@ -16,7 +16,7 @@ fn bench_kzg_proof(c: &mut Criterion) {
         let random_blob: Vec<u8> = (0..10000).map(|_| rng.gen_range(32..=126) as u8).collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_eval_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         let index =
             rand::thread_rng().gen_range(0..input_poly.len_underlying_blob_field_elements());
@@ -30,7 +30,7 @@ fn bench_kzg_proof(c: &mut Criterion) {
         let random_blob: Vec<u8> = (0..30000).map(|_| rng.gen_range(32..=126) as u8).collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_eval_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         let index =
             rand::thread_rng().gen_range(0..input_poly.len_underlying_blob_field_elements());
@@ -44,7 +44,7 @@ fn bench_kzg_proof(c: &mut Criterion) {
         let random_blob: Vec<u8> = (0..50000).map(|_| rng.gen_range(32..=126) as u8).collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_eval_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         let index =
             rand::thread_rng().gen_range(0..input_poly.len_underlying_blob_field_elements());

benches/bench_kzg_verify.rs

@@ -16,7 +16,7 @@ fn bench_kzg_verify(c: &mut Criterion) {
         let random_blob: Vec<u8> = (0..10000).map(|_| rng.gen_range(32..=126) as u8).collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_eval_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         let index =
             rand::thread_rng().gen_range(0..input_poly.len_underlying_blob_field_elements());
@@ -33,7 +33,7 @@ fn bench_kzg_verify(c: &mut Criterion) {
         let random_blob: Vec<u8> = (0..30000).map(|_| rng.gen_range(32..=126) as u8).collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_eval_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         let index =
             rand::thread_rng().gen_range(0..input_poly.len_underlying_blob_field_elements());
@@ -50,7 +50,7 @@ fn bench_kzg_verify(c: &mut Criterion) {
         let random_blob: Vec<u8> = (0..50000).map(|_| rng.gen_range(32..=126) as u8).collect();
         let input = Blob::from_raw_data(&random_blob);
         let input_poly = input.to_polynomial_eval_form();
-        kzg.data_setup_custom(1, input.len().try_into().unwrap())
+        kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap())
             .unwrap();
         let index =
             rand::thread_rng().gen_range(0..input_poly.len_underlying_blob_field_elements());

src/errors.rs

@@ -14,7 +14,6 @@ pub enum PolynomialError {
     /// Error related to Fast Fourier Transform (FFT) operations with a descriptive message.
     #[error("FFT error: {0}")]
     FFTError(String),
-
     /// A generic error with a descriptive message.
     #[error("generic error: {0}")]
     GenericError(String),

src/helpers.rs

@@ -155,14 +155,14 @@ pub fn to_fr_array(data: &[u8]) -> Vec<Fr> {
 ///
 /// # Arguments
 /// * `data_fr` - Slice of field elements to convert to bytes
-/// * `max_data_size` - Maximum allowed size in bytes for the output buffer
+/// * `max_output_size` - Maximum allowed size in bytes for the output buffer
 ///
 /// # Returns
 /// * `Vec<u8>` - Byte array containing the encoded field elements, truncated if needed
 ///
 /// # Details
 /// - Each field element is converted to BYTES_PER_FIELD_ELEMENT bytes
-/// - Output is truncated to max_data_size if total bytes would exceed it
+/// - Output is truncated to max_output_size if total bytes would exceed it
 ///
 /// # Example
 /// ```
@@ -175,7 +175,7 @@ pub fn to_fr_array(data: &[u8]) -> Vec<Fr> {
 ///                  131072,
 ///                  ).unwrap();
 /// let input = Blob::from_raw_data(b"random data for blob");
-/// kzg.calculate_roots_of_unity(input.len().try_into().unwrap()).unwrap();
+/// kzg.calculate_and_store_roots_of_unity(input.len().try_into().unwrap()).unwrap();
 /// ```
 pub fn to_byte_array(data_fr: &[Fr], max_output_size: usize) -> Vec<u8> {
     // Calculate the number of field elements in input

src/kzg.rs

@@ -43,16 +43,8 @@ pub struct KZG {
     // in coefficient form. To commit against a polynomial in evaluation form, we need to transform
     // the SRS points to lagrange form using IFFT.
     g1: Vec<G1Affine>,
-    params: Params,
     srs_order: u64,
     expanded_roots_of_unity: Vec<Fr>,
-}
-
-#[derive(Debug, PartialEq, Clone)]
-struct Params {
-    chunk_length: u64,
-    num_chunks: u64,
-    max_fft_width: u64,
     completed_setup: bool,
 }
 
@@ -74,14 +66,9 @@ impl KZG {
 
         Ok(Self {
             g1: g1_points,
-            params: Params {
-                chunk_length: 0,
-                num_chunks: 0,
-                max_fft_width: 0,
-                completed_setup: false,
-            },
             srs_order: srs_order.into(),
             expanded_roots_of_unity: vec![],
+            completed_setup: false,
         })
     }
 
@@ -100,31 +87,11 @@ impl KZG {
     /// # Details
     /// - Generates roots of unity needed for FFT operations
     /// - Calculates KZG operational parameters for commitment scheme
-    ///
-    /// # Example
     /// ```
-    /// use ark_std::One;
-    /// use rust_kzg_bn254::helpers::to_byte_array;
-    /// use ark_bn254::Fr;
-    ///
-    /// let elements = vec![Fr::one(), Fr::one(), Fr::one()];
-    /// let max_size = 64;
-    /// let bytes = to_byte_array(&elements, max_size);
-    /// assert_eq!(bytes.len(), 64);
-    /// // bytes will contain up to max_size bytes from the encoded elements
-    /// ```
-    fn calculate_roots_of_unity_standalone(
+    fn calculate_roots_of_unity(
         length_of_data_after_padding: u64,
         srs_order: u64,
-    ) -> Result<(Params, Vec<Fr>), KzgError> {
-        // Initialize parameters
-        let mut params = Params {
-            num_chunks: 0_u64,
-            chunk_length: 0_u64,
-            max_fft_width: 0_u64,
-            completed_setup: false,
-        };
-
+    ) -> Result<Vec<Fr>, KzgError> {
         // Calculate log2 of the next power of two of the length of data after padding
         let log2_of_evals = (length_of_data_after_padding
             .div_ceil(32)
@@ -137,9 +104,6 @@ impl KZG {
                 )
             })?;
 
-        // Set the maximum FFT width
-        params.max_fft_width = 1_u64 << log2_of_evals;
-
         // Check if the length of data after padding is valid with respect to the SRS order
         if length_of_data_after_padding
             .div_ceil(BYTES_PER_FIELD_ELEMENT as u64)
@@ -166,103 +130,45 @@ impl KZG {
         // Remove the last element to avoid duplication
         expanded_roots_of_unity.truncate(expanded_roots_of_unity.len() - 1);
 
-        // Mark the setup as completed
-        params.completed_setup = true;
-
         // Return the parameters and the expanded roots of unity
-        Ok((params, expanded_roots_of_unity))
+        Ok(expanded_roots_of_unity)
     }
 
-    /// Similar to [Kzg::data_setup_mins], but mainly used for setting up Kzg
-    /// for testing purposes. Used to specify the number of chunks and chunk
-    /// length. These parameters are then used to calculate the FFT params
-    /// required for FFT operations.
-    pub fn data_setup_custom(
-        &mut self,
-        num_of_nodes: u64,
-        padded_input_data_size: u64,
-    ) -> Result<(), KzgError> {
-        let floor = u64::try_from(BYTES_PER_FIELD_ELEMENT)
-            .map_err(|e| KzgError::SerializationError(e.to_string()))?;
-        let len_of_data_in_elements = padded_input_data_size.div_ceil(floor);
-        let min_num_chunks = len_of_data_in_elements.div_ceil(num_of_nodes);
-        self.data_setup_mins(min_num_chunks, num_of_nodes)
-    }
-
-    /// Used to specify the number of chunks and chunk length.
-    /// These parameters are then used to calculate the FFT params required for
-    /// FFT operations.
-    pub fn data_setup_mins(
-        &mut self,
-        min_chunk_length: u64,
-        min_num_chunks: u64,
-    ) -> Result<(), KzgError> {
-        let mut params = Params {
-            num_chunks: min_num_chunks.next_power_of_two(),
-            chunk_length: min_chunk_length.next_power_of_two(),
-            max_fft_width: 0_u64,
-            completed_setup: false,
-        };
-
-        let number_of_evaluations = params.chunk_length * params.num_chunks;
-        let mut log2_of_evals = number_of_evaluations
-            .to_f64()
-            .ok_or_else(|| {
-                KzgError::GenericError("Failed to convert number_of_evaluations to f64".to_string())
-            })?
-            .log2()
-            .to_u8()
-            .ok_or_else(|| {
-                KzgError::GenericError("Failed to convert number_of_evaluations to u8".to_string())
-            })?;
-        params.max_fft_width = 1_u64 << log2_of_evals;
-
-        if params.chunk_length == 1 {
-            log2_of_evals = (2 * params.num_chunks)
-                .to_f64()
-                .ok_or_else(|| {
-                    KzgError::GenericError("Failed to convert num_chunks to f64".to_string())
-                })?
-                .log2()
-                .to_u8()
-                .ok_or_else(|| {
-                    KzgError::GenericError("Failed to convert num_chunks to u8".to_string())
-                })?;
-        }
-
-        if params.chunk_length * params.num_chunks >= self.srs_order {
-            return Err(KzgError::SerializationError(
-                "the supplied encoding parameters are not valid with respect to the SRS."
-                    .to_string(),
-            ));
-        }
-
-        let primitive_roots_of_unity = Self::get_primitive_roots_of_unity()?;
-        let found_root_of_unity = primitive_roots_of_unity
-            .get(log2_of_evals.to_usize().ok_or_else(|| {
-                KzgError::GenericError("Failed to convert log2_of_evals to usize".to_string())
-            })?)
-            .ok_or_else(|| KzgError::GenericError("Root of unity not found".to_string()))?;
-        let mut expanded_roots_of_unity = Self::expand_root_of_unity(found_root_of_unity);
-        expanded_roots_of_unity.truncate(expanded_roots_of_unity.len() - 1);
-
-        params.completed_setup = true;
-        self.params = params;
-        self.expanded_roots_of_unity = expanded_roots_of_unity;
-
-        Ok(())
-    }
-
-    pub fn calculate_roots_of_unity(
+    /// Calculates the roots of unities and assigns it to the struct
+    ///
+    /// # Arguments
+    /// * `length_of_data_after_padding` - Length of the blob data after padding in bytes.
+    ///
+    /// # Returns
+    /// * `Result<(), KzgError>`
+    ///
+    /// # Details
+    /// - Generates roots of unity needed for FFT operations
+    ///
+    /// # Example
+    /// ```
+    /// use rust_kzg_bn254::kzg::KZG;
+    /// use rust_kzg_bn254::blob::Blob;
+    /// use ark_std::One;
+    /// use ark_bn254::Fr;
+    ///
+    /// let mut kzg = KZG::setup(
+    ///        "tests/test-files/mainnet-data/g1.131072.point",
+    ///        "",
+    ///        "tests/test-files/mainnet-data/g2.point.powerOf2",
+    ///        268435456,
+    ///        131072,
+    ///    ).unwrap();
+    /// let input_blob = Blob::from_raw_data(b"test blob data");
+    /// kzg.calculate_and_store_roots_of_unity(input_blob.len().try_into().unwrap()).unwrap();
+    /// ```
+    pub fn calculate_and_store_roots_of_unity(
         &mut self,
         length_of_data_after_padding: u64,
     ) -> Result<(), KzgError> {
-        let (params, roots_of_unity) = Self::calculate_roots_of_unity_standalone(
-            length_of_data_after_padding,
-            self.srs_order,
-        )?;
-        self.params = params;
-        self.params.completed_setup = true;
+        let roots_of_unity =
+            Self::calculate_roots_of_unity(length_of_data_after_padding, self.srs_order)?;
+        self.completed_setup = true;
         self.expanded_roots_of_unity = roots_of_unity;
         Ok(())
     }
@@ -532,7 +438,7 @@ impl KZG {
         polynomial: &PolynomialEvalForm,
         z_fr: &Fr,
     ) -> Result<G1Affine, KzgError> {
-        if !self.params.completed_setup {
+        if !self.completed_setup {
             return Err(KzgError::GenericError(
                 "setup is not complete, run the data_setup functions".to_string(),
             ));
@@ -623,7 +529,7 @@ impl KZG {
         polynomial: &PolynomialEvalForm,
         z_fr: &Fr,
     ) -> Result<G1Affine, KzgError> {
-        if !self.params.completed_setup {
+        if !self.completed_setup {
             return Err(KzgError::GenericError(
                 "setup is not complete, run one of the setup functions".to_string(),
             ));
@@ -937,8 +843,7 @@ impl KZG {
         let blob_size = polynomial.len_underlying_blob_bytes();
 
         // Step 2: Calculate roots of unity for the given blob size and SRS order
-        let (_, roots_of_unity) =
-            Self::calculate_roots_of_unity_standalone(blob_size as u64, srs_order)?;
+        let roots_of_unity = Self::calculate_roots_of_unity(blob_size as u64, srs_order)?;
 
         // Step 3: Ensure the polynomial length matches the domain length
         if polynomial.len() != roots_of_unity.len() {