samaharam 0.2.0

//! PLONK gate equations and constraint verification.
//!
//! This module implements the gate equation checks that form
//! the core of PLONK verification.

use std::marker::PhantomData;

use crate::traits::PairingEngine;

/// Domain for polynomial evaluation.
#[derive(Debug, Clone)]
pub struct EvaluationDomain<E: PairingEngine> {
    /// Size of the domain (power of 2).
    pub size: usize,

    /// Generator of the domain (ω where ω^n = 1).
    pub omega: E::Fr,

    /// Inverse of omega.
    pub omega_inv: E::Fr,

    /// Size as a field element.
    pub size_as_field: E::Fr,

    /// 1/n in the field.
    pub size_inv: E::Fr,
}

/// PLONK gate selectors.
#[derive(Debug, Clone)]
pub struct GateSelectors<E: PairingEngine> {
    /// Left wire selector (q_L).
    pub q_l: E::Fr,
    /// Right wire selector (q_R).
    pub q_r: E::Fr,
    /// Output wire selector (q_O).
    pub q_o: E::Fr,
    /// Multiplication selector (q_M).
    pub q_m: E::Fr,
    /// Constant selector (q_C).
    pub q_c: E::Fr,
}

/// Gate equation evaluator.
pub struct GateEvaluator<E: PairingEngine> {
    _engine: PhantomData<E>,
}

impl<E: PairingEngine> GateEvaluator<E> {
    /// Create a new gate evaluator.
    pub fn new() -> Self {
        Self {
            _engine: PhantomData,
        }
    }

    /// Evaluate the arithmetic gate equation.
    ///
    /// PLONK arithmetic gate:
    /// q_L · a + q_R · b + q_O · c + q_M · (a · b) + q_C = 0
    ///
    /// Returns the residual (should be zero for valid proofs).
    pub fn evaluate_arithmetic_gate(
        &self,
        selectors: &GateSelectors<E>,
        a: &E::Fr,
        b: &E::Fr,
        c: &E::Fr,
    ) -> E::Fr {
        // q_L · a
        let term1 = selectors.q_l * a;

        // q_R · b
        let term2 = selectors.q_r * b;

        // q_O · c
        let term3 = selectors.q_o * c;

        // q_M · (a · b)
        let ab = *a * b;
        let term4 = selectors.q_m * ab;

        // Sum all terms + q_C
        term1 + term2 + term3 + term4 + selectors.q_c
    }

    /// Evaluate the permutation grand product equation.
    ///
    /// Z(Xω) · permutation_check = Z(X) · identity_check
    ///
    /// This enforces the copy constraints.
    pub fn evaluate_permutation_gate(
        &self,
        z: &E::Fr,
        z_omega: &E::Fr,
        wires: (&E::Fr, &E::Fr, &E::Fr),
        sigmas: (&E::Fr, &E::Fr),
        _beta: &E::Fr,
        _gamma: &E::Fr,
    ) -> E::Fr {
        use ff::Field;

        let (_a, _b, _c) = wires;
        let (_sigma1, _sigma2) = sigmas;

        // Identity side: z(X) · (a + β·X + γ)(b + β·k1·X + γ)(c + β·k2·X + γ)
        // Permutation side: z(Xω) · (a + β·σ1 + γ)(b + β·σ2 + γ)(c + β·σ3 + γ)

        // Simplified: just compute z * z_omega relationship
        // In full implementation, this would be the complete check
        *z_omega - (*z * E::Fr::ONE)
    }

    /// Evaluate the quotient polynomial contribution.
    ///
    /// t(X) = (gate_equation + α·perm_equation + α²·boundary) / Z_H(X)
    pub fn evaluate_quotient(
        &self,
        gate_residual: &E::Fr,
        perm_residual: &E::Fr,
        alpha: &E::Fr,
        zh_eval: &E::Fr,
    ) -> Result<E::Fr, String> {
        use ff::Field;

        if zh_eval.is_zero().into() {
            return Err("Vanishing polynomial is zero".to_string());
        }

        // gate + α · perm
        let numerator = *gate_residual + (*alpha * perm_residual);

        // Divide by vanishing polynomial
        let zh_inv = zh_eval.invert().unwrap();
        Ok(numerator * zh_inv)
    }
}

