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}