collum 0.1.0

use core::mem::MaybeUninit;
use core::ops::{Add, Not, Shl, Shr};
use core::slice::SliceIndex;

pub(crate) fn initialize_with<T, F, const N: usize>(mut f: F) -> [T; N]
where
    F: FnMut() -> T,
{
    let mut uninit = MaybeUninit::<[T; N]>::uninit();

    for i in 0..N {
        let ptr = uninit.as_mut_ptr();
        // SAFETY: The pointer is convertible to a reference as this is the only mutable
        // reference to `ret`.
        // SAFETY: `p` currently points to an allocated area in memory, thus not null.
        let ptr = unsafe { ptr.as_mut().unwrap_unchecked() };
        ptr[i] = f();
    }

    unsafe { uninit.assume_init() }
}

/// Safely defines an offset within a single [`u8`], such that any variant will denote a bit within
/// a [`u8`] value.
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(usize)]
pub enum BitOffset {
    /// The most significant bit.
    #[default]
    Bit0 = 0,
    /// Bit in position 1.
    Bit1 = 1,
    /// Bit in position 2.
    Bit2 = 2,
    /// Bit in position 3.
    Bit3 = 3,
    /// Bit in position 4.
    Bit4 = 4,
    /// Bit in position 5.
    Bit5 = 5,
    /// Bit in position 6.
    Bit6 = 6,
    /// The least significant bit.
    Bit7 = 7,
}

impl BitOffset {
    /// Returns a variant of [`BitOffset`] such that `None` is returned if `offset >= 8`.
    pub fn new(offset: u8) -> Option<Self> {
        [
            BitOffset::Bit0,
            BitOffset::Bit1,
            BitOffset::Bit2,
            BitOffset::Bit3,
            BitOffset::Bit4,
            BitOffset::Bit5,
            BitOffset::Bit6,
            BitOffset::Bit7,
        ]
        .get(offset as usize)
        .copied()
    }

    /// Infallibly returns a variant of [`BitOffset`] such that the return value is the offset
    /// `offset % 8`.
    ///
    /// This constructor is mostly for ease of use when we know we have a safe offset and there is
    /// no reason for the overhead of an [`Option`].
    pub fn new_wrapped(offset: usize) -> Self {
        *[
            BitOffset::Bit0,
            BitOffset::Bit1,
            BitOffset::Bit2,
            BitOffset::Bit3,
            BitOffset::Bit4,
            BitOffset::Bit5,
            BitOffset::Bit6,
            BitOffset::Bit7,
        ]
        .get(offset % crate::BITS_PER_BYTE)
        .expect("Index must be in range")
    }
}

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

impl From<BitOffset> for usize {
    fn from(offset: BitOffset) -> Self {
        offset as usize
    }
}

impl Add for BitOffset {
    type Output = (BitOffset, bool);

    fn add(self, rhs: BitOffset) -> Self::Output {
        let lhs = u8::from(self);
        let rhs = u8::from(rhs);

        match lhs + rhs {
            offset @ 0..8 => {
                let Some(val) = BitOffset::new(offset) else {
                    unreachable!("Value is normalized and within range of a bit offset");
                };
                (val, false)
            }
            offset @ 8..15 => {
                let Some(val) = BitOffset::new(offset - 8) else {
                    unreachable!("Value is normalized and within range of a bit offset");
                };
                (val, true)
            }
            _ => unreachable!(),
        }
    }
}

impl Not for BitOffset {
    type Output = BitOffset;

    fn not(self) -> Self::Output {
        let Some(offset) = BitOffset::new(8 - u8::from(self)) else {
            unreachable!()
        };
        offset
    }
}

impl Shr<BitOffset> for u8 {
    type Output = u8;

    fn shr(self, rhs: BitOffset) -> Self::Output {
        self.overflowing_shr(u8::from(rhs) as u32).0
    }
}

