arcium-core-utils 0.8.6

use primitives::algebra::elliptic_curve::NumCoordinates;
use typenum::Unsigned;

use crate::{
    circuit::{
        errors::{CircuitError, ConversionError},
        gate::Gate,
        preprocessing::AES_S_BOX_N_NETWORK_ROUNDS,
        AlgebraicType,
        BitShareBinaryOp,
        BitShareUnaryOp,
        FieldShareBinaryOp,
        FieldShareUnaryOp,
        FieldType,
        GateIndex,
        Input,
        PlaintextBitConstraint,
        PointPlaintextBinaryOp,
        PointShareBinaryOp,
        PointShareUnaryOp,
        ShareOrPlaintext,
        MAX_EXPR_DEPTH,
        MAX_EXPR_NODES,
    },
    config::MpcConfig,
    key_recovery::{MXE_KEY_RECOVERY_D, MXE_KEY_RECOVERY_N},
};

/// A circuit composed of a sequence of gates, input and output identifiers.
///
/// Each circuit gate contains additional information about its output characteristics, refer to
///  the ` GateExt ` struct for more details.
///
/// The circuit is always valid because the gate addition is validated.
#[derive(Default, PartialEq, Debug, Clone)]
pub struct Circuit<C: MpcConfig> {
    /// The circuit gates.
    pub(super) gates: Vec<GateExt<C>>,
    /// The input gates in order of definition
    pub(super) inputs: Vec<GateIndex>,
    /// The output gates in order of definition
    pub(super) outputs: Vec<GateIndex>,
}

/// A circuit gate together with additional information about its output.
/// The additional information is automatically deduced when the gate is added to the circuit.
#[derive(Clone, Debug, PartialEq)]
pub struct GateExt<C: MpcConfig> {
    pub gate: Gate<C>,
    pub output: GateOutput,
    pub level: GateLevel,
}

/// Gate output characteristics like algebraic type, visibility, and batch size
#[derive(PartialEq, Copy, Clone, Debug)]
pub struct GateOutput {
    pub(super) algebraic_type: AlgebraicType,
    pub(super) form: ShareOrPlaintext,
    pub(super) batch_size: u32,
    /// Whether the wire carries daBits rather than plain elements of `algebraic_type`.
    ///
    /// A daBit is a bit authenticated in both backends at once, so it is a different Rust type
    /// from a share of the same field and needs a different task. The two are indistinguishable
    /// by `algebraic_type` and `form` alone, which is why this is tracked separately.
    pub(super) is_dabit: bool,
}

/// The level of a gate in a circuit. All gates with the same level can be executed in parallel as
/// they do not depend on each other.
///
/// The gate level is a pair of integers: the first one is the communication round (level), and the
/// second one is a relative level within the same communication round. The gate communication round
/// is the number of communications rounds which have passed after the gate execution.
#[derive(Copy, Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Default)]
pub struct GateLevel {
    comm_level: usize,
    level: usize, // relative counter for ordering gates within a multiplicative level
}

impl<C: MpcConfig> GateExt<C> {
    pub fn new(gate: Gate<C>, output: GateOutput, level: GateLevel) -> Self {
        Self {
            gate,
            output,
            level,
        }
    }
}

impl GateOutput {
    pub fn get_type(&self) -> AlgebraicType {
        self.algebraic_type
    }

    pub fn get_field_type(&self) -> Result<FieldType, ConversionError> {
        FieldType::try_from(self.algebraic_type)
    }

    pub fn get_field_type_unchecked(&self) -> FieldType {
        self.get_field_type().unwrap()
    }

    pub fn get_form(&self) -> ShareOrPlaintext {
        self.form
    }

    /// Whether this wire carries daBits. See [`GateOutput::is_dabit`].
    pub fn is_dabit(&self) -> bool {
        self.is_dabit
    }

    pub fn get_batch_size(&self) -> u32 {
        self.batch_size
    }

    fn is_field(&self) -> bool {
        FieldType::try_from(self.algebraic_type).is_ok()
    }

    fn is_bit(&self) -> bool {
        self.algebraic_type == AlgebraicType::Bit
    }

    fn is_point(&self) -> bool {
        self.algebraic_type == AlgebraicType::Point
    }

    fn is_base_field(&self) -> bool {
        self.algebraic_type == AlgebraicType::BaseField
    }

    fn is_scalar_field(&self) -> bool {
        self.algebraic_type == AlgebraicType::ScalarField
    }

    fn is_share(&self) -> bool {
        self.form == ShareOrPlaintext::Share
    }

    pub fn is_plaintext(&self) -> bool {
        self.form == ShareOrPlaintext::Plaintext
    }

    fn with_new_form(self, form: ShareOrPlaintext) -> Self {
        let mut res = self;
        res.form = form;
        res
    }

    fn with_new_type(self, algebraic_type: AlgebraicType) -> Self {
        let mut res = self;
        res.algebraic_type = algebraic_type;
        res
    }

    fn with_new_batch_size(self, batch_size: u32) -> Self {
        let mut res = self;
        res.batch_size = batch_size;
        res
    }

    /// For the gates that read one component out of a daBit: the result is a plain share.
    fn without_dabit(self) -> Self {
        let mut res = self;
        res.is_dabit = false;
        res
    }
}

impl GateLevel {
    fn next(&self, comm_rounds: usize) -> GateLevel {
        if comm_rounds > 0 {
            GateLevel {
                comm_level: self.comm_level + comm_rounds,
                level: 0,
            }
        } else {
            GateLevel {
                comm_level: self.comm_level,
                level: self.level + 1,
            }
        }
    }

    pub fn comm_level(&self) -> usize {
        self.comm_level
    }
}

