samaharam 0.1.0

//! Adapter for jf-plonk proofs (Jellyfish).
//!
//! Supports proofs generated by jf-plonk (https://github.com/EspressoSystems/jellyfish)
//! as used by EspressoSystems projects.

use super::external::{AdapterError, ExternalProof, ProofMetadata};
use super::parsing;
use crate::backend::bn254::Bn254;
use crate::crypto::AccumulatorInstance;
use group::Curve;
use halo2curves::bn256::{Fr, G1Affine, G1};

/// PLONK proof from jf-plonk (Jellyfish).
///
/// jf-plonk uses UltraPlonk with custom gates, making it more
/// complex than standard PLONK.
#[derive(Debug, Clone)]
pub struct JfPlonkProof {
    /// Wire polynomial commitments (typically 5 for UltraPlonk).
    pub wire_commits: Vec<G1Affine>,
    /// Permutation polynomial commitment.
    pub prod_perm_poly_commit: G1Affine,
    /// Split quotient polynomial commitments.
    pub split_quot_poly_commits: Vec<G1Affine>,
    /// Opening proof at evaluation point.
    pub opening_proof: G1Affine,
    /// Shifted opening proof (for permutation).
    pub shifted_opening_proof: G1Affine,
    /// Polynomial evaluations at zeta.
    pub poly_evals: Vec<Fr>,
    /// Public inputs.
    pub public_inputs: Vec<Fr>,
}

impl JfPlonkProof {
    /// Create from components.
    pub fn new(
        wire_commits: Vec<G1Affine>,
        prod_perm_poly_commit: G1Affine,
        split_quot_poly_commits: Vec<G1Affine>,
        opening_proof: G1Affine,
        shifted_opening_proof: G1Affine,
        poly_evals: Vec<Fr>,
        public_inputs: Vec<Fr>,
    ) -> Self {
        Self {
            wire_commits,
            prod_perm_poly_commit,
            split_quot_poly_commits,
            opening_proof,
            shifted_opening_proof,
            poly_evals,
            public_inputs,
        }
    }

    /// Create a mock proof for testing or as placeholder.
    pub fn mock(num_wires: usize, public_inputs: Vec<Fr>) -> Self {
        let g1 = G1::generator().to_affine();
        let g1_2 = (G1::generator() * Fr::from(2u64)).to_affine();

        Self {
            wire_commits: vec![g1; num_wires],
            prod_perm_poly_commit: g1_2,
            split_quot_poly_commits: vec![g1, g1, g1],
            opening_proof: g1,
            shifted_opening_proof: g1_2,
            poly_evals: vec![Fr::from(1u64); num_wires + 1], // +1 for permutation eval
            public_inputs,
        }
    }

    /// Parse from jf-plonk serialized format.
    ///
    /// jf-plonk proof layout (typical):
    /// - Wire commits: variable (N * 64 bytes)
    /// - Perm commit: 64 bytes
    /// - Quotient commits: 3 * 64 bytes
    /// - Opening: 64 bytes
    /// - Shifted opening: 64 bytes
    /// - Evaluations: variable (M * 32 bytes)
    pub fn from_bytes(data: &[u8], num_wires: usize) -> Result<Self, AdapterError> {
        const G1_SIZE: usize = 64;
        const FR_SIZE: usize = 32;

        let min_size = G1_SIZE * (num_wires + 5); // wires + perm + 3 quotients + opening + shifted
        if data.len() < min_size {
            return Err(AdapterError::InvalidFormat(format!(
                "Data too short for jf-plonk: {} < {}",
                data.len(),
                min_size
            )));
        }

        let mut offset = 0;

        // Parse wire commitments
        let mut wire_commits = Vec::with_capacity(num_wires);
        for i in 0..num_wires {
            let wire = parsing::parse_g1_uncompressed(&data[offset..offset + G1_SIZE])
                .map_err(|e| AdapterError::InvalidPoint(format!("wire[{}]: {}", i, e)))?;
            wire_commits.push(wire);
            offset += G1_SIZE;
        }

        // Parse permutation commitment
        let prod_perm_poly_commit = parsing::parse_g1_uncompressed(&data[offset..offset + G1_SIZE])
            .map_err(|e| AdapterError::InvalidPoint(format!("perm: {}", e)))?;
        offset += G1_SIZE;

        // Parse quotient commitments (typically 3)
        let mut split_quot_poly_commits = Vec::with_capacity(3);
        for i in 0..3 {
            let quot = parsing::parse_g1_uncompressed(&data[offset..offset + G1_SIZE])
                .map_err(|e| AdapterError::InvalidPoint(format!("quot[{}]: {}", i, e)))?;
            split_quot_poly_commits.push(quot);
            offset += G1_SIZE;
        }

        // Parse opening proof
        let opening_proof = parsing::parse_g1_uncompressed(&data[offset..offset + G1_SIZE])
            .map_err(|e| AdapterError::InvalidPoint(format!("opening: {}", e)))?;
        offset += G1_SIZE;

        // Parse shifted opening proof
        let shifted_opening_proof = parsing::parse_g1_uncompressed(&data[offset..offset + G1_SIZE])
            .map_err(|e| AdapterError::InvalidPoint(format!("shifted: {}", e)))?;
        offset += G1_SIZE;

        // Parse evaluations (remaining bytes / 32)
        let mut poly_evals = Vec::new();
        while offset + FR_SIZE <= data.len() {
            let eval = parsing::parse_fr_bytes(&data[offset..offset + FR_SIZE])
                .map_err(|e| AdapterError::ParseError(format!("eval: {}", e)))?;
            poly_evals.push(eval);
            offset += FR_SIZE;
        }

        Ok(Self {
            wire_commits,
            prod_perm_poly_commit,
            split_quot_poly_commits,
            opening_proof,
            shifted_opening_proof,
            poly_evals,
            public_inputs: vec![],
        })
    }



