arcium-core-utils 0.7.3

use std::hash::Hash;

use primitives::algebra::{
    elliptic_curve::{Point, Scalar},
    BoxedUint,
};
use serde::{Deserialize, Serialize};
use wincode::{SchemaRead, SchemaWrite};

use crate::{
    circuit::{
        errors::CircuitError,
        AlgebraicType,
        BatchSize,
        BitPlaintextBinaryOp,
        BitPlaintextUnaryOp,
        BitShareBinaryOp,
        BitShareUnaryOp,
        Constant,
        FieldPlaintextBinaryOp,
        FieldPlaintextUnaryOp,
        FieldShareBinaryOp,
        FieldShareUnaryOp,
        FieldType,
        GateIndex,
        Input,
        PointPlaintextBinaryOp,
        PointPlaintextUnaryOp,
        PointShareBinaryOp,
        PointShareUnaryOp,
        Slice,
    },
    config::MpcConfig,
};

/// Gate operations, where the operation arguments correspond to _wire_ label.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, SchemaRead, SchemaWrite)]
#[serde(bound(
    serialize = "Scalar<C::Curve>: Serialize, Point<C::Curve>: Serialize",
    deserialize = "Scalar<C::Curve>: Deserialize<'de>, Point<C::Curve>: Deserialize<'de>"
))]
#[repr(C)]
pub enum Gate<C: MpcConfig> {
    /// Input a wire
    Input(Input),
    /// Input a constant value
    Constant(Constant<C>),
    /// Generate random shares
    Random {
        algebraic_type: AlgebraicType,
        batch_size: BatchSize,
    },
    /// Field share unary operations
    FieldShareUnaryOp {
        x: GateIndex,
        op: FieldShareUnaryOp,
    },
    /// Field share binary operations, where the second wire may be a plaintext.
    FieldShareBinaryOp {
        x: GateIndex,
        y: GateIndex,
        op: FieldShareBinaryOp,
    },
    BatchSummation {
        x: GateIndex,
    },
    BitShareUnaryOp {
        x: GateIndex,
        op: BitShareUnaryOp,
    },
    BitShareBinaryOp {
        x: GateIndex,
        y: GateIndex,
        op: BitShareBinaryOp,
    },
    /// Operations with elliptic curve points
    PointShareUnaryOp {
        p: GateIndex,
        op: PointShareUnaryOp,
    },
    PointShareBinaryOp {
        p: GateIndex,
        y: GateIndex,
        op: PointShareBinaryOp,
    },
    /// Field plaintext unary operations
    FieldPlaintextUnaryOp {
        x: GateIndex,
        op: FieldPlaintextUnaryOp,
    },
    /// Field plaintext binary operations
    FieldPlaintextBinaryOp {
        x: GateIndex,
        y: GateIndex,
        op: FieldPlaintextBinaryOp,
    },
    BitPlaintextUnaryOp {
        x: GateIndex,
        op: BitPlaintextUnaryOp,
    },
    BitPlaintextBinaryOp {
        x: GateIndex,
        y: GateIndex,
        op: BitPlaintextBinaryOp,
    },
    PointPlaintextUnaryOp {
        p: GateIndex,
        op: PointPlaintextUnaryOp,
    },
    PointPlaintextBinaryOp {
        p: GateIndex,
        y: GateIndex,
        op: PointPlaintextBinaryOp,
    },
    /// Request a daBit
    DaBit {
        field_type: FieldType,
        batch_size: BatchSize,
    },
    GetDaBitFieldShare {
        x: GateIndex,
    },
    GetDaBitSharedBit {
        x: GateIndex,
    },
    /// Base field exponentiation operation
    BaseFieldPow {
        x: GateIndex,
        exp: BoxedUint,
    },
    /// Bit plaintext conversion operations
    BitPlaintextToField {
        x: GateIndex,
        field_type: FieldType,
    },
    FieldPlaintextToBit {
        x: GateIndex,
    },
    /// Get a slice of elements from a batched wire
    ExtractFromBatch {
        x: GateIndex,
        slice: Slice,
    },
    CollectToBatch {
        wires: Vec<GateIndex>,
    },
    PointFromPlaintextCoordinates {
        wires: Vec<GateIndex>,
    },
    PlaintextPointToCoordinates {
        point: GateIndex,
    },
    PlaintextKeccakF1600 {
        x: GateIndex,
    },
    CompressPlaintextPoint {
        point: GateIndex,
    },
    KeyRecoveryPlaintextComputeErrors {
        d_minus_one: GateIndex,
        syndromes: GateIndex,
    },
    AesGcmKeyStream {
        round_keys: GateIndex,
        iv: GateIndex,
        n_ciphertext_blocks: u32,
    },
    #[cfg(any(test, feature = "dev"))]
    AesKeySchedule {
        key: GateIndex,
    },
    GhashPowersOfH {
        h: GateIndex,
        n_ciphertext_blocks: u32,
    },
    Ghash {
        x: GateIndex,
        powers_of_h: GateIndex,
    },
}

