arcium-core-utils 0.6.0

use std::ops::{Index, IndexMut};

use derive_more::derive::{Add, AddAssign, Sub, SubAssign};
use primitives::{algebra::elliptic_curve::Curve, correlated_randomness::bundler::BundleConsumer};

use crate::{
    circuit::{
        AlgebraicType,
        BitShareBinaryOp,
        Circuit,
        FieldShareBinaryOp,
        FieldShareUnaryOp,
        FieldType,
        Gate,
        GateExt,
        Input,
        PointShareBinaryOp,
        PointShareUnaryOp,
        ShareOrPlaintext,
    },
    preprocessing::iterator::PreprocessingIterator,
};

/// Field specific preprocessing requirements for a circuit.
#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, Add, AddAssign, Sub, SubAssign)]
pub struct FieldCircuitPreprocessing {
    pub singlets: usize,
    pub triples: usize,
    pub dabits: usize,
}

/// Preprocessing requirements for a circuit.
#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, Add, AddAssign, Sub, SubAssign)]
pub struct CircuitPreprocessing {
    pub bit_singlets: usize,
    pub bit_triples: usize,
    pub base_field: FieldCircuitPreprocessing,
    pub scalar: FieldCircuitPreprocessing,
    pub mersenne107: FieldCircuitPreprocessing,
}

impl Index<FieldType> for CircuitPreprocessing {
    type Output = FieldCircuitPreprocessing;

    fn index(&self, index: FieldType) -> &Self::Output {
        match index {
            FieldType::BaseField => &self.base_field,
            FieldType::ScalarField => &self.scalar,
            FieldType::Mersenne107 => &self.mersenne107,
        }
    }
}

impl IndexMut<FieldType> for CircuitPreprocessing {
    fn index_mut(&mut self, index: FieldType) -> &mut Self::Output {
        match index {
            FieldType::BaseField => &mut self.base_field,
            FieldType::ScalarField => &mut self.scalar,
            FieldType::Mersenne107 => &mut self.mersenne107,
        }
    }
}

impl<C: Curve> BundleConsumer for Circuit<C> {
    type Iterator = PreprocessingIterator<C>;

    fn required_preprocessing(&self) -> CircuitPreprocessing {
        let mut circuit_preprocessing = CircuitPreprocessing::default();
        for gate in self.iter_gates_ext() {
            self.add_to_required_preprocessing(gate, &mut circuit_preprocessing);
        }
        circuit_preprocessing
    }
}

