samaharam 0.2.0

//! KZG Polynomial Commitment Scheme.
//!
//! Implements Kate-Zaverucha-Goldberg commitments for PLONK proofs.

use std::marker::PhantomData;

use crate::traits::PairingEngine;

/// A KZG commitment (a single G1 point).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KzgCommitment<E: PairingEngine> {
    /// The commitment point in G1.
    pub point: E::G1Affine,
}

/// A KZG opening proof.
#[derive(Debug, Clone)]
pub struct KzgProof<E: PairingEngine> {
    /// The quotient polynomial commitment.
    pub quotient: E::G1Affine,
}

/// Structured Reference String for KZG.
#[derive(Debug, Clone)]
pub struct KzgSrs<E: PairingEngine> {
    /// Powers of tau in G1: [τ^0]₁, [τ^1]₁, ..., [τ^n]₁
    pub powers_of_tau_g1: Vec<E::G1Affine>,

    /// [τ]₂ in G2 for pairing checks
    pub tau_g2: E::G2Affine,

    /// [1]₂ in G2 (generator)
    pub g2_generator: E::G2Affine,
}

impl<E: PairingEngine> KzgSrs<E> {
    /// Validate the SRS has consistent tau powers.
    ///
    /// Checks: e([τ^i]₁, [1]₂) = e([τ^(i-1)]₁, [τ]₂) for each i > 0
    /// This ensures the powers are correctly formed.
    ///
    /// Returns Ok(()) if valid, Err with description if invalid.
    pub fn validate(&self) -> Result<(), String> {
        use group::prime::PrimeCurveAffine;

        if self.powers_of_tau_g1.is_empty() {
            return Err("SRS has no G1 powers".to_string());
        }

        // Check G1[0] is the generator
        if self.powers_of_tau_g1[0].is_identity().into() {
            return Err("SRS G1[0] is identity, should be generator".to_string());
        }

        // Check G2 generator is not identity
        if self.g2_generator.is_identity().into() {
            return Err("SRS G2 generator is identity".to_string());
        }

        // Check tau_g2 is not identity
        if self.tau_g2.is_identity().into() {
            return Err("SRS tau_g2 is identity".to_string());
        }

        // For each consecutive pair (i, i+1), verify:
        // e([τ^i]₁, [τ]₂) == e([τ^(i+1)]₁, [1]₂)
        // 
        // This ensures τ^(i+1) = τ * τ^i
        for i in 0..(self.powers_of_tau_g1.len().saturating_sub(1)) {
            let g1_i = &self.powers_of_tau_g1[i];
            let g1_i_plus_1 = &self.powers_of_tau_g1[i + 1];

            // Skip if either point is identity (shouldn't happen in valid SRS)
            if g1_i.is_identity().into() {
                return Err(format!("SRS G1[{}] is identity", i));
            }

            // Check pairing: e(G1[i], tau_g2) == e(G1[i+1], g2_gen)
            let lhs = E::pairing(g1_i, &self.tau_g2);
            let rhs = E::pairing(g1_i_plus_1, &self.g2_generator);

            if lhs != rhs {
                return Err(format!(
                    "SRS tau consistency check failed at index {}: e([τ^{}]₁, [τ]₂) ≠ e([τ^{}]₁, [1]₂)",
                    i, i, i + 1
                ));
            }
        }

        Ok(())
    }

    /// Validate only basic structure (no pairing checks).
    /// 
    /// This is faster but less thorough than full validation.
    pub fn validate_structure(&self) -> Result<(), String> {
        use group::prime::PrimeCurveAffine;

        if self.powers_of_tau_g1.is_empty() {
            return Err("SRS has no G1 powers".to_string());
        }

        if self.powers_of_tau_g1[0].is_identity().into() {
            return Err("SRS G1[0] is identity".to_string());
        }

        if self.g2_generator.is_identity().into() {
            return Err("G2 generator is identity".to_string());
        }

        if self.tau_g2.is_identity().into() {
            return Err("tau_g2 is identity".to_string());
        }

        Ok(())
    }
}


