bit-buf 0.1.7

use crate::{BitBuf, Error, Result, StorageMut};

impl<S: StorageMut> BitBuf<S> {
  /// Write a [`u8`] in BE-bit-order at `byte_offset` without performing bound checks
  ///
  /// # Safety
  ///
  /// * This is UB if [`byte_offset >= self.bytes().len()`][Self::bytes]
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if [`byte_offset >= self.bytes().len()`][Self::bytes]
  #[inline(always)]
  pub unsafe fn write_u8_be_aligned_full_at_unchecked(
    &mut self,
    byte_offset: usize,
    v: u8,
  ) -> &mut Self {
    let bytes = self.bytes_mut();

    debug_assert!(
      byte_offset < bytes.len(),
      "BitBuf::write_u8_be_aligned_full_at_unchecked: index out of bounds! len is {}, offset is {}",
      bytes.len(),
      byte_offset,
    );

    unsafe { *bytes.get_unchecked_mut(byte_offset) = v };
    self
  }

  /// Write a [`u8`] in BE-bit-order without performing bound checks, advancing the internal cursor
  ///
  /// # Safety
  ///
  /// * The internal cursor must be byte-aligned ([`self.is_aligned()`][Self::is_aligned])
  /// * This is UB if the internal cursor points past the end of storage (<code>[self.byte_pos()][Self::byte_pos] >= [self.bytes().len()][Self::bytes]</code>)
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if the internal cursor is not byte-aligned ([`!self.is_aligned()`][Self::is_aligned])
  /// * Panics in debug mode if the internal cursor points past the end of storage (<code>[self.byte_pos()][Self::byte_pos] >= [self.bytes().len()][Self::bytes]</code>)
  #[inline(always)]
  pub unsafe fn write_u8_be_aligned_full_unchecked(&mut self, v: u8) -> &mut Self {
    debug_assert!(
      self.is_aligned(),
      "BitBuf::write_u8_be_aligned_full_unchecked called at unaligned bit position: {}",
      self.pos(),
    );

    unsafe { self.write_u8_be_aligned_full_at_unchecked(self.byte_pos(), v) };
    self.advance_bytes(1);
    self
  }

  /// Write a [`u8`] in BE-bit-order at `offset` without performing bound checks
  ///
  /// # Safety
  ///
  /// * This is UB if <code>(offset / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  /// * This is UB if <code>(offset % 8 != 0) && (offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if <code>(offset / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  /// * Panics in debug mode if <code>(offset % 8 != 0) && (offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub unsafe fn write_u8_be_full_at_unchecked(&mut self, offset: usize, v: u8) -> &mut Self {
    let byte_idx = offset / 8;
    let shift = offset % 8;

    let bytes = self.bytes_mut();

    debug_assert!(
      byte_idx < bytes.len(),
      "BitBuf::write_u8_be_full_at_unchecked: index out of bounds! len is {}, byte_idx is {}",
      bytes.len(),
      byte_idx,
    );

    if shift == 0 {
      unsafe { self.write_u8_be_aligned_full_at_unchecked(byte_idx, v) };
    } else {
      let next_idx = byte_idx + 1;

      debug_assert!(
        next_idx < bytes.len(),
        "BitBuf::write_u8_be_full_at_unchecked: lookahead index out of bounds! len is {}, lookahead byte_idx is {}",
        bytes.len(),
        next_idx,
      );

      // first byte, higher part of value
      unsafe {
        // mask to keep high bits of the current byte
        let mask = !0 << (8 - shift);
        let b = bytes.get_unchecked(byte_idx) & mask;
        // high bits of the value, shifted down
        let new = v >> shift;
        *bytes.get_unchecked_mut(byte_idx) = b | new;
      }

      // second byte, lower part of value
      unsafe {
        // mask to keep low bits of the current byte
        let mask = !0 >> shift;
        let b = bytes.get_unchecked(next_idx) & mask;
        // low bits of the value, shifted up
        let new = v << (8 - shift);
        *bytes.get_unchecked_mut(next_idx) = b | new;
      }
    }

