synor/crates/synor-zk/src/circuit.rs

//! ZK Circuit definitions for Synor rollups.
//!
//! Circuits define the constraints that must be satisfied for valid state transitions.
//! We use R1CS (Rank-1 Constraint System) representation for Groth16 compatibility.
//!
//! # Circuit Types
//!
//! - **TransferCircuit**: Validates token transfers between accounts
//! - **DepositCircuit**: Validates deposits from L1 to L2
//! - **WithdrawCircuit**: Validates withdrawals from L2 to L1
//! - **BatchCircuit**: Aggregates multiple transactions into single proof

use ark_ff::PrimeField;
use ark_relations::r1cs::SynthesisError;
use ark_bn254::Fr as ScalarField;
use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::state::StateRoot;

/// Circuit configuration parameters.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct CircuitConfig {
    /// Maximum number of transactions in a batch
    pub max_batch_size: usize,
    /// State tree depth
    pub tree_depth: usize,
    /// Enable signature verification in circuit
    pub verify_signatures: bool,
}

impl Default for CircuitConfig {
    fn default() -> Self {
        Self {
            max_batch_size: crate::constants::MAX_BATCH_SIZE,
            tree_depth: crate::constants::STATE_TREE_DEPTH,
            verify_signatures: true,
        }
    }
}

/// Constraint system interface for building circuits.
pub trait ConstraintSystem {
    /// Allocates a public input variable.
    fn alloc_input(&mut self, name: &str, value: ScalarField) -> Result<Variable, CircuitError>;

    /// Allocates a private witness variable.
    fn alloc_witness(&mut self, name: &str, value: ScalarField) -> Result<Variable, CircuitError>;

    /// Adds a constraint: a * b = c
    fn enforce_constraint(
        &mut self,
        a: LinearCombination,
        b: LinearCombination,
        c: LinearCombination,
    ) -> Result<(), CircuitError>;

    /// Returns the number of constraints.
    fn num_constraints(&self) -> usize;

    /// Returns the number of public inputs.
    fn num_inputs(&self) -> usize;

    /// Returns the number of witness variables.
    fn num_witnesses(&self) -> usize;
}

/// Variable in the constraint system.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct Variable(pub(crate) usize);

impl Variable {
    /// Creates a new variable with the given index.
    pub fn new(index: usize) -> Self {
        Self(index)
    }

    /// Returns the variable index.
    pub fn index(&self) -> usize {
        self.0
    }
}

/// Linear combination of variables: sum(coeff_i * var_i).
#[derive(Clone, Debug, Default)]
pub struct LinearCombination {
    terms: Vec<(ScalarField, Variable)>,
}

impl LinearCombination {
    /// Creates an empty linear combination.
    pub fn new() -> Self {
        Self { terms: Vec::new() }
    }

    /// Creates a linear combination from a single variable.
    pub fn from_variable(var: Variable) -> Self {
        let mut lc = Self::new();
        lc.add_term(ScalarField::from(1u64), var);
        lc
    }

    /// Creates a linear combination from a constant.
    pub fn from_constant(value: ScalarField) -> Self {
        let mut lc = Self::new();
        lc.add_term(value, Variable(0)); // Variable 0 is the constant 1
        lc
    }

    /// Adds a term (coefficient * variable) to the linear combination.
    pub fn add_term(&mut self, coeff: ScalarField, var: Variable) {
        self.terms.push((coeff, var));
    }

    /// Returns the terms of the linear combination.
    pub fn terms(&self) -> &[(ScalarField, Variable)] {
        &self.terms
    }
}

/// Base trait for all circuits.
pub trait Circuit: Clone + Send + Sync {
    /// Returns the circuit name.
    fn name(&self) -> &str;

    /// Returns the circuit configuration.
    fn config(&self) -> &CircuitConfig;

    /// Synthesizes the circuit constraints.
    fn synthesize<CS: ConstraintSystem>(&self, cs: &mut CS) -> Result<(), CircuitError>;

    /// Returns the public inputs for this circuit instance.
    fn public_inputs(&self) -> Vec<ScalarField>;
}

