bits_io/buf/
chain.rs

1use crate::prelude::BitBufMut;
2
3use super::bit_buf::BitBuf;
4
5/// A `Chain` sequences to `BitBuf`s.
6///
7/// `Chain` is an adaptor that links to underlying buffers and provides a continuous view across
8/// both.
9pub struct Chain<T, U> {
10    a: T,
11    b: U,
12}
13
14impl<T, U> Chain<T, U> {
15    pub fn new(a: T, b: U) -> Self {
16        Self { a, b }
17    }
18
19    pub fn first_ref(&self) -> &T {
20        &self.a
21    }
22
23    pub fn last_ref(&self) -> &U {
24        &self.b
25    }
26
27    pub fn first_mut(&mut self) -> &mut T {
28        &mut self.a
29    }
30
31    pub fn last_mut(&mut self) -> &mut U {
32        &mut self.b
33    }
34
35    pub fn into_inner(self) -> (T, U) {
36        (self.a, self.b)
37    }
38}
39
40impl<T, U> BitBuf for Chain<T, U>
41where
42    T: BitBuf,
43    U: BitBuf,
44{
45    fn advance_bits(&mut self, mut count: usize) {
46        let a_rem = self.a.remaining_bits();
47
48        if a_rem != 0 {
49            if a_rem >= count {
50                self.a.advance_bits(count);
51                return;
52            }
53
54            // Consume what is left of a
55            self.a.advance_bits(a_rem);
56
57            count -= a_rem;
58        }
59        self.b.advance_bits(count);
60    }
61
62    fn remaining_bits(&self) -> usize {
63        self.a
64            .remaining_bits()
65            .saturating_add(self.b.remaining_bits())
66    }
67
68    fn chunk_bits(&self) -> &crate::prelude::BitSlice {
69        if self.a.has_remaining_bits() {
70            self.a.chunk_bits()
71        } else {
72            self.b.chunk_bits()
73        }
74    }
75
76    fn chunk_bytes(&self) -> &[u8] {
77        if self.a.has_remaining_bytes() {
78            self.a.chunk_bytes()
79        } else {
80            self.b.chunk_bytes()
81        }
82    }
83
84    fn byte_aligned(&self) -> bool {
85        self.a.byte_aligned() && self.b.byte_aligned()
86    }
87}
88
89impl<T, U> BitBufMut for Chain<T, U>
90where
91    T: BitBufMut,
92    U: BitBufMut,
93{
94    fn advance_mut_bits(&mut self, mut count: usize) {
95        let a_rem = self.a.remaining_mut_bits();
96
97        if a_rem != 0 {
98            if a_rem >= count {
99                self.a.advance_mut_bits(count);
100                return;
101            }
102
103            // Consume what's left in a
104            self.a.advance_mut_bits(a_rem);
105
106            count -= a_rem;
107        }
108        self.b.advance_mut_bits(count);
109    }
110
111    fn chunk_mut_bits(&mut self) -> &mut crate::prelude::BitSlice {
112        if self.a.has_remaining_mut_bits() {
113            self.a.chunk_mut_bits()
114        } else {
115            self.b.chunk_mut_bits()
116        }
117    }
118
119    fn chunk_mut_bytes(&mut self) -> &mut bytes::buf::UninitSlice {
120        if self.a.has_reminaing_mut_bytes() {
121            self.a.chunk_mut_bytes()
122        } else {
123            self.b.chunk_mut_bytes()
124        }
125    }
126
127    fn remaining_mut_bits(&self) -> usize {
128        self.a
129            .remaining_mut_bits()
130            .saturating_add(self.b.remaining_mut_bits())
131    }
132
133    fn byte_aligned_mut(&self) -> bool {
134        self.a.byte_aligned_mut() && self.b.byte_aligned_mut()
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use crate::prelude::*;
141
142    use super::*;
143
144    #[test]
145    fn test_bit_buf_chain() {
146        let left = Bits::from(bits![1, 1, 1, 1, 0, 0, 0, 0]);
147        let right = Bits::from(bits![1, 0, 1, 0, 1, 0, 1, 0]);
148
149        let mut chain = left.chain(right);
150
151        let mut data = [0u8; 2];
152        chain.copy_to_slice_bytes(&mut data);
153
154        assert_eq!(data, [0b11110000, 0b10101010]);
155    }
156
157    #[test]
158    fn test_bit_buf_mut_chain() {
159        let mut left = [0u8; 1];
160        let mut right = [0u8; 1];
161
162        let mut chain = (&mut left[..]).chain_mut(&mut right[..]);
163        chain.put_u16::<NetworkOrder>(42).unwrap();
164        assert_eq!(&left[..], [0]);
165        assert_eq!(&right[..], [42]);
166    }
167}