ashv2 6.0.2

//! A three bit number that wraps around.

use core::fmt::Display;
use core::num::NonZero;
use core::ops::{Add, AddAssign};

const MASK: u8 = 0b0000_0111;
const UNUSED_BITS: u8 = !MASK;

/// A three bit unsigned integer which wraps on adding.
// The inner `NonZero<u8>` type is used in conjunction with `#[repr(transparent)]`
// to allow niche optimizations of this type when used in e.g. an `Option.
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[repr(transparent)]
pub struct WrappingU3(NonZero<u8>);

impl WrappingU3 {
    /// The zero value.
    pub const ZERO: Self = Self::from_u8_lossy(0);

    /// Creates a new three bit number.
    #[must_use]
    pub const fn from_u8_lossy(n: u8) -> Self {
        #[expect(unsafe_code)]
        // SAFETY: We create a three bit number by applying `MASK` to `n`.
        // Finally, we OR the result with `UNUSED_BITS`, which makes the number non-zero.
        Self(unsafe { NonZero::new_unchecked(n & MASK | UNUSED_BITS) })
    }

    /// Returns the number as an u8.
    #[must_use]
    pub const fn as_u8(self) -> u8 {
        self.0.get() & MASK
    }
}

impl Add for WrappingU3 {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        self.add(rhs.as_u8())
    }
}

impl Add<u8> for WrappingU3 {
    type Output = Self;

    fn add(self, rhs: u8) -> Self::Output {
        Self::from_u8_lossy(self.as_u8().wrapping_add(rhs))
    }
}

impl AddAssign for WrappingU3 {
    fn add_assign(&mut self, rhs: Self) {
        self.add_assign(rhs.as_u8());
    }
}

impl AddAssign<u8> for WrappingU3 {
    fn add_assign(&mut self, rhs: u8) {
        *self = Self::from_u8_lossy(self.as_u8().wrapping_add(rhs));
    }
}

impl Default for WrappingU3 {
    fn default() -> Self {
        Self::ZERO
    }
}

impl Display for WrappingU3 {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.as_u8())
    }
}

impl PartialEq<u8> for WrappingU3 {
    fn eq(&self, other: &u8) -> bool {
        self.as_u8() == *other
    }
}

#[cfg(test)]
mod tests {
    use super::WrappingU3;

    #[test]
    fn test_new() {
        for n in u8::MIN..=u8::MAX {
            let number = WrappingU3::from_u8_lossy(n);
            assert_eq!(number.as_u8(), n % 8);
        }
    }

    #[test]
    fn test_const_zero() {
        assert_eq!(WrappingU3::ZERO.as_u8(), 0);
    }

    #[test]
    fn test_as_u8() {
        for n in u8::MIN..=u8::MAX {
            let number = WrappingU3::from_u8_lossy(n);
            assert_eq!(number.as_u8(), n % 8);
        }
    }

    #[test]
    fn test_add_self() {
        for n in 0..=u8::MAX {
            for rhs in 0..=u8::MAX {
                let number = WrappingU3::from_u8_lossy(n) + WrappingU3::from_u8_lossy(rhs);
                assert_eq!(number.as_u8(), n.wrapping_add(rhs) % 8);
            }
        }
    }

    #[test]
    fn test_add_u8() {
        for n in 0..=u8::MAX {
            for rhs in 0..=u8::MAX {
                let number = WrappingU3::from_u8_lossy(n) + rhs;
                assert_eq!(number.as_u8(), n.wrapping_add(rhs) % 8);
            }
        }
    }

    #[test]
    fn test_add_assign_self() {
        for n in 0..=u8::MAX {
            for rhs in 0..=u8::MAX {
                let mut number = WrappingU3::from_u8_lossy(n);
                number += WrappingU3::from_u8_lossy(rhs);
                assert_eq!(number.as_u8(), n.wrapping_add(rhs) % 8);
            }
        }
    }

    #[test]
    fn test_add_assign_u8() {
        for n in 0..=u8::MAX {
            for rhs in 0..=u8::MAX {
                let mut number = WrappingU3::from_u8_lossy(n);
                number += rhs;
                assert_eq!(number.as_u8(), n.wrapping_add(rhs) % 8);
            }
        }
    }
}