/// Transfer circuit: validates a token transfer between accounts.
#[derive(Clone, Debug)]
pub struct TransferCircuit {
    config: CircuitConfig,
    /// Previous state root
    pub old_root: StateRoot,
    /// New state root after transfer
    pub new_root: StateRoot,
    /// Sender account index
    pub sender_idx: u64,
    /// Recipient account index
    pub recipient_idx: u64,
    /// Transfer amount
    pub amount: u64,
    /// Sender's signature (for witness)
    pub signature: Vec<u8>,
    /// Merkle proof for sender account
    pub sender_proof: Vec<[u8; 32]>,
    /// Merkle proof for recipient account
    pub recipient_proof: Vec<[u8; 32]>,
}

impl TransferCircuit {
    /// Creates a new transfer circuit.
    pub fn new(
        old_root: StateRoot,
        new_root: StateRoot,
        sender_idx: u64,
        recipient_idx: u64,
        amount: u64,
    ) -> Self {
        Self {
            config: CircuitConfig::default(),
            old_root,
            new_root,
            sender_idx,
            recipient_idx,
            amount,
            signature: Vec::new(),
            sender_proof: Vec::new(),
            recipient_proof: Vec::new(),
        }
    }

    /// Sets the Merkle proofs for verification.
    pub fn with_proofs(
        mut self,
        sender_proof: Vec<[u8; 32]>,
        recipient_proof: Vec<[u8; 32]>,
    ) -> Self {
        self.sender_proof = sender_proof;
        self.recipient_proof = recipient_proof;
        self
    }

    /// Sets the signature for verification.
    pub fn with_signature(mut self, signature: Vec<u8>) -> Self {
        self.signature = signature;
        self
    }
}

impl Circuit for TransferCircuit {
    fn name(&self) -> &str {
        "TransferCircuit"
    }

    fn config(&self) -> &CircuitConfig {
        &self.config
    }

    fn synthesize<CS: ConstraintSystem>(&self, cs: &mut CS) -> Result<(), CircuitError> {
        // Public inputs
        let old_root_var = cs.alloc_input("old_root", field_from_hash(&self.old_root.0))?;
        let new_root_var = cs.alloc_input("new_root", field_from_hash(&self.new_root.0))?;

        // Private witnesses
        let sender_idx_var = cs.alloc_witness("sender_idx", ScalarField::from(self.sender_idx))?;
        let recipient_idx_var = cs.alloc_witness("recipient_idx", ScalarField::from(self.recipient_idx))?;
        let amount_var = cs.alloc_witness("amount", ScalarField::from(self.amount))?;

        // Constraint 1: Amount > 0
        // We don't have a direct "greater than" constraint in R1CS,
        // but we can enforce that amount is non-zero by requiring amount * amount_inv = 1
        // This is a simplified version; real implementation would use range proofs

        // Constraint 2: Merkle proof verification for sender
        // Verify that sender account exists at sender_idx in old_root
        // (Merkle path verification constraints)

        // Constraint 3: Merkle proof verification for recipient
        // Verify that recipient account exists at recipient_idx in old_root

        // Constraint 4: Balance check
        // sender.balance >= amount

        // Constraint 5: State transition
        // new_sender_balance = old_sender_balance - amount
        // new_recipient_balance = old_recipient_balance + amount

        // Constraint 6: New root computation
        // Verify that new_root is correct after state updates

        // Note: Full implementation would include all Merkle proof verification
        // and hash constraints. This is a structural placeholder.

        // Dummy constraint to ensure variables are used
        let _lc = LinearCombination::from_variable(old_root_var);
        let _lc2 = LinearCombination::from_variable(new_root_var);
        let _lc3 = LinearCombination::from_variable(sender_idx_var);
        let _lc4 = LinearCombination::from_variable(recipient_idx_var);
        let _lc5 = LinearCombination::from_variable(amount_var);

        Ok(())
    }