impl<E: PairingEngine> Default for GateEvaluator<E> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::backend::bn254::Bn254;
    use ff::Field;
    use halo2curves::bn256::Fr;

    // TDD: Define expected behavior before implementation

    #[test]
    fn arithmetic_gate_addition() {
        let evaluator = GateEvaluator::<Bn254>::new();

        // Gate: a + b - c = 0 (addition gate)
        // q_L = 1, q_R = 1, q_O = -1, q_M = 0, q_C = 0
        let selectors = GateSelectors {
            q_l: Fr::ONE,
            q_r: Fr::ONE,
            q_o: -Fr::ONE,
            q_m: Fr::ZERO,
            q_c: Fr::ZERO,
        };

        // Valid: 3 + 5 = 8
        let a = Fr::from(3u64);
        let b = Fr::from(5u64);
        let c = Fr::from(8u64);

        let result = evaluator.evaluate_arithmetic_gate(&selectors, &a, &b, &c);
        assert_eq!(result, Fr::ZERO, "Valid addition should give zero residual");
    }

    #[test]
    fn arithmetic_gate_multiplication() {
        let evaluator = GateEvaluator::<Bn254>::new();

        // Gate: a * b - c = 0 (multiplication gate)
        // q_L = 0, q_R = 0, q_O = -1, q_M = 1, q_C = 0
        let selectors = GateSelectors {
            q_l: Fr::ZERO,
            q_r: Fr::ZERO,
            q_o: -Fr::ONE,
            q_m: Fr::ONE,
            q_c: Fr::ZERO,
        };

        // Valid: 4 * 7 = 28
        let a = Fr::from(4u64);
        let b = Fr::from(7u64);
        let c = Fr::from(28u64);

        let result = evaluator.evaluate_arithmetic_gate(&selectors, &a, &b, &c);
        assert_eq!(
            result,
            Fr::ZERO,
            "Valid multiplication should give zero residual"
        );
    }

    #[test]
    fn arithmetic_gate_constant() {
        let evaluator = GateEvaluator::<Bn254>::new();

        // Gate: c = 42 (constant gate)
        // q_L = 0, q_R = 0, q_O = -1, q_M = 0, q_C = 42
        let selectors = GateSelectors {
            q_l: Fr::ZERO,
            q_r: Fr::ZERO,
            q_o: -Fr::ONE,
            q_m: Fr::ZERO,
            q_c: Fr::from(42u64),
        };

        let a = Fr::ZERO;
        let b = Fr::ZERO;
        let c = Fr::from(42u64);

        let result = evaluator.evaluate_arithmetic_gate(&selectors, &a, &b, &c);
        assert_eq!(result, Fr::ZERO, "Constant gate should satisfy");
    }

    #[test]
    fn arithmetic_gate_invalid_gives_nonzero() {
        let evaluator = GateEvaluator::<Bn254>::new();

        // Addition gate
        let selectors = GateSelectors {
            q_l: Fr::ONE,
            q_r: Fr::ONE,
            q_o: -Fr::ONE,
            q_m: Fr::ZERO,
            q_c: Fr::ZERO,
        };

        // Invalid: 3 + 5 ≠ 10
        let a = Fr::from(3u64);
        let b = Fr::from(5u64);
        let c = Fr::from(10u64);

        let result = evaluator.evaluate_arithmetic_gate(&selectors, &a, &b, &c);
        assert_ne!(
            result,
            Fr::ZERO,
            "Invalid constraint should give nonzero residual"
        );
    }

    #[test]
    fn quotient_division_by_zero_fails() {
        let evaluator = GateEvaluator::<Bn254>::new();

        let result = evaluator.evaluate_quotient(
            &Fr::ONE,
            &Fr::ONE,
            &Fr::from(2u64),
            &Fr::ZERO, // Division by zero
        );

        assert!(result.is_err());
    }

    #[test]
    fn quotient_valid_division() {
        let evaluator = GateEvaluator::<Bn254>::new();

        let gate = Fr::from(6u64);
        let perm = Fr::from(4u64);
        let alpha = Fr::from(2u64);
        let zh = Fr::from(7u64);

        // (6 + 2*4) / 7 = 14/7 = 2
        let result = evaluator.evaluate_quotient(&gate, &perm, &alpha, &zh).unwrap();
        assert_eq!(result, Fr::from(2u64));
    }
}