arcium-core-utils 0.7.3

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

use derive_more::derive::{Add, AddAssign, Sub, SubAssign};
use primitives::correlated_randomness::bundler::BundleConsumer;
use serde::{Deserialize, Serialize};
use wincode::{SchemaRead, SchemaWrite};

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

/// Number of network rounds needed to compute the AES S-box (currently using the algorithm of [Boyar and Peralta](https://eprint.iacr.org/2011/332.pdf)).
/// The achievable minimum is 4 but this would require to concatenate columns in a way that we have
/// to clone them first. In our batching, however, we try to minimize cloning and thus end up
/// performing the 34 ANDs in 8 rounds.
pub(crate) const AES_S_BOX_N_NETWORK_ROUNDS: usize = 8;

/// Number of bit triples needed to compute the AES S-box (currently using the algorithm of [Boyar and Peralta](https://eprint.iacr.org/2011/332.pdf)).
pub(crate) const AES_S_BOX_N_TRIPLES: usize = 34;

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

impl FieldCircuitPreprocessing {
    /// Component-wise maximum of two requirement/position vectors.
    pub fn componentwise_max(self, other: Self) -> Self {
        Self {
            singlets: self.singlets.max(other.singlets),
            triples: self.triples.max(other.triples),
            dabits: self.dabits.max(other.dabits),
        }
    }
}

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

impl CircuitPreprocessing {
    /// Component-wise maximum of two position vectors — used to agree on a resync target across
    /// parties: take the per-stream maximum of everyone's [`StreamBundler::positions`], then feed
    /// it to [`StreamBundler::resync`].
    ///
    /// [`StreamBundler::positions`]: crate::preprocessing::bundler::StreamBundler::positions
    /// [`StreamBundler::resync`]: crate::preprocessing::bundler::StreamBundler::resync
    pub fn componentwise_max(self, other: Self) -> Self {
        Self {
            bit_singlets: self.bit_singlets.max(other.bit_singlets),
            bit_triples: self.bit_triples.max(other.bit_triples),
            base_field: self.base_field.componentwise_max(other.base_field),
            scalar: self.scalar.componentwise_max(other.scalar),
            mpc_field: self.mpc_field.componentwise_max(other.mpc_field),
        }
    }
}

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::MpcField => &self.mpc_field,
        }
    }
}

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::MpcField => &mut self.mpc_field,
        }
    }
}

impl<C: MpcConfig> 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: MpcConfig> 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::MpcField => {
                    circuit_preprocessing.mpc_field.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::PointFromPlaintextCoordinates { .. }
            | Gate::PlaintextPointToCoordinates { .. }
            | Gate::PlaintextKeccakF1600 { .. }
            | Gate::CompressPlaintextPoint { .. }
            | Gate::KeyRecoveryPlaintextComputeErrors { .. }
            | Gate::Ghash { .. } => (),
            #[cfg(any(test, feature = "dev"))]
            Gate::AesKeySchedule { key, .. } => {
                let key_length = self.gate_ext(*key).map(|g| g.output.batch_size).ok();
                circuit_preprocessing.bit_triples += key_length
                    .and_then(|len| n_triples_aes_key_schedule(len as usize))
                    .expect("Something went wrong with Circuit::add_to_required_preprocessing for Gate::AesKeySchedule")
            }
            Gate::AesGcmKeyStream {
                round_keys,
                n_ciphertext_blocks,
                ..
            } => {
                let round_keys_length =
                    self.gate_ext(*round_keys).map(|g| g.output.batch_size).ok();
                circuit_preprocessing.bit_triples += round_keys_length
                    .and_then(|len| {
                        n_triples_aes_gcm_key_stream(len as usize, *n_ciphertext_blocks)
                    })
                    .expect("Something went wrong with Circuit::add_to_required_preprocessing for Gate::AesGcmKeyStream")
            }
            Gate::GhashPowersOfH {
                n_ciphertext_blocks,
                ..
            } => {
                circuit_preprocessing.bit_triples +=
                    (*n_ciphertext_blocks as usize - 1) * n_triples_gf2_128_multiply()
            }
        };
    }
}

pub fn n_triples_aes_key_schedule(security_level: usize) -> Option<usize> {
    let n_sub_bytes_calls = match security_level {
        128 => Some(10),
        192 => Some(8),
        256 => Some(13),
        _ => None,
    };
    // sub_bytes is called on vectors of 4 bytes
    n_sub_bytes_calls.map(|n_calls| 4 * AES_S_BOX_N_TRIPLES * n_calls)
}

pub fn n_triples_aes_gcm_key_stream(
    round_keys_length: usize,
    n_ciphertext_blocks: u32,
) -> Option<usize> {
    // round keys length must be 11 * 128, 13 * 128 or 15 * 128
    // for AES-128, AES-192 and AES-256 respectively
    let n_rounds = match round_keys_length {
        1408 => Some(10),
        1664 => Some(12),
        1920 => Some(14),
        _ => None,
    };
    // For all but the first round, sub_bytes is computed on slices of length
    // 1+n_ciphertext_blocks. In the first round, the first 12 S-boxes (inverses)
    // are computed on slices of length 1.
    n_rounds.map(|n| {
        AES_S_BOX_N_TRIPLES
            * (12
                + 4 * (1 + n_ciphertext_blocks as usize)
                + (n - 1) * 16 * (1 + n_ciphertext_blocks as usize))
    })
}

pub fn n_triples_gf2_128_multiply() -> usize {
    // 128 shift-and-add steps × 128 bitands
    128 * 128
}

#[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,
            },
            mpc_field: 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,
            },
            mpc_field: 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.mpc_field.dabits, 0);
        assert_eq!(c.mpc_field.singlets, 3);
        assert_eq!(c.mpc_field.triples, 2);
        assert_eq!(c.scalar.dabits, 3);
        assert_eq!(c.base_field.dabits, 5);
    }
}