impl<C: Curve> Circuit<C> {
    /// Updates the circuit preprocessing structure with the requirements of this gate.
    pub fn add_to_required_preprocessing(
        &self,
        gate: &GateExt<C>,
        circuit_preprocessing: &mut CircuitPreprocessing,
    ) {
        let batch_size = gate.output.get_batch_size() as usize;
        match &gate.gate {
            Gate::Input(Input::SecretPlaintext { algebraic_type, .. })
            | Gate::Random { algebraic_type, .. } => match algebraic_type {
                AlgebraicType::ScalarField | AlgebraicType::Point => {
                    circuit_preprocessing.scalar.singlets += batch_size;
                }
                AlgebraicType::BaseField => {
                    circuit_preprocessing.base_field.singlets += batch_size;
                }
                AlgebraicType::Bit => {
                    circuit_preprocessing.bit_singlets += batch_size;
                }
                AlgebraicType::Mersenne107 => {
                    circuit_preprocessing.mersenne107.singlets += batch_size;
                }
            },
            Gate::FieldShareUnaryOp { op, .. } => {
                let field_type = gate.output.get_field_type_unchecked();
                match op {
                    FieldShareUnaryOp::MulInverse | FieldShareUnaryOp::IsZero => {
                        circuit_preprocessing[field_type].triples += batch_size;
                        circuit_preprocessing[field_type].singlets += batch_size;
                    }
                    FieldShareUnaryOp::Open | FieldShareUnaryOp::Neg => (),
                }
            }
            Gate::FieldShareBinaryOp { op, y, .. } => match op {
                FieldShareBinaryOp::Mul => {
                    let field_type = gate.output.get_field_type_unchecked();
                    if self.gate_output_unchecked(*y).get_form() == ShareOrPlaintext::Share {
                        circuit_preprocessing[field_type].triples += batch_size;
                    }
                }
                FieldShareBinaryOp::Add => (),
            },
            Gate::PointShareUnaryOp { op, .. } => match op {
                PointShareUnaryOp::IsZero => {
                    circuit_preprocessing.scalar.triples += batch_size;
                    circuit_preprocessing.scalar.singlets += batch_size;
                }
                PointShareUnaryOp::Open | PointShareUnaryOp::Neg => (),
            },
            Gate::PointShareBinaryOp { op, p, y, .. } => match op {
                PointShareBinaryOp::ScalarMul => {
                    if self.gate_output_unchecked(*p).get_form() == ShareOrPlaintext::Share
                        && self.gate_output_unchecked(*y).get_form() == ShareOrPlaintext::Share
                    {
                        circuit_preprocessing.scalar.triples += batch_size;
                    }
                }
                PointShareBinaryOp::Add => (),
            },
            Gate::BitShareBinaryOp { op, y, .. } => match op {
                BitShareBinaryOp::And | BitShareBinaryOp::Or => {
                    if self.gate_output_unchecked(*y).get_form() == ShareOrPlaintext::Share {
                        circuit_preprocessing.bit_triples += batch_size;
                    }
                }
                BitShareBinaryOp::Xor => (),
            },
            Gate::BaseFieldPow { .. } => {
                unimplemented!("Removed from the Bundler, need to choose exponent to set it back.")
            }
            Gate::DaBit { field_type, .. } => {
                circuit_preprocessing[*field_type].dabits += batch_size
            }

            Gate::Input(_)
            | Gate::Constant { .. }
            | Gate::BatchSummation { .. }
            | Gate::BitShareUnaryOp { .. }
            | Gate::FieldPlaintextUnaryOp { .. }
            | Gate::FieldPlaintextBinaryOp { .. }
            | Gate::BitPlaintextUnaryOp { .. }
            | Gate::BitPlaintextBinaryOp { .. }
            | Gate::PointPlaintextUnaryOp { .. }
            | Gate::PointPlaintextBinaryOp { .. }
            | Gate::GetDaBitFieldShare { .. }
            | Gate::GetDaBitSharedBit { .. }
            | Gate::BitPlaintextToField { .. }
            | Gate::FieldPlaintextToBit { .. }
            | Gate::ExtractFromBatch { .. }
            | Gate::CollectToBatch { .. }
            | Gate::PointFromPlaintextExtendedEdwards { .. }
            | Gate::PlaintextPointToExtendedEdwards { .. }
            | Gate::PlaintextKeccakF1600 { .. }
            | Gate::CompressPlaintextPoint { .. }
            | Gate::KeyRecoveryPlaintextComputeErrors { .. } => (),
        };
    }
}

#[cfg(test)]
mod tests {
    use crate::circuit::preprocessing::{CircuitPreprocessing, FieldCircuitPreprocessing};

    #[test]
    fn test_circuit_preprocessing_add() {
        let a = CircuitPreprocessing {
            bit_singlets: 0,
            bit_triples: 1,
            base_field: FieldCircuitPreprocessing {
                singlets: 3,
                triples: 4,
                dabits: 2,
            },
            scalar: FieldCircuitPreprocessing {
                singlets: 1,
                triples: 2,
                dabits: 1,
            },
            mersenne107: FieldCircuitPreprocessing {
                singlets: 0,
                triples: 0,
                dabits: 0,
            },
        };
        let b = CircuitPreprocessing {
            bit_singlets: 3,
            bit_triples: 4,
            base_field: FieldCircuitPreprocessing {
                singlets: 0,
                triples: 5,
                dabits: 3,
            },
            scalar: FieldCircuitPreprocessing {
                singlets: 2,
                triples: 3,
                dabits: 2,
            },
            mersenne107: FieldCircuitPreprocessing {
                singlets: 3,
                triples: 2,
                dabits: 0,
            },
        };

        let c = a + b;

        assert_eq!(c.scalar.singlets, 3);
        assert_eq!(c.scalar.triples, 5);
        assert_eq!(c.base_field.singlets, 3);
        assert_eq!(c.base_field.triples, 9);
        assert_eq!(c.bit_singlets, 3);
        assert_eq!(c.bit_triples, 5);
        assert_eq!(c.mersenne107.dabits, 0);
        assert_eq!(c.mersenne107.singlets, 3);
        assert_eq!(c.mersenne107.triples, 2);
        assert_eq!(c.scalar.dabits, 3);
        assert_eq!(c.base_field.dabits, 5);
    }
}