impl<C: MpcConfig> Gate<C> {
    /// Check whether the gate is an input gate.
    pub fn is_input(&self) -> bool {
        matches!(self, Gate::Input { .. })
    }

    /// Returns the indices of gate inputs.
    pub fn get_inputs(&self) -> Vec<GateIndex> {
        match &self {
            Gate::Input(_) | Gate::Random { .. } | Gate::Constant(_) | Gate::DaBit { .. } => {
                Vec::new()
            }

            Gate::FieldShareUnaryOp { x, .. }
            | Gate::BatchSummation { x, .. }
            | Gate::BitShareUnaryOp { x, .. }
            | Gate::PointShareUnaryOp { p: x, .. }
            | Gate::FieldPlaintextUnaryOp { x, .. }
            | Gate::BitPlaintextUnaryOp { x, .. }
            | Gate::PointPlaintextUnaryOp { p: x, .. }
            | Gate::GetDaBitFieldShare { x, .. }
            | Gate::GetDaBitSharedBit { x, .. }
            | Gate::BaseFieldPow { x, .. }
            | Gate::BitPlaintextToField { x, .. }
            | Gate::FieldPlaintextToBit { x, .. }
            | Gate::ExtractFromBatch { x, .. }
            | Gate::PlaintextPointToCoordinates { point: x, .. }
            | Gate::CompressPlaintextPoint { point: x, .. }
            | Gate::PlaintextKeccakF1600 { x }
            | Gate::GhashPowersOfH { h: x, .. } => {
                vec![*x]
            }
            #[cfg(any(test, feature = "dev"))]
            Gate::AesKeySchedule { key } => {
                vec![*key]
            }

            Gate::FieldShareBinaryOp { x, y, .. }
            | Gate::BitShareBinaryOp { x, y, .. }
            | Gate::PointShareBinaryOp { p: x, y, .. }
            | Gate::FieldPlaintextBinaryOp { x, y, .. }
            | Gate::BitPlaintextBinaryOp { x, y, .. }
            | Gate::PointPlaintextBinaryOp { p: x, y, .. }
            | Gate::KeyRecoveryPlaintextComputeErrors {
                d_minus_one: x,
                syndromes: y,
                ..
            }
            | Gate::AesGcmKeyStream {
                round_keys: x,
                iv: y,
                ..
            }
            | Gate::Ghash { x, powers_of_h: y } => {
                vec![*x, *y]
            }

            Gate::CollectToBatch { wires, .. }
            | Gate::PointFromPlaintextCoordinates { wires, .. } => wires.clone(),
        }
    }

    /// Maps inplace gate inputs using the given function.
    pub fn map_inputs<F: FnMut(GateIndex) -> GateIndex>(mut self, mut f: F) -> Self {
        match &mut self {
            Gate::Input(_) | Gate::Random { .. } | Gate::Constant(_) | Gate::DaBit { .. } => (),

            Gate::FieldShareUnaryOp { x, .. }
            | Gate::BatchSummation { x, .. }
            | Gate::BitShareUnaryOp { x, .. }
            | Gate::PointShareUnaryOp { p: x, .. }
            | Gate::FieldPlaintextUnaryOp { x, .. }
            | Gate::BitPlaintextUnaryOp { x, .. }
            | Gate::PointPlaintextUnaryOp { p: x, .. }
            | Gate::GetDaBitFieldShare { x, .. }
            | Gate::GetDaBitSharedBit { x, .. }
            | Gate::BaseFieldPow { x, .. }
            | Gate::BitPlaintextToField { x, .. }
            | Gate::FieldPlaintextToBit { x, .. }
            | Gate::ExtractFromBatch { x, .. }
            | Gate::PlaintextPointToCoordinates { point: x, .. }
            | Gate::CompressPlaintextPoint { point: x, .. }
            | Gate::PlaintextKeccakF1600 { x }
            | Gate::GhashPowersOfH { h: x, .. } => {
                *x = f(*x);
            }

            #[cfg(any(test, feature = "dev"))]
            Gate::AesKeySchedule { key } => {
                *key = f(*key);
            }

            Gate::FieldShareBinaryOp { x, y, .. }
            | Gate::BitShareBinaryOp { x, y, .. }
            | Gate::PointShareBinaryOp { p: x, y, .. }
            | Gate::FieldPlaintextBinaryOp { x, y, .. }
            | Gate::BitPlaintextBinaryOp { x, y, .. }
            | Gate::PointPlaintextBinaryOp { p: x, y, .. }
            | Gate::KeyRecoveryPlaintextComputeErrors {
                d_minus_one: x,
                syndromes: y,
                ..
            }
            | Gate::AesGcmKeyStream {
                round_keys: x,
                iv: y,
                ..
            }
            | Gate::Ghash { x, powers_of_h: y } => {
                *x = f(*x);
                *y = f(*y);
            }

            Gate::CollectToBatch { wires, .. }
            | Gate::PointFromPlaintextCoordinates { wires, .. } => {
                wires.iter_mut().for_each(|x| *x = f(*x))
            }
        };

        self
    }

