commonware_codec/

codec.rs

//! Core traits for encoding and decoding.

use crate::error::Error;
use bytes::{Buf, BufMut, Bytes, BytesMut};

/// Trait for types with a known, fixed encoded size.
///
/// Implementing this trait signifies that the encoded representation of this type *always* has the
/// same byte length, regardless of the specific value.
///
/// This automatically provides an implementation of [EncodeSize].
pub trait FixedSize {
    /// The size of the encoded value (in bytes).
    const SIZE: usize;
}

/// Trait for types that can provide their encoded size in bytes.
///
/// This must be implemented by all encodable types. For types implementing [FixedSize], this
/// trait is implemented automatically. For variable-size types, this requires calculating the size
/// based on the value.
pub trait EncodeSize {
    /// Returns the encoded size of this value (in bytes).
    fn encode_size(&self) -> usize;

    /// Returns the encoded size excluding bytes passed to [`BufsMut::push`]
    /// during [`Write::write_bufs`]. Used to size the working buffer for inline
    /// writes. Override alongside [`Write::write_bufs`] for types where large
    /// [`Bytes`] fields go via push; failing to do so will over-allocate.
    fn encode_inline_size(&self) -> usize {
        self.encode_size()
    }
}

// Automatically implement `EncodeSize` for types that are `FixedSize`.
impl<T: FixedSize> EncodeSize for T {
    fn encode_size(&self) -> usize {
        Self::SIZE
    }
}

/// Trait for types that can be written (encoded) to a byte buffer.
pub trait Write {
    /// Writes the binary representation of `self` to the provided buffer `buf`.
    ///
    /// Implementations should panic if the buffer doesn't have enough capacity.
    fn write(&self, buf: &mut impl BufMut);

    /// Writes to a [`BufsMut`], allowing existing [`Bytes`] chunks to be
    /// appended via [`BufsMut::push`] instead of written inline. Must encode
    /// to the same format as [`Write::write`]. Defaults to [`Write::write`].
    fn write_bufs(&self, buf: &mut impl BufsMut) {
        self.write(buf);
    }
}

/// Trait for types that can be read (decoded) from a byte buffer.
pub trait Read: Sized {
    /// The `Cfg` type parameter allows passing configuration during the read process. This is
    /// crucial for safely decoding untrusted data, for example, by providing size limits for
    /// collections or strings.
    ///
    /// Use `Cfg = ()` if no configuration is needed for a specific type.
    type Cfg: Clone + Send + Sync + 'static;

    /// Reads a value from the buffer using the provided configuration `cfg`.
    ///
    /// Implementations should consume the exact number of bytes required from `buf` to reconstruct
    /// the value.
    ///
    /// Returns [Error] if decoding fails due to invalid data, insufficient bytes in the buffer,
    /// or violation of constraints imposed by the `cfg`.
    fn read_cfg(buf: &mut impl Buf, cfg: &Self::Cfg) -> Result<Self, Error>;
}

/// Trait combining [Write] and [EncodeSize] for types that can be fully encoded.
///
/// This trait provides the convenience [Encode::encode] method which handles
/// buffer allocation, writing, and size assertion in one go.
pub trait Encode: Write + EncodeSize {
    /// Encodes `self` into a new [Bytes] buffer.
    ///
    /// This method calculates the required size using [EncodeSize::encode_size], allocates a
    /// buffer of that exact capacity, writes the value using [Write::write], and performs a
    /// sanity check assertion.
    ///
    /// # Panics
    ///
    /// Panics if `encode_size()` does not return the same number of bytes actually written by
    /// `write()`
    fn encode(&self) -> Bytes {
        self.encode_mut().freeze()
    }

    /// Encodes `self` into a new [BytesMut] buffer.
    ///
    /// This method calculates the required size using [EncodeSize::encode_size], allocates a
    /// buffer of that exact capacity, writes the value using [Write::write], and performs a
    /// sanity check assertion.
    ///
    /// # Panics
    ///
    /// Panics if `encode_size()` does not return the same number of bytes actually written by
    /// `write()`
    fn encode_mut(&self) -> BytesMut {
        let len = self.encode_size();
        let mut buffer = BytesMut::with_capacity(len);
        self.write(&mut buffer);
        assert_eq!(buffer.len(), len, "write() did not write expected bytes");
        buffer
    }
}

// Automatically implement `Encode` for types that implement `Write` and `EncodeSize`.
impl<T: Write + EncodeSize> Encode for T {}

/// Convenience trait combining `Encode` with thread-safety bounds.
///
/// Represents types that can be fully encoded and safely shared across threads.
pub trait EncodeShared: Encode + Send + Sync {}

// Automatically implement `EncodeShared` for types that meet all bounds.
impl<T: Encode + Send + Sync> EncodeShared for T {}

