arcium-core-utils 0.7.3

/// A serialization optimized representation of a circuit.
use primitives::algebra::elliptic_curve::{Point, Scalar};
use serde::{Deserialize, Serialize};

use super::{errors::CircuitError, Circuit, Gate, GateIndex};
use crate::config::{name_tag, tag_name, MpcConfig};

/// Version of the serialized circuit layout. Bump on any change to the byte format so stale
/// blobs fail with [`CircuitError::UnsupportedFormatVersion`] instead of a parse error.
const CIRCUIT_FORMAT_VERSION: u8 = 1;

/// Serialization/deserialization optimized representation of a circuit.
#[derive(Serialize, Deserialize, Default)]
#[serde(bound(
    serialize = "Scalar<C::Curve>: Serialize, Point<C::Curve>: Serialize",
    deserialize = "Scalar<C::Curve>: Deserialize<'de>, Point<C::Curve>: Deserialize<'de>"
))]
#[repr(C)]
struct CompressedCircuit<C: MpcConfig> {
    /// Serialization format version (see [`CIRCUIT_FORMAT_VERSION`]). First in the layout so a
    /// stale blob fails fast.
    pub format_version: u8,
    /// Identifies the curve the circuit was built for (see [`name_tag`]).
    pub curve_tag: [u8; 32],
    /// Identifies the MPC field the circuit was built for (see [`name_tag`]). The tags precede
    /// the gates so a config mismatch fails before they are parsed.
    pub mpc_field_tag: [u8; 32],
    /// The circuit operations.
    pub ops: Vec<Gate<C>>,
    /// The output gates in order of definition
    pub output_gates: Vec<GateIndex>,
}

impl<C: MpcConfig> From<&Circuit<C>> for CompressedCircuit<C> {
    fn from(value: &Circuit<C>) -> Self {
        CompressedCircuit {
            format_version: CIRCUIT_FORMAT_VERSION,
            curve_tag: name_tag::<C::Curve>(),
            mpc_field_tag: name_tag::<C::Field>(),
            ops: value.iter_gates().cloned().collect(),
            output_gates: value.iter_output_indices().copied().collect(),
        }
    }
}

impl<C: MpcConfig> TryFrom<CompressedCircuit<C>> for Circuit<C> {
    type Error = CircuitError<C>;

    fn try_from(circuit: CompressedCircuit<C>) -> Result<Self, Self::Error> {
        if circuit.format_version != CIRCUIT_FORMAT_VERSION {
            return Err(CircuitError::UnsupportedFormatVersion {
                expected: CIRCUIT_FORMAT_VERSION,
                found: circuit.format_version,
            });
        }
        let expected_curve = name_tag::<C::Curve>();
        if circuit.curve_tag != expected_curve {
            return Err(CircuitError::CurveMismatch {
                expected: tag_name(&expected_curve),
                found: tag_name(&circuit.curve_tag),
            });
        }
        let expected_field = name_tag::<C::Field>();
        if circuit.mpc_field_tag != expected_field {
            return Err(CircuitError::MpcFieldMismatch {
                expected: tag_name(&expected_field),
                found: tag_name(&circuit.mpc_field_tag),
            });
        }

        let mut res = Self {
            gates: Vec::with_capacity(circuit.ops.len()),
            inputs: Vec::new(),
            outputs: Vec::with_capacity(circuit.output_gates.len()),
        };

        for gate in circuit.ops.into_iter() {
            res.add_gate(gate)?;
        }

        for index in circuit.output_gates.into_iter() {
            res.add_output(index)?;
        }

        Ok(res)
    }
}

mod bincode {

    use serde::{Deserialize, Deserializer, Serialize, Serializer};

    use super::{Circuit, CompressedCircuit};
    use crate::config::MpcConfig;

    impl<C: MpcConfig> Serialize for Circuit<C> {
        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where
            S: Serializer,
        {
            let circuit_serde: CompressedCircuit<C> = self.into();
            circuit_serde.serialize(serializer)
        }
    }

    impl<'de, C: MpcConfig> Deserialize<'de> for Circuit<C> {
        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where
            D: Deserializer<'de>,
        {
            let circuit_serde = CompressedCircuit::<C>::deserialize(deserializer)
                .map_err(serde::de::Error::custom)?;
            let circuit = circuit_serde.try_into();
            circuit.map_err(serde::de::Error::custom)
        }
    }
}

mod wincode {
    use core::{
        mem::{self, MaybeUninit},
        ptr,
    };