    /// Parse from custom jf-plonk format.
    ///
    /// This format uses 5 wires in its UltraPlonk configuration.
    pub fn from_jf_bytes(data: &[u8]) -> Result<Self, AdapterError> {
        Self::from_bytes(data, 5)
    }

    /// Set public inputs after parsing.
    pub fn with_public_inputs(mut self, inputs: Vec<Fr>) -> Self {
        self.public_inputs = inputs;
        self
    }
}

impl ExternalProof<Bn254> for JfPlonkProof {
    fn to_accumulator_instances(&self) -> Result<Vec<AccumulatorInstance<Bn254>>, AdapterError> {
        if self.wire_commits.is_empty() {
            return Err(AdapterError::InvalidFormat("No wire commitments".to_string()));
        }

        let mut instances = Vec::new();

        // Evaluation point (would come from transcript in real impl)
        let zeta = Fr::from(7u64);

        // Create instance for each wire polynomial
        for (i, wire_commit) in self.wire_commits.iter().enumerate() {
            let eval = self.poly_evals.get(i)
                .copied()
                .ok_or_else(|| AdapterError::InvalidFormat(
                    format!("Missing evaluation for wire {}", i)
                ))?;
            instances.push(AccumulatorInstance {
                commitment: *wire_commit,
                evaluation: eval,
                point: zeta,
                quotient: self.opening_proof,
            });
        }

        // Add permutation polynomial instance
        let perm_eval = self.poly_evals.get(self.wire_commits.len())
            .copied()
            .ok_or_else(|| AdapterError::InvalidFormat(
                "Missing permutation evaluation".to_string()
            ))?;
        instances.push(AccumulatorInstance {
            commitment: self.prod_perm_poly_commit,
            evaluation: perm_eval,
            point: zeta,
            quotient: self.shifted_opening_proof,
        });

        Ok(instances)
    }

    fn public_inputs(&self) -> &[Fr] {
        &self.public_inputs
    }

    fn metadata(&self) -> ProofMetadata {
        ProofMetadata {
            system: "jf-plonk",
            proof_type: "ultraplonk",
            curve: "bn254",
            num_public_inputs: self.public_inputs.len(),
        }
    }

    fn validate_format(&self) -> Result<(), AdapterError> {
        use group::prime::PrimeCurveAffine;

        if self.wire_commits.is_empty() {
            return Err(AdapterError::InvalidFormat("No wire commitments".to_string()));
        }

        // Validate wire commitments
        for (i, comm) in self.wire_commits.iter().enumerate() {
            if comm.is_identity().into() {
                return Err(AdapterError::InvalidPoint(format!(
                    "wire_commits[{}] is identity",
                    i
                )));
            }
        }

        // Validate permutation commitment
        if self.prod_perm_poly_commit.is_identity().into() {
            return Err(AdapterError::InvalidPoint(
                "prod_perm_poly_commit is identity".to_string(),
            ));
        }

        // Validate opening proofs
        if self.opening_proof.is_identity().into() {
            return Err(AdapterError::InvalidPoint("opening_proof is identity".to_string()));
        }

        // Check evals match wire count (at minimum)
        if self.poly_evals.len() < self.wire_commits.len() {
            return Err(AdapterError::InvalidFormat(format!(
                "Not enough evaluations: {} < {}",
                self.poly_evals.len(),
                self.wire_commits.len()
            )));
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn jf_plonk_proof_creation() {
        let proof = JfPlonkProof::mock(5, vec![Fr::from(42u64)]);

        assert_eq!(proof.wire_commits.len(), 5);
        assert_eq!(proof.public_inputs().len(), 1);
        assert_eq!(proof.metadata().system, "jf-plonk");
        assert_eq!(proof.metadata().proof_type, "ultraplonk");
    }

    #[test]
    fn jf_plonk_to_accumulator() {
        let proof = JfPlonkProof::mock(5, vec![Fr::from(100u64)]);

        let instances = proof.to_accumulator_instances().unwrap();
        // 5 wires + 1 permutation = 6 instances
        assert_eq!(instances.len(), 6);
    }

    #[test]
    fn jf_plonk_validate_format() {
        let proof = JfPlonkProof::mock(3, vec![]);
        assert!(proof.validate_format().is_ok());
    }

    #[test]
    fn jf_plonk_validate_empty_wires_fails() {
        let proof = JfPlonkProof::new(
            vec![],
            G1::generator().to_affine(),
            vec![],
            G1::generator().to_affine(),
            G1::generator().to_affine(),
            vec![],
            vec![],
        );

        assert!(proof.validate_format().is_err());
    }

    #[test]
    fn jf_plonk_with_public_inputs() {
        let proof = JfPlonkProof::mock(3, vec![])
            .with_public_inputs(vec![Fr::from(1u64), Fr::from(2u64)]);

        assert_eq!(proof.public_inputs.len(), 2);
    }
}