impl<C: MpcConfig> Circuit<C> {
    pub fn new() -> Self {
        Self::default()
    }

    /// Tries to add a gate to the circuit.
    ///
    /// This function validates the gate before adding the gate or fails otherwise.
    pub fn add_gate(&mut self, gate: Gate<C>) -> Result<GateIndex, CircuitError<C>> {
        self.validate_gate(&gate)?;

        let index = self.nb_gates();
        if index == GateIndex::MAX {
            return Err(CircuitError::CircuitTooBig);
        }

        let gate_output = self.comp_gate_output(&gate);
        let level = self.comp_gate_level(&gate);

        if gate.is_input() {
            self.inputs.push(index);
        }
        self.gates.push(GateExt::new(gate, gate_output?, level));

        Ok(index)
    }

    /// Tries to set a gate as circuit output.
    ///
    /// This function fails if there is no gate with the given index.
    pub fn add_output(&mut self, index: GateIndex) -> Result<(), CircuitError<C>> {
        if index < self.nb_gates() {
            self.outputs.push(index);
            Ok(())
        } else {
            Err(CircuitError::GateIndexOutOfBounds(index, self.nb_gates()))
        }
    }

    pub fn nb_gates(&self) -> GateIndex {
        self.gates.len() as GateIndex
    }

    pub fn nb_inputs(&self) -> GateIndex {
        self.inputs.len() as GateIndex
    }

    pub fn nb_outputs(&self) -> GateIndex {
        self.outputs.len() as GateIndex
    }

    /// Consumes the circuit and returns the list of gates.
    pub fn into_gates(self) -> Vec<GateExt<C>> {
        self.gates
    }

    pub fn iter_gates_ext(
        &self,
    ) -> impl ExactSizeIterator<Item = &GateExt<C>> + DoubleEndedIterator {
        self.gates.iter()
    }

    pub fn iter_gates(&self) -> impl ExactSizeIterator<Item = &Gate<C>> + DoubleEndedIterator {
        self.gates.iter().map(|g| &g.gate)
    }

    pub fn iter_output_indices(&self) -> impl ExactSizeIterator<Item = &GateIndex> {
        self.outputs.iter()
    }

    pub fn iter_input_indices(&self) -> impl ExactSizeIterator<Item = &GateIndex> {
        self.inputs.iter()
    }

    pub fn gate_ext(&self, index: GateIndex) -> Result<&GateExt<C>, CircuitError<C>> {
        if index < self.nb_gates() {
            Ok(&self.gates[index as usize])
        } else {
            Err(CircuitError::GateIndexOutOfBounds(index, self.nb_gates()))
        }
    }

    pub fn gate_ext_unchecked(&self, index: GateIndex) -> &GateExt<C> {
        &self.gates[index as usize]
    }

    pub fn gate(&self, index: GateIndex) -> Result<&Gate<C>, CircuitError<C>> {
        self.gate_ext(index).map(|g| &g.gate)
    }

    pub fn gate_unchecked(&self, index: GateIndex) -> &Gate<C> {
        &self.gate_ext_unchecked(index).gate
    }

    pub fn gate_output(&self, index: GateIndex) -> Result<GateOutput, CircuitError<C>> {
        self.gate_ext(index).map(|g| g.output)
    }

    pub fn gate_output_unchecked(&self, index: GateIndex) -> GateOutput {
        self.gate_ext_unchecked(index).output
    }

    pub fn gate_level(&self, index: GateIndex) -> Result<GateLevel, CircuitError<C>> {
        self.gate_ext(index).map(|g| g.level)
    }

    pub fn gate_level_unchecked(&self, index: GateIndex) -> GateLevel {
        self.gate_ext_unchecked(index).level
    }
}

macro_rules! check_algebraic_type {
    ($exp_type:expr, $found_type:expr) => {
        if $exp_type != $found_type {
            return Err(CircuitError::InvalidGateAlgebraicType {
                expected: $exp_type,
                found: $found_type,
            });
        }
    };
}

impl<C: MpcConfig> Circuit<C> {
    /// Opens and outputs as scalar field a given gate
    pub fn open_and_output_scalar(&mut self, x: GateIndex) -> Result<(), CircuitError<C>> {
        check_algebraic_type!(AlgebraicType::ScalarField, self.gate_output(x)?.get_type());
        let opening_index = self.add_gate(Gate::FieldShareUnaryOp {
            x,
            op: FieldShareUnaryOp::Open,
        })?;
        self.add_output(opening_index)
    }

    /// Opens and outputs as base field a given gate
    pub fn open_and_output_base_field(&mut self, x: GateIndex) -> Result<(), CircuitError<C>> {
        check_algebraic_type!(AlgebraicType::BaseField, self.gate_output(x)?.get_type());
        let opening_index = self.add_gate(Gate::FieldShareUnaryOp {
            x,
            op: FieldShareUnaryOp::Open,
        })?;
        self.add_output(opening_index)
    }

    /// Opens and outputs a given gate as the MPC field
    pub fn open_and_output_mpc_field(&mut self, x: GateIndex) -> Result<(), CircuitError<C>> {
        check_algebraic_type!(AlgebraicType::MpcField, self.gate_output(x)?.get_type());
        let opening_index = self.add_gate(Gate::FieldShareUnaryOp {
            x,
            op: FieldShareUnaryOp::Open,
        })?;
        self.add_output(opening_index)
    }

    /// Opens and outputs as point a given gate
    pub fn open_and_output_point(&mut self, p: GateIndex) -> Result<(), CircuitError<C>> {
        check_algebraic_type!(AlgebraicType::Point, self.gate_output(p)?.get_type());
        let opening_index = self.add_gate(Gate::PointShareUnaryOp {
            p,
            op: PointShareUnaryOp::Open,
        })?;
        self.add_output(opening_index)
    }