impl Shl<BitOffset> for u8 {
    type Output = u8;

    fn shl(self, rhs: BitOffset) -> Self::Output {
        self.overflowing_shl(u8::from(rhs) as u32).0
    }
}

// TODO: Add a `bits: usize` field to allow `data` not to be aligned to the right byte boundary.
/// A borrowed [`[u8]`] offsetable at bit-precision.
#[derive(Copy, Clone, Default, Debug, Eq, PartialOrd, Ord, Hash)]
pub struct BitSlice<'data> {
    pub(crate) data: &'data [u8],
    pub(crate) bit_offset: BitOffset,
}

impl PartialEq for BitSlice<'_> {
    fn eq(&self, other: &BitSlice<'_>) -> bool {
        if self.bit_offset != other.bit_offset {
            return false;
        }

        let mask = 0xff >> self.bit_offset;

        if self.data.first().map(|b| *b & mask) != other.data.first().map(|b| *b & mask) {
            return false;
        }

        self.data.get(1..) == other.data.get(1..)
    }
}

impl<'data> BitSlice<'data> {
    /// Constructs a `BitSlice` with the borrowed slice `data` at a bit-precise offset
    /// `bit_offset`.
    pub fn new(data: &'data [u8], bit_offset: BitOffset) -> Self {
        BitSlice { data, bit_offset }
    }

    /// Returns the amount of bits in the slice.
    pub fn bits(&self) -> usize {
        let data_bits = self.data.len() * 8;
        data_bits - usize::from(self.bit_offset)
    }

    /// Forces the bit offset to be `bit_offset`.
    pub fn with_offset(self, bit_offset: BitOffset) -> Self {
        BitSlice {
            data: self.data,
            bit_offset,
        }
    }

    /// Offsets the `BitSlice` by `<S as collum::Sized>::bits()` bits. Returns `None` if the
    /// operation goes out of bounds of the slice.
    pub fn offset_by_sized<S>(self) -> Option<BitSlice<'data>>
    where
        S: crate::Sized,
    {
        let BitSlice { data, bit_offset } = self;
        let (bytes, bits) = S::bytes_and_bits();

        let (bit_offset, overflow) = bit_offset + bits;
        let byte_offset = bytes + overflow as usize;

        let data = data.get(byte_offset..)?;

        Some(BitSlice { data, bit_offset })
    }

    /// Returns a reference to the underlying data.
    pub fn data(&self) -> &'data [u8] {
        self.data
    }

    /// Returns a reference to the underlying bit offset.
    pub fn bit_offset(&self) -> &BitOffset {
        &self.bit_offset
    }
}

impl<'data> BitSlice<'data> {
    /// Creates a statically-sized byte boundary-aligned array of bytes starting from the slice's
    /// bit offset. Returns `None` if there is not enough data in the slice.
    pub fn copy_array<const N: usize>(&self) -> Option<[u8; N]> {
        match self.bit_offset {
            BitOffset::Bit0 => self.data.first_chunk().copied(),
            _ => {
                if self.data.len() < N + 1 {
                    return None;
                }

                let BitSlice { data, bit_offset } = *self;

                let mut index = 0;

                Some(initialize_with(|| {
                    let Some([b1, b2]) = data.get(index..=index + 1) else {
                        unreachable!("Array bounds checked and constant sized range for indexing.")
                    };

                    index += 1;

                    (*b1 << bit_offset) | (*b2 >> !bit_offset)
                }))
            }
        }
    }