/// KZG commitment scheme operations.
pub struct KzgScheme<E: PairingEngine> {
    srs: KzgSrs<E>,
    _engine: PhantomData<E>,
}

impl<E: PairingEngine> KzgScheme<E> {
    /// Create a new KZG scheme with the given SRS.
    pub fn new(srs: KzgSrs<E>) -> Self {
        Self {
            srs,
            _engine: PhantomData,
        }
    }

    /// Commit to a polynomial.
    ///
    /// commitment = Σ coeff_i * [τ^i]₁
    pub fn commit(&self, coefficients: &[E::Fr]) -> Result<KzgCommitment<E>, String> {
        use group::Curve;

        if coefficients.len() > self.srs.powers_of_tau_g1.len() {
            return Err("Polynomial degree exceeds SRS size".to_string());
        }

        // MSM: multi-scalar multiplication
        let result = self.msm(
            &self.srs.powers_of_tau_g1[..coefficients.len()],
            coefficients,
        );

        Ok(KzgCommitment {
            point: result.to_affine(),
        })
    }

    /// Verify a KZG opening proof.
    ///
    /// Checks: e(C - [v]₁, [1]₂) = e(π, [τ]₂ - [z]₂)  
    /// Rearranged: e(C - [v]₁, [1]₂) · e(-π, [τ]₂) · e(π·z, [1]₂) = 1
    /// Simplified: e(C - [v]₁ + π·z, [1]₂) · e(-π, [τ]₂) = 1
    ///
    /// For efficiency, we check: e(C - [v]₁ + π·z, [1]₂) == e(π, [τ]₂)
    pub fn verify(
        &self,
        commitment: &KzgCommitment<E>,
        point: &E::Fr,
        value: &E::Fr,
        proof: &KzgProof<E>,
    ) -> bool {
        use group::{Curve, Group};
        use group::prime::PrimeCurveAffine;

        // Security Patch: Reject identity point for commitment
        // Prevents bypassing verification with empty proofs (0 - 0 = 0)
        if commitment.point.is_identity().into() {
            return false;
        }

        // Compute [v]₁ = v · G1
        let v_g1 = self.scalar_mul_g1(&E::G1::generator().to_affine(), value);

        // Compute C - [v]₁
        let comm_g1: E::G1 = commitment.point.into();
        let c_minus_v = comm_g1 - v_g1;

        // Compute π · z (the evaluation point scaled proof)
        let quot_g1: E::G1 = proof.quotient.into();
        let pi_times_z = quot_g1 * *point;

        // LHS of pairing: C - [v]₁ + π·z
        let lhs_point = (c_minus_v + pi_times_z).to_affine();

        // Pairing check: e(C - [v]₁ + π·z, [1]₂) == e(π, [τ]₂)
        let lhs = E::pairing(&lhs_point, &self.srs.g2_generator);
        let rhs = E::pairing(&proof.quotient, &self.srs.tau_g2);

        lhs == rhs
    }

