1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105

#![no_std]
use core::{
    mem::size_of,
    num::NonZeroUsize,
    ops::{BitOrAssign, RangeFrom},
};
use nom::{
    error::{make_error, ContextError, ErrorKind, ParseError},
    IResult, InputIter, InputLength, Needed, Slice,
};
use num_traits::{PrimInt, Signed, WrappingNeg};

// TODO: when stable, use const NonZeroUsize::new(1).unwrap()
const NEED_ONE: nom::Needed = Needed::Size(unsafe { NonZeroUsize::new_unchecked(1) });

/// Maximum LEB128-encoded size of an integer type
const fn leb128_size<T>() -> usize {
    let bits = size_of::<T>() * 8;
    (bits + 6) / 7 // equivalent to ceil(bits/7) w/o floats
}

macro_rules! impl_generic_leb128 {
    ($fn_name:ident, $int_ty:ident, $post:tt, $int_name:expr) => {
        #[doc="Recognizes an LEB128-encoded number that fits in a `"]
        #[doc=$int_name]
        #[doc="`."]
        #[inline]
        pub fn $fn_name<I, E>(input: I) -> IResult<I, $int_ty, E>
        where
            I: Clone + Slice<RangeFrom<usize>> + InputIter<Item = u8> + InputLength,
            E: ParseError<I> + ContextError<I>,
        {
            let mut res = 0;
            let mut shift = 0;

            for (pos, byte) in input.iter_indices() {
                if (byte & 0x80) == 0 {
                    res |= (byte as $int_ty) << shift;
                    $post(&mut res, shift, byte);
                    return Ok((input.slice(pos + 1..), res));
                } else if pos == leb128_size::<$int_ty>() - 1 {
                    return Err(nom::Err::Error(E::add_context(
                        input.clone(),
                        concat!("LEB128 integer is too big to fit in ", $int_name),
                        make_error(input, ErrorKind::TooLarge),
                    )));
                } else {
                    res |= ((byte & 0x7F) as $int_ty) << shift;
                }
                shift += 7;
            }

            Err(nom::Err::Incomplete(NEED_ONE))
        }
    };
    ($fn_name:ident, $int_ty:ident, $post:tt) => {
        impl_generic_leb128!($fn_name, $int_ty, $post, stringify!($int_ty));
    };
}

macro_rules! impl_unsigned_leb128 {
    ($fn_name:ident, $int_ty:ident) => {
        impl_generic_leb128!($fn_name, $int_ty, (|_, _, _| {}));
    };
}

impl_unsigned_leb128!(leb128_u8, u8);
impl_unsigned_leb128!(leb128_u16, u16);
impl_unsigned_leb128!(leb128_u32, u32);
impl_unsigned_leb128!(leb128_u64, u64);
impl_unsigned_leb128!(leb128_u128, u128);
impl_unsigned_leb128!(leb128_usize, usize);

#[inline]
fn sign_extend<T: BitOrAssign + PrimInt + Signed + WrappingNeg>(
    res: &mut T,
    shift: usize,
    byte: u8,
) {
    // leb128_generic skips the last shift update for efficiency on unsigned ints

    if (shift < size_of::<T>() * 8 - 7) && ((byte & 0x40) != 0) {
        // sign extend
        *res |= (T::one() << (shift + 7)).wrapping_neg()
    }
}

macro_rules! impl_signed_leb128 {
    ($fn_name:ident, $int_ty:ident) => {
        impl_generic_leb128!($fn_name, $int_ty, sign_extend);
    };
}

impl_signed_leb128!(leb128_i8, i8);
impl_signed_leb128!(leb128_i16, i16);
impl_signed_leb128!(leb128_i32, i32);
impl_signed_leb128!(leb128_i64, i64);
impl_signed_leb128!(leb128_i128, i128);
impl_signed_leb128!(leb128_isize, isize);

#[cfg(any(test, fuzzing))]
mod tests;

#[cfg(fuzzing)]
pub use tests::{write_signed_leb128, write_unsigned_leb128};