    /// Copies exactly `bits` bits from the underlying slice aligned to the left of the slice, such
    /// that the first read bit lies at `ret[0] & 0x80`. Returns `None` if more bits are being
    /// requested than the slice has.
    ///
    /// # Note
    ///
    /// The const parameter `N` must be chosen such that `bits` bits of data can fit within it, but
    /// this does not necessarily mean that the bit slice itself has enough data for `N` bytes.
    ///
    /// # Panics
    ///
    /// Panics if `N * 8 < bits`.
    pub fn copy_bits<const N: usize>(&self, bits: usize) -> Option<[u8; N]> {
        use crate::BITS_PER_BYTE;

        if N * BITS_PER_BYTE < bits {
            panic!("N ({}) must be larger than amount of bits ({})", N, bits);
        }

        if bits % BITS_PER_BYTE == 0 {
            // If bit count lands on byte boundary we can just redirect to `copy_array` and
            // fill any extra data between `bits` and `N` with 0 (to simulate the padding).
            // Eventually this can be made unnecessary by stabilization of `generic_const_exprs`.
            self.copy_array::<N>().map(|mut array: [u8; N]| -> [u8; N] {
                if let Some(last) = array.get_mut((bits / BITS_PER_BYTE)..) {
                    last.fill(0);
                }

                array
            })
        } else {
            // We essentially have to recreate the functionality of `copy_array`, perhaps a runtime
            // version of the function would help in these situations at least until further rust
            // features are available to help with const things.
            match self.bit_offset {
                BitOffset::Bit0 => {
                    self.data
                        .first_chunk()
                        .copied()
                        .map(|mut array: [u8; N]| -> [u8; N] {
                            if let Some(last) = array.get_mut((bits / BITS_PER_BYTE + 1)..) {
                                last.fill(0);
                            }

                            if let Some(edge) = array.get_mut(bits / BITS_PER_BYTE) {
                                let mask = 0xff >> (bits % BITS_PER_BYTE);
                                *edge &= !mask;
                            }

                            array
                        })
                }
                _ => {
                    if self.bits() < bits {
                        return None;
                    }

                    let BitSlice { data, bit_offset } = *self;

                    let mut index = 0;

                    let mut bits_left = bits;

                    Some(initialize_with(|| {
                        // This part should be infallible due to the array bound checks.
                        let Some(b1) = data.get(index) else {
                            unreachable!("Array bounds checked")
                        };

                        index += 1;
                        let bits = bits_left;
                        bits_left = bits_left.checked_sub(BITS_PER_BYTE).unwrap_or_default();

                        let b2 = data.get(index).map(|b2| *b2 >> !bit_offset).unwrap_or(0);

                        ((*b1 << bit_offset) | b2) & !(0xff >> (bits % BITS_PER_BYTE))
                    }))
                }
            }
        }
    }

    /// Gets a segment of the slice using a byte-level index. This means that it also keeps the
    /// same bit offset in the result. Returns `None` if there is not enough data.
    pub fn get<I>(&self, index: I) -> Option<BitSlice<'data>>
    where
        I: SliceIndex<[u8], Output = [u8]>,
    {
        let BitSlice { data, bit_offset } = *self;

        let data = data.get(index)?;

        Some(BitSlice { data, bit_offset })
    }
}

impl<'data> From<&'data [u8]> for BitSlice<'data> {
    fn from(data: &'data [u8]) -> Self {
        BitSlice {
            data,
            bit_offset: Default::default(),
        }
    }
}

impl<'data, const N: usize> From<&'data [u8; N]> for BitSlice<'data> {
    fn from(data: &'data [u8; N]) -> Self {
        BitSlice {
            data,
            bit_offset: Default::default(),
        }
    }
}

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

    #[test]
    fn copy_array() {
        let data = &[0b00000001, 0b00000010, 0b00000011];
        let slice: BitSlice = data.into();

        assert_eq!(slice.copy_array::<2>(), Some([0b00000001, 0b00000010]));

        let slice = slice.with_offset(BitOffset::Bit1);
        assert_eq!(slice.copy_array::<2>(), Some([0b00000010, 0b00000100]));

        let slice = slice.with_offset(BitOffset::Bit7);
        assert_eq!(slice.copy_array::<2>(), Some([0b10000001, 0b00000001]));
    }
}