third_party/rust/scroll/src/leb128.rs

   1 use crate::ctx::TryFromCtx;
   2 use crate::error;
   3 use crate::Pread;
   4 use core::convert::{AsRef, From};
   5 use core::result;
   6 use core::u8;
   7
   8 #[derive(Debug, PartialEq, Copy, Clone)]
   9 /// An unsigned leb128 integer
  10 pub struct Uleb128 {
  11     value: u64,
  12     count: usize,
  13 }
  14
  15 impl Uleb128 {
  16     #[inline]
  17     /// Return how many bytes this Uleb128 takes up in memory
  18     pub fn size(&self) -> usize {
  19         self.count
  20     }
  21     #[inline]
  22     /// Read a variable length u64 from `bytes` at `offset`
  23     pub fn read(bytes: &[u8], offset: &mut usize) -> error::Result<u64> {
  24         let tmp = bytes.pread::<Uleb128>(*offset)?;
  25         *offset += tmp.size();
  26         Ok(tmp.into())
  27     }
  28 }
  29
  30 impl AsRef<u64> for Uleb128 {
  31     fn as_ref(&self) -> &u64 {
  32         &self.value
  33     }
  34 }
  35
  36 impl From<Uleb128> for u64 {
  37     #[inline]
  38     fn from(uleb128: Uleb128) -> u64 {
  39         uleb128.value
  40     }
  41 }
  42
  43 #[derive(Debug, PartialEq, Copy, Clone)]
  44 /// An signed leb128 integer
  45 pub struct Sleb128 {
  46     value: i64,
  47     count: usize,
  48 }
  49
  50 impl Sleb128 {
  51     #[inline]
  52     /// Return how many bytes this Sleb128 takes up in memory
  53     pub fn size(&self) -> usize {
  54         self.count
  55     }
  56     #[inline]
  57     /// Read a variable length i64 from `bytes` at `offset`
  58     pub fn read(bytes: &[u8], offset: &mut usize) -> error::Result<i64> {
  59         let tmp = bytes.pread::<Sleb128>(*offset)?;
  60         *offset += tmp.size();
  61         Ok(tmp.into())
  62     }
  63 }
  64
  65 impl AsRef<i64> for Sleb128 {
  66     fn as_ref(&self) -> &i64 {
  67         &self.value
  68     }
  69 }
  70
  71 impl From<Sleb128> for i64 {
  72     #[inline]
  73     fn from(sleb128: Sleb128) -> i64 {
  74         sleb128.value
  75     }
  76 }
  77
  78 // Below implementation heavily adapted from: https://github.com/fitzgen/leb128
  79 const CONTINUATION_BIT: u8 = 1 << 7;
  80 const SIGN_BIT: u8 = 1 << 6;
  81
  82 #[inline]
  83 fn mask_continuation(byte: u8) -> u8 {
  84     byte & !CONTINUATION_BIT
  85 }
  86
  87 // #[inline]
  88 // fn mask_continuation_u64(val: u64) -> u8 {
  89 //     let byte = val & (u8::MAX as u64);
  90 //     mask_continuation(byte as u8)
  91 // }
  92
  93 impl<'a> TryFromCtx<'a> for Uleb128 {
  94     type Error = error::Error;
  95     #[inline]
  96     fn try_from_ctx(src: &'a [u8], _ctx: ()) -> result::Result<(Self, usize), Self::Error> {
  97         let mut result = 0;
  98         let mut shift = 0;
  99         let mut count = 0;
 100         loop {
 101             let byte: u8 = src.pread(count)?;
 102
 103             if shift == 63 && byte != 0x00 && byte != 0x01 {
 104                 return Err(error::Error::BadInput {
 105                     size: src.len(),
 106                     msg: "failed to parse",
 107                 });
 108             }
 109
 110             let low_bits = u64::from(mask_continuation(byte));
 111             result |= low_bits << shift;
 112
 113             count += 1;
 114             shift += 7;
 115
 116             if byte & CONTINUATION_BIT == 0 {
 117                 return Ok((
 118                     Uleb128 {
 119                         value: result,
 120                         count,
 121                     },
 122                     count,
 123                 ));
 124             }
 125         }
 126     }
 127 }
 128
 129 impl<'a> TryFromCtx<'a> for Sleb128 {
 130     type Error = error::Error;
 131     #[inline]
 132     fn try_from_ctx(src: &'a [u8], _ctx: ()) -> result::Result<(Self, usize), Self::Error> {
 133         let o = 0;
 134         let offset = &mut 0;
 135         let mut result = 0;
 136         let mut shift = 0;
 137         let size = 64;
 138         let mut byte: u8;
 139         loop {
 140             byte = src.gread(offset)?;
 141
 142             if shift == 63 && byte != 0x00 && byte != 0x7f {
 143                 return Err(error::Error::BadInput {
 144                     size: src.len(),
 145                     msg: "failed to parse",
 146                 });
 147             }
 148
 149             let low_bits = i64::from(mask_continuation(byte));
 150             result |= low_bits << shift;
 151             shift += 7;
 152
 153             if byte & CONTINUATION_BIT == 0 {
 154                 break;
 155             }
 156         }
 157
 158         if shift < size && (SIGN_BIT & byte) == SIGN_BIT {
 159             // Sign extend the result.
 160             result |= !0 << shift;
 161         }
 162         let count = *offset - o;
 163         Ok((
 164             Sleb128 {
 165                 value: result,
 166                 count,
 167             },
 168             count,
 169         ))
 170     }
 171 }
 172
 173 #[cfg(test)]
 174 mod tests {
 175     use super::super::LE;
 176     use super::{Sleb128, Uleb128};
 177
 178     const CONTINUATION_BIT: u8 = 1 << 7;
 179     //const SIGN_BIT: u8 = 1 << 6;
 180
 181     #[test]
 182     fn uleb_size() {
 183         use super::super::Pread;
 184         let buf = [2u8 | CONTINUATION_BIT, 1];
 185         let bytes = &buf[..];
 186         let num = bytes.pread::<Uleb128>(0).unwrap();
 187         println!("num: {:?}", &num);
 188         assert_eq!(130u64, num.into());
 189         assert_eq!(num.size(), 2);
 190
 191         let buf = [0x00, 0x01];
 192         let bytes = &buf[..];
 193         let num = bytes.pread::<Uleb128>(0).unwrap();
 194         println!("num: {:?}", &num);
 195         assert_eq!(0u64, num.into());
 196         assert_eq!(num.size(), 1);
 197
 198         let buf = [0x21];
 199         let bytes = &buf[..];
 200         let num = bytes.pread::<Uleb128>(0).unwrap();
 201         println!("num: {:?}", &num);
 202         assert_eq!(0x21u64, num.into());
 203         assert_eq!(num.size(), 1);
 204     }
 205
 206     #[test]
 207     fn uleb128() {
 208         use super::super::Pread;
 209         let buf = [2u8 | CONTINUATION_BIT, 1];
 210         let bytes = &buf[..];
 211         let num = bytes.pread::<Uleb128>(0).expect("Should read Uleb128");
 212         assert_eq!(130u64, num.into());
 213         assert_eq!(
 214             386,
 215             bytes.pread_with::<u16>(0, LE).expect("Should read number")
 216         );
 217     }
 218
 219     #[test]
 220     fn uleb128_overflow() {
 221         use super::super::Pread;
 222         let buf = [
 223             2u8 | CONTINUATION_BIT,
 224             2 | CONTINUATION_BIT,
 225             2 | CONTINUATION_BIT,
 226             2 | CONTINUATION_BIT,
 227             2 | CONTINUATION_BIT,
 228             2 | CONTINUATION_BIT,
 229             2 | CONTINUATION_BIT,
 230             2 | CONTINUATION_BIT,
 231             2 | CONTINUATION_BIT,
 232             2 | CONTINUATION_BIT,
 233             1,
 234         ];
 235         let bytes = &buf[..];
 236         assert!(bytes.pread::<Uleb128>(0).is_err());
 237     }
 238
 239     #[test]
 240     fn sleb128() {
 241         use super::super::Pread;
 242         let bytes = [0x7fu8 | CONTINUATION_BIT, 0x7e];
 243         let num: i64 = bytes
 244             .pread::<Sleb128>(0)
 245             .expect("Should read Sleb128")
 246             .into();
 247         assert_eq!(-129, num);
 248     }
 249 }