1use crate::prelude::BitBufMut;
2
3use super::bit_buf::BitBuf;
4
5pub 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 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 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}