    /// Tries to replace the gate inputs with the given ones.
    ///
    /// This function returns an error if the number of given inputs does not match the number of
    /// gate inputs.
    pub fn try_replace_inputs(mut self, inputs: Vec<GateIndex>) -> Result<Self, CircuitError<C>> {
        if inputs.len() != self.get_inputs().len() {
            return Err(CircuitError::InvalidGateInputCount {
                expected: self.get_inputs().len(),
                found: inputs.len(),
            });
        }

        match &mut self {
            Gate::Input(_) | Gate::Random { .. } | Gate::Constant(_) | Gate::DaBit { .. } => (),

            Gate::FieldShareUnaryOp { x, .. }
            | Gate::BatchSummation { x, .. }
            | Gate::BitShareUnaryOp { x, .. }
            | Gate::PointShareUnaryOp { p: x, .. }
            | Gate::FieldPlaintextUnaryOp { x, .. }
            | Gate::BitPlaintextUnaryOp { x, .. }
            | Gate::PointPlaintextUnaryOp { p: x, .. }
            | Gate::GetDaBitFieldShare { x, .. }
            | Gate::GetDaBitSharedBit { x, .. }
            | Gate::BaseFieldPow { x, .. }
            | Gate::BitPlaintextToField { x, .. }
            | Gate::FieldPlaintextToBit { x, .. }
            | Gate::ExtractFromBatch { x, .. }
            | Gate::PlaintextPointToCoordinates { point: x, .. }
            | Gate::CompressPlaintextPoint { point: x, .. }
            | Gate::PlaintextKeccakF1600 { x }
            | Gate::GhashPowersOfH { h: x, .. } => {
                *x = inputs[0];
            }
            #[cfg(any(test, feature = "dev"))]
            Gate::AesKeySchedule { key } => {
                *key = inputs[0];
            }

            Gate::FieldShareBinaryOp { x, y, .. }
            | Gate::BitShareBinaryOp { x, y, .. }
            | Gate::PointShareBinaryOp { p: x, y, .. }
            | Gate::FieldPlaintextBinaryOp { x, y, .. }
            | Gate::BitPlaintextBinaryOp { x, y, .. }
            | Gate::PointPlaintextBinaryOp { p: x, y, .. }
            | Gate::KeyRecoveryPlaintextComputeErrors {
                d_minus_one: x,
                syndromes: y,
                ..
            }
            | Gate::AesGcmKeyStream {
                round_keys: x,
                iv: y,
                ..
            }
            | Gate::Ghash { x, powers_of_h: y } => {
                *x = inputs[0];
                *y = inputs[1];
            }

            Gate::CollectToBatch { wires, .. }
            | Gate::PointFromPlaintextCoordinates { wires, .. } => *wires = inputs,
        };

        Ok(self)
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use super::*;
    use crate::{circuit::FieldShareBinaryOp, config::DefaultConfig as C};

    #[test]
    fn test_ser_gate() {
        let scalar_gate: Gate<C> = Gate::FieldShareBinaryOp {
            x: 1,
            y: 3,
            op: FieldShareBinaryOp::Add,
        };
        let point_gate: Gate<C> = Gate::PointShareBinaryOp {
            p: 1,
            y: 3,
            op: PointShareBinaryOp::Add,
        };

        let scalar_gate_ser = bincode::serialize(&scalar_gate).unwrap();
        let point_gate_ser = bincode::serialize(&point_gate).unwrap();

        let scalar_gate_de: Gate<C> = bincode::deserialize(&scalar_gate_ser).unwrap();
        let point_gate_de: Gate<C> = bincode::deserialize(&point_gate_ser).unwrap();

        assert_eq!(scalar_gate, scalar_gate_de);
        assert_eq!(point_gate, point_gate_de);
        let set = HashSet::from([scalar_gate, scalar_gate_de, point_gate, point_gate_de]);
        assert_eq!(set.len(), 2)
    }
}