bits_io/buf/
limit.rs

1use crate::prelude::BitBufMut;
2
3pub struct Limit<T> {
4    inner: T,
5    limit: usize,
6}
7
8impl<T> Limit<T> {
9    pub fn new(inner: T, limit: usize) -> Self {
10        Self { inner, limit }
11    }
12
13    pub fn into_inner(self) -> T {
14        self.inner
15    }
16
17    pub fn get_ref(&self) -> &T {
18        &self.inner
19    }
20
21    pub fn get_mut(&mut self) -> &mut T {
22        &mut self.inner
23    }
24
25    pub fn limit(&self) -> usize {
26        self.limit
27    }
28
29    pub fn set_limit(&mut self, limit: usize) {
30        self.limit = limit;
31    }
32}
33
34impl<T: BitBufMut> BitBufMut for Limit<T> {
35    fn advance_mut_bits(&mut self, count: usize) {
36        assert!(count <= self.limit);
37        self.inner.advance_mut_bits(count);
38        self.limit -= count;
39    }
40
41    fn chunk_mut_bits(&mut self) -> &mut crate::prelude::BitSlice {
42        let chunk = self.inner.chunk_mut_bits();
43        let end = std::cmp::min(chunk.len(), self.limit);
44        &mut chunk[..end]
45    }
46
47    fn chunk_mut_bytes(&mut self) -> &mut bytes::buf::UninitSlice {
48        assert!(self.byte_aligned_mut());
49        let chunk = self.inner.chunk_mut_bytes();
50        let byte_limit = self.limit / 8;
51        let end = std::cmp::min(chunk.len(), byte_limit);
52        &mut chunk[..end]
53    }
54
55    fn remaining_mut_bits(&self) -> usize {
56        std::cmp::min(self.inner.remaining_mut_bits(), self.limit)
57    }
58
59    fn byte_aligned_mut(&self) -> bool {
60        self.inner.byte_aligned_mut() && self.limit % 8 == 0
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    #[test]
69    fn test_limit() {
70        let data = &mut [0u8; 10];
71
72        let lim = data.limit_bytes(2);
73        assert_eq!(lim.remaining_mut_bytes(), 2);
74        let lim = data.limit_bits(16);
75        assert_eq!(lim.remaining_mut_bits(), 16);
76    }
77}