samaharam 0.2.0

//! Proof data structure with TypeState.

use std::marker::PhantomData;

use super::state::{Batched, Pending, ProofState, Verified};
use crate::error::Error;
use crate::registry::VkId;
use crate::traits::PairingEngine;

/// A zero-knowledge proof with compile-time state tracking.
///
/// The type parameter `S` tracks the proof's lifecycle state:
/// - `Pending`: Just created, needs verification
/// - `Verified`: Verified against VK, can be submitted
/// - `Batched`: In aggregation queue
/// - `Aggregated`: Final aggregated proof
///
/// # Example
///
/// ```rust,ignore
/// let pending = Proof::<Bn254, Pending>::new(data, inputs, vk_id);
/// let verified = pending.verify(&registry)?; // Pending -> Verified
/// let batched = verified.submit();           // Verified -> Batched
/// ```
#[derive(Debug)]
pub struct Proof<E: PairingEngine, S: ProofState> {
    /// Raw proof data (serialized proof bytes)
    data: Vec<u8>,

    /// Public inputs for this proof
    public_inputs: Vec<E::Fr>,

    /// ID of the verification key used for this proof
    vk_id: VkId,

    /// Zero-sized state marker
    _state: PhantomData<S>,
}

impl<E: PairingEngine, S: ProofState> Proof<E, S> {
    /// Get the verification key ID for this proof.
    pub fn vk_id(&self) -> VkId {
        self.vk_id
    }

    /// Get the public inputs for this proof.
    pub fn public_inputs(&self) -> &[E::Fr] {
        &self.public_inputs
    }

    /// Get the raw proof data.
    pub fn data(&self) -> &[u8] {
        &self.data
    }
}

impl<E: PairingEngine> Proof<E, Pending> {
    /// Create a new pending proof.
    ///
    /// # Arguments
    ///
    /// * `data` - Serialized proof bytes
    /// * `public_inputs` - Public inputs for verification
    /// * `vk_id` - ID of the registered verification key
    pub fn new(data: Vec<u8>, public_inputs: Vec<E::Fr>, vk_id: VkId) -> Self {
        Self {
            data,
            public_inputs,
            vk_id,
            _state: PhantomData,
        }
    }

    /// Verify this proof against its registered VK.
    ///
    /// Transitions the proof from `Pending` to `Verified` state.
    ///
    /// # Verification Steps
    ///
    /// 1. Checks that the VK exists in the registry
    /// 2. Validates public input count against VK limits
    /// 3. Deserializes proof data into PlonkProof
    /// 4. Runs cryptographic verification via PlonkVerifier
    ///
    /// # Errors
    ///
    /// - `Error::UnknownVk` if the VK is not registered
    /// - `Error::PublicInputMismatch` if input count exceeds limit
    /// - `Error::VerificationFailed` if proof is invalid or cryptographic check fails
    pub fn verify(
        self,
        registry: &crate::registry::VkRegistry<E>,
    ) -> Result<Proof<E, Verified>, Error> {
        use crate::crypto::{PlonkProof, PlonkVerifier};

        // 1. Check VK exists in registry
        let registered_vk = registry.require(self.vk_id)?;

        // 2. Validate public input count
        if self.public_inputs.len() > registered_vk.vk.num_public_inputs {
            return Err(Error::PublicInputMismatch {
                expected: registered_vk.vk.num_public_inputs,
                got: self.public_inputs.len(),
            });
        }

        // 3. Validate proof data is non-empty
        if self.data.is_empty() {
            return Err(Error::VerificationFailed("empty proof data".into()));
        }

        // 4. Deserialize proof
        let plonk_proof = PlonkProof::<E>::from_bytes(&self.data).map_err(|e| {
            Error::VerificationFailed(format!("proof deserialization failed: {}", e))
        })?;

        // 5. Cryptographic verification via PlonkVerifier
        let verifier = PlonkVerifier::new(registered_vk.vk.clone());
        let is_valid = verifier
            .verify(&plonk_proof, &self.public_inputs)
            .map_err(|e| Error::VerificationFailed(format!("verification error: {}", e)))?;

        if !is_valid {
            return Err(Error::VerificationFailed("pairing check failed".into()));
        }

        // 6. Transition to Verified state
        Ok(Proof {
            data: self.data,
            public_inputs: self.public_inputs,
            vk_id: self.vk_id,
            _state: PhantomData,
        })
    }
}

impl<E: PairingEngine> Proof<E, Verified> {
    /// Create a verified proof directly (for testing/internal use).
    ///
    /// # Safety
    ///
    /// This bypasses verification - only use for testing.
    #[allow(dead_code)]
    pub(crate) fn new_verified(data: Vec<u8>, public_inputs: Vec<E::Fr>, vk_id: VkId) -> Self {
        Self {
            data,
            public_inputs,
            vk_id,
            _state: PhantomData,
        }
    }

    /// Submit this verified proof for aggregation.
    ///
    /// Transitions the proof from `Verified` to `Batched` state.
    pub fn submit(self) -> Proof<E, Batched> {
        Proof {
            data: self.data,
            public_inputs: self.public_inputs,
            vk_id: self.vk_id,
            _state: PhantomData,
        }
    }
}

