use core::{
convert::Infallible,
marker::PhantomData,
num::NonZeroUsize,
};
use qubit_codec::Codec;
use crate::{
Leb128DecodeError,
Leb128DecodePolicy,
NonStrict,
};
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct Leb128Codec<T, P = NonStrict> {
marker: PhantomData<fn() -> (T, P)>,
}
macro_rules! impl_unsigned_leb128_codec {
($ty:ty) => {
impl<P> Leb128Codec<$ty, P>
where
P: Leb128DecodePolicy,
{
pub const MIN_UNITS_PER_VALUE: usize = 1;
pub const MAX_UNITS_PER_VALUE: usize =
(<$ty>::BITS as usize).div_ceil(7);
#[inline(always)]
pub unsafe fn decode_unchecked(
input: &[u8],
index: usize,
) -> Result<($ty, core::num::NonZeroUsize), Leb128DecodeError> {
debug_assert!(
input.len().saturating_sub(index)
>= Self::MIN_UNITS_PER_VALUE
);
let (value, consumed) = unsafe {
read_uleb_unchecked::<P>(
input,
index,
<$ty>::BITS,
Self::MAX_UNITS_PER_VALUE,
)?
};
Ok((value as $ty, consumed))
}
#[inline(always)]
pub unsafe fn encode_unchecked(
value: $ty,
output: &mut [u8],
index: usize,
) -> usize {
unsafe { write_uleb_unchecked(output, index, value as u128) }
}
}
unsafe impl<P> Codec for Leb128Codec<$ty, P>
where
P: Leb128DecodePolicy,
{
type Value = $ty;
type Unit = u8;
type DecodeError = Leb128DecodeError;
type EncodeError = Infallible;
#[inline(always)]
fn min_units_per_value(&self) -> core::num::NonZeroUsize {
core::num::NonZeroUsize::MIN
}
#[inline(always)]
fn max_units_per_value(&self) -> core::num::NonZeroUsize {
unsafe {
core::num::NonZeroUsize::new_unchecked(
Self::MAX_UNITS_PER_VALUE,
)
}
}
#[inline(always)]
unsafe fn decode_unchecked(
&self,
input: &[u8],
index: usize,
) -> Result<($ty, core::num::NonZeroUsize), Self::DecodeError> {
debug_assert!(
input.len().saturating_sub(index)
>= Self::MIN_UNITS_PER_VALUE
);
unsafe { Self::decode_unchecked(input, index) }
}
#[inline(always)]
unsafe fn encode_unchecked(
&self,
value: &$ty,
output: &mut [u8],
index: usize,
) -> Result<usize, Self::EncodeError> {
debug_assert!(
output.len().saturating_sub(index)
>= Self::MAX_UNITS_PER_VALUE
);
Ok(unsafe { Self::encode_unchecked(*value, output, index) })
}
}
};
}
macro_rules! impl_signed_leb128_codec {
($ty:ty) => {
impl<P> Leb128Codec<$ty, P>
where
P: Leb128DecodePolicy,
{
pub const MIN_UNITS_PER_VALUE: usize = 1;
pub const MAX_UNITS_PER_VALUE: usize =
(<$ty>::BITS as usize).div_ceil(7);
#[inline(always)]
pub unsafe fn decode_unchecked(
input: &[u8],
index: usize,
) -> Result<($ty, core::num::NonZeroUsize), Leb128DecodeError> {
debug_assert!(
input.len().saturating_sub(index)
>= Self::MIN_UNITS_PER_VALUE
);
let (value, consumed) = unsafe {
read_sleb_unchecked::<P>(
input,
index,
<$ty>::BITS,
Self::MAX_UNITS_PER_VALUE,
)?
};
Ok((value as $ty, consumed))
}
#[inline(always)]
pub unsafe fn encode_unchecked(
value: $ty,
output: &mut [u8],
index: usize,
) -> usize {
unsafe { write_sleb_unchecked(output, index, value as i128) }
}
}
unsafe impl<P> Codec for Leb128Codec<$ty, P>
where
P: Leb128DecodePolicy,
{
type Value = $ty;
type Unit = u8;
type DecodeError = Leb128DecodeError;
type EncodeError = Infallible;
#[inline(always)]
fn min_units_per_value(&self) -> core::num::NonZeroUsize {
core::num::NonZeroUsize::MIN
}
#[inline(always)]
fn max_units_per_value(&self) -> core::num::NonZeroUsize {
unsafe {
core::num::NonZeroUsize::new_unchecked(
Self::MAX_UNITS_PER_VALUE,
)
}
}
#[inline(always)]
unsafe fn decode_unchecked(
&self,
input: &[u8],
index: usize,
) -> Result<($ty, core::num::NonZeroUsize), Self::DecodeError> {
debug_assert!(
input.len().saturating_sub(index)
>= Self::MIN_UNITS_PER_VALUE
);
unsafe { Self::decode_unchecked(input, index) }
}
#[inline(always)]
unsafe fn encode_unchecked(
&self,
value: &$ty,
output: &mut [u8],
index: usize,
) -> Result<usize, Self::EncodeError> {
debug_assert!(
output.len().saturating_sub(index)
>= Self::MAX_UNITS_PER_VALUE
);
Ok(unsafe { Self::encode_unchecked(*value, output, index) })
}
}
};
}
impl_unsigned_leb128_codec!(u8);
impl_unsigned_leb128_codec!(u16);
impl_unsigned_leb128_codec!(u32);
impl_unsigned_leb128_codec!(u64);
impl_unsigned_leb128_codec!(u128);
impl_unsigned_leb128_codec!(usize);
impl_signed_leb128_codec!(i8);
impl_signed_leb128_codec!(i16);
impl_signed_leb128_codec!(i32);
impl_signed_leb128_codec!(i64);
impl_signed_leb128_codec!(i128);
impl_signed_leb128_codec!(isize);
#[inline]
unsafe fn read_uleb_unchecked<P>(
input: &[u8],
index: usize,
bits: u32,
max_bytes: usize,
) -> Result<(u128, core::num::NonZeroUsize), Leb128DecodeError>
where
P: Leb128DecodePolicy,
{
debug_assert!(input.len().saturating_sub(index) >= 1);
let available = input.len().saturating_sub(index).min(max_bytes);
match unsafe {
read_uleb_prefix_unchecked::<P>(
input, index, bits, max_bytes, available,
)
} {
Ok(Some((value, consumed))) => {
debug_assert!(consumed > 0);
let consumed =
unsafe { core::num::NonZeroUsize::new_unchecked(consumed) };
Ok((value, consumed))
}
Ok(None) => {
debug_assert!(
available < usize::MAX,
"available byte count overflowed"
);
let required =
unsafe { NonZeroUsize::new_unchecked(available + 1) };
Err(Leb128DecodeError::incomplete(index, required, available))
}
Err(error) => Err(error),
}
}
#[inline]
unsafe fn read_uleb_prefix_unchecked<P>(
input: &[u8],
index: usize,
bits: u32,
max_bytes: usize,
available: usize,
) -> Result<Option<(u128, usize)>, Leb128DecodeError>
where
P: Leb128DecodePolicy,
{
debug_assert!(
available <= max_bytes,
"available bytes exceed LEB128 maximum width"
);
let mut value = 0u128;
let mut shift = 0u32;
let base = unsafe { input.as_ptr().add(index) };
for offset in 0..available {
let byte = unsafe { *base.add(offset) };
let payload = u128::from(byte & 0x7F);
value |= payload << shift;
if byte & 0x80 == 0 {
if offset == max_bytes - 1
&& !unsigned_final_payload_fits(byte, bits, offset)
{
return Err(malformed_decode_error(
index,
index + offset,
offset + 1,
));
}
let consumed = offset + 1;
if P::STRICT && !has_canonical_uleb_len(value, consumed) {
return Err(noncanonical_decode_error(index, consumed));
}
return Ok(Some((value, consumed)));
}
shift += 7;
}
if available < max_bytes {
return Ok(None);
}
Err(malformed_decode_error(
index,
index + max_bytes - 1,
max_bytes,
))
}
#[inline]
unsafe fn read_sleb_unchecked<P>(
input: &[u8],
index: usize,
bits: u32,
max_bytes: usize,
) -> Result<(i128, core::num::NonZeroUsize), Leb128DecodeError>
where
P: Leb128DecodePolicy,
{
debug_assert!(input.len().saturating_sub(index) >= 1);
let available = input.len().saturating_sub(index).min(max_bytes);
match unsafe {
read_sleb_prefix_unchecked::<P>(
input, index, bits, max_bytes, available,
)
} {
Ok(Some((value, consumed))) => {
debug_assert!(consumed > 0);
let consumed =
unsafe { core::num::NonZeroUsize::new_unchecked(consumed) };
Ok((value, consumed))
}
Ok(None) => {
debug_assert!(
available < usize::MAX,
"available byte count overflowed"
);
let required =
unsafe { NonZeroUsize::new_unchecked(available + 1) };
Err(Leb128DecodeError::incomplete(index, required, available))
}
Err(error) => Err(error),
}
}
#[inline]
unsafe fn read_sleb_prefix_unchecked<P>(
input: &[u8],
index: usize,
bits: u32,
max_bytes: usize,
available: usize,
) -> Result<Option<(i128, usize)>, Leb128DecodeError>
where
P: Leb128DecodePolicy,
{
debug_assert!(
available <= max_bytes,
"available bytes exceed LEB128 maximum width"
);
let mut value = 0i128;
let mut shift = 0u32;
let base = unsafe { input.as_ptr().add(index) };
for offset in 0..available {
let byte = unsafe { *base.add(offset) };
let payload = i128::from(byte & 0x7F);
value |= payload << shift;
if byte & 0x80 == 0 {
if offset == max_bytes - 1
&& !signed_final_payload_fits(byte, bits, offset)
{
return Err(malformed_decode_error(
index,
index + offset,
offset + 1,
));
}
if byte & 0x40 != 0 && shift + 7 < i128::BITS {
value |= (!0i128) << (shift + 7);
}
let consumed = offset + 1;
if P::STRICT && !has_canonical_sleb_len(value, consumed) {
return Err(noncanonical_decode_error(index, consumed));
}
return Ok(Some((value, consumed)));
}
shift += 7;
}
if available < max_bytes {
return Ok(None);
}
Err(malformed_decode_error(
index,
index + max_bytes - 1,
max_bytes,
))
}
#[cold]
fn malformed_decode_error(
start_index: usize,
error_index: usize,
consumed: usize,
) -> Leb128DecodeError {
debug_assert!(consumed > 0, "malformed LEB128 errors must consume bytes");
let consumed = unsafe { NonZeroUsize::new_unchecked(consumed) };
Leb128DecodeError::malformed(start_index, error_index, consumed)
}
#[cold]
fn noncanonical_decode_error(
index: usize,
consumed: usize,
) -> Leb128DecodeError {
debug_assert!(
consumed > 0,
"non-canonical LEB128 errors must consume bytes"
);
let consumed = unsafe { NonZeroUsize::new_unchecked(consumed) };
Leb128DecodeError::noncanonical(index, consumed)
}
#[must_use]
#[inline(always)]
fn unsigned_final_payload_fits(byte: u8, bits: u32, offset: usize) -> bool {
let used_bits = bits - offset as u32 * 7;
byte >> used_bits == 0
}
#[must_use]
#[inline(always)]
fn signed_final_payload_fits(byte: u8, bits: u32, offset: usize) -> bool {
let used_bits = bits - offset as u32 * 7;
let payload = byte & 0x7F;
let used_mask = ((1u16 << used_bits) - 1) as u8;
let unused_mask = 0x7F & !used_mask;
let sign_bit = 1u8 << (used_bits - 1);
if payload & sign_bit == 0 {
payload & unused_mask == 0
} else {
payload & unused_mask == unused_mask
}
}
#[must_use]
#[inline(always)]
fn has_canonical_uleb_len(value: u128, actual_len: usize) -> bool {
canonical_uleb_len(value) == actual_len
}
#[must_use]
#[inline(always)]
fn has_canonical_sleb_len(value: i128, actual_len: usize) -> bool {
canonical_sleb_len(value) == actual_len
}
#[must_use]
#[inline]
fn canonical_uleb_len(mut value: u128) -> usize {
let mut len = 1;
while value >= 0x80 {
value >>= 7;
len += 1;
}
len
}
#[must_use]
#[inline]
fn canonical_sleb_len(mut value: i128) -> usize {
let mut len = 0;
loop {
let byte = (value & 0x7F) as u8;
let sign_bit_set = byte & 0x40 != 0;
value >>= 7;
len += 1;
if (value == 0 && !sign_bit_set) || (value == -1 && sign_bit_set) {
return len;
}
}
}
unsafe fn write_uleb_unchecked(
output: &mut [u8],
index: usize,
mut value: u128,
) -> usize {
let mut offset = 0;
loop {
let mut byte = (value & 0x7F) as u8;
value >>= 7;
if value != 0 {
byte |= 0x80;
}
unsafe {
*output.as_mut_ptr().add(index + offset) = byte;
}
offset += 1;
if value == 0 {
return offset;
}
}
}
unsafe fn write_sleb_unchecked(
output: &mut [u8],
index: usize,
mut value: i128,
) -> usize {
let mut offset = 0;
loop {
let mut byte = (value & 0x7F) as u8;
let sign_bit_set = byte & 0x40 != 0;
value >>= 7;
let done =
(value == 0 && !sign_bit_set) || (value == -1 && sign_bit_set);
if !done {
byte |= 0x80;
}
unsafe {
*output.as_mut_ptr().add(index + offset) = byte;
}
offset += 1;
if done {
return offset;
}
}
}