fray/

bitfields.rs

use crate::{BitContainer, BitContainerFor};

mod fields;
pub use fields::{Field, FieldType};
pub mod bitorder;
use bitorder::BitOrder;
mod private;

/// Provides all the convenience methods for interacting with a **bitfield structure**.
///
/// This trait is:
/// - **sealed** (cannot be implemented manually),
/// - intended to be obtained via the `bitfield` macro,
/// - the main API end-users will work with.
///
/// # Core methods
///
/// - [`new`](Self::new):
///   Construct a new instance with an empty [`BitContainer`].
///
/// - [`into_inner`](Self::into_inner):
///   Consume `self` and return the underlying raw **inner value**
///   ([`BitFieldImpl::Container::Inner`](BitContainer::Inner)).
///
/// - [`get`](Self::get) / [`try_get`](Self::try_get):
///   Read the value of a field.
///
/// - [`set`](Self::set) / [`try_set`](Self::try_set):
///   Write the value of a field.
///
/// - [`with`](Self::with) / [`try_with`](Self::try_with):
///   Write the value of a field and return `&mut Self` for chaining.
///
/// # How it works
///
/// Internally, the macro implements [`BitFieldImpl`] for your type, and
/// [`BitField`] is automatically implemented via a generic blanket impl
/// over all types that implement [`BitFieldImpl`].
pub trait BitField: private::Sealed + BitFieldImpl {
    ///   Consume `self` and return the underlying raw **inner value**
    ///   ([`BitField::Container::Inner`](BitContainer::Inner)).
    #[inline]
    fn into_inner(self) -> <Self::Container as BitContainer>::Inner {
        self.into().into_inner()
    }

    /// Construct a new instance with an empty [`BitContainer`].
    #[inline]
    fn new() -> Self {
        Self::Container::empty().into()
    }

    #[doc(hidden)]
    #[inline]
    fn _set<F>(&mut self, value: F::BitsType, _: private::Token)
    where
        F: Field<Self>,
        Self::Container: BitContainerFor<F::BitsType>,
    {
        debug_assert!(F::OFFSET + F::SIZE <= Self::Container::SIZE);
        self.as_mut().store(value, F::OFFSET, F::SIZE);
    }

    #[doc(hidden)]
    #[inline]
    fn _get<F>(&self, _: private::Token) -> F::BitsType
    where
        F: Field<Self>,
        Self::Container: BitContainerFor<F::BitsType>,
    {
        debug_assert!(F::OFFSET + F::SIZE <= Self::Container::SIZE);
        self.as_ref().retrieve(F::OFFSET, F::SIZE)
    }

    /// Read the value of field `F`. Returns `F::Type` directly.
    ///
    /// Requires an infallible conversion from `F::BitsType` to `F::Type`.
    #[inline]
    fn get<F>(&self) -> F::Type
    where
        F: Field<Self>,
        Self::Container: BitContainerFor<F::BitsType>,
        F::BitsType: Into<F::Type>,
    {
        self._get::<F>(private::Token).into()
    }

    /// Read the value of field `F`.
    ///
    /// Returns a `Result` containing:
    /// - `F::Type` if the conversion succeeds,
    /// - an error (`TryInto<F::Type>::Error`) if it fails.
    #[inline]
    fn try_get<F>(&self) -> Result<F::Type, <F::BitsType as TryInto<F::Type>>::Error>
    where
        F: Field<Self>,
        Self::Container: BitContainerFor<F::BitsType>,
        F::BitsType: TryInto<F::Type>,
    {
        self._get::<F>(private::Token).try_into()
    }

    /// Write the value of field `F`.
    ///
    /// Requires an infallible conversion from `F::Type` to `F::BitsType`.
    #[inline]
    fn set<F>(&mut self, value: F::Type)
    where
        F: Field<Self>,
        Self::Container: BitContainerFor<F::BitsType>,
        F::Type: Into<F::BitsType>,
    {
        self._set::<F>(value.into(), private::Token);
    }

    /// Write the value of field `F`.
    ///
    /// Returns a `Result` containing:
    /// - `()` if the conversion succeeds,
    /// - an error (`TryInto<F::BitsType>::Error`) if it fails.
    #[inline]
    fn try_set<F>(&mut self, value: F::Type) -> Result<(), <F::Type as TryInto<F::BitsType>>::Error>
    where
        F: Field<Self>,
        Self::Container: BitContainerFor<F::BitsType>,
        F::Type: TryInto<F::BitsType>,
    {
        self._set::<F>(value.try_into()?, private::Token);
        Ok(())
    }

    /// Write the value of field `F` and return `&mut Self` for chaining.
    ///
    /// Requires an infallible conversion from `F::Type` to `F::BitsType`.
    #[inline]
    fn with<F>(&mut self, value: F::Type) -> &mut Self
    where
        F: Field<Self>,
        Self::Container: BitContainerFor<F::BitsType>,
        F::Type: Into<F::BitsType>,
    {
        self.set::<F>(value);
        self
    }

