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

use crate::private::{
    checks::private::Sealed,
    PopBits,
    PushBits,
};

/// A bit buffer that allows to pop bits from it.
pub struct PopBuffer<T> {
    bytes: T,
}

impl<T> PopBuffer<T> {
    /// Creates a new pop buffer from the given bytes.
    #[inline]
    pub(super) fn from_bytes(bytes: T) -> Self {
        Self { bytes }
    }
}

impl Sealed for PopBuffer<u8> {}

impl PopBits for PopBuffer<u8> {
    #[inline]
    fn pop_bits(&mut self, amount: u32) -> u8 {
        let Self { bytes } = self;
        let orig_ones = bytes.count_ones();
        debug_assert!(1 <= amount && amount <= 8);
        let res = *bytes & ((0x01_u16.wrapping_shl(amount)).wrapping_sub(1) as u8);
        *bytes = bytes.checked_shr(amount).unwrap_or(0);
        debug_assert_eq!(res.count_ones() + bytes.count_ones(), orig_ones);
        res
    }
}

macro_rules! impl_pop_bits {
    ( $($type:ty),+ ) => {
        $(
            impl Sealed for PopBuffer<$type> {}

            impl PopBits for PopBuffer<$type> {
                #[inline]
                fn pop_bits(&mut self, amount: u32) -> u8 {
                    let Self { bytes } = self;
                    let orig_ones = bytes.count_ones();
                    debug_assert!(1 <= amount && amount <= 8);
                    let bitmask = 0xFF >> (8 - amount);
                    let res = (*bytes & bitmask) as u8;
                    *bytes = bytes.checked_shr(amount).unwrap_or(0);
                    debug_assert_eq!(res.count_ones() + bytes.count_ones(), orig_ones);
                    res
                }
            }
        )+
    };
}
impl_pop_bits!(u16, u32, u64, u128);

/// A bit buffer that allows to push bits onto it.
pub struct PushBuffer<T> {
    bytes: T,
}

impl<T> PushBuffer<T> {
    /// Returns the underlying bytes of the push buffer.
    #[inline]
    pub(super) fn into_bytes(self) -> T {
        self.bytes
    }
}

macro_rules! impl_push_bits {
    ( $($type:ty),+ ) => {
        $(
            impl Sealed for PushBuffer<$type> {}

            impl Default for PushBuffer<$type> {
                #[inline]
                fn default() -> Self {
                    Self { bytes: <$type as Default>::default() }
                }
            }

            impl PushBits for PushBuffer<$type> {
                #[inline]
                fn push_bits(&mut self, amount: u32, bits: u8) {
                    let Self { bytes } = self;
                    let orig_ones = bytes.count_ones();
                    debug_assert!(1 <= amount && amount <= 8);
                    let bitmask = 0xFF >> (8 - amount as u8);
                    *bytes = bytes.wrapping_shl(amount) | ((bits & bitmask) as $type);
                    debug_assert_eq!((bits & bitmask).count_ones() + orig_ones, bytes.count_ones());
                }
            }
        )+
    }
}
impl_push_bits!(u8, u16, u32, u64, u128);