    /// Opens and outputs as bit a given gate
    pub fn open_and_output_bit(&mut self, x: GateIndex) -> Result<(), CircuitError<C>> {
        check_algebraic_type!(AlgebraicType::Bit, self.gate_output(x)?.get_type());
        let opening_index = self.add_gate(Gate::BitShareUnaryOp {
            x,
            op: BitShareUnaryOp::Open,
        })?;
        self.add_output(opening_index)
    }
}

impl<C: MpcConfig> Circuit<C> {
    /// Validates a gate.
    ///
    /// Checks that:
    ///     - gate inputs are present in the circuit
    ///     - gate input types correspond to the specification
    ///     - gate input batch sizes are compatible
    ///     - gate parameters are valid
    fn validate_gate(&self, gate: &Gate<C>) -> Result<(), CircuitError<C>> {
        macro_rules! check_op {
            ($msg:expr, $gate:expr => $($val1:expr, $op:tt, $val2:expr);+$(;)?) => {
                $(if !($val1 $op $val2) {
                    return Err(CircuitError::InvalidGate(
                        $gate.clone(),
                        format!("{}: {:?} {:?} {:?}", $msg, $val1, stringify!($op), $val2),
                    ));
                })+
            };
        }

        macro_rules! check_gate_properties {
            ($gate:expr, $($func:ident),* $(,)?) => {
                $(if !($gate.output.$func()) {
                    return Err(CircuitError::InvalidGate(
                        $gate.gate.clone(),
                        format!("{:?} fails - {}", $gate.output, stringify!($func)))); }
                )*
            };
        }

        match gate {
            Gate::Input(input) => {
                check_op!(
                    "input batch size must be non-zero",
                    gate =>
                    0, <, input.batch_size();
                );
            }
            Gate::Constant(constant) => {
                check_op!(
                    "constant batch size must be non-zero",
                    gate =>
                     0, <, constant.batch_size()?;
                );
            }
            Gate::Random { batch_size, .. } => {
                check_op!(
                    "random batch size must be non-zero",
                    gate =>
                     0, <, *batch_size;
                );
            }
            Gate::FieldShareUnaryOp { x, .. } => {
                check_gate_properties!(self.gate_ext(*x)?, is_field, is_share);
            }
            Gate::FieldShareBinaryOp { x, y, .. } => {
                let (gx, gy) = (self.gate_ext(*x)?, self.gate_ext(*y)?);
                check_gate_properties!(gx, is_field, is_share);
                check_gate_properties!(gy, is_field);
                check_op!(
                    "inputs must have same batch-size and field type",
                    gate =>
                    gx.output.batch_size, ==, gy.output.batch_size;
                    gx.output.algebraic_type, ==, gy.output.algebraic_type
                );
            }
            Gate::BatchSummation { x } => {
                self.gate_ext(*x)?;
            }
            Gate::BitShareUnaryOp { x, .. } => {
                check_gate_properties!(self.gate_ext(*x)?, is_bit, is_share);
            }
            Gate::BitShareBinaryOp { x, y, .. } => {
                let (gx, gy) = (self.gate_ext(*x)?, self.gate_ext(*y)?);
                check_gate_properties!(gx, is_bit, is_share);
                check_gate_properties!(gy, is_bit);
                check_op!(
                    "inputs must have same batch-size",
                    gate =>
                    gx.output.batch_size, ==, gy.output.batch_size
                );
            }
            Gate::PointShareUnaryOp { p: x, .. } => {
                check_gate_properties!(self.gate_ext(*x)?, is_point, is_share);
            }
            Gate::PointShareBinaryOp { p: x, y, op } => {
                let (gx, gy) = (self.gate_ext(*x)?, self.gate_ext(*y)?);
                if gx.output.is_plaintext() && gy.output.is_plaintext() {
                    return Err(CircuitError::InvalidGate(
                        gate.clone(),
                        "at least one input must be share".to_string(),
                    ));
                }
                check_gate_properties!(gx, is_point);
                match op {
                    PointShareBinaryOp::Add => {
                        check_gate_properties!(gy, is_point);
                    }
                    PointShareBinaryOp::ScalarMul => {
                        check_gate_properties!(gy, is_scalar_field);
                    }
                };
                check_op!(
                    "inputs must have same batch-size",
                    gate =>
                    gx.output.batch_size, ==, gy.output.batch_size
                );
            }
            Gate::FieldPlaintextUnaryOp { x, .. } => {
                check_gate_properties!(self.gate_ext(*x)?, is_field, is_plaintext);
            }
            Gate::FieldPlaintextBinaryOp { x, y, .. } => {
                let (gx, gy) = (self.gate_ext(*x)?, self.gate_ext(*y)?);
                check_gate_properties!(gx, is_field, is_plaintext);
                check_gate_properties!(gy, is_field, is_plaintext);
                check_op!(
                    "inputs must have same field type",
                    gate =>
                    gx.output.algebraic_type, ==, gy.output.algebraic_type;
                    gx.output.batch_size, ==, gy.output.batch_size
                );
            }
            Gate::BitPlaintextUnaryOp { x, .. } => {
                check_gate_properties!(self.gate_ext(*x)?, is_bit, is_plaintext);
            }
            Gate::BitPlaintextBinaryOp { x, y, .. } => {
                let (gx, gy) = (self.gate_ext(*x)?, self.gate_ext(*y)?);
                check_gate_properties!(gx, is_bit, is_plaintext);
                check_gate_properties!(gy, is_bit, is_plaintext);
                check_op!(
                    "inputs must have same batch-size",
                    gate =>
                    gx.output.batch_size, ==, gy.output.batch_size
                );
            }
            Gate::PointPlaintextUnaryOp { p: x, .. } => {
                check_gate_properties!(self.gate_ext(*x)?, is_point, is_plaintext);
            }
            Gate::PointPlaintextBinaryOp { p: x, y, op } => {
                let (gx, gy) = (self.gate_ext(*x)?, self.gate_ext(*y)?);
                check_gate_properties!(gx, is_point, is_plaintext);
                match op {
                    PointPlaintextBinaryOp::Add => {
                        check_gate_properties!(gy, is_point, is_plaintext);
                    }
                    PointPlaintextBinaryOp::ScalarMul => {
                        check_gate_properties!(gy, is_scalar_field, is_plaintext);
                    }
                }
                check_op!(
                    "inputs must have same batch-size",
                    gate =>
                    gx.output.batch_size, ==, gy.output.batch_size
                );
            }
            Gate::DaBit { batch_size, .. } => {
                check_op!(
                    "input batch size must be non-zero",
                    gate =>
                    0, <, *batch_size
                );
            }
            Gate::GetDaBitFieldShare { x, .. } => {
                // By convention, we suppose that the `DaBit` output is a field element
                check_gate_properties!(self.gate_ext(*x)?, is_field, is_share);
            }
            Gate::GetDaBitSharedBit { x, .. } => {
                // By convention, we suppose that the `DaBit` output is a field element
                check_gate_properties!(self.gate_ext(*x)?, is_field, is_share);
            }
            Gate::BaseFieldPow { x, .. } => {
                check_gate_properties!(self.gate_ext(*x)?, is_base_field, is_share);
            }
            Gate::BitPlaintextToField { x, .. } => {
                check_gate_properties!(self.gate_ext(*x)?, is_bit, is_plaintext);
            }
            Gate::FieldPlaintextToBit { x, .. } => {
                check_gate_properties!(self.gate_ext(*x)?, is_field, is_plaintext);
            }
            Gate::ExtractFromBatch { x, slice, .. } => {
                let gx = self.gate_ext(*x)?;
                if slice.is_empty() {
                    return Err(CircuitError::InvalidGate(
                        gate.clone(),
                        format!("slice must be non-empty: {slice:?}"),
                    ));
                }
                if slice.get_indices()
                    .into_iter()
                    .max()
                    .expect("non-empty slice expected") // never fails as we check that the slice is non-empty
                    >= gx.output.batch_size
                {
                    return Err(CircuitError::InvalidGate(
                        gate.clone(),
                        format!("slice indices out-of-range: {slice:?}"),
                    ));
                }
            }
            Gate::CollectToBatch { wires } => {
                check_op!("expected at least one input", gate => 0, <, wires.len());
                let first = self.gate_ext(wires[0])?.output;
                for x in wires.iter().skip(1) {
                    let gx = self.gate_ext(*x)?.output;
                    check_op!(
                        "all inputs must have the same type",
                        gate =>
                        first.algebraic_type, ==, gx.algebraic_type;
                        first.form, ==, gx.form;
                        first.is_dabit, ==, gx.is_dabit
                    );
                }
            }
            Gate::GatherFromBatches { parts } => {
                check_op!("expected at least one part", gate => 0, <, parts.len());
                let first = self.gate_output(parts[0].0)?;
                for (x, slice) in parts {
                    let gx = self.gate_output(*x)?;
                    check_op!(
                        "all parts must have the same type",
                        gate =>
                        first.algebraic_type, ==, gx.algebraic_type;
                        first.form, ==, gx.form;
                        first.is_dabit, ==, gx.is_dabit
                    );
                    if slice.is_empty() {
                        return Err(CircuitError::InvalidGate(
                            gate.clone(),
                            format!("slice must be non-empty: {slice:?}"),
                        ));
                    }
                    if slice
                        .get_indices()
                        .into_iter()
                        .max()
                        .expect("non-empty slice expected") // never fails, the slice is non-empty
                        >= gx.batch_size
                    {
                        return Err(CircuitError::InvalidGate(
                            gate.clone(),
                            format!("slice indices out-of-range: {slice:?}"),
                        ));
                    }
                }
            }
            Gate::PointFromPlaintextCoordinates { wires } => {
                check_op!(
                    "expected one input per point coordinate",
                    gate => wires.len(), ==, NumCoordinates::<C::Curve>::USIZE
                );
                for x in wires {
                    let gx = self.gate_ext(*x)?;
                    check_gate_properties!(gx, is_base_field, is_plaintext);
                    check_op!("expected batch-size 1", gate => gx.output.batch_size, ==, 1);
                }
            }
            Gate::PlaintextPointToCoordinates { point: x, .. } => {
                let gx = self.gate_ext(*x)?;
                check_gate_properties!(gx, is_point, is_plaintext);
                check_op!("expected batch-size 1", gate => gx.output.batch_size, ==, 1);
            }
            Gate::PlaintextKeccakF1600 { x } => {
                let gx = self.gate_ext(*x)?;
                check_gate_properties!(gx, is_bit, is_plaintext);
                check_op!("expected batch-size 1600", gate => gx.output.batch_size, ==, 1600);
            }
            Gate::CompressPlaintextPoint { point: x, .. } => {
                let gx = self.gate_ext(*x)?;
                check_gate_properties!(gx, is_point, is_plaintext);
                check_op!("expected batch-size 1", gate => gx.output.batch_size, ==, 1);
            }
            Gate::KeyRecoveryPlaintextComputeErrors {
                d_minus_one,
                syndromes,
            } => {
                let g1 = self.gate_ext(*d_minus_one)?;
                let g2 = self.gate_ext(*syndromes)?;
                check_gate_properties!(g1, is_base_field, is_plaintext);
                check_gate_properties!(g2, is_base_field, is_plaintext);

                check_op!("expected batch-size 1", gate => g1.output.batch_size, ==, 1);
                // TODO: Check that the batch size of `g2` is correct.
                check_op!(format!("expected batch-size {}", MXE_KEY_RECOVERY_D - 1),
                    gate => g2.output.batch_size, ==, MXE_KEY_RECOVERY_D as u32 - 1);
            }
            Gate::AesKeySchedule { key } => {
                let g = self.gate_ext(*key)?;
                check_gate_properties!(g, is_bit, is_share);
                let key_length = g.output.batch_size;
                check_op!("key_length: expected 128, 192 or 256", gate => matches!(key_length, 128 | 192 | 256), ==, true);
            }
            Gate::AesGcmKeyStream {
                round_keys,
                iv,
                n_ciphertext_blocks,
            } => {
                let g1 = self.gate_ext(*round_keys)?;
                let g2 = self.gate_ext(*iv)?;
                check_gate_properties!(g1, is_bit, is_share);
                check_gate_properties!(g2, is_bit, is_plaintext);
                let is_multiple = g1.output.batch_size.is_multiple_of(128);
                check_op!("round_keys: expected batch-size multiple of 128", gate => is_multiple, ==, true);
                let n_round_key_blocks = g1.output.batch_size / 128;
                // 11, 13 and 15 correspond to AES-128, AES-192 and AES-256 respectively
                check_op!("round_keys: expected 11, 13 or 15 128-bit round_keys", gate => matches!(n_round_key_blocks, 11 | 13 | 15), ==, true);
                check_op!("IV: expected batch-size 96", gate => g2.output.batch_size, ==, 96);
                // The batch size of a gate input/ouput is a u32, and since this task outputs
                // (1+n_ciphertext_blocks) * 128 bit shares we require n_ciphertext_blocks to be
                // less than 2^24.
                check_op!("n_ciphertext_blocks: expected at most 2^24 - 1", gate => *n_ciphertext_blocks, <, 1u32 << 24);
            }
            Gate::GhashPowersOfH {
                h,
                n_ciphertext_blocks,
            } => {
                let g = self.gate_ext(*h)?;
                check_gate_properties!(g, is_bit, is_share);
                let len = g.output.batch_size;
                check_op!("batch size of h must be 128", gate => len, ==, 128);
                check_op!("n_ciphertext_blocks must be positive", gate => *n_ciphertext_blocks, >, 0);
            }
            Gate::Ghash { x, powers_of_h } => {
                let g1 = self.gate_ext(*x)?;
                let g2 = self.gate_ext(*powers_of_h)?;
                check_gate_properties!(g1, is_bit, is_plaintext);
                check_gate_properties!(g2, is_bit, is_share);
                let x_len = g1.output.batch_size;
                let powers_of_h_len = g2.output.batch_size;
                check_op!("batch size of x must be equal to batch size of powers_of_h", gate => x_len, ==, powers_of_h_len);
                check_op!("x must be non-empty", gate => x_len, !=, 0);
                check_op!("batch size of x must be a multiple of 128", gate => x_len.is_multiple_of(128), ==, true);
            }
            Gate::ConstrainPlaintextBits { x, clauses, .. } => {
                let gx = self.gate_ext(*x)?;
                check_gate_properties!(gx, is_bit, is_plaintext);
                let batch_size = gx.output.batch_size;
                check_op!("expected at least one clause", gate => clauses.is_empty(), ==, false);

                // Resolves how wide a wire is, for the operand length checks below. `None` for a
                // wire that is out of range or not a plaintext bit batch; the per-wire checks
                // report those properly.
                let wire_bits = |wire: GateIndex| -> Option<u32> {
                    let g = self.gate_ext(wire).ok()?;
                    (g.output.is_bit() && g.output.is_plaintext()).then_some(g.output.batch_size)
                };

                for (i, clause) in clauses.iter().enumerate() {
                    check_op!(format!("clause {i} is empty"), gate => clause.is_empty(), ==, false);

                    // Whichever clause ends up satisfied is the only one pinning the bits down, so
                    // coverage has to hold clause by clause rather than across the disjunction.
                    let mut covered = vec![false; batch_size as usize];
                    for constraint in clause.constraints() {
                        for operand in constraint.operands() {
                            check_op!(
                                format!("clause {i}: expression nests deeper than {MAX_EXPR_DEPTH}"),
                                gate => operand.depth(), <=, MAX_EXPR_DEPTH
                            );
                            check_op!(
                                format!("clause {i}: expression has more than {MAX_EXPR_NODES} nodes"),
                                gate => operand.node_count(), <=, MAX_EXPR_NODES
                            );
                            for slice in operand.slices() {
                                check_op!(format!("clause {i}: slices must be non-empty"), gate => slice.is_empty(), ==, false);
                                for index in slice.get_indices() {
                                    check_op!(
                                        format!("clause {i}: slice index out-of-range: {slice:?}"),
                                        gate => index, <, batch_size
                                    );
                                }
                            }
                            for wire in operand.wires() {
                                let g = self.gate_ext(wire)?;
                                check_gate_properties!(g, is_bit, is_plaintext);
                            }
                        }

                        match constraint {
                            PlaintextBitConstraint::Signature {
                                scheme,
                                signature,
                                message,
                                public_key,
                            } => {
                                check_op!(
                                    format!("clause {i}: expected a {}-bit {scheme:?} signature", scheme.signature_bits()),
                                    gate => signature.static_len(&wire_bits), ==, Some(scheme.signature_bits())
                                );
                                check_op!(
                                    format!("clause {i}: expected a {}-bit {scheme:?} public key", scheme.public_key_bits()),
                                    gate => public_key.static_len(&wire_bits), ==, Some(scheme.public_key_bits())
                                );
                                // The message is signed as bytes, so a partial trailing byte could
                                // never verify.
                                let message_bits = message.static_len(&wire_bits);
                                check_op!(
                                    format!("clause {i}: signed message must be a whole number of bytes"),
                                    gate => message_bits.is_none_or(|bits| bits.is_multiple_of(8)), ==, true
                                );
                            }
                            PlaintextBitConstraint::Equality { bits, expected } => {
                                let (lhs, rhs) =
                                    (bits.static_len(&wire_bits), expected.static_len(&wire_bits));
                                // Only enforced when both widths are known; an operand of
                                // as-yet-unknown width is the case `static_len` leaves room for.
                                check_op!(
                                    format!("clause {i}: the two sides of an equality must be the same length"),
                                    gate => lhs.zip(rhs).is_none_or(|(lhs, rhs)| lhs == rhs), ==, true
                                );
                            }
                            PlaintextBitConstraint::Comparison { lhs, rhs, .. } => {
                                let (lhs, rhs) =
                                    (lhs.static_len(&wire_bits), rhs.static_len(&wire_bits));
                                check_op!(
                                    format!("clause {i}: the two sides of a comparison must be the same length"),
                                    gate => lhs.zip(rhs).is_none_or(|(lhs, rhs)| lhs == rhs), ==, true
                                );
                                // The ordering is over bytes, so a partial trailing byte has no
                                // place in it.
                                check_op!(
                                    format!("clause {i}: a comparison's operands must be a whole number of bytes"),
                                    gate => lhs.into_iter().chain(rhs).all(|bits| bits.is_multiple_of(8)), ==, true
                                );
                            }
                        }

                        // Every operand was range-checked above; only the ones this constraint
                        // *pins* count towards coverage. See `covering_operands`.
                        for operand in constraint.covering_operands() {
                            for slice in operand.slices() {
                                for index in slice.get_indices() {
                                    covered[index as usize] = true;
                                }
                            }
                        }
                    }
                    let uncovered = covered.iter().filter(|c| !**c).count();
                    check_op!(
                        format!("clause {i}: every input bit must be covered by a constraint, {uncovered} are not"),
                        gate => uncovered, ==, 0
                    );
                }
            }
        }

        Ok(())
    }

