rust_kzg_bn254_prover/
kzg.rs

1use ark_bn254::{Fr, G1Affine, G1Projective};
2use ark_ec::{CurveGroup, VariableBaseMSM};
3use ark_poly::{EvaluationDomain, GeneralEvaluationDomain};
4use ark_std::{ops::Div, Zero};
5use num_traits::ToPrimitive;
6use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
7use rust_kzg_bn254_primitives::{
8    blob::Blob,
9    errors::KzgError,
10    helpers,
11    polynomial::{PolynomialCoeffForm, PolynomialEvalForm},
12};
13
14use crate::srs::SRS;
15
16/// Main interesting struct of the rust-kzg-bn254 crate.
17/// [Kzg] is a struct that holds the SRS points in monomial form, and
18/// provides methods for committing to a blob, (either via a [Blob] itself,
19/// or a [PolynomialCoeffForm] or [PolynomialEvalForm]), and generating and
20/// verifying proofs.
21///
22/// The [Blob] and [PolynomialCoeffForm]/[PolynomialEvalForm] structs are mostly
23/// <https://en.wikipedia.org/wiki/Passive_data_structure> with
24/// constructor and few helper methods.
25#[derive(Debug, PartialEq, Clone)]
26pub struct KZG {
27    expanded_roots_of_unity: Vec<Fr>,
28}
29
30impl Default for KZG {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36impl KZG {
37    pub fn new() -> Self {
38        Self {
39            expanded_roots_of_unity: vec![],
40        }
41    }
42
43    /// Calculates the roots of unities and assigns it to the struct
44    ///
45    /// # Arguments
46    /// * `length_of_data_after_padding` - Length of the blob data after padding in bytes.
47    ///
48    /// # Returns
49    /// * `Result<(), KzgError>`
50    ///
51    /// # Details
52    /// - Generates roots of unity needed for FFT operations
53    ///
54    /// # Example
55    /// ```
56    /// use rust_kzg_bn254_prover::kzg::KZG;
57    /// use rust_kzg_bn254_primitives::blob::Blob;
58    /// use ark_std::One;
59    /// use ark_bn254::Fr;
60    ///
61    /// let mut kzg = KZG::new();
62    /// let input_blob = Blob::from_raw_data(b"test blob data");
63    /// kzg.calculate_and_store_roots_of_unity(input_blob.len().try_into().unwrap()).unwrap();
64    /// ```
65    pub fn calculate_and_store_roots_of_unity(
66        &mut self,
67        length_of_data_after_padding: u64,
68    ) -> Result<(), KzgError> {
69        let roots_of_unity = helpers::calculate_roots_of_unity(length_of_data_after_padding)?;
70        self.expanded_roots_of_unity = roots_of_unity;
71        Ok(())
72    }
73
74    pub fn get_roots_of_unities(&self) -> Vec<Fr> {
75        self.expanded_roots_of_unity.clone()
76    }
77
78    /// helper function to get the
79    pub fn get_nth_root_of_unity(&self, i: usize) -> Option<&Fr> {
80        self.expanded_roots_of_unity.get(i)
81    }
82
83    /// Commit the polynomial with the srs values loaded into [Kzg].
84    pub fn commit_eval_form(
85        &self,
86        polynomial: &PolynomialEvalForm,
87        srs: &SRS,
88    ) -> Result<G1Affine, KzgError> {
89        if polynomial.len() > srs.g1.len() {
90            return Err(KzgError::SerializationError(
91                "polynomial length is not correct".to_string(),
92            ));
93        }
94
95        // When the polynomial is in evaluation form, use IFFT to transform monomial srs
96        // points to lagrange form.
97        let bases = self.g1_ifft(polynomial.len(), srs)?;
98
99        match G1Projective::msm(&bases, polynomial.evaluations()) {
100            Ok(res) => Ok(res.into_affine()),
101            Err(err) => Err(KzgError::CommitError(err.to_string())),
102        }
103    }
104
105    /// Commit the polynomial with the srs values loaded into [Kzg].
106    pub fn commit_coeff_form(
107        &self,
108        polynomial: &PolynomialCoeffForm,
109        srs: &SRS,
110    ) -> Result<G1Affine, KzgError> {
111        if polynomial.len() > srs.g1.len() {
112            return Err(KzgError::SerializationError(
113                "polynomial length is not correct".to_string(),
114            ));
115        }
116        // When the polynomial is in coefficient form, use the original srs points (in
117        // monomial form).
118        let bases = srs.g1[..polynomial.len()].to_vec();
119
120        match G1Projective::msm(&bases, polynomial.coeffs()) {
121            Ok(res) => Ok(res.into_affine()),
122            Err(err) => Err(KzgError::CommitError(err.to_string())),
123        }
124    }
125
126    /// Helper function for `compute_kzg_proof()` and `compute_blob_kzg_proof()`
127    fn compute_proof_impl(
128        &self,
129        polynomial: &PolynomialEvalForm,
130        z_fr: &Fr,
131        srs: &SRS,
132    ) -> Result<G1Affine, KzgError> {
133        // Verify polynomial length matches that of the roots of unity
134        if polynomial.len() != self.expanded_roots_of_unity.len() {
135            return Err(KzgError::GenericError(
136                "inconsistent length between blob and root of unities".to_string(),
137            ));
138        }
139
140        let eval_fr = polynomial.evaluations();
141        // Pre-allocate vector for shifted polynomial p(x) - y
142        let mut poly_shift: Vec<Fr> = Vec::with_capacity(eval_fr.len());
143
144        // Evaluate polynomial at the point z
145        // This gives us y = p(z)
146        let y_fr = helpers::evaluate_polynomial_in_evaluation_form(polynomial, z_fr)?;
147
148        // Compute p(x) - y for each evaluation point
149        // This is the numerator of the quotient polynomial
150        for fr in eval_fr {
151            poly_shift.push(*fr - y_fr);
152        }
153
154        // Compute denominator polynomial (x - z) at each root of unity
155        let mut denom_poly = Vec::<Fr>::with_capacity(self.expanded_roots_of_unity.len());
156        for root_of_unity in self.expanded_roots_of_unity.iter().take(eval_fr.len()) {
157            denom_poly.push(*root_of_unity - z_fr);
158        }
159
160        // Pre-allocate vector for quotient polynomial evaluations
161        let mut quotient_poly = Vec::<Fr>::with_capacity(self.expanded_roots_of_unity.len());
162
163        // Compute quotient polynomial q(x) = (p(x) - y)/(x - z) at each root of unity
164        for i in 0..self.expanded_roots_of_unity.len() {
165            if denom_poly[i].is_zero() {
166                // Special case: when x = z, use L'Hôpital's rule
167                // Compute the derivative evaluation instead
168                quotient_poly.push(self.compute_quotient_eval_on_domain(z_fr, eval_fr, &y_fr));
169            } else {
170                // Normal case: direct polynomial division
171                quotient_poly.push(poly_shift[i].div(denom_poly[i]));
172            }
173        }
174
175        let quotient_poly_eval_form = PolynomialEvalForm::new(quotient_poly);
176        self.commit_eval_form(&quotient_poly_eval_form, srs)
177    }
178
179    /// commit to a [Blob], by transforming it into a [PolynomialEvalForm] and
180    /// then calling [Kzg::commit_eval_form].
181    pub fn commit_blob(&self, blob: &Blob, srs: &SRS) -> Result<G1Affine, KzgError> {
182        let polynomial = blob.to_polynomial_eval_form();
183        self.commit_eval_form(&polynomial, srs)
184    }
185
186    pub fn compute_proof_with_known_z_fr_index(
187        &self,
188        polynomial: &PolynomialEvalForm,
189        index: u64,
190        srs: &SRS,
191    ) -> Result<G1Affine, KzgError> {
192        // Convert u64 index to usize for array indexing
193        let usized_index = index.to_usize().ok_or(KzgError::GenericError(
194            "Index conversion to usize failed".to_string(),
195        ))?;
196
197        // Get the root of unity at the specified index
198        let z_fr = self
199            .get_nth_root_of_unity(usized_index)
200            .ok_or_else(|| KzgError::GenericError("Root of unity not found".to_string()))?;
201
202        // Compute the KZG proof at the selected root of unity
203        // This delegates to the main proof computation function
204        // using our selected evaluation point
205        self.compute_proof(polynomial, z_fr, srs)
206    }
207
208    /// Compute a kzg proof from a polynomial in evaluation form.
209    /// We don't currently support proofs for polynomials in coefficient form,
210    /// but one can take the FFT of the polynomial in coefficient form to
211    /// get the polynomial in evaluation form. This is available via the
212    /// method [PolynomialCoeffForm::to_eval_form].
213    /// TODO(anupsv): Accept bytes instead of Fr element. Ref: https://github.com/Layr-Labs/rust-kzg-bn254/issues/29
214    pub fn compute_proof(
215        &self,
216        polynomial: &PolynomialEvalForm,
217        z_fr: &Fr,
218        srs: &SRS,
219    ) -> Result<G1Affine, KzgError> {
220        // Verify that polynomial length matches roots of unity length
221        if polynomial.len() != self.expanded_roots_of_unity.len() {
222            return Err(KzgError::GenericError(
223                "inconsistent length between blob and root of unities".to_string(),
224            ));
225        }
226
227        // Call the implementation to compute the actual proof
228        // This will:
229        // 1. Evaluate polynomial at z
230        // 2. Compute quotient polynomial q(x) = (p(x) - p(z)) / (x - z)
231        // 3. Generate KZG proof as commitment to q(x)
232        self.compute_proof_impl(polynomial, z_fr, srs)
233    }
234
235    /// refer to DA for more context
236    pub fn compute_quotient_eval_on_domain(&self, z_fr: &Fr, eval_fr: &[Fr], value_fr: &Fr) -> Fr {
237        let mut quotient = Fr::zero();
238        let mut fi: Fr = Fr::zero();
239        let mut numerator: Fr = Fr::zero();
240        let mut denominator: Fr = Fr::zero();
241        let mut temp: Fr = Fr::zero();
242
243        self.expanded_roots_of_unity
244            .iter()
245            .enumerate()
246            .for_each(|(i, omega_i)| {
247                if *omega_i == *z_fr {
248                    return;
249                }
250                fi = eval_fr[i] - value_fr;
251                numerator = fi * omega_i;
252                denominator = z_fr - omega_i;
253                denominator *= z_fr;
254                temp = numerator.div(denominator);
255                quotient += temp;
256            });
257
258        quotient
259    }
260
261    /// function to compute the inverse FFT
262    pub fn g1_ifft(&self, length: usize, srs: &SRS) -> Result<Vec<G1Affine>, KzgError> {
263        // is not power of 2
264        if !length.is_power_of_two() {
265            return Err(KzgError::FFTError(
266                "length provided is not a power of 2".to_string(),
267            ));
268        }
269
270        let points_projective: Vec<G1Projective> = srs.g1[..length]
271            .par_iter()
272            .map(|&p| G1Projective::from(p))
273            .collect();
274        let ifft_result: Vec<_> = GeneralEvaluationDomain::<Fr>::new(length)
275            .ok_or(KzgError::FFTError(
276                "Could not perform IFFT due to domain consturction error".to_string(),
277            ))?
278            .ifft(&points_projective)
279            .par_iter()
280            .map(|p| p.into_affine())
281            .collect();
282
283        Ok(ifft_result)
284    }
285
286    /// TODO(anupsv): Match 4844 specs w.r.t to the inputs. Ref: https://github.com/Layr-Labs/rust-kzg-bn254/issues/30
287    pub fn compute_blob_proof(
288        &self,
289        blob: &Blob,
290        commitment: &G1Affine,
291        srs: &SRS,
292    ) -> Result<G1Affine, KzgError> {
293        // Validate that the commitment is a valid point on the G1 curve
294        // This prevents potential invalid curve attacks
295        if !commitment.is_on_curve() || !commitment.is_in_correct_subgroup_assuming_on_curve() {
296            return Err(KzgError::NotOnCurveError(
297                "commitment not on curve".to_string(),
298            ));
299        }
300
301        // Convert the blob to a polynomial in evaluation form
302        // This is necessary because KZG proofs work with polynomials
303        let blob_poly = blob.to_polynomial_eval_form();
304
305        // Compute the evaluation challenge using Fiat-Shamir heuristic
306        // This challenge determines the point at which we evaluate the polynomial
307        let evaluation_challenge = helpers::compute_challenge(blob, commitment)?;
308
309        // Compute the actual KZG proof using the polynomial and evaluation point
310        // This creates a proof that the polynomial evaluates to a specific value at the challenge point
311        // The proof is a single G1 point that can be used to verify the evaluation
312        self.compute_proof_impl(&blob_poly, &evaluation_challenge, srs)
313    }
314}
rust_kzg_bn254_prover/kzg.rs

rust_kzg_bn254_prover/
kzg.rs
