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
106
107
108
109
110

use bitvec::mem::BitMemory;
use bitvec::prelude::*;
use std::fmt;

#[derive(Default)]
pub struct BitVecReader {
    bs: BitVec<Msb0, u8>,
    offset: usize,
}

impl BitVecReader {
    pub fn new(data: Vec<u8>) -> Self {
        Self {
            bs: BitVec::from_vec(data),
            offset: 0,
        }
    }

    #[inline(always)]
    pub fn get(&mut self) -> bool {
        let val = self.bs.get(self.offset).unwrap();
        self.offset += 1;

        *val
    }

    #[inline(always)]
    pub fn get_n<T: BitMemory>(&mut self, n: usize) -> T {
        let val = self.bs[self.offset..self.offset + n].load_be::<T>();
        self.offset += n;

        val
    }

    // bitstring.py implementation: https://github.com/scott-griffiths/bitstring/blob/master/bitstring.py#L1706
    #[inline(always)]
    pub fn get_ue(&mut self) -> u64 {
        let oldpos = self.offset;
        let mut pos = self.offset;

        loop {
            match self.bs.get(pos) {
                Some(val) => {
                    if !val {
                        pos += 1;
                    } else {
                        break;
                    }
                }
                None => panic!("Out of bounds index: {}", pos),
            }
        }

        let leading_zeroes = pos - oldpos;
        let mut code_num = (1 << leading_zeroes) - 1;

        if leading_zeroes > 0 {
            if pos + leading_zeroes + 1 > self.bs.len() {
                panic!("Out of bounds attempt");
            }

            code_num += self.bs[pos + 1..pos + leading_zeroes + 1].load_be::<u64>();
            pos += leading_zeroes + 1;
        } else {
            assert_eq!(code_num, 0);
            pos += 1;
        }

        self.offset = pos;

        code_num
    }

    // bitstring.py implementation: https://github.com/scott-griffiths/bitstring/blob/master/bitstring.py#L1767
    #[inline(always)]
    pub fn get_se(&mut self) -> i64 {
        let code_num = self.get_ue();

        let m = ((code_num + 1) as f64 / 2.0).floor() as u64;

        if code_num % 2 == 0 {
            -(m as i64)
        } else {
            m as i64
        }
    }

    pub fn is_aligned(&self) -> bool {
        self.offset % 8 == 0
    }

    pub fn available(&self) -> usize {
        self.bs.len() - self.offset
    }

    pub fn skip_n(&mut self, n: usize) {
        self.offset += n;
    }
}

impl fmt::Debug for BitVecReader {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "BitVecReader: {{offset: {}, len: {}}}",
            self.offset,
            self.bs.len()
        )
    }
}