    /// Multi-scalar multiplication using Pippenger's bucket method.
    ///
    /// Complexity: O(n/c + 2^c) vs O(n) for naive method.
    /// For typical sizes (n ~ 2^16), this provides ~10x speedup.
    fn msm(&self, bases: &[E::G1Affine], scalars: &[E::Fr]) -> E::G1 {
        use ff::PrimeField;
        use group::Group;

        let n = bases.len();
        if n == 0 {
            return E::G1::identity();
        }

        // For small n, use naive method
        if n < 32 {
            return self.msm_naive(bases, scalars);
        }

        // Pippenger's algorithm
        // Choose optimal window size: c ≈ log2(n) / 2
        let c = (64 - (n as u64).leading_zeros()) as usize / 2;
        let c = c.clamp(4, 16); // Clamp to reasonable range

        let num_windows = 256_usize.div_ceil(c); // ceil(256/c)
        let num_buckets = 1 << c; // 2^c buckets per window

        let mut result = E::G1::identity();

        // Process each window
        for w in (0..num_windows).rev() {
            // Double the result c times (shift)
            for _ in 0..c {
                result = result.double();
            }

            // Buckets for this window
            let mut buckets: Vec<E::G1> = vec![E::G1::identity(); num_buckets];

            // Assign points to buckets based on scalar window
            for (base, scalar) in bases.iter().zip(scalars.iter()) {
                // Extract c-bit window from scalar
                let repr = scalar.to_repr();
                let bytes: &[u8] = repr.as_ref();

                let bit_offset = w * c;
                let bucket_idx = Self::extract_window(bytes, bit_offset, c);

                if bucket_idx > 0 {
                    let base_proj: E::G1 = (*base).into();
                    buckets[bucket_idx] += base_proj;
                }
            }

            // Accumulate buckets: sum_{i=1}^{2^c-1} i * buckets[i]
            // Using running sum for efficiency
            let mut running_sum = E::G1::identity();
            let mut window_sum = E::G1::identity();

            for i in (1..num_buckets).rev() {
                running_sum += buckets[i];
                window_sum += running_sum;
            }

            result += window_sum;
        }

        result
    }

    /// Extract a c-bit window from scalar bytes at given bit offset.
    fn extract_window(bytes: &[u8], bit_offset: usize, c: usize) -> usize {
        let byte_offset = bit_offset / 8;
        let bit_shift = bit_offset % 8;

        // Read up to 3 bytes to handle window crossing byte boundaries
        let mut val: u32 = 0;
        for i in 0..3 {
            if byte_offset + i < bytes.len() {
                val |= (bytes[byte_offset + i] as u32) << (8 * i);
            }
        }

        let mask = (1u32 << c) - 1;
        ((val >> bit_shift) & mask) as usize
    }

    /// Naive O(n) MSM for small inputs.
    fn msm_naive(&self, bases: &[E::G1Affine], scalars: &[E::Fr]) -> E::G1 {
        use group::Group;

        let mut result = E::G1::identity();
        for (base, scalar) in bases.iter().zip(scalars.iter()) {
            let base_proj: E::G1 = (*base).into();
            result += base_proj * *scalar;
        }
        result
    }

    /// Scalar multiplication in G1.
    fn scalar_mul_g1(&self, base: &E::G1Affine, scalar: &E::Fr) -> E::G1 {
        let base_proj: E::G1 = (*base).into();
        base_proj * *scalar
    }