    /// Write the value of field `F`.
    ///
    /// Returns a `Result` containing:
    /// - `&mut Self` if the conversion succeeds allowing chaning,
    /// - an error (`TryInto<F::BitsType>::Error`) if it fails.
    #[inline]
    fn try_with<F>(
        &mut self,
        value: F::Type,
    ) -> Result<&mut Self, <F::Type as TryInto<F::BitsType>>::Error>
    where
        F: Field<Self>,
        Self::Container: BitContainerFor<F::BitsType>,
        F::Type: TryInto<F::BitsType>,
    {
        self.try_set::<F>(value)?;
        Ok(self)
    }
}

impl<T: BitFieldImpl + private::Sealed> BitField for T {}

/// Internal trait required by [`BitField`].
///
/// `BitFieldImpl` is **not meant to be implemented manually**.
/// It is automatically implemented by the `bitfield` macro.
///
/// # Usage
///
/// End-users generally do **not** interact with `BitFieldImpl` directly.
/// Instead, they use the convenience methods provided by [`BitField`].
pub trait BitFieldImpl
where
    Self::Container: Into<Self>,
    Self: Into<Self::Container>,
    Self: AsRef<Self::Container>,
    Self: AsMut<Self::Container>,
{
    /// The underlying storage type for the bitfield.
    ///
    /// This type defines where the raw bits are stored.
    /// It must implement [`BitContainer`].
    type Container: BitContainer;

    /// Indicates the bit numbering order used by the bitfield.
    ///
    /// This associated type has no functional impact; it exists purely for
    /// **semantic clarity**.
    ///
    /// Set automatically by the [`bitorder`](crate::bitfield) attribute
    /// (`LSB0` by default).
    type BitOrder: BitOrder;
}

impl<T: BitFieldImpl> private::Sealed for T {}

#[cfg(test)]
mod tests {
    use crate::iterable::BitIterableContainer;

    use super::*;

    #[allow(dead_code)]
    #[allow(clippy::upper_case_acronyms)]
    #[allow(missing_debug_implementations)]
    mod dns_flags {
        use core::fmt::Debug;

        use super::*;
        use crate::{bitorder::LSB0, iterable::BitIterableContainer};

        // https://datatracker.ietf.org/doc/html/rfc1035#section-4.1.1
        // #[bitfield(repr(u16), bitorder(msb0))]
        // pub struct DNSFlags {
        //     QR: bool,
        //     OPCODE: OpCode,
        //     AA: bool,
        //     TC: bool,
        //     RD: bool,
        //     RA: bool,
        //     #[bits(3)]
        //     _z: (),
        //     RCODE: Rcode,
        // }
        pub mod fields {
            macro_rules! create_field {
                ($name:ident, $bf:ty, $t: ty, $bit_t:ty, $offset:expr) => {
                    pub enum $name {}

                    impl crate::Field<$bf> for $name {
                        type Type = $t;
                        type BitsType = $bit_t;
                        const OFFSET: usize = $offset;
                    }
                };
            }
            create_field!(QR, super::DNSFlags, bool, bool, 0);
            create_field!(OPCODE, super::DNSFlags, super::OpCode, u8, 1);
            create_field!(AA, super::DNSFlags, bool, bool, 5);
            create_field!(TC, super::DNSFlags, bool, bool, 6);
            create_field!(RD, super::DNSFlags, bool, bool, 7);
            create_field!(RA, super::DNSFlags, bool, bool, 8);
            // offset = 8 + 1 (RA size) + 3 (zeroes)
            create_field!(RCODE, super::DNSFlags, super::Rcode, u8, 12);
        }

        #[repr(u8)]
        #[derive(Debug, Clone, Copy, PartialEq, Eq)]
        pub enum OpCode {
            Query = 0,
            IQuery = 1,
            Status = 2,
        }
        impl FieldType for OpCode {
            const SIZE: usize = 4;

            type BitsType = u8;
        }
        impl TryFrom<u8> for OpCode {
            type Error = ();

            fn try_from(value: u8) -> Result<Self, Self::Error> {
                Ok(match value {
                    0 => Self::Query,
                    1 => Self::IQuery,
                    2 => Self::Status,
                    _ => return Err(()),
                })
            }
        }

        impl From<OpCode> for u8 {
            fn from(value: OpCode) -> Self {
                value as u8
            }
        }

        #[repr(u8)]
        #[derive(Debug, Clone, Copy, PartialEq, Eq)]
        pub enum Rcode {
            NoError = 0,
            FormatError = 1,
            ServerFailure = 2,
            NameError = 3,
            NotImplemented = 4,
            Refused = 5,
        }
        impl FieldType for Rcode {
            const SIZE: usize = 4;

            type BitsType = u8;
        }

        impl TryFrom<u8> for Rcode {
            type Error = ();

            fn try_from(value: u8) -> Result<Self, Self::Error> {
                Ok(match value {
                    0 => Self::NoError,
                    1 => Self::FormatError,
                    2 => Self::ServerFailure,
                    3 => Self::NameError,
                    4 => Self::NotImplemented,
                    5 => Self::Refused,
                    _ => return Err(()),
                })
            }
        }

        #[derive(Clone, Copy)]
        pub struct DNSFlags(BitIterableContainer<u16>);