    use ::wincode::{
        containers,
        io::{Reader, Writer},
        len::BincodeLen,
        ReadResult,
        SchemaRead,
        SchemaWrite,
        TypeMeta,
        WriteResult,
    };

    use super::*;
    pub type BincodeLenU32 = BincodeLen<{ 2 << 32 }>;

    // TODO: optimize so that we dont need to cast into CircuitSerde twice
    impl<C: MpcConfig> SchemaWrite for Circuit<C> {
        type Src = Self;

        const TYPE_META: TypeMeta = <CompressedCircuit<C> as SchemaWrite>::TYPE_META;

        fn size_of(src: &Self::Src) -> WriteResult<usize> {
            let circuit_serde: CompressedCircuit<C> = src.into();
            <CompressedCircuit<C> as SchemaWrite>::size_of(&circuit_serde)
        }

        fn write(writer: &mut impl Writer, src: &Self::Src) -> WriteResult<()> {
            let circuit_serde: CompressedCircuit<C> = src.into();
            <CompressedCircuit<C> as SchemaWrite>::write(writer, &circuit_serde)
        }
    }

    impl<C: MpcConfig> ::wincode::SchemaWrite for CompressedCircuit<C> {
        type Src = Self;
        #[allow(clippy::arithmetic_side_effects)]
        const TYPE_META: TypeMeta = if let (
            TypeMeta::Static {
                size: v,
                zero_copy: zc_v,
            },
            TypeMeta::Static {
                size: t,
                zero_copy: zc_t,
            },
            TypeMeta::Static {
                size: a,
                zero_copy: zc_a,
            },
            TypeMeta::Static {
                size: b,
                zero_copy: zc_b,
            },
        ) = (
            <u8 as SchemaWrite>::TYPE_META,
            <[u8; 32] as SchemaWrite>::TYPE_META,
            <containers::Vec<Gate<C>, BincodeLenU32> as SchemaWrite>::TYPE_META,
            <Vec<GateIndex> as SchemaWrite>::TYPE_META,
        ) {
            let serialized_size = v + 2 * t + a + b;
            let no_padding = serialized_size == size_of::<Self>();
            TypeMeta::Static {
                size: serialized_size,
                zero_copy: no_padding && zc_v && zc_t && zc_a && zc_b,
            }
        } else {
            TypeMeta::Dynamic
        };
        #[inline]
        fn size_of(src: &Self::Src) -> WriteResult<usize> {
            if let TypeMeta::Static { size, .. } = <Self as SchemaWrite>::TYPE_META {
                return Ok(size);
            }
            let mut total = 0usize;
            total += <u8 as SchemaWrite>::size_of(&src.format_version)?;
            total += <[u8; 32] as SchemaWrite>::size_of(&src.curve_tag)?;
            total += <[u8; 32] as SchemaWrite>::size_of(&src.mpc_field_tag)?;
            total += <containers::Vec<Gate<C>, BincodeLenU32> as SchemaWrite>::size_of(&src.ops)?;
            total += <Vec<GateIndex> as SchemaWrite>::size_of(&src.output_gates)?;
            Ok(total)
        }
        #[inline]
        fn write(writer: &mut impl Writer, src: &Self::Src) -> WriteResult<()> {
            // Macro to reduce duplication in field writing
            macro_rules! write_fields {
                ($writer:expr) => {{
                    <u8 as SchemaWrite>::write($writer, &src.format_version)?;
                    <[u8; 32] as SchemaWrite>::write($writer, &src.curve_tag)?;
                    <[u8; 32] as SchemaWrite>::write($writer, &src.mpc_field_tag)?;
                    <containers::Vec<Gate<C>, BincodeLenU32> as SchemaWrite>::write(
                        $writer, &src.ops,
                    )?;
                    <Vec<GateIndex> as SchemaWrite>::write($writer, &src.output_gates)?;
                }};
            }

            match <Self as SchemaWrite>::TYPE_META {
                TypeMeta::Static { size, .. } => {
                    let writer = &mut unsafe { writer.as_trusted_for(size) }?;
                    write_fields!(writer);
                    writer.finish()?;
                }
                TypeMeta::Dynamic => {
                    write_fields!(writer);
                }
            }
            Ok(())
        }
    }

    impl<'de, C: MpcConfig> SchemaRead<'de> for Circuit<C> {
        type Dst = Self;
        const TYPE_META: TypeMeta = <CompressedCircuit<C> as SchemaRead>::TYPE_META;