    fn public_inputs(&self) -> Vec<ScalarField> {
        vec![
            field_from_hash(&self.old_root.0),
            field_from_hash(&self.new_root.0),
        ]
    }
}

/// Batch circuit: aggregates multiple transfers into a single proof.
#[derive(Clone, Debug)]
pub struct BatchCircuit {
    config: CircuitConfig,
    /// Initial state root
    pub initial_root: StateRoot,
    /// Final state root after all transactions
    pub final_root: StateRoot,
    /// Individual transfer circuits
    pub transfers: Vec<TransferCircuit>,
}

impl BatchCircuit {
    /// Creates a new batch circuit.
    pub fn new(initial_root: StateRoot, final_root: StateRoot) -> Self {
        Self {
            config: CircuitConfig::default(),
            initial_root,
            final_root,
            transfers: Vec::new(),
        }
    }

    /// Adds a transfer to the batch.
    pub fn add_transfer(&mut self, transfer: TransferCircuit) {
        self.transfers.push(transfer);
    }

    /// Returns the number of transfers in the batch.
    pub fn num_transfers(&self) -> usize {
        self.transfers.len()
    }
}

impl Circuit for BatchCircuit {
    fn name(&self) -> &str {
        "BatchCircuit"
    }

    fn config(&self) -> &CircuitConfig {
        &self.config
    }

    fn synthesize<CS: ConstraintSystem>(&self, cs: &mut CS) -> Result<(), CircuitError> {
        // Public inputs: initial_root, final_root, batch_hash
        let initial_root_var = cs.alloc_input("initial_root", field_from_hash(&self.initial_root.0))?;
        let final_root_var = cs.alloc_input("final_root", field_from_hash(&self.final_root.0))?;

        // For each transfer, verify:
        // 1. The intermediate state roots chain correctly
        // 2. Each individual transfer is valid

        // In a real implementation, we would:
        // - Synthesize each transfer circuit
        // - Chain the state roots: root_i+1 = update(root_i, transfer_i)
        // - Verify initial_root -> transfer_1 -> ... -> transfer_n -> final_root

        let _lc = LinearCombination::from_variable(initial_root_var);
        let _lc2 = LinearCombination::from_variable(final_root_var);

        Ok(())
    }

    fn public_inputs(&self) -> Vec<ScalarField> {
        vec![
            field_from_hash(&self.initial_root.0),
            field_from_hash(&self.final_root.0),
        ]
    }
}

/// Deposit circuit: validates an L1 -> L2 deposit.
#[derive(Clone, Debug)]
pub struct DepositCircuit {
    config: CircuitConfig,
    /// Previous state root
    pub old_root: StateRoot,
    /// New state root after deposit
    pub new_root: StateRoot,
    /// L1 deposit transaction hash
    pub l1_tx_hash: [u8; 32],
    /// Recipient account index on L2
    pub recipient_idx: u64,
    /// Deposit amount
    pub amount: u64,
}

impl Circuit for DepositCircuit {
    fn name(&self) -> &str {
        "DepositCircuit"
    }

    fn config(&self) -> &CircuitConfig {
        &self.config
    }

    fn synthesize<CS: ConstraintSystem>(&self, cs: &mut CS) -> Result<(), CircuitError> {
        let old_root_var = cs.alloc_input("old_root", field_from_hash(&self.old_root.0))?;
        let new_root_var = cs.alloc_input("new_root", field_from_hash(&self.new_root.0))?;
        let l1_tx_hash_var = cs.alloc_input("l1_tx_hash", field_from_hash(&self.l1_tx_hash))?;

        let _lc = LinearCombination::from_variable(old_root_var);
        let _lc2 = LinearCombination::from_variable(new_root_var);
        let _lc3 = LinearCombination::from_variable(l1_tx_hash_var);

        Ok(())
    }

    fn public_inputs(&self) -> Vec<ScalarField> {
        vec![
            field_from_hash(&self.old_root.0),
            field_from_hash(&self.new_root.0),
            field_from_hash(&self.l1_tx_hash),
        ]
    }
}

