arcium-core-utils 0.8.5

use primitives::correlated_randomness::bundler::BundleConsumer;

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

/// 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;

/// Number of GF(2) lanes the semi-honest online-phase packs into one machine word for
/// [`Gate::AesGcmKeyStream`]/[`Gate::GhashPowersOfH`] (`online_phase::tasks::field::packed_bit::
/// PACKED_BIT_LEN`, duplicated here since `core-utils` doesn't depend on `online-phase`).
pub(crate) const SEMI_HONEST_PACKED_LANES: usize = 64;

/// Preprocessing requirements for a circuit, one count per stream.
pub type CircuitPreprocessing = PerStream<usize>;

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[PreprocessingKind::ScalarSinglets] += batch_size;
                }
                AlgebraicType::BaseField => {
                    circuit_preprocessing[PreprocessingKind::BaseFieldSinglets] += batch_size;
                }
                AlgebraicType::Bit => {
                    circuit_preprocessing[PreprocessingKind::BitSinglets] += batch_size;
                }
                AlgebraicType::MpcField => {
                    circuit_preprocessing[PreprocessingKind::MpcFieldSinglets] += batch_size;
                }
            },
            Gate::FieldShareUnaryOp { op, .. } => {
                let field_type = gate.output.get_field_type_unchecked();
                match op {
                    FieldShareUnaryOp::MulInverse | FieldShareUnaryOp::IsZero => {
                        circuit_preprocessing
                            [PreprocessingKind::of(field_type, CorrelationKind::Triples)] +=
                            batch_size;
                        circuit_preprocessing
                            [PreprocessingKind::of(field_type, CorrelationKind::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
                            [PreprocessingKind::of(field_type, CorrelationKind::Triples)] +=
                            batch_size;
                    }
                }
                FieldShareBinaryOp::Add => (),
            },
            Gate::PointShareUnaryOp { op, .. } => match op {
                PointShareUnaryOp::IsZero => {
                    circuit_preprocessing[PreprocessingKind::ScalarTriples] += batch_size;
                    circuit_preprocessing[PreprocessingKind::ScalarSinglets] += 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[PreprocessingKind::ScalarTriples] += 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[PreprocessingKind::BitTriples] += 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
                    [PreprocessingKind::of(*field_type, CorrelationKind::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 { .. }
            | Gate::ConstrainPlaintextBits { .. } => (),
            Gate::AesKeySchedule { key, .. } => {
                let key_length = self.gate_ext_unchecked(*key).output.batch_size;
                circuit_preprocessing[PreprocessingKind::BitTriples] += n_triples_aes_key_schedule(key_length 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_unchecked(*round_keys).output.batch_size;
                // No threat-model parameter here, so this must cover whichever model lowers the
                // gate: the semi-honest packed task recomputes the IV's sub_bytes in every packed
                // lane (a fixed surcharge), and the malicious task leaves that excess unconsumed.
                circuit_preprocessing[PreprocessingKind::BitTriples] +=
                    n_triples_aes_gcm_key_stream(round_keys_length as usize, *n_ciphertext_blocks)
                        .map(|n| n + AES_S_BOX_N_TRIPLES * 12 * (SEMI_HONEST_PACKED_LANES - 1))
                        .expect("Something went wrong with Circuit::add_to_required_preprocessing for Gate::AesGcmKeyStream")
            }
            Gate::GhashPowersOfH {
                n_ciphertext_blocks,
                ..
            } => {
                circuit_preprocessing[PreprocessingKind::BitTriples] +=
                    (*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::preprocessing::{PerStream, PreprocessingKind};

    #[test]
    fn test_circuit_preprocessing_add() {
        // [bit s/t, base s/t/d, scalar s/t/d, mpc s/t/d]
        let a = PerStream([0, 1, 3, 4, 2, 1, 2, 1, 0, 0, 0]);
        let b = PerStream([3, 4, 0, 5, 3, 2, 3, 2, 3, 2, 0]);

        let c = a + b;

        assert_eq!(c[PreprocessingKind::ScalarSinglets], 3);
        assert_eq!(c[PreprocessingKind::ScalarTriples], 5);
        assert_eq!(c[PreprocessingKind::BaseFieldSinglets], 3);
        assert_eq!(c[PreprocessingKind::BaseFieldTriples], 9);
        assert_eq!(c[PreprocessingKind::BitSinglets], 3);
        assert_eq!(c[PreprocessingKind::BitTriples], 5);
        assert_eq!(c[PreprocessingKind::MpcFieldDaBits], 0);
        assert_eq!(c[PreprocessingKind::MpcFieldSinglets], 3);
        assert_eq!(c[PreprocessingKind::MpcFieldTriples], 2);
        assert_eq!(c[PreprocessingKind::ScalarDaBits], 3);
        assert_eq!(c[PreprocessingKind::BaseFieldDaBits], 5);
    }
}