        fn read(reader: &mut impl Reader<'de>, dst: &mut MaybeUninit<Self::Dst>) -> ReadResult<()> {
            let mut circuit_serde = MaybeUninit::new(CompressedCircuit::<C>::default());
            <CompressedCircuit<C> as SchemaRead>::read(reader, &mut circuit_serde)?;
            let circuit_serde = unsafe { circuit_serde.assume_init() };
            let circuit: Self::Dst = circuit_serde.try_into().map_err(|e| match e {
                CircuitError::UnsupportedFormatVersion { .. } => ::wincode::ReadError::Custom(
                    "Circuit was serialized with an unsupported format version",
                ),
                CircuitError::CurveMismatch { .. } => ::wincode::ReadError::Custom(
                    "Circuit was built for a different curve than this runtime's config",
                ),
                CircuitError::MpcFieldMismatch { .. } => ::wincode::ReadError::Custom(
                    "Circuit was built for a different MPC field than this runtime's config",
                ),
                _ => {
                    ::wincode::ReadError::Custom("Invalid cast from CircuitSerde to Circuit struct")
                }
            })?;
            dst.write(circuit);
            Ok(())
        }
    }

    impl<'de, C: MpcConfig> SchemaRead<'de> for CompressedCircuit<C> {
        type Dst = Self;
        #[allow(clippy::arithmetic_side_effects)]
        const TYPE_META: TypeMeta = if let (
            TypeMeta::Static {
                size: v,
                zero_copy: zc_v,
            },
            TypeMeta::Static {
                size: t,
                zero_copy: zc_t,
            },
            TypeMeta::Static {
                size: a,
                zero_copy: zc_a,
            },
            TypeMeta::Static {
                size: b,
                zero_copy: zc_b,
            },
        ) = (
            <u8 as SchemaRead<'de>>::TYPE_META,
            <[u8; 32] as SchemaRead<'de>>::TYPE_META,
            <containers::Vec<Gate<C>, BincodeLenU32> as SchemaRead<'de>>::TYPE_META,
            <Vec<GateIndex> as SchemaRead<'de>>::TYPE_META,
        ) {
            let serialized_size = v + 2 * t + a + b;
            let no_padding = serialized_size == size_of::<Self>();
            TypeMeta::Static {
                size: serialized_size,
                zero_copy: no_padding && zc_v && zc_t && zc_a && zc_b,
            }
        } else {
            TypeMeta::Dynamic
        };
        #[inline]
        fn read(reader: &mut impl Reader<'de>, dst: &mut MaybeUninit<Self::Dst>) -> ReadResult<()> {
            struct DropGuard<C: MpcConfig> {
                init_count: u8,
                dst_ptr: *mut CompressedCircuit<C>,
            }
            impl<C: MpcConfig> Drop for DropGuard<C> {
                #[cold]
                fn drop(&mut self) {
                    let dst_ptr = self.dst_ptr;
                    let init_count = self.init_count;
                    match init_count {
                        0 => {}
                        1u8 => unsafe {
                            ptr::drop_in_place(&raw mut (*dst_ptr).ops);
                        },
                        _ => unreachable!("init_count out of bounds"),
                    }
                }
            }
            // Macro to reduce duplication in field reading
            macro_rules! read_fields {
                ($reader:expr, $dst_ptr:expr, $guard:expr) => {{
                    let init_count = &mut $guard.init_count;
                    // The version byte and tags are Copy: no drop tracking needed.
                    <u8 as SchemaRead<'de>>::read($reader, unsafe {
                        &mut *(&raw mut (*$dst_ptr).format_version).cast::<MaybeUninit<_>>()
                    })?;
                    <[u8; 32] as SchemaRead<'de>>::read($reader, unsafe {
                        &mut *(&raw mut (*$dst_ptr).curve_tag).cast::<MaybeUninit<_>>()
                    })?;
                    <[u8; 32] as SchemaRead<'de>>::read($reader, unsafe {
                        &mut *(&raw mut (*$dst_ptr).mpc_field_tag).cast::<MaybeUninit<_>>()
                    })?;
                    <wincode::containers::Vec<Gate<C>, BincodeLenU32> as SchemaRead<'de>>::read(
                        $reader,
                        unsafe { &mut *(&raw mut (*$dst_ptr).ops).cast::<MaybeUninit<_>>() },
                    )?;
                    *init_count += 1;
                    <Vec<GateIndex> as SchemaRead<'de>>::read($reader, unsafe {
                        &mut *(&raw mut (*$dst_ptr).output_gates).cast::<MaybeUninit<_>>()
                    })?;
                    mem::forget($guard);
                }};
            }

            let dst_ptr = dst.as_mut_ptr();
            let mut guard = DropGuard {
                init_count: 0,
                dst_ptr,
            };

            match <Self as SchemaRead<'de>>::TYPE_META {
                TypeMeta::Static { size, .. } => {
                    read_fields!(&mut unsafe { reader.as_trusted_for(size) }?, dst_ptr, guard);
                }
                TypeMeta::Dynamic => {
                    read_fields!(reader, dst_ptr, guard);
                }
            }
            Ok(())
        }
    }
}

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

    #[test]
    fn valid_circuit() {
        let mut circuit = Circuit::<C>::new();
        let input_gate1 = circuit
            .add_gate(Gate::Input(Input::SecretPlaintext {
                inputer: 0,
                algebraic_type: AlgebraicType::ScalarField,
                batch_size: 1,
            }))
            .unwrap();
        assert_eq!(input_gate1, 0);
        let input_gate2 = circuit
            .add_gate(Gate::Input(Input::SecretPlaintext {
                inputer: 1,
                algebraic_type: AlgebraicType::ScalarField,
                batch_size: 1,
            }))
            .unwrap();
        assert_eq!(input_gate2, 1);
        let add_gate = circuit
            .add_gate(Gate::FieldShareBinaryOp {
                x: input_gate1,
                y: input_gate2,
                op: FieldShareBinaryOp::Add,
            })
            .unwrap();
        assert_eq!(add_gate, 2);
        circuit.add_output(add_gate).unwrap();
        assert_eq!(
            circuit.iter_output_indices().copied().collect::<Vec<_>>(),
            vec![2]
        );
    }

    #[test]
    fn test_ser_circuit_bincode() {
        let circuit = create_add_tree_circuit(18);
        let serialized = ::bincode::serialize(&circuit).unwrap();
        let circuit_de: Circuit<C> = ::bincode::deserialize(&serialized).unwrap();

        assert_eq!(circuit, circuit_de);
    }

    #[test]
    fn test_ser_circuit_wincode() {
        let circuit = create_add_tree_circuit(18);

        // Wincode roundtrip
        let serialized = ::wincode::serialize(&circuit).expect("Serialization failed");
        let deserialized: Circuit<C> =
            ::wincode::deserialize(&serialized).expect("Deserialization failed");

        assert_eq!(circuit, deserialized);
    }

    /// A circuit serialized for one MPC field must not deserialize under a config with another,
    /// in either format.
    #[test]
    fn test_cross_field_rejection() {
        use crate::config::Gf2_128Config;

        let circuit = create_add_tree_circuit::<C>(2);

        let bin = ::bincode::serialize(&circuit).unwrap();
        let err = ::bincode::deserialize::<Circuit<Gf2_128Config>>(&bin).unwrap_err();
        assert!(
            err.to_string().contains("MPC field"),
            "unexpected error: {err}"
        );

        let win = ::wincode::serialize(&circuit).unwrap();
        assert!(::wincode::deserialize::<Circuit<Gf2_128Config>>(&win).is_err());

        // Same-config roundtrips still pass the tag check.
        assert_eq!(::bincode::deserialize::<Circuit<C>>(&bin).unwrap(), circuit);
    }

    /// A blob with an unknown format version is rejected before anything else is parsed, and a
    /// corrupted curve tag is rejected before the field tag.
    #[test]
    fn test_version_and_curve_rejection() {
        let circuit = create_add_tree_circuit::<C>(2);
        let bin = ::bincode::serialize(&circuit).unwrap();

        // Layout: format_version (1 byte) | curve_tag (32) | mpc_field_tag (32) | gates...
        let mut wrong_version = bin.clone();
        wrong_version[0] ^= 0xFF;
        let err = ::bincode::deserialize::<Circuit<C>>(&wrong_version).unwrap_err();
        assert!(
            err.to_string().contains("format version"),
            "unexpected error: {err}"
        );

        let mut wrong_curve = bin;
        wrong_curve[1] ^= 0xFF;
        let err = ::bincode::deserialize::<Circuit<C>>(&wrong_curve).unwrap_err();
        assert!(err.to_string().contains("curve"), "unexpected error: {err}");
    }

    /// The circuit id commits to the MPC field: identical shapes under different configs get
    /// different ids.
    #[test]
    fn test_circuit_id_commits_to_field() {
        use crate::{circuit::circuit_id::CircuitId, config::Gf2_128Config};

        let id_default = CircuitId::of(&create_add_tree_circuit::<C>(2));
        let id_gf2_128 = CircuitId::of(&create_add_tree_circuit::<Gf2_128Config>(2));
        assert_ne!(id_default.as_bytes(), id_gf2_128.as_bytes());
    }
}