ashv2 1.2.8

//! 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 NON_ZERO_BIT: u8 = 0b0000_1000;

/// A three bit number.
#[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 `NON_ZERO_BIT`, which makes the number non-zero.
        Self(unsafe { NonZero::new_unchecked(n & MASK | NON_ZERO_BIT) })
    }

    /// 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::from_u8_lossy(self.as_u8().wrapping_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 = Self::from_u8_lossy(self.as_u8().wrapping_add(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
    }
}

impl From<WrappingU3> for u8 {
    fn from(value: WrappingU3) -> Self {
        value.as_u8()
    }
}

impl TryFrom<u8> for WrappingU3 {
    type Error = u8;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        let u3 = Self::from_u8_lossy(value);

        if u3.as_u8() == value {
            Ok(u3)
        } else {
            Err(value)
        }
    }
}

#[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);
            }
        }
    }
}