/// Trait combining [Read] with a check for remaining bytes.
///
/// Ensures that *all* bytes from the input buffer were consumed during decoding.
pub trait Decode: Read {
    /// Decodes a value from `buf` using `cfg`, ensuring the entire buffer is consumed.
    ///
    /// Returns [Error] if decoding fails via [Read::read_cfg] or if there are leftover bytes in
    /// `buf` after reading.
    fn decode_cfg(mut buf: impl Buf, cfg: &Self::Cfg) -> Result<Self, Error> {
        let result = Self::read_cfg(&mut buf, cfg)?;

        // Check that the buffer is fully consumed.
        let remaining = buf.remaining();
        if remaining > 0 {
            return Err(Error::ExtraData(remaining));
        }

        Ok(result)
    }
}

// Automatically implement `Decode` for types that implement `Read`.
impl<T: Read> Decode for T {}

/// Convenience trait combining [Encode] and [Decode].
///
/// Represents types that can be both fully encoded and decoded.
pub trait Codec: Encode + Decode {}

/// Automatically implement `Codec` for types that implement `Encode` and `Decode`.
impl<T: Encode + Decode> Codec for T {}

/// Convenience trait for [FixedSize] types that can be encoded directly into a fixed-size array.
pub trait EncodeFixed: Write + FixedSize {
    /// Encodes `self` into a fixed-size byte array `[u8; N]`.
    ///
    /// # Panics
    ///
    /// Panics if `N` is not equal to `<Self as FixedSize>::SIZE`.
    /// Also panics if the `write()` implementation does not write exactly `N` bytes.
    fn encode_fixed<const N: usize>(&self) -> [u8; N] {
        // Ideally this is a compile-time check, but we can't do that in the current Rust version
        // without adding a new generic parameter to the trait.
        assert_eq!(
            N,
            Self::SIZE,
            "Can't encode {} bytes into {} bytes",
            Self::SIZE,
            N
        );

        let mut array = [0u8; N];
        let mut buf = &mut array[..];
        self.write(&mut buf);
        assert_eq!(buf.len(), 0);
        array
    }
}

// Automatically implement `EncodeFixed` for types that implement `Write` and `FixedSize`.
impl<T: Write + FixedSize> EncodeFixed for T {}

/// Convenience trait combining `FixedSize` and `Codec`.
///
/// Represents types that can be both fully encoded and decoded from a fixed-size byte sequence.
pub trait CodecFixed: Codec + FixedSize {}

// Automatically implement `CodecFixed` for types that implement `Codec` and `FixedSize`.
impl<T: Codec + FixedSize> CodecFixed for T {}

/// Convenience trait combining `Codec` with thread-safety bounds.
///
/// Represents types that can be fully encoded/decoded and safely shared across threads.
pub trait CodecShared: Codec + Send + Sync {}

// Automatically implement `CodecShared` for types that meet all bounds.
impl<T: Codec + Send + Sync> CodecShared for T {}

/// Convenience trait combining `CodecFixed` with thread-safety bounds and unit config.
///
/// Represents fixed-size types that can be fully encoded/decoded, require no configuration,
/// and can be safely shared across threads.
pub trait CodecFixedShared: CodecFixed<Cfg = ()> + Send + Sync {}

// Automatically implement `CodecFixedShared` for types that meet all bounds.
impl<T: CodecFixed<Cfg = ()> + Send + Sync> CodecFixedShared for T {}

/// A [`BufMut`] that can also append pre-existing [`Bytes`] chunks.
pub trait BufsMut: BufMut {
    /// Appends a [`Bytes`] chunk instead of writing its contents inline into
    /// the destination buffer.
    fn push(&mut self, bytes: impl Into<Bytes>);
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        extensions::{DecodeExt, ReadExt},
        Error,
    };
    use bytes::Bytes;

    #[test]
    fn test_insufficient_buffer() {
        let mut reader = Bytes::from_static(&[0x01, 0x02]);
        assert!(matches!(u32::read(&mut reader), Err(Error::EndOfBuffer)));
    }

    #[test]
    fn test_extra_data() {
        let encoded = Bytes::from_static(&[0x01, 0x02]);
        assert!(matches!(u8::decode(encoded), Err(Error::ExtraData(1))));
    }

    #[test]
    fn test_encode_fixed() {
        let value = 42u32;
        let encoded: [u8; 4] = value.encode_fixed();
        let decoded = <u32>::decode(&encoded[..]).unwrap();
        assert_eq!(value, decoded);
    }

    #[test]
    #[should_panic(expected = "Can't encode 4 bytes into 5 bytes")]
    fn test_encode_fixed_panic() {
        let _: [u8; 5] = 42u32.encode_fixed();
    }
}