bitkit 3.2.3

//! The [`Bits<T>`] newtype — the primary entry point of the crate.

use crate::{BitError, IntoBitRange};

/// Width-typed wrapper around a primitive unsigned integer.
///
/// All bit-level operations are methods on this type. Per-width inherent
/// `impl` blocks are generated by an internal macro; from the outside it
/// looks like a single generic surface.
///
/// ```
/// use bitkit::Bits;
/// let x = Bits::<u32>::new(0b1011_0000);
/// assert_eq!(x.isolate_lowest_set_bit(), Bits::new(0b0001_0000));
/// ```
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
#[repr(transparent)]
pub struct Bits<T>(pub T);

impl<T: core::fmt::Debug> core::fmt::Debug for Bits<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "Bits({:?})", self.0)
    }
}

impl<T> Bits<T> {
    /// Wraps a raw value.
    #[inline] pub const fn new(value: T) -> Self { Self(value) }
}

impl<T: Copy> Bits<T> {
    /// Returns the underlying value.
    #[inline] pub const fn get(self) -> T { self.0 }
}

// ---------------- per-width method macro ----------------------------------

macro_rules! impl_bits {
    ($ty:ty) => {
        impl Bits<$ty> {
            /// Number of bits.
            pub const BITS: u32 = <$ty>::BITS;

            /// Bit `index` set? `Err` if `index >= BITS`.
            #[inline] pub const fn has_bit(self, index: u32) -> Result<bool, BitError> {
                if index < Self::BITS { Ok((self.0 & ((1 as $ty) << index)) != 0) } else { Err(BitError::IndexOutOfRange) }
            }
            /// Set bit `index`. `Err` if `index >= BITS`.
            #[inline] pub const fn set_bit(self, index: u32) -> Result<Self, BitError> {
                if index < Self::BITS { Ok(Self(self.0 | ((1 as $ty) << index))) } else { Err(BitError::IndexOutOfRange) }
            }
            /// Clear bit `index`. `Err` if `index >= BITS`.
            #[inline] pub const fn clear_bit(self, index: u32) -> Result<Self, BitError> {
                if index < Self::BITS { Ok(Self(self.0 & !((1 as $ty) << index))) } else { Err(BitError::IndexOutOfRange) }
            }
            /// Toggle bit `index`. `Err` if `index >= BITS`.
            #[inline] pub const fn toggle_bit(self, index: u32) -> Result<Self, BitError> {
                if index < Self::BITS { Ok(Self(self.0 ^ ((1 as $ty) << index))) } else { Err(BitError::IndexOutOfRange) }
            }

            /// `has_bit` with index normalized modulo `BITS`.
            #[inline] pub const fn has_bit_wrapping(self, index: u32) -> bool { (self.0 & ((1 as $ty) << (index % Self::BITS))) != 0 }
            /// `set_bit` with index normalized modulo `BITS`.
            #[inline] pub const fn set_bit_wrapping(self, index: u32) -> Self { Self(self.0 | ((1 as $ty) << (index % Self::BITS))) }
            /// `clear_bit` with index normalized modulo `BITS`.
            #[inline] pub const fn clear_bit_wrapping(self, index: u32) -> Self { Self(self.0 & !((1 as $ty) << (index % Self::BITS))) }
            /// `toggle_bit` with index normalized modulo `BITS`.
            #[inline] pub const fn toggle_bit_wrapping(self, index: u32) -> Self { Self(self.0 ^ ((1 as $ty) << (index % Self::BITS))) }

            /// `x & -x` — isolates the lowest set bit, or `0`.
            #[inline] pub const fn isolate_lowest_set_bit(self) -> Self { Self(self.0 & self.0.wrapping_neg()) }
            /// `x & (x - 1)` — clears the lowest set bit.
            #[inline] pub const fn clear_lowest_set_bit(self) -> Self { Self(self.0 & self.0.wrapping_sub(1)) }
            /// `!x & -(!x)` — isolates the lowest zero bit, or `0`.
            #[inline] pub const fn isolate_lowest_zero_bit(self) -> Self { let n = !self.0; Self(n & n.wrapping_neg()) }
            /// `x | (x + 1)` — sets the lowest zero bit.
            #[inline] pub const fn set_lowest_zero_bit(self) -> Self { Self(self.0 | self.0.wrapping_add(1)) }
            /// `x ^ (x - 1)` — mask through the lowest set bit (BMI1 `BLSMSK`).
            #[inline] pub const fn mask_through_lowest_set_bit(self) -> Self { Self(self.0 ^ self.0.wrapping_sub(1)) }
            /// Reverse the order of bits. Hardware: ARM `RBIT`, x86 SWAR.
            #[inline] pub const fn reverse_bits(self) -> Self { Self(self.0.reverse_bits()) }
            /// Keep the low `n` bits, zero the rest (BMI2 `BZHI` semantics).
            /// `Err(IndexOutOfRange)` if `n > BITS`.
            #[inline] pub const fn zero_high_bits(self, n: u32) -> Result<Self, BitError> {
                match Self::low_mask(n) { Ok(m) => Ok(Self(self.0 & m.0)), Err(e) => Err(e) }
            }

            /// Set bits.
            #[inline] pub const fn count_ones(self) -> u32 { self.0.count_ones() }
            /// Cleared bits.
            #[inline] pub const fn count_zeros(self) -> u32 { self.0.count_zeros() }
            /// Leading zeros.
            #[inline] pub const fn leading_zeros(self) -> u32 { self.0.leading_zeros() }
            /// Trailing zeros.
            #[inline] pub const fn trailing_zeros(self) -> u32 { self.0.trailing_zeros() }

            /// True if exactly one bit set.
            #[inline] pub const fn is_power_of_two(self) -> bool { self.0 != 0 && (self.0 & self.0.wrapping_sub(1)) == 0 }
            /// Smallest power of two `>= self`. `0` → [`BitError::UnderflowFromZero`].
            #[inline]
            pub const fn next_power_of_two(self) -> Result<Self, BitError> {
                if self.0 == 0 { return Err(BitError::UnderflowFromZero); }
                match <$ty>::checked_next_power_of_two(self.0) { Some(p) => Ok(Self(p)), None => Err(BitError::Overflow) }
            }
            /// Largest power of two `<= self`. `0` → [`BitError::UnderflowFromZero`].
            #[inline]
            pub const fn previous_power_of_two(self) -> Result<Self, BitError> {
                if self.0 == 0 { Err(BitError::UnderflowFromZero) }
                else { Ok(Self((1 as $ty) << (Self::BITS - 1 - self.0.leading_zeros()))) }
            }

            /// Mask with the low `width` bits set.
            #[inline]
            pub const fn low_mask(width: u32) -> Result<Self, BitError> {
                if width > Self::BITS { Err(BitError::IndexOutOfRange) }
                else if width == Self::BITS { Ok(Self(!(0 as $ty))) }
                else { Ok(Self(((1 as $ty) << width) - 1)) }
            }
            /// Mask with the high `width` bits set.
            #[inline]
            pub const fn high_mask(width: u32) -> Result<Self, BitError> {
                if width > Self::BITS { Err(BitError::IndexOutOfRange) }
                else if width == 0 { Ok(Self(0 as $ty)) }
                else if width == Self::BITS { Ok(Self(!(0 as $ty))) }
                else { Ok(Self((!(0 as $ty)) << (Self::BITS - width))) }
            }
            /// Mask with bits in `range` set (`Range<u32>` or `(u32, u32)`).
            #[inline]
            pub fn range_mask(range: impl IntoBitRange) -> Result<Self, BitError> {
                let (s, e) = range.into_bit_range();
                Self::range_mask_indices(s, e)
            }
            /// Const-callable form of [`range_mask`](Self::range_mask).
            #[inline]
            pub const fn range_mask_indices(start: u32, end: u32) -> Result<Self, BitError> {
                if end > Self::BITS || start > end { return Err(BitError::InvalidRange); }
                if start == end { return Ok(Self(0 as $ty)); }
                let w = end - start;
                let low: $ty = if w == Self::BITS { !(0 as $ty) } else { ((1 as $ty) << w) - 1 };
                Ok(Self(low << start))
            }

            /// Extract bit field at `range`, right-aligned.
            #[inline]
            pub fn extract(self, range: impl IntoBitRange) -> Result<Self, BitError> {
                let (s, e) = range.into_bit_range();
                self.extract_indices(s, e)
            }
            /// Const-callable form of [`extract`](Self::extract).
            #[inline]
            pub const fn extract_indices(self, start: u32, end: u32) -> Result<Self, BitError> {
                if end > Self::BITS || start > end { return Err(BitError::InvalidRange); }
                if start == end { return Ok(Self(0 as $ty)); }
                let w = end - start;
                let m: $ty = if w == Self::BITS { !(0 as $ty) } else { ((1 as $ty) << w) - 1 };
                Ok(Self((self.0 >> start) & m))
            }
            /// Insert `value` into bit field at `range`. Fails if `value` doesn't fit.
            #[inline]
            pub fn insert(self, range: impl IntoBitRange, value: $ty) -> Result<Self, BitError> {
                let (s, e) = range.into_bit_range();
                self.insert_indices(s, e, value)
            }
            /// Const-callable form of [`insert`](Self::insert).
            #[inline]
            pub const fn insert_indices(self, start: u32, end: u32, value: $ty) -> Result<Self, BitError> {
                if end > Self::BITS || start > end { return Err(BitError::InvalidRange); }
                if start == end {
                    return if value != 0 { Err(BitError::FieldValueTooLarge) } else { Ok(self) };
                }
                let w = end - start;
                let vm: $ty = if w == Self::BITS { !(0 as $ty) } else { ((1 as $ty) << w) - 1 };
                if value & !vm != 0 { return Err(BitError::FieldValueTooLarge); }
                let fm = vm << start;
                Ok(Self((self.0 & !fm) | (value << start)))
            }
            /// Replace bit field at `range` with `value`, silently truncating to field width.
            #[inline]
            pub fn replace(self, range: impl IntoBitRange, value: $ty) -> Result<Self, BitError> {
                let (s, e) = range.into_bit_range();
                self.replace_indices(s, e, value)
            }
            /// Const-callable form of [`replace`](Self::replace).
            #[inline]
            pub const fn replace_indices(self, start: u32, end: u32, value: $ty) -> Result<Self, BitError> {
                if end > Self::BITS || start > end { return Err(BitError::InvalidRange); }
                if start == end { return Ok(self); }
                let w = end - start;
                let vm: $ty = if w == Self::BITS { !(0 as $ty) } else { ((1 as $ty) << w) - 1 };
                let fm = vm << start;
                Ok(Self((self.0 & !fm) | ((value & vm) << start)))
            }

            /// Iterator over set-bit indexes (ascending).
            #[inline] pub const fn set_bits(self) -> SetBits<$ty> { SetBits(self.0) }
            /// Iterator over isolated set-bit values (ascending).
            #[inline] pub const fn set_bit_values(self) -> SetBitValues<$ty> { SetBitValues(self.0) }
            /// Iterator over zero-bit indexes (ascending).
            #[inline] pub const fn zero_bits(self) -> ZeroBits<$ty> { ZeroBits(self.0) }
            /// Iterator over submasks of `self` (descending; includes `self` and `0`).
            #[inline] pub const fn submasks(self) -> Submasks<$ty> { Submasks { mask: self.0, current: Some(self.0) } }
            /// Iterator over proper submasks of `self` (excludes `self`).
            #[inline]
            pub const fn proper_submasks(self) -> ProperSubmasks<$ty> {
                ProperSubmasks { inner: Submasks { mask: self.0, current: Some(self.0) }, skipped: false }
            }
        }

        impl From<$ty> for Bits<$ty> { #[inline] fn from(v: $ty) -> Self { Self(v) } }
        impl From<Bits<$ty>> for $ty { #[inline] fn from(b: Bits<$ty>) -> Self { b.0 } }
        impl core::ops::BitAnd for Bits<$ty> { type Output = Self; #[inline] fn bitand(self, r: Self) -> Self { Self(self.0 & r.0) } }
        impl core::ops::BitOr  for Bits<$ty> { type Output = Self; #[inline] fn bitor (self, r: Self) -> Self { Self(self.0 | r.0) } }
        impl core::ops::BitXor for Bits<$ty> { type Output = Self; #[inline] fn bitxor(self, r: Self) -> Self { Self(self.0 ^ r.0) } }
        impl core::ops::Not    for Bits<$ty> { type Output = Self; #[inline] fn not   (self)         -> Self { Self(!self.0) } }
        impl core::ops::BitAndAssign for Bits<$ty> { #[inline] fn bitand_assign(&mut self, r: Self) { self.0 &= r.0; } }
        impl core::ops::BitOrAssign  for Bits<$ty> { #[inline] fn bitor_assign (&mut self, r: Self) { self.0 |= r.0; } }
        impl core::ops::BitXorAssign for Bits<$ty> { #[inline] fn bitxor_assign(&mut self, r: Self) { self.0 ^= r.0; } }
        impl core::fmt::Binary   for Bits<$ty> { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { core::fmt::Binary  ::fmt(&self.0, f) } }
        impl core::fmt::LowerHex for Bits<$ty> { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { core::fmt::LowerHex::fmt(&self.0, f) } }
        impl core::fmt::UpperHex for Bits<$ty> { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { core::fmt::UpperHex::fmt(&self.0, f) } }
    };
}

impl_bits!(u8);
impl_bits!(u16);
impl_bits!(u32);
impl_bits!(u64);
impl_bits!(u128);
impl_bits!(usize);

// ---------------- gather / scatter (hardware-accelerated) ----------------

use crate::accel;

impl Bits<u32> {
    /// Pack the bits of `self` selected by `mask` into the low bits of the
    /// result (PEXT semantics). One instruction on x86_64 BMI2; SWAR fallback otherwise.
    ///
    /// ```
    /// # use bitkit::Bits;
    /// let v = Bits::<u32>::new(0b1011_0101);
    /// assert_eq!(v.gather(Bits::new(0b1001_0101)).get(), 0b1111);
    /// ```
    #[inline] pub fn gather(self, mask: Self) -> Self { Self(accel::pext_u32(self.0, mask.0)) }
    /// Scatter the low `popcount(mask)` bits of `self` into the positions
    /// of `mask` (PDEP semantics). Inverse of [`gather`](Self::gather).
    ///
    /// ```
    /// # use bitkit::Bits;
    /// assert_eq!(Bits::<u32>::new(0b1111).scatter(Bits::new(0b1001_0101)).get(), 0b1001_0101);
    /// ```
    #[inline] pub fn scatter(self, mask: Self) -> Self { Self(accel::pdep_u32(self.0, mask.0)) }
}

impl Bits<u64> {
    /// PEXT for `u64`. See [`Bits::<u32>::gather`].
    #[inline] pub fn gather(self, mask: Self) -> Self { Self(accel::pext_u64(self.0, mask.0)) }
    /// PDEP for `u64`. See [`Bits::<u32>::scatter`].
    #[inline] pub fn scatter(self, mask: Self) -> Self { Self(accel::pdep_u64(self.0, mask.0)) }
}

// ---------------- iterators -----------------------------------------------
//
// The `for x in iter` form goes through `next()` which Rust cannot always
// fold across the loop boundary; for hot inner loops, prefer
// `iter.for_each(|x| ...)`, overridden below for a tight bit-manipulation
// loop with no `Option<T>` round-trip.

/// Iterator over set-bit indexes (ascending). See [`Bits::set_bits`].
#[derive(Clone, Debug)] pub struct SetBits<T>(T);
/// Iterator over isolated set-bit values (ascending). See [`Bits::set_bit_values`].
#[derive(Clone, Debug)] pub struct SetBitValues<T>(T);
/// Iterator over zero-bit indexes (ascending). See [`Bits::zero_bits`].
#[derive(Clone, Debug)] pub struct ZeroBits<T>(T);
/// Iterator over submasks of a mask (descending). See [`Bits::submasks`].
#[derive(Clone, Debug)] pub struct Submasks<T> { mask: T, current: Option<T> }
/// Iterator over proper submasks. See [`Bits::proper_submasks`].
#[derive(Clone, Debug)] pub struct ProperSubmasks<T> { inner: Submasks<T>, skipped: bool }

macro_rules! impl_iters {
    ($ty:ty) => {
        impl Iterator for SetBits<$ty> {
            type Item = u32;
            #[inline]
            fn next(&mut self) -> Option<u32> {
                if self.0 == 0 { None } else { let i = self.0.trailing_zeros(); self.0 &= self.0.wrapping_sub(1); Some(i) }
            }
            #[inline]
            fn size_hint(&self) -> (usize, Option<usize>) { let c = self.0.count_ones() as usize; (c, Some(c)) }
            #[inline]
            fn for_each<F: FnMut(u32)>(self, mut f: F) {
                let mut v = self.0;
                while v != 0 { let i = v.trailing_zeros(); v &= v.wrapping_sub(1); f(i); }
            }
        }
        impl ExactSizeIterator for SetBits<$ty> {}
        impl core::iter::FusedIterator for SetBits<$ty> {}

        impl Iterator for SetBitValues<$ty> {
            type Item = $ty;
            #[inline]
            fn next(&mut self) -> Option<$ty> {
                if self.0 == 0 { None } else { let b = self.0 & self.0.wrapping_neg(); self.0 ^= b; Some(b) }
            }
            #[inline]
            fn size_hint(&self) -> (usize, Option<usize>) { let c = self.0.count_ones() as usize; (c, Some(c)) }
            #[inline]
            fn for_each<F: FnMut($ty)>(self, mut f: F) {
                let mut v = self.0;
                while v != 0 { let b = v & v.wrapping_neg(); v ^= b; f(b); }
            }
        }
        impl ExactSizeIterator for SetBitValues<$ty> {}
        impl core::iter::FusedIterator for SetBitValues<$ty> {}

        impl Iterator for ZeroBits<$ty> {
            type Item = u32;
            #[inline]
            fn next(&mut self) -> Option<u32> {
                if self.0 == !(0 as $ty) { None }
                else { let i = (!self.0).trailing_zeros(); self.0 |= (1 as $ty) << i; Some(i) }
            }
            #[inline]
            fn size_hint(&self) -> (usize, Option<usize>) { let c = self.0.count_zeros() as usize; (c, Some(c)) }
            #[inline]
            fn for_each<F: FnMut(u32)>(self, mut f: F) {
                let mut v = self.0;
                while v != !(0 as $ty) { let i = (!v).trailing_zeros(); v |= (1 as $ty) << i; f(i); }
            }
        }
        impl ExactSizeIterator for ZeroBits<$ty> {}
        impl core::iter::FusedIterator for ZeroBits<$ty> {}

        impl Iterator for Submasks<$ty> {
            type Item = Bits<$ty>;
            #[inline]
            fn next(&mut self) -> Option<Bits<$ty>> {
                let c = self.current?;
                self.current = if c == 0 { None } else { Some(c.wrapping_sub(1) & self.mask) };
                Some(Bits(c))
            }
            #[inline]
            fn for_each<F: FnMut(Bits<$ty>)>(self, mut f: F) {
                if let Some(start) = self.current {
                    let (mut c, mask) = (start, self.mask);
                    loop { f(Bits(c)); if c == 0 { break; } c = c.wrapping_sub(1) & mask; }
                }
            }
        }
        impl core::iter::FusedIterator for Submasks<$ty> {}

        impl Iterator for ProperSubmasks<$ty> {
            type Item = Bits<$ty>;
            #[inline]
            fn next(&mut self) -> Option<Bits<$ty>> {
                if !self.skipped { self.skipped = true; let _ = self.inner.next(); }
                self.inner.next()
            }
            #[inline]
            fn for_each<F: FnMut(Bits<$ty>)>(mut self, f: F) {
                if !self.skipped { let _ = self.inner.next(); }
                self.inner.for_each(f);
            }
        }
        impl core::iter::FusedIterator for ProperSubmasks<$ty> {}
    };
}
impl_iters!(u8); impl_iters!(u16); impl_iters!(u32); impl_iters!(u64); impl_iters!(u128); impl_iters!(usize);