    /// Scalar multiplication in G2.
    ///
    /// Note: This is currently unused in the verification flow since KZG
    /// verification only requires G1 operations. The G2 points ([τ]G₂, G₂)
    /// come directly from the SRS without scalar multiplication.
    #[allow(dead_code)]
    fn scalar_mul_g2(&self, base: &E::G2Affine, _scalar: &E::Fr) -> E::G2Affine {
        // G2 scalar multiplication is not needed for KZG verification.
        // The pairing check uses fixed G2 points from SRS.
        // If needed, extend PairingEngine trait with G2 projective type.
        *base
    }
}

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

    fn mock_srs(size: usize) -> KzgSrs<Bn254> {
        // Generate a mock SRS with random tau
        let tau = Fr::random(OsRng);

        let g1_gen = G1::generator();
        let g2_gen = halo2curves::bn256::G2::generator();

        // Powers of tau
        let mut powers_of_tau_g1 = Vec::with_capacity(size);
        let mut current = Fr::ONE;
        for _ in 0..size {
            powers_of_tau_g1.push((g1_gen * current).to_affine());
            current *= tau;
        }

        KzgSrs {
            powers_of_tau_g1,
            tau_g2: (g2_gen * tau).to_affine(),
            g2_generator: g2_gen.to_affine(),
        }
    }

    // TDD: RED - These tests define expected behavior

    #[test]
    fn kzg_commit_to_constant_polynomial() {
        let srs = mock_srs(16);
        let scheme = KzgScheme::<Bn254>::new(srs.clone());

        // Constant polynomial p(x) = 5
        let coeffs = vec![Fr::from(5u64)];
        let commitment = scheme.commit(&coeffs).unwrap();

        // C = 5 * [1]₁ = 5 * G1
        let expected = (G1::generator() * Fr::from(5u64)).to_affine();
        assert_eq!(commitment.point, expected);
    }

    #[test]
    fn kzg_commit_to_linear_polynomial() {
        let srs = mock_srs(16);
        let scheme = KzgScheme::<Bn254>::new(srs.clone());

        // Linear polynomial p(x) = 3 + 7x
        let coeffs = vec![Fr::from(3u64), Fr::from(7u64)];
        let commitment = scheme.commit(&coeffs).unwrap();

        // C = 3*[1]₁ + 7*[τ]₁
        let tau_1: G1 = srs.powers_of_tau_g1[1].into();
        let expected = (G1::generator() * Fr::from(3u64) + tau_1 * Fr::from(7u64)).to_affine();

        assert_eq!(commitment.point, expected);
    }

    #[test]
    fn kzg_commit_rejects_oversized_polynomial() {
        let srs = mock_srs(4);
        let scheme = KzgScheme::<Bn254>::new(srs);

        // Polynomial with 5 coefficients, but SRS only supports degree 3
        let coeffs = vec![Fr::ONE; 5];
        let result = scheme.commit(&coeffs);

        assert!(result.is_err());
    }

    #[test]
    fn kzg_commitment_is_deterministic() {
        let srs = mock_srs(16);
        let scheme = KzgScheme::<Bn254>::new(srs);

        let coeffs = vec![Fr::from(42u64), Fr::from(17u64)];

        let c1 = scheme.commit(&coeffs).unwrap();
        let c2 = scheme.commit(&coeffs).unwrap();

        assert_eq!(c1, c2);
    }

    #[test]
    fn kzg_commitment_is_homomorphic() {
        let srs = mock_srs(16);
        let scheme = KzgScheme::<Bn254>::new(srs);

        let p1 = vec![Fr::from(3u64), Fr::from(5u64)];
        let p2 = vec![Fr::from(7u64), Fr::from(11u64)];
        let p_sum = vec![Fr::from(10u64), Fr::from(16u64)]; // p1 + p2

        let c1 = scheme.commit(&p1).unwrap();
        let c2 = scheme.commit(&p2).unwrap();
        let c_sum = scheme.commit(&p_sum).unwrap();

        // C(p1 + p2) = C(p1) + C(p2)
        let c1_proj: G1 = c1.point.into();
        let c2_proj: G1 = c2.point.into();
        let computed_sum = (c1_proj + c2_proj).to_affine();

        assert_eq!(c_sum.point, computed_sum);
    }

    // TDD: KZG Verify tests

    #[test]
    fn kzg_verify_correct_opening() {
        let srs = mock_srs(16);
        let scheme = KzgScheme::<Bn254>::new(srs.clone());

        // Polynomial p(x) = 5 (constant)
        let coeffs = vec![Fr::from(5u64)];
        let commitment = scheme.commit(&coeffs).unwrap();

        // p(z) = 5 for any z, so the quotient is 0
        let point = Fr::from(7u64);
        let value = Fr::from(5u64);

        // Quotient polynomial q(x) = (p(x) - v) / (x - z) = 0
        let proof = KzgProof {
            quotient: G1::identity().to_affine(),
        };

        let result = scheme.verify(&commitment, &point, &value, &proof);
        assert!(result, "Valid opening should verify");
    }

    #[test]
    fn kzg_srs_has_correct_structure() {
        let srs = mock_srs(8);

        // SRS should have the requested number of G1 points
        assert_eq!(srs.powers_of_tau_g1.len(), 8);

        // First element should be the generator
        assert_eq!(srs.powers_of_tau_g1[0], G1::generator().to_affine());
    }
}