arcium-primitives 0.4.2

use std::mem::MaybeUninit;

use derive_more::derive::{AsMut, AsRef, Deref, DerefMut, From};
use serde::{Deserialize, Serialize};
use wincode::{io::Writer, ReadResult, SchemaRead, SchemaWrite, WriteResult};

// TODO: Remove once v0.7.0 is released

#[cfg(not(target_pointer_width = "64"))]
compile_error!("this crate builds on 64-bit platforms only");

/// A wrapper around `crypto_bigint::BoxedUint` that implements `Serialize` and `Deserialize`.
#[derive(
    Debug, Clone, Deref, DerefMut, PartialEq, Eq, Hash, AsRef, AsMut, From, PartialOrd, Ord,
)]
#[as_ref(forward)]
#[from(forward)]
#[repr(transparent)]
pub struct BoxedUint(crypto_bigint::BoxedUint);

impl Serialize for BoxedUint {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let bytes = self.0.to_le_bytes();
        bytes.serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for BoxedUint {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let bytes = Vec::<u8>::deserialize(deserializer)?;
        let bits_precision = bytes.len() * 8;
        let boxed_uint = crypto_bigint::BoxedUint::from_le_slice(&bytes, bits_precision as u32)
            .map_err(serde::de::Error::custom)?;
        Ok(BoxedUint(boxed_uint))
    }
}

impl SchemaWrite for BoxedUint {
    type Src = Self;

    fn size_of(src: &Self::Src) -> WriteResult<usize> {
        let limbs = src
            .0
            .as_limbs()
            .iter()
            .map(|limb| limb.0)
            .collect::<Box<[u64]>>();
        <Box<[u64]> as SchemaWrite>::size_of(&limbs)
    }

    fn write(writer: &mut impl Writer, src: &Self::Src) -> WriteResult<()> {
        let limbs = src
            .0
            .as_limbs()
            .iter()
            .map(|limb| limb.0)
            .collect::<Box<[u64]>>();
        <Box<[u64]> as SchemaWrite>::write(writer, &limbs)
    }
}

impl<'de> SchemaRead<'de> for BoxedUint {
    type Dst = Self;

    fn read(
        reader: &mut impl wincode::io::Reader<'de>,
        dst: &mut MaybeUninit<Self::Dst>,
    ) -> ReadResult<()> {
        // SAFETY: Wrappers maintain the same memory layout as their inner type, even without
        // #[repr(transparent)]. Reference: https://github.com/rust-lang/unsafe-code-guidelines/issues/34#issuecomment-431114929
        // Additionally, we only build on 64-bit platforms, so u64 limbs are valid.
        let dst_u64 = unsafe { &mut *dst.as_mut_ptr().cast::<MaybeUninit<Box<[u64]>>>() };
        <Box<[u64]> as SchemaRead>::read(reader, dst_u64)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    // Roundtrip test BoxedUint serialization
    use bincode;
    use crypto_bigint::U256;

    use super::BoxedUint;
    #[test]
    fn test_boxed_uint_bincode() {
        let original = BoxedUint(crypto_bigint::BoxedUint::from(U256::from_u64(123456789)));
        let serialized = bincode::serialize(&original).expect("Serialization failed");
        let deserialized: BoxedUint =
            bincode::deserialize(&serialized).expect("Deserialization failed");
        assert_eq!(original, deserialized);
    }

    #[test]
    fn test_boxed_uint_wincode() {
        let original = BoxedUint(crypto_bigint::BoxedUint::from(U256::from_u64(123456789)));
        let serialized = wincode::serialize(&original).expect("Serialization failed");
        let deserialized: BoxedUint =
            wincode::deserialize(&serialized).expect("Deserialization failed");
        assert_eq!(original, deserialized);
    }
}