    self
  }

  /// Write a [`u8`] in BE-bit-order without performing bound checks, advancing the internal cursor
  ///
  /// # Safety
  ///
  /// * This is UB if <code>(self.pos() / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  /// * This is UB if <code>(self.pos() % 8 != 0) && (self.pos() / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if <code>(self.pos() / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  /// * Panics in debug mode if <code>(self.pos() % 8 != 0) && (self.pos() / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub unsafe fn write_u8_be_full_unchecked(&mut self, v: u8) -> &mut Self {
    unsafe { self.write_u8_be_full_at_unchecked(self.pos(), v) };
    self.advance_bytes(1);
    self
  }

  /// Write a BE-bit-order [`u8`] at `offset` while performing bound checks
  ///
  /// # Errors
  ///
  /// * Returns [`Error::OutOfBounds`]
  ///   * if <code>(offset / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  ///   * if <code>(offset % 8 != 0) && (offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub fn try_write_u8_be_full_at(&mut self, offset: usize, v: u8) -> Result<&mut Self> {
    let byte_idx = offset / 8;
    let shift = offset % 8;

    let len = self.bytes().len();

    if byte_idx >= len {
      return Err(Error::OutOfBounds);
    }

    if shift != 0 && byte_idx + 1 >= len {
      return Err(Error::OutOfBounds);
    }

    Ok(unsafe { self.write_u8_be_full_at_unchecked(offset, v) })
  }

  /// Write a BE-bit-order [`u8`] while performing bound checks, advancing the internal cursor
  ///
  /// # Errors
  ///
  /// * Returns [`Error::OutOfBounds`]
  ///   * if <code>(self.pos() / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  ///   * if <code>(self.pos() % 8 != 0) && (self.pos() / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub fn try_write_u8_be_full(&mut self, v: u8) -> Result<&mut Self> {
    self.try_write_u8_be_full_at(self.pos(), v)?;
    self.advance_bytes(1);
    Ok(self)
  }

  /// Write a BE-bit-order [`u8`] at `offset`, panicking on out of bounds
  ///
  /// # Panics
  ///
  /// * Panics if <code>(offset / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  /// * Panics if <code>(offset % 8 != 0) && (offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub fn write_u8_be_full_at(&mut self, offset: usize, v: u8) -> &mut Self {
    self
      .try_write_u8_be_full_at(offset, v)
      .expect("BitBuf::write_u8_be_full_at out of bounds")
  }

  /// Write a BE-bit-order [`u8`], panicking on out of bounds, advancing the internal cursor
  ///
  /// # Panics
  ///
  /// * Panics if <code>(self.pos() / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  /// * Panics if <code>(self.pos() % 8 != 0) && (self.pos() / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub fn write_u8_be_full(&mut self, v: u8) -> &mut Self {
    self
      .try_write_u8_be_full(v)
      .expect("BitBuf::write_u8_be_full out of bounds")
  }

  /// Write a [`u16`] in [BE-bit, BE-byte] order at `byte_offset` without performing bound checks
  ///
  /// # Safety
  ///
  /// * This is UB if [`byte_offset + 1 >= self.bytes().len()`][Self::bytes]
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if [`byte_offset + 1 >= self.bytes().len()`][Self::bytes]
  #[inline(always)]
  pub unsafe fn write_u16_be_aligned_full_at_unchecked(
    &mut self,
    byte_offset: usize,
    v: u16,
  ) -> &mut Self {
    let bytes = self.bytes_mut();

    debug_assert!(
      byte_offset + 1 < bytes.len(),
      "BitBuf::write_u16_be_aligned_full_at_unchecked: index out of bounds! len is {}, offset is {}",
      bytes.len(),
      byte_offset + 1,
    );

    let [high, low] = v.to_be_bytes();

    unsafe {
      *bytes.get_unchecked_mut(byte_offset) = high;
      *bytes.get_unchecked_mut(byte_offset + 1) = low;
    }

    self
  }

  /// Write a [`u16`] in [BE-bit, BE-byte] order without performing bound checks, advancing the internal cursor
  ///
  /// # Safety
  ///
  /// * The internal cursor must be byte-aligned ([`self.is_aligned()`][Self::is_aligned])
  /// * This is UB if the internal cursor points past the end of storage (<code>[self.byte_pos()][Self::byte_pos] + 1 >= [self.bytes().len()][Self::bytes]</code>)
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if the internal cursor is not byte-aligned ([`!self.is_aligned()`][Self::is_aligned])
  /// * Panics in debug mode if the internal cursor points past the end of storage (<code>[self.byte_pos()][Self::byte_pos] + 1 >= [self.bytes().len()][Self::bytes]</code>)
  #[inline(always)]
  pub unsafe fn write_u16_be_aligned_full_unchecked(&mut self, v: u16) -> &mut Self {
    debug_assert!(
      self.is_aligned(),
      "BitBuf::write_u16_be_aligned_full_unchecked called at unaligned bit position: {}",
      self.pos(),
    );

    unsafe { self.write_u16_be_aligned_full_at_unchecked(self.byte_pos(), v) };
    self.advance_bytes(2);
    self
  }

  /// Write a [`u16`] in [BE-bit, BE-byte] order at `offset` without performing bound checks
  ///
  /// # Safety
  ///
  /// * This is UB if <code>(offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  /// * This is UB if <code>(offset % 8 != 0) && (offset / 8) + 2 >= [`self.bytes().len()`][Self::bytes]</code>
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if <code>(offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  /// * Panics in debug mode if <code>(offset % 8 != 0) && (offset / 8) + 2 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub unsafe fn write_u16_be_full_at_unchecked(&mut self, offset: usize, v: u16) -> &mut Self {
    let first_idx = offset / 8;
    let shift = offset % 8;

    let bytes = self.bytes_mut();

    debug_assert!(
      first_idx + 1 < bytes.len(),
      "BitBuf::write_u16_be_full_at_unchecked: index out of bounds! len is {}, byte_idx is {}",
      bytes.len(),
      first_idx + 1,
    );

    if shift == 0 {
      unsafe { self.write_u16_be_aligned_full_at_unchecked(first_idx, v) };
    } else {
      let second_idx = first_idx + 1;
      let third_idx = first_idx + 2;

      debug_assert!(
        third_idx < bytes.len(),
        "BitBuf::write_u16_be_full_at_unchecked: lookahead index out of bounds! len is {}, lookahead byte_idx is {}",
        bytes.len(),
        third_idx,
      );

      let [high, low] = v.to_be_bytes();

      // first byte, higher part of value
      unsafe {
        // mask to keep high bits of the current byte
        let mask = !0 << (8 - shift);
        let b = bytes.get_unchecked(first_idx) & mask;
        // high bits of the value, shifted down
        let new = high >> shift;
        *bytes.get_unchecked_mut(first_idx) = b | new;
      }

      // second byte, middle part of value
      unsafe {
        // no need to grab byte since all 8 bits will be overwritten
        // high middle bits of the value, shifted up
        let b_high = high << (8 - shift);
        // low middle bits of the value, shifted down
        let b_low = low >> shift;
        *bytes.get_unchecked_mut(second_idx) = b_high | b_low;
      }

      // third byte, lower part of value
      unsafe {
        // mask to keep low bits of the current byte
        let mask = !0 >> shift;
        let b = bytes.get_unchecked(third_idx) & mask;
        // low bits of the value, shifted up
        let new = low << (8 - shift);
        *bytes.get_unchecked_mut(third_idx) = b | new;
      }
    }

    self
  }

  /// Write a [`u16`] in [BE-bit, BE-byte] order without performing bound checks, advancing the internal cursor
  ///
  /// # Safety
  ///
  /// * This is UB if <code>(self.pos() / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  /// * This is UB if <code>(self.pos() % 8 != 0) && (self.pos() / 8) + 2 >= [`self.bytes().len()`][Self::bytes]</code>
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if <code>(self.pos() / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  /// * Panics in debug mode if <code>(self.pos() % 8 != 0) && (self.pos() / 8) + 2 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub unsafe fn write_u16_be_full_unchecked(&mut self, v: u16) -> &mut Self {
    unsafe { self.write_u16_be_full_at_unchecked(self.pos(), v) };
    self.advance_bytes(2);
    self
  }

  /// Write a [BE-bit, BE-byte] order [`u16`] at `offset` while performing bound checks
  ///
  /// # Errors
  ///
  /// * Returns [`Error::OutOfBounds`]
  ///   * if <code>(offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  ///   * if <code>(offset % 8 != 0) && (offset / 8) + 2 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub fn try_write_u16_be_full_at(&mut self, offset: usize, v: u16) -> Result<&mut Self> {
    let first_idx = offset / 8;
    let shift = offset % 8;

    let len = self.bytes().len();

    if first_idx + 1 >= len {
      return Err(Error::OutOfBounds);
    }

    if shift != 0 && first_idx + 2 >= len {
      return Err(Error::OutOfBounds);
    }

    Ok(unsafe { self.write_u16_be_full_at_unchecked(offset, v) })
  }

  /// Write a [BE-bit, BE-byte] order [`u16`] while performing bound checks, advancing the internal cursor
  ///
  /// # Errors
  ///
  /// * Returns [`Error::OutOfBounds`]
  ///   * if <code>(self.pos() / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  ///   * if <code>(self.pos() % 8 != 0) && (self.pos() / 8) + 2 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub fn try_write_u16_be_full(&mut self, v: u16) -> Result<&mut Self> {
    self.try_write_u16_be_full_at(self.pos(), v)?;
    self.advance_bytes(2);
    Ok(self)
  }

  /// Write a [BE-bit, BE-byte] order [`u16`] at `offset`, panicking on out of bounds
  ///
  /// # Panics
  ///
  /// * Panics if <code>(offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  /// * Panics if <code>(offset % 8 != 0) && (offset / 8) + 2 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub fn write_u16_be_full_at(&mut self, offset: usize, v: u16) -> &mut Self {
    self
      .try_write_u16_be_full_at(offset, v)
      .expect("BitBuf::write_u16_be_full_at out of bounds")
  }

  /// Write a [BE-bit, BE-byte] order [`u16`], panicking on out of bounds, advancing the internal cursor
  ///
  /// # Panics
  ///
  /// * Panics if <code>(self.pos() / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  /// * Panics if <code>(self.pos() % 8 != 0) && (self.pos() / 8) + 2 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub fn write_u16_be_full(&mut self, v: u16) -> &mut Self {
    self
      .try_write_u16_be_full(v)
      .expect("BitBuf::write_u16_be_full out of bounds")
  }
}