        impl From<DNSFlags> for <<DNSFlags as BitFieldImpl>::Container as BitContainer>::Inner {
            fn from(value: DNSFlags) -> Self {
                value.into_inner()
            }
        }

        impl From<BitIterableContainer<u16>> for DNSFlags {
            fn from(value: BitIterableContainer<u16>) -> Self {
                Self(value)
            }
        }

        impl From<DNSFlags> for BitIterableContainer<u16> {
            fn from(value: DNSFlags) -> Self {
                value.0
            }
        }

        impl AsRef<BitIterableContainer<u16>> for DNSFlags {
            fn as_ref(&self) -> &BitIterableContainer<u16> {
                &self.0
            }
        }

        impl AsMut<BitIterableContainer<u16>> for DNSFlags {
            fn as_mut(&mut self) -> &mut BitIterableContainer<u16> {
                &mut self.0
            }
        }

        impl BitFieldImpl for DNSFlags {
            type Container = BitIterableContainer<u16>;
            type BitOrder = LSB0;
        }

        impl Debug for DNSFlags {
            fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
                use crate::debug::PrettyResult;
                f.debug_struct("DNSFlags")
                    .field("QR", &PrettyResult::from(&self.try_get::<QR>()))
                    .field("OPCODE", &PrettyResult::from(&self.try_get::<OPCODE>()))
                    .field("AA", &PrettyResult::from(&self.try_get::<AA>()))
                    .field("TC", &PrettyResult::from(&self.try_get::<TC>()))
                    .field("RD", &PrettyResult::from(&self.try_get::<RD>()))
                    .field("RA", &PrettyResult::from(&self.try_get::<RA>()))
                    .field("RCODE", &PrettyResult::from(&self.try_get::<RCODE>()))
                    .finish()
            }
        }
    }
    use dns_flags::{DNSFlags, OpCode, fields::*};

    #[test]
    fn set_bool_basic() {
        let mut dns_flags = DNSFlags::new();

        dns_flags.set::<QR>(true);
        assert_eq!(dns_flags.into_inner(), 0b1u16);

        dns_flags.set::<QR>(true);
        assert_eq!(dns_flags.into_inner(), 0b1u16);

        dns_flags.set::<QR>(false);
        assert_eq!(dns_flags.into_inner(), 0b0u16);

        dns_flags.set::<QR>(false);
        assert_eq!(dns_flags.into_inner(), 0b0u16);

        dns_flags.set::<RA>(true);
        assert_eq!(dns_flags.into_inner(), 0b1_0000_0000u16);

        dns_flags.set::<RA>(true);
        assert_eq!(dns_flags.into_inner(), 0b1_0000_0000u16);

        dns_flags.set::<RA>(false);
        assert_eq!(dns_flags.into_inner(), 0b0u16);

        dns_flags.set::<RA>(false);
        assert_eq!(dns_flags.into_inner(), 0b0u16);
    }

    #[test]
    fn set_bool_dont_overlap() {
        let mut dns_flags = DNSFlags::new();
        dns_flags.set::<TC>(true);
        dns_flags.set::<RD>(true);
        dns_flags.set::<RA>(true);
        assert_eq!(dns_flags.into_inner(), 0b1_1100_0000u16);

        dns_flags.set::<RD>(false);
        assert_eq!(dns_flags.into_inner(), 0b1_0100_0000u16);

        dns_flags.set::<RD>(true);
        assert_eq!(dns_flags.into_inner(), 0b1_1100_0000u16);
    }

    #[test]
    fn get_bool() {
        let dns_flags = DNSFlags::from(BitIterableContainer::from(0b1u16));
        assert!(dns_flags.get::<QR>());

        let dns_flags = DNSFlags::new();
        assert!(!dns_flags.get::<RA>());
        assert!(!dns_flags.get::<QR>());

        let dns_flags = DNSFlags::from(BitIterableContainer::from(0b1_0000_0000u16));
        assert!(dns_flags.get::<RA>());
    }

    #[test]
    fn set_custom_type() {
        let mut dns_flags = DNSFlags::new();

        dns_flags.set::<OPCODE>(OpCode::IQuery);
        assert_eq!(dns_flags.into_inner(), 0b10u16);

        dns_flags.set::<OPCODE>(OpCode::Status);
        assert_eq!(dns_flags.into_inner(), 0b100u16);

        dns_flags.set::<OPCODE>(OpCode::Query);
        assert_eq!(dns_flags.into_inner(), 0b0u16);
    }

    #[test]
    fn get_custom_type() {
        let dns_flags = DNSFlags::from(BitIterableContainer::from(0b10u16));
        assert_eq!(dns_flags.try_get::<OPCODE>(), Ok(OpCode::IQuery));

        let dns_flags = DNSFlags::from(BitIterableContainer::from(0b100u16));
        assert_eq!(dns_flags.try_get::<OPCODE>(), Ok(OpCode::Status));

        let dns_flags = DNSFlags::new();
        assert_eq!(dns_flags.try_get::<OPCODE>(), Ok(OpCode::Query));
    }
}