    /// Computes the output type of gate.
    ///
    /// **Note: ** This function can panic if the gate is not valid.
    fn comp_gate_output(&self, gate: &Gate<C>) -> Result<GateOutput, CircuitError<C>> {
        let r = match gate {
            Gate::Input(input_type) => GateOutput {
                batch_size: input_type.batch_size(),
                algebraic_type: input_type.algebraic_type(),
                form: input_type.share_or_plaintext(),
                is_dabit: false,
            },

            Gate::Constant(const_type) => GateOutput {
                batch_size: const_type.batch_size()?,
                algebraic_type: const_type.algebraic_type(),
                form: ShareOrPlaintext::Plaintext,
                is_dabit: false,
            },

            Gate::Random {
                algebraic_type,
                batch_size,
            } => GateOutput {
                batch_size: *batch_size,
                algebraic_type: *algebraic_type,
                form: ShareOrPlaintext::Share,
                is_dabit: false,
            },

            Gate::FieldShareUnaryOp { x, op } => match op {
                FieldShareUnaryOp::Neg | FieldShareUnaryOp::MulInverse => {
                    self.gate_output_unchecked(*x)
                }
                FieldShareUnaryOp::Open | FieldShareUnaryOp::IsZero => self
                    .gate_output_unchecked(*x)
                    .with_new_form(ShareOrPlaintext::Plaintext),
            },

            Gate::FieldShareBinaryOp { x, .. }
            | Gate::BitShareBinaryOp { x, .. }
            | Gate::FieldPlaintextUnaryOp { x, .. }
            | Gate::FieldPlaintextBinaryOp { x, .. }
            | Gate::BitPlaintextUnaryOp { x, .. }
            | Gate::BitPlaintextBinaryOp { x, .. }
            | Gate::PointPlaintextUnaryOp { p: x, .. }
            | Gate::PointPlaintextBinaryOp { p: x, .. }
            | Gate::BaseFieldPow { x, .. } => self.gate_output_unchecked(*x),

            Gate::BatchSummation { x, .. } => self.gate_output_unchecked(*x).with_new_batch_size(1),

            Gate::PointShareBinaryOp { p: x, .. } => self
                .gate_output_unchecked(*x)
                .with_new_form(ShareOrPlaintext::Share),

            Gate::BitShareUnaryOp { x, op } => match op {
                BitShareUnaryOp::Not => self.gate_output_unchecked(*x),
                BitShareUnaryOp::Open => self
                    .gate_output_unchecked(*x)
                    .with_new_form(ShareOrPlaintext::Plaintext),
            },

            Gate::PointShareUnaryOp { p: x, op } => match op {
                PointShareUnaryOp::Neg => self.gate_output_unchecked(*x),
                PointShareUnaryOp::Open => self
                    .gate_output_unchecked(*x)
                    .with_new_form(ShareOrPlaintext::Plaintext),
                PointShareUnaryOp::IsZero => self
                    .gate_output_unchecked(*x)
                    .with_new_form(ShareOrPlaintext::Plaintext)
                    .with_new_type(AlgebraicType::ScalarField),
            },

            Gate::DaBit {
                field_type,
                batch_size,
            } => GateOutput {
                batch_size: *batch_size,
                algebraic_type: AlgebraicType::from(*field_type),
                form: ShareOrPlaintext::Share,
                is_dabit: true,
            },

            Gate::GetDaBitFieldShare { x, .. } => self.gate_output_unchecked(*x).without_dabit(),

            Gate::GetDaBitSharedBit { x, .. } => self
                .gate_output_unchecked(*x)
                .with_new_type(AlgebraicType::Bit)
                .without_dabit(),

            Gate::BitPlaintextToField { x, field_type } => self
                .gate_output_unchecked(*x)
                .with_new_type(AlgebraicType::from(*field_type)),

            Gate::FieldPlaintextToBit { x } => self
                .gate_output_unchecked(*x)
                .with_new_type(AlgebraicType::Bit),

            Gate::ExtractFromBatch { x, slice } => self
                .gate_output_unchecked(*x)
                .with_new_batch_size(slice.len()),

            Gate::CollectToBatch { wires, .. } => {
                let batch_size = wires
                    .iter()
                    .map(|x| self.gate_output_unchecked(*x).batch_size)
                    .sum();
                self.gate_output_unchecked(wires[0])
                    .with_new_batch_size(batch_size)
            }

            Gate::GatherFromBatches { parts } => {
                let batch_size = parts.iter().map(|(_, slice)| slice.len()).sum();
                self.gate_output_unchecked(parts[0].0)
                    .with_new_batch_size(batch_size)
            }

            Gate::PointFromPlaintextCoordinates { .. } => GateOutput {
                algebraic_type: AlgebraicType::Point,
                form: ShareOrPlaintext::Plaintext,
                batch_size: 1,
                is_dabit: false,
            },
            Gate::PlaintextPointToCoordinates { .. } => GateOutput {
                algebraic_type: AlgebraicType::BaseField,
                form: ShareOrPlaintext::Plaintext,
                batch_size: NumCoordinates::<C::Curve>::U32,
                is_dabit: false,
            },
            Gate::PlaintextKeccakF1600 { .. } => GateOutput {
                algebraic_type: AlgebraicType::Bit,
                form: ShareOrPlaintext::Plaintext,
                batch_size: 1600,
                is_dabit: false,
            },
            Gate::CompressPlaintextPoint { .. } => GateOutput {
                algebraic_type: AlgebraicType::Bit,
                form: ShareOrPlaintext::Plaintext,
                batch_size: 256,
                is_dabit: false,
            },
            Gate::KeyRecoveryPlaintextComputeErrors { .. } => GateOutput {
                algebraic_type: AlgebraicType::BaseField,
                form: ShareOrPlaintext::Plaintext,
                batch_size: MXE_KEY_RECOVERY_N as u32,
                is_dabit: false,
            },
            Gate::AesKeySchedule { key } => {
                let key_length = self.gate_output_unchecked(*key).batch_size;
                let n_round_key_bits = match key_length {
                    128 => 11 * 128,
                    192 => 13 * 128,
                    256 => 15 * 128,
                    _ => {
                        // the length of key was validated in Circuit::validate_gate
                        panic!("key length {key_length} does not match any of 128, 192 or 256");
                    }
                };
                GateOutput {
                    algebraic_type: AlgebraicType::Bit,
                    form: ShareOrPlaintext::Share,
                    batch_size: n_round_key_bits,
                    is_dabit: false,
                }
            }
            Gate::AesGcmKeyStream {
                n_ciphertext_blocks,
                ..
            } => GateOutput {
                algebraic_type: AlgebraicType::Bit,
                form: ShareOrPlaintext::Share,
                batch_size: (1 + n_ciphertext_blocks) * 128,
                is_dabit: false,
            },
            Gate::GhashPowersOfH {
                n_ciphertext_blocks,
                ..
            } => GateOutput {
                algebraic_type: AlgebraicType::Bit,
                form: ShareOrPlaintext::Share,
                batch_size: n_ciphertext_blocks * 128,
                is_dabit: false,
            },
            Gate::Ghash { .. } => GateOutput {
                algebraic_type: AlgebraicType::Bit,
                form: ShareOrPlaintext::Share,
                batch_size: 128,
                is_dabit: false,
            },
            // The agreed bits, plus a trailing success bit.
            Gate::ConstrainPlaintextBits { x, .. } => {
                let output = self.gate_output_unchecked(*x);
                let batch_size = output.batch_size;
                output.with_new_batch_size(batch_size + 1)
            }
        };

        Ok(r)
    }