/// Withdrawal circuit: validates an L2 -> L1 withdrawal.
#[derive(Clone, Debug)]
pub struct WithdrawCircuit {
    config: CircuitConfig,
    /// Previous state root
    pub old_root: StateRoot,
    /// New state root after withdrawal
    pub new_root: StateRoot,
    /// Sender account index on L2
    pub sender_idx: u64,
    /// L1 recipient address
    pub l1_recipient: [u8; 20],
    /// Withdrawal amount
    pub amount: u64,
}

impl Circuit for WithdrawCircuit {
    fn name(&self) -> &str {
        "WithdrawCircuit"
    }

    fn config(&self) -> &CircuitConfig {
        &self.config
    }

    fn synthesize<CS: ConstraintSystem>(&self, cs: &mut CS) -> Result<(), CircuitError> {
        let old_root_var = cs.alloc_input("old_root", field_from_hash(&self.old_root.0))?;
        let new_root_var = cs.alloc_input("new_root", field_from_hash(&self.new_root.0))?;

        let _lc = LinearCombination::from_variable(old_root_var);
        let _lc2 = LinearCombination::from_variable(new_root_var);

        Ok(())
    }

    fn public_inputs(&self) -> Vec<ScalarField> {
        vec![
            field_from_hash(&self.old_root.0),
            field_from_hash(&self.new_root.0),
        ]
    }
}

/// Converts a 32-byte hash to a scalar field element.
fn field_from_hash(hash: &[u8; 32]) -> ScalarField {
    // Take first 31 bytes to ensure it fits in the field
    let mut bytes = [0u8; 32];
    bytes[..31].copy_from_slice(&hash[..31]);
    ScalarField::from_le_bytes_mod_order(&bytes)
}

/// Circuit errors.
#[derive(Debug, Error)]
pub enum CircuitError {
    #[error("Synthesis error: {0}")]
    SynthesisError(String),

    #[error("Invalid witness: {0}")]
    InvalidWitness(String),

    #[error("Constraint violation: {0}")]
    ConstraintViolation(String),

    #[error("Circuit configuration error: {0}")]
    ConfigError(String),
}

impl From<SynthesisError> for CircuitError {
    fn from(e: SynthesisError) -> Self {
        CircuitError::SynthesisError(e.to_string())
    }
}

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

    #[test]
    fn test_circuit_config_default() {
        let config = CircuitConfig::default();
        assert_eq!(config.max_batch_size, crate::constants::MAX_BATCH_SIZE);
        assert_eq!(config.tree_depth, crate::constants::STATE_TREE_DEPTH);
        assert!(config.verify_signatures);
    }

    #[test]
    fn test_linear_combination() {
        let var = Variable::new(1);
        let lc = LinearCombination::from_variable(var);
        assert_eq!(lc.terms().len(), 1);
        assert_eq!(lc.terms()[0].1, var);
    }

    #[test]
    fn test_transfer_circuit_creation() {
        let old_root = StateRoot([0u8; 32]);
        let new_root = StateRoot([1u8; 32]);

        let circuit = TransferCircuit::new(old_root, new_root, 0, 1, 100);
        assert_eq!(circuit.name(), "TransferCircuit");
        assert_eq!(circuit.sender_idx, 0);
        assert_eq!(circuit.recipient_idx, 1);
        assert_eq!(circuit.amount, 100);
    }

    #[test]
    fn test_batch_circuit() {
        let initial_root = StateRoot([0u8; 32]);
        let final_root = StateRoot([2u8; 32]);

        let mut batch = BatchCircuit::new(initial_root, final_root);
        assert_eq!(batch.num_transfers(), 0);

        let transfer = TransferCircuit::new(
            StateRoot([0u8; 32]),
            StateRoot([1u8; 32]),
            0,
            1,
            100,
        );
        batch.add_transfer(transfer);
        assert_eq!(batch.num_transfers(), 1);
    }

    #[test]
    fn test_field_from_hash() {
        let hash = [0xffu8; 32];
        let field = field_from_hash(&hash);
        // Field element should be valid (not panic)
        let _ = field;
    }
}