1use std::cmp;
2
3fn get_bytes_u16(x: u16) -> [u8; 2] {
4 let mut buf = [0 as u8; 2];
5 buf[0] = (x >> 8) as u8;
6 buf[1] = (x & 0xFF) as u8;
7 buf
8}
9
10fn get_bytes_u32(x: u32) -> [u8; 4] {
11 let mut buf = [0 as u8; 4];
12 buf[0] = ((x >> 24) & 0xFF) as u8;
13 buf[1] = ((x >> 16) & 0xFF) as u8;
14 buf[2] = ((x >> 8) & 0xFF) as u8;
15 buf[3] = (x & 0xFF) as u8;
16 buf
17}
18
19fn get_bytes_u64(x: u64) -> [u8; 8] {
20 let mut buf = [0 as u8; 8];
21 buf[0] = ((x >> 56) & 0xFF) as u8;
22 buf[1] = ((x >> 48) & 0xFF) as u8;
23 buf[2] = ((x >> 40) & 0xFF) as u8;
24 buf[3] = ((x >> 32) & 0xFF) as u8;
25 buf[4] = ((x >> 24) & 0xFF) as u8;
26 buf[5] = ((x >> 16) & 0xFF) as u8;
27 buf[6] = ((x >> 8) & 0xFF) as u8;
28 buf[7] = (x & 0xFF) as u8;
29 buf
30}
31
32pub struct ByteBufMut<'a> {
33 pos: usize,
34 buf: &'a mut [u8],
35}
36
37impl<'a> ByteBufMut<'a> {
38
39 pub fn wrap(buf: &'a mut [u8]) -> ByteBufMut<'a> {
40 ByteBufMut { pos: 0, buf }
41 }
42
43 pub fn put_bytes(&mut self, bytes: &[u8]) -> usize {
44 let pos = self.pos;
45 for (i, &x) in bytes.iter().enumerate() {
46 self.pos = pos + i;
47 if self.pos >= self.buf.len() {
48 return i;
49 }
50 self.buf[self.pos] = x;
51 }
52 self.pos += 1;
53 bytes.len()
54 }
55
56 pub fn put_u8(&mut self, x: u8) -> usize {
57 self.put_bytes(&[x])
58 }
59
60 pub fn put_u16(&mut self, x: u16) -> usize {
61 self.put_bytes(&get_bytes_u16(x))
62 }
63
64 pub fn put_u32(&mut self, x: u32) -> usize {
65 self.put_bytes(&get_bytes_u32(x))
66 }
67
68 pub fn put_u64(&mut self, x: u64) -> usize {
69 self.put_bytes(&get_bytes_u64(x))
70 }
71
72 pub fn pos(&self) -> usize { self.pos }
73}
74
75pub struct ByteBuf<'a> {
76 pos: usize,
77 buf: &'a [u8],
78}
79
80impl<'a> ByteBuf<'a> {
81
82 pub fn wrap(buf: &'a [u8]) -> ByteBuf<'a> {
83 ByteBuf { pos: 0, buf }
84 }
85
86 pub fn get_bytes(&mut self, n: usize) -> &[u8] {
87 let at = self.pos;
88 let to = cmp::min(self.buf.len(), self.pos + n);
89 self.pos += n;
90 &self.buf[at..to]
91 }
92
93 pub fn get_u8(&mut self) -> Option<u8> {
94 let n = 1;
95 let bytes = self.get_bytes(n);
96 if bytes.len() == n {
97 let x = bytes[0];
98 Some(x)
99 } else {
100 None
101 }
102 }
103
104 pub fn get_u16(&mut self) -> Option<u16> {
105 let n = 2;
106 let bytes = self.get_bytes(n);
107 if bytes.len() == n {
108 let x = ((bytes[0] as u16) << 8) + bytes[1] as u16;
109 Some(x)
110 } else {
111 None
112 }
113 }
114
115 pub fn get_u32(&mut self) -> Option<u32> {
116 let n = 4;
117 let bytes = self.get_bytes(n);
118 if bytes.len() == n {
119 let x = ((bytes[0] as u32) << 24) +
120 ((bytes[1] as u32) << 16) +
121 ((bytes[2] as u32) << 8) +
122 bytes[3] as u32;
123 Some(x)
124 } else {
125 None
126 }
127 }
128
129 pub fn get_u64(&mut self) -> Option<u64> {
130 let n = 8;
131 let bytes = self.get_bytes(n);
132 if bytes.len() == n {
133 let x = ((bytes[0] as u64) << 56) +
134 ((bytes[1] as u64) << 48) +
135 ((bytes[2] as u64) << 40) +
136 ((bytes[3] as u64) << 32) +
137 ((bytes[4] as u64) << 24) +
138 ((bytes[5] as u64) << 16) +
139 ((bytes[6] as u64) << 8) +
140 bytes[7] as u64;
141 Some(x)
142 } else {
143 None
144 }
145 }
146
147 pub fn pos(&self) -> usize { self.pos }
148}
149
150#[cfg(test)]
151mod tests {
152 use super::*;
153
154 #[test]
155 fn test_write_u8() {
156 let x: u8 = 0xAB;
157 let mut buf = [0 as u8; 1];
158 let mut bb = ByteBufMut::wrap(&mut buf);
159 assert_eq!(bb.put_u8(x), 1);
160 assert_eq!(bb.pos, 1);
161 assert_eq!([x], buf);
162 }
163
164 #[test]
165 fn test_write_u16() {
166 let x: u16 = 0xABCD;
167 let mut buf = [0 as u8; 2];
168 let mut bb = ByteBufMut::wrap(&mut buf);
169 assert_eq!(bb.put_u16(x), 2);
170 assert_eq!(bb.pos, 2);
171 assert_eq!(buf, [0xAB, 0xCD]);
172 }
173
174 #[test]
175 fn test_write_u32() {
176 let x: u32 = 0xAABBCCDD;
177 let mut buf = [0 as u8; 4];
178 let mut bb = ByteBufMut::wrap(&mut buf);
179 assert_eq!(bb.put_u32(x), 4);
180 assert_eq!(bb.pos, 4);
181 assert_eq!(buf, [0xAA, 0xBB, 0xCC, 0xDD]);
182 }
183
184 #[test]
185 fn test_write_u64() {
186 let x: u64 = 0xAABBCCDDEEFFABCD;
187 let mut buf = [0 as u8; 8];
188 let mut bb = ByteBufMut::wrap(&mut buf);
189 assert_eq!(bb.put_u64(x), 8);
190 assert_eq!(bb.pos, 8);
191 assert_eq!(buf, [0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0xAB, 0xCD]);
192 }
193
194 #[test]
195 fn test_write_all() {
196 let a: u8 = 0x01;
197 let b: u16 = 0x0002;
198 let c: u32 = 0x00000003;
199 let d: u64 = 0x0000000000000004;
200 let mut buf = [0 as u8; 15];
201 {
202 let mut bb = ByteBufMut::wrap(&mut buf);
203 assert_eq!(bb.put_u8(a), 1);
204 assert_eq!(bb.pos, 1);
205 assert_eq!(bb.put_u16(b), 2);
206 assert_eq!(bb.pos, 3);
207 assert_eq!(bb.put_u32(c), 4);
208 assert_eq!(bb.pos, 7);
209 assert_eq!(bb.put_u64(d), 8);
210 assert_eq!(bb.pos, 15);
211 }
212 assert_eq!(buf, [0x01, 0x00, 0x02, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04]);
213 }
214
215 #[test]
216 fn test_read_u8() {
217 let x = 0xAA as u8;
218 let buf = [x];
219 let mut bb = ByteBuf::wrap(&buf);
220 assert_eq!(bb.get_u8(), Some(x));
221 assert_eq!(bb.pos, 1);
222 }
223
224 #[test]
225 fn test_read_u16() {
226 let x = 0xAABB as u16;
227 let buf = [0xAA, 0xBB];
228 let mut bb = ByteBuf::wrap(&buf);
229 assert_eq!(bb.get_u16(), Some(x));
230 assert_eq!(bb.pos, 2);
231 }
232
233 #[test]
234 fn test_read_u32() {
235 let x = 0xAABBCCDD as u32;
236 let buf = [0xAA, 0xBB, 0xCC, 0xDD];
237 let mut bb = ByteBuf::wrap(&buf);
238 assert_eq!(bb.get_u32(), Some(x));
239 assert_eq!(bb.pos, 4);
240 }
241
242 #[test]
243 fn test_read_u64() {
244 let x = 0xAABBCCDDEEFFABCD as u64;
245 let buf = [0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0xAB, 0xCD];
246 let mut bb = ByteBuf::wrap(&buf);
247 assert_eq!(bb.get_u64(), Some(x));
248 assert_eq!(bb.pos, 8);
249 }
250
251 #[test]
252 fn test_read_all() {
253 let buf = [0x01, 0x00, 0x02, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04];
254 let mut bb = ByteBuf::wrap(&buf);
255 assert_eq!(bb.get_u8(), Some(0x01));
256 assert_eq!(bb.pos, 1);
257 assert_eq!(bb.get_u16(), Some(0x0002));
258 assert_eq!(bb.pos, 3);
259 assert_eq!(bb.get_u32(), Some(0x00000003));
260 assert_eq!(bb.pos, 7);
261 assert_eq!(bb.get_u64(), Some(0x0000000000000004));
262 assert_eq!(bb.pos, 15);
263 }
264
265 #[test]
266 fn test_write_all_read_all() {
267 let mut buf = [0 as u8; 1024];
268
269 {
270 let mut bb = ByteBufMut::wrap(&mut buf);
271 bb.put_u8(1 as u8);
272 bb.put_u16(2 as u16);
273 bb.put_u32(3 as u32);
274 bb.put_u64(4 as u64);
275 }
276
277 let exp: [u8; 15] = [1, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4];
278 assert_eq!(exp, buf[0..15]);
279
280 let mut bb = ByteBuf::wrap(&buf);
281 assert_eq!(1 as u8, bb.get_u8().unwrap());
282 assert_eq!(bb.pos, 1);
283 assert_eq!(2 as u16, bb.get_u16().unwrap());
284 assert_eq!(bb.pos, 3);
285 assert_eq!(3 as u32, bb.get_u32().unwrap());
286 assert_eq!(bb.pos, 7);
287 assert_eq!(4 as u64, bb.get_u64().unwrap());
288 assert_eq!(bb.pos, 15);
289 }
290
291 #[test]
292 fn test_write_slice() {
293 let mut buf = [0 as u8; 64];
294 let mut bb = ByteBufMut::wrap(&mut buf);
295
296 let a = b"hello there";
297 let b = b"sup?";
298 bb.put_bytes(a);
299 bb.put_bytes(b);
300 let p = bb.pos;
301
302 assert_eq!(b"hello theresup?\0", &buf[0..(a.len() + b.len() + 1)]);
303 assert_eq!(p, a.len() + b.len());
304 }
305
306 #[test]
307 fn test_read_slice() {
308 let mut buf: &[u8] = b"what is up?!...\0";
309 let mut bb = ByteBuf::wrap(&mut buf);
310
311 assert_eq!(b"what ", bb.get_bytes(5));
312 assert_eq!(bb.pos, 5);
313 }
314}