    /// Computes the number of rounds required to evaluate the gate.
    ///
    /// **Note: ** This function can panic if the gate is not valid.
    fn comp_gate_comm_rounds(&self, gate: &Gate<C>) -> usize {
        match gate {
            Gate::Input(input_type) => match input_type {
                Input::SecretPlaintext { .. } => 1,
                _ => 0,
            },

            Gate::Constant(_) | Gate::Random { .. } => 0,

            Gate::FieldShareUnaryOp { op, .. } => match op {
                FieldShareUnaryOp::Neg => 0,
                FieldShareUnaryOp::MulInverse => 2,
                FieldShareUnaryOp::Open => 1,
                FieldShareUnaryOp::IsZero => 2,
            },
            Gate::FieldShareBinaryOp { op, y, .. } => {
                match (op, self.gate_output_unchecked(*y).form) {
                    (FieldShareBinaryOp::Mul, ShareOrPlaintext::Share) => 1,
                    (FieldShareBinaryOp::Mul, ShareOrPlaintext::Plaintext)
                    | (FieldShareBinaryOp::Add, _) => 0,
                }
            }
            Gate::BatchSummation { .. } => 0,
            Gate::BitShareUnaryOp { op, .. } => match op {
                BitShareUnaryOp::Not => 0,
                BitShareUnaryOp::Open => 1,
            },
            Gate::BitShareBinaryOp { op, y, .. } => {
                match (op, self.gate_output_unchecked(*y).form) {
                    (BitShareBinaryOp::Xor, _) => 0,
                    (_, ShareOrPlaintext::Share) => 1,
                    (_, ShareOrPlaintext::Plaintext) => 0,
                }
            }
            Gate::PointShareUnaryOp { op, .. } => match op {
                PointShareUnaryOp::Neg => 0,
                PointShareUnaryOp::Open => 1,
                PointShareUnaryOp::IsZero => 2,
            },
            Gate::PointShareBinaryOp { y, op, .. } => {
                match (op, self.gate_output_unchecked(*y).form) {
                    (PointShareBinaryOp::Add, _) => 0,
                    (PointShareBinaryOp::ScalarMul, ShareOrPlaintext::Share) => 1,
                    (PointShareBinaryOp::ScalarMul, ShareOrPlaintext::Plaintext) => 0,
                }
            }

            Gate::BaseFieldPow { .. } => 2,

            Gate::FieldPlaintextUnaryOp { .. }
            | Gate::FieldPlaintextBinaryOp { .. }
            | Gate::BitPlaintextUnaryOp { .. }
            | Gate::BitPlaintextBinaryOp { .. }
            | Gate::PointPlaintextUnaryOp { .. }
            | Gate::PointPlaintextBinaryOp { .. }
            | Gate::DaBit { .. }
            | Gate::GetDaBitFieldShare { .. }
            | Gate::GetDaBitSharedBit { .. }
            | Gate::BitPlaintextToField { .. }
            | Gate::FieldPlaintextToBit { .. }
            | Gate::ExtractFromBatch { .. }
            | Gate::CollectToBatch { .. }
            | Gate::GatherFromBatches { .. }
            | Gate::PointFromPlaintextCoordinates { .. }
            | Gate::PlaintextPointToCoordinates { .. }
            | Gate::PlaintextKeccakF1600 { .. }
            | Gate::CompressPlaintextPoint { .. }
            | Gate::KeyRecoveryPlaintextComputeErrors { .. }
            | Gate::Ghash { .. } => 0,

            // Two: the broadcast in which peers exchange what they fetched, and the one in which
            // they exchange a digest of what they reconciled it to. Constraint checking and
            // selection in between are local, and so is comparing the digests.
            Gate::ConstrainPlaintextBits { .. } => 2,

            Gate::AesKeySchedule { key, .. } => {
                let key_length = self.gate_ext(*key).map(|g| g.output.batch_size);
                let n_sub_bytes_calls = match key_length {
                    Ok(128) => 10,
                    Ok(192) => 8,
                    Ok(256) => 13,
                    _ => panic!("Invalid key length"),
                };
                n_sub_bytes_calls * AES_S_BOX_N_NETWORK_ROUNDS
            }

            Gate::AesGcmKeyStream { round_keys, .. } => {
                let round_keys_length = self.gate_ext(*round_keys).map(|g| g.output.batch_size);
                // 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 {
                    Ok(1408) => 10,
                    Ok(1664) => 12,
                    Ok(1920) => 14,
                    _ => panic!("Invalid round keys length"),
                };
                // in the first round we do a small optimization though this incurs two calls
                // to the S-box
                (1 + n_rounds) * AES_S_BOX_N_NETWORK_ROUNDS
            }
            Gate::GhashPowersOfH {
                n_ciphertext_blocks,
                ..
            } => {
                if *n_ciphertext_blocks <= 1 {
                    0
                } else {
                    (n_ciphertext_blocks - 1).ilog2() as usize + 1
                }
            }
        }
    }

    /// Computes the communication level of the gate.
    ///
    /// **Note: ** This function can panic if the gate is not valid.
    fn comp_gate_level(&self, gate: &Gate<C>) -> GateLevel {
        let comm_rounds = self.comp_gate_comm_rounds(gate);
        match gate
            .get_inputs()
            .iter()
            .map(|pred| self.gate_level_unchecked(*pred))
            .max()
        {
            None => GateLevel::default(),
            Some(preds_level) => preds_level.next(comm_rounds),
        }
    }

    #[cfg(feature = "dev")]
    pub fn remove_secret_shared_inputs(&mut self) {
        let mut inputs_to_remove = std::collections::HashSet::new();
        for (idx, gate) in self.gates.iter_mut().enumerate() {
            if let &Gate::Input(Input::Share {
                algebraic_type,
                batch_size,
            }) = &gate.gate
            {
                gate.gate = Gate::Random {
                    algebraic_type,
                    batch_size,
                };
                inputs_to_remove.insert(idx as GateIndex);
            }
        }
        self.inputs.retain(|gate| !inputs_to_remove.contains(gate));
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        circuit::{AlgebraicType, Circuit, FieldShareBinaryOp, Gate, Input},
        config::DefaultConfig as C,
    };

    #[test]
    fn test_circuit_new() {
        let mut circuit = Circuit::<C>::new();

        let x = circuit
            .add_gate(Gate::Input(Input::SecretPlaintext {
                inputer: 0,
                algebraic_type: AlgebraicType::MpcField,
                batch_size: 3,
            }))
            .unwrap();

        let y = circuit
            .add_gate(Gate::Input(Input::SecretPlaintext {
                inputer: 0,
                algebraic_type: AlgebraicType::MpcField,
                batch_size: 3,
            }))
            .unwrap();

        let z = circuit
            .add_gate(Gate::FieldShareBinaryOp {
                x,
                y,
                op: FieldShareBinaryOp::Mul,
            })
            .unwrap();

        circuit.add_output(z).unwrap();

        assert_eq!(circuit.nb_inputs(), 2);
        assert_eq!(circuit.nb_gates(), 2 + 1);
        assert_eq!(circuit.nb_outputs(), 1);
    }
}