spawn-zk-snarks 0.1.5

// Copyright (c) 2023 spawn-zk-snarks developers
//
// Licensed under the MIT License
// <LICENSE-MIT or http://opensource.org/licenses/MIT>

use ark_bn254::{Bn254, Fr};
use ark_groth16::{
    PreparedVerifyingKey, Proof as ArkProof, ProvingKey, VerifyingKey,
};
use ark_snark::SNARK;
use ark_groth16::Groth16;
use ark_ec::pairing::Pairing;
use ark_serialize::CanonicalSerialize;
use ark_std::rand::thread_rng;
use thiserror::Error;
use crate::circuits::{CircuitError, ArithmeticCircuit};

/// Represents all possible errors that can occur during Groth16 operations.
/// 
/// This enum provides a comprehensive set of errors that might occur during
/// setup, proving, verification, and other related operations.
#[derive(Error, Debug)]
pub enum Groth16Error {
    /// Occurs when the setup phase fails, typically due to invalid parameters
    /// or insufficient randomness.
    #[error("Setup failed")]
    SetupError,

    /// Occurs when proof generation fails, usually due to invalid witness
    /// or constraint system issues.
    #[error("Proving failed")]
    ProvingError,

    /// Occurs when proof verification fails, indicating either an invalid proof
    /// or mismatched public inputs.
    #[error("Verification failed")]
    VerificationError,

    /// Wraps an error from the underlying circuit implementation.
    /// This can occur during constraint generation or witness computation.
    #[error("Circuit error: {0}")]
    CircuitError(#[from] CircuitError),

    /// Occurs when serialization or deserialization of proofs, keys,
    /// or other data structures fails.
    #[error("Serialization error")]
    SerializationError,
}

/// Contains all necessary parameters for the Groth16 proving system.
/// 
/// This structure holds the proving key, verifying key, and prepared verifying key
/// needed for generating and verifying zero-knowledge proofs.
#[derive(Debug)]
pub struct Groth16Setup<E: Pairing> {
    /// The proving key used to generate proofs. This contains the secret parameters
    /// and should be kept private in some applications.
    pub proving_key: ProvingKey<E>,

    /// The verification key used to verify proofs. This can be made public
    /// and distributed to verifiers.
    pub verifying_key: VerifyingKey<E>,

    /// A preprocessed version of the verification key that allows for
    /// more efficient proof verification.
    pub prepared_verifying_key: PreparedVerifyingKey<E>,
}

impl Groth16Setup<Bn254> {
    /// Creates a new Groth16 setup for a circuit with the specified parameters.
    /// 
    /// This function generates the proving and verification keys necessary for
    /// creating and verifying zero-knowledge proofs.
    /// 
    /// # Arguments
    /// * `num_constraints` - Number of constraints in the arithmetic circuit
    /// * `num_variables` - Number of variables used in the circuit
    /// 
    /// # Returns
    /// * `Result<Self, Groth16Error>` - The setup parameters or an error
    /// 
    /// # Example
    /// ```
    /// # use spawn_zk_snarks::Groth16Setup;
    /// let setup = Groth16Setup::new(3, 2)?;
    /// ```
    pub fn new(num_constraints: usize, num_variables: usize) -> Result<Self, Groth16Error> {
        let circuit: ArithmeticCircuit<Fr> = ArithmeticCircuit::new(num_constraints, num_variables, 1);
        let rng = &mut thread_rng();
        
        let (pk, vk) = Groth16::<Bn254>::circuit_specific_setup(
            circuit,
            rng,
        ).map_err(|_| Groth16Error::SetupError)?;

        let pvk = match Groth16::<Bn254>::process_vk(&vk) {
            Ok(pvk) => pvk,
            Err(_) => return Err(Groth16Error::SetupError),
        };

        Ok(Self {
            proving_key: pk,
            verifying_key: vk,
            prepared_verifying_key: pvk,
        })
    }

    /// Generates a zero-knowledge proof using the provided inputs and witness.
    /// 
    /// # Arguments
    /// * `inputs` - Public inputs to the circuit
    /// * `witness` - Private witness values
    /// 
    /// # Returns
    /// * `Result<ArkProof<Bn254>, Groth16Error>` - The generated proof or an error
    /// 
    /// # Example
    /// ```
    /// # use spawn_zk_snarks::{Groth16Setup, Fr};
    /// # let setup = Groth16Setup::new(3, 2)?;
    /// let proof = setup.prove(&inputs, &witness)?;
    /// ```
    pub fn prove(&self, inputs: &[Fr], witness: &[Fr]) -> Result<ArkProof<Bn254>, Groth16Error> {
        let rng = &mut thread_rng();
        
        let circuit = ArithmeticCircuit {
            num_constraints: inputs.len(),
            num_variables: witness.len(),
            num_inputs: 1,
            witness: witness.to_vec(),
        };

        Groth16::<Bn254>::create_random_proof_with_reduction(
            circuit,
            &self.proving_key,
            rng,
        ).map_err(|_| Groth16Error::ProvingError)
    }

    /// Verifies a zero-knowledge proof against the provided public inputs.
    /// 
    /// # Arguments
    /// * `proof` - The proof to verify
    /// * `public_inputs` - Public inputs used in the proof
    /// 
    /// # Returns
    /// * `Result<bool, Groth16Error>` - Whether the proof is valid
    /// 
    /// # Example
    /// ```
    /// # use spawn_zk_snarks::Groth16Setup;
    /// # let setup = Groth16Setup::new(3, 2)?;
    /// let is_valid = setup.verify(&proof, &public_inputs)?;
    /// assert!(is_valid);
    /// ```
    pub fn verify(
        &self,
        proof: &ArkProof<Bn254>,
        public_inputs: &[Fr],
    ) -> Result<bool, Groth16Error> {
        <Groth16<Bn254> as SNARK<Fr>>::verify_with_processed_vk(
            &self.prepared_verifying_key,
            public_inputs,
            proof,
        ).map_err(|_| Groth16Error::VerificationError)
    }

    /// Converts a proof to EVM-compatible format for on-chain verification.
    /// 
    /// # Arguments
    /// * `proof` - The proof to convert
    /// 
    /// # Returns
    /// * `Result<Vec<u8>, Groth16Error>` - The serialized proof bytes
    pub fn proof_to_evm_format(proof: &ArkProof<Bn254>) -> Result<Vec<u8>, Groth16Error> {
        let mut bytes = Vec::new();
        proof.serialize_compressed(&mut bytes)
            .map_err(|_| Groth16Error::SerializationError)?;
        Ok(bytes)
    }

    /// Estimates the gas cost for verifying a proof on Ethereum.
    /// 
    /// # Arguments
    /// * `proof_size` - Size of the proof in bytes
    /// * `num_inputs` - Number of public inputs
    /// 
    /// # Returns
    /// * `u64` - Estimated gas cost in units
    pub fn estimate_verification_gas(proof_size: usize, num_inputs: usize) -> u64 {
        let base_cost = 150_000;
        let input_cost = num_inputs as u64 * 1_000;
        let data_cost = proof_size as u64 * 16;
        base_cost + input_cost + data_cost
    }
}