impl<E: PairingEngine> Proof<E, super::state::Aggregated> {
    /// Create a new aggregated proof.
    pub(crate) fn new_aggregated(data: Vec<u8>, public_inputs: Vec<E::Fr>) -> Self {
        Self {
            data,
            public_inputs,
            vk_id: VkId::new(0), // Aggregated proofs don't have a single VK
            _state: PhantomData,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::backend::bn254::Bn254;
    use crate::crypto::VerificationKey;
    use crate::registry::VkRegistry;
    use group::{Curve, Group};
    use halo2curves::bn256::{Fr, G1, G2};
    use rand::rngs::OsRng;

    fn mock_vk(num_public_inputs: usize) -> VerificationKey<Bn254> {
        VerificationKey {
            num_public_inputs,
            domain_size: 1024,
            selector_commitments: vec![
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
            ],
            permutation_commitments: vec![G1::random(OsRng).to_affine()],
            x_g2: G2::random(OsRng).to_affine(),
            g2_generator: G2::generator().to_affine(),
        }
    }

    /// Generate valid mock proof data that passes PlonkProof::from_bytes
    fn mock_proof_data() -> Vec<u8> {
        use crate::crypto::{PlonkProof, ProofEvaluations};
        use ff::Field;

        let proof = PlonkProof::<Bn254> {
            wire_commitments: [
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
            ],
            z_commitment: G1::random(OsRng).to_affine(),
            t_commitments: vec![
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
            ],
            opening_proof: G1::random(OsRng).to_affine(),
            shifted_opening_proof: G1::random(OsRng).to_affine(),
            evaluations: ProofEvaluations {
                a_eval: Fr::random(OsRng),
                b_eval: Fr::random(OsRng),
                c_eval: Fr::random(OsRng),
                s1_eval: Fr::random(OsRng),
                s2_eval: Fr::random(OsRng),
                z_shifted_eval: Fr::random(OsRng),
            },
        };

        proof.to_bytes()
    }

    #[test]
    fn proof_is_generic_over_state() {
        // This is a compile-time check that Proof<E, S> works with different states
        fn _accepts_pending<E: PairingEngine>(_: &Proof<E, Pending>) {}
        fn _accepts_verified<E: PairingEngine>(_: &Proof<E, Verified>) {}
        fn _accepts_batched<E: PairingEngine>(_: &Proof<E, Batched>) {}

        // The following would fail to compile (good!):
        // fn mixed<E: PairingEngine>(p: Proof<E, Pending>) { _accepts_verified(&p); }
    }

    #[test]
    fn state_transitions_consume_proof() {
        // This documents that verify() and submit() take ownership
        // preventing double-submission or re-verification
        fn _verify_consumes<E: PairingEngine>(p: Proof<E, Pending>) {
            // p.verify(registry) consumes p
            // p would not be usable after this
            let _ = p;
        }

        fn _submit_consumes<E: PairingEngine>(p: Proof<E, Verified>) {
            // p.submit() consumes p
            let _ = p;
        }
    }

    #[test]
    fn verify_rejects_unknown_vk() {
        let registry = VkRegistry::<Bn254>::new();
        // Note: VkId(999) is not registered
        let proof =
            Proof::<Bn254, Pending>::new(mock_proof_data(), vec![Fr::from(1u64)], VkId::new(999));

        let result = proof.verify(&registry);
        assert!(matches!(result, Err(Error::UnknownVk(_))));
    }

    #[test]
    fn verify_rejects_too_many_public_inputs() {
        let registry = VkRegistry::<Bn254>::new();
        let vk_id = registry.register("test", mock_vk(2)); // max 2 inputs

        let proof = Proof::<Bn254, Pending>::new(
            mock_proof_data(),
            vec![Fr::from(1u64), Fr::from(2u64), Fr::from(3u64)], // 3 inputs
            vk_id,
        );

        let result = proof.verify(&registry);
        assert!(matches!(
            result,
            Err(Error::PublicInputMismatch {
                expected: 2,
                got: 3
            })
        ));
    }

    #[test]
    fn verify_rejects_empty_proof() {
        let registry = VkRegistry::<Bn254>::new();
        let vk_id = registry.register("test", mock_vk(5));

        let proof = Proof::<Bn254, Pending>::new(
            vec![], // Empty proof data
            vec![Fr::from(1u64)],
            vk_id,
        );

        let result = proof.verify(&registry);
        assert!(matches!(result, Err(Error::VerificationFailed(_))));
    }

    #[test]
    fn verify_accepts_valid_proof() {
        let registry = VkRegistry::<Bn254>::new();
        // VK num_public_inputs must match the actual inputs provided
        let vk_id = registry.register("test", mock_vk(2)); // exactly 2 inputs

        let proof = Proof::<Bn254, Pending>::new(
            mock_proof_data(),                    // Valid proof data
            vec![Fr::from(1u64), Fr::from(2u64)], // 2 inputs matching VK
            vk_id,
        );

        let verified = proof.verify(&registry).expect("should verify");
        assert_eq!(verified.vk_id(), vk_id);
        assert_eq!(verified.public_inputs().len(), 2);
    }

    #[test]
    fn verify_allows_exact_max_public_inputs() {
        let registry = VkRegistry::<Bn254>::new();
        let vk_id = registry.register("test", mock_vk(3)); // max 3 inputs

        let proof = Proof::<Bn254, Pending>::new(
            mock_proof_data(),
            vec![Fr::from(1u64), Fr::from(2u64), Fr::from(3u64)], // Exactly 3
            vk_id,
        );

        let result = proof.verify(&registry);
        assert!(result.is_ok()); // Should succeed with exact count
    }
}