1/*
2 * Copyright (c) 2023.
3 *
4 * This software is free software;
5 *
6 * You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
7 */
8
9use core::mem::size_of;
10
11enum Mode {
12 // Big endian
13 BE,
14 // Little Endian
15 LE
16}
17
18static ERROR_MSG: &str = "No more space";
19
20/// Encapsulates a simple Byte writer with
21/// support for Endian aware writes
22pub struct ZByteWriter<'a> {
23 buffer: &'a mut [u8],
24 position: usize
25}
26
27impl<'a> ZByteWriter<'a> {
28 /// Write bytes from the buf into the bytestream
29 /// and return how many bytes were written
30 ///
31 /// # Arguments
32 /// - `buf`: The bytes to be written to the bytestream
33 ///
34 /// # Returns
35 /// - `Ok(usize)` - Number of bytes written
36 /// This number may be less than `buf.len()` if the length of the buffer is greater
37 /// than the internal bytestream length
38 ///
39 /// If you want to be sure that all bytes were written, see [`write_all`](Self::write_all)
40 ///
41 #[inline]
42 pub fn write(&mut self, buf: &[u8]) -> Result<usize, &'static str> {
43 let min = buf.len().min(self.bytes_left());
44 // write
45 self.buffer[self.position..self.position + min].copy_from_slice(&buf[0..min]);
46 self.position += min;
47
48 Ok(min)
49 }
50 /// Write all bytes from `buf` into the bytestream and return
51 /// and panic if not all bytes were written to the bytestream
52 ///
53 /// # Arguments
54 /// - `buf`: The bytes to be written into the bytestream
55 ///
56 ///# Returns
57 /// - `Ok(())`: Indicates all bytes were written into the bytestream
58 /// - `Err(&static str)`: In case all the bytes could not be written
59 /// to the stream
60 pub fn write_all(&mut self, buf: &[u8]) -> Result<(), &'static str> {
61 let size = self.write(buf)?;
62
63 if size != buf.len() {
64 return Err("Could not write the whole buffer");
65 }
66 Ok(())
67 }
68 /// Create a new bytestream writer
69 /// Bytes are written from the start to the end and not assumptions
70 /// are made of the nature of the underlying stream
71 ///
72 /// # Arguments
73 pub fn new(data: &'a mut [u8]) -> ZByteWriter<'a> {
74 ZByteWriter {
75 buffer: data,
76 position: 0
77 }
78 }
79 /// Return number of unwritten bytes in this stream
80 ///
81 /// # Example
82 /// ```
83 /// use zune_core::bytestream::ZByteWriter;
84 /// let mut storage = [0;10];
85 ///
86 /// let writer = ZByteWriter::new(&mut storage);
87 /// assert_eq!(writer.bytes_left(),10); // no bytes were written
88 /// ```
89 pub const fn bytes_left(&self) -> usize {
90 self.buffer.len().saturating_sub(self.position)
91 }
92
93 /// Return the number of bytes the writer has written
94 ///
95 /// ```
96 /// use zune_core::bytestream::ZByteWriter;
97 /// let mut stream = ZByteWriter::new(&mut []);
98 /// assert_eq!(stream.position(),0);
99 /// ```
100 pub const fn position(&self) -> usize {
101 self.position
102 }
103
104 /// Write a single byte into the bytestream or error out
105 /// if there is not enough space
106 ///
107 /// # Example
108 /// ```
109 /// use zune_core::bytestream::ZByteWriter;
110 /// let mut buf = [0;10];
111 /// let mut stream = ZByteWriter::new(&mut buf);
112 /// assert!(stream.write_u8_err(34).is_ok());
113 /// ```
114 /// No space
115 /// ```
116 /// use zune_core::bytestream::ZByteWriter;
117 /// let mut stream = ZByteWriter::new(&mut []);
118 /// assert!(stream.write_u8_err(32).is_err());
119 /// ```
120 ///
121 pub fn write_u8_err(&mut self, byte: u8) -> Result<(), &'static str> {
122 match self.buffer.get_mut(self.position) {
123 Some(m_byte) => {
124 self.position += 1;
125 *m_byte = byte;
126
127 Ok(())
128 }
129 None => Err(ERROR_MSG)
130 }
131 }
132
133 /// Write a single byte in the stream or don't write
134 /// anything if the buffer is full and cannot support the byte read
135 ///
136 /// Should be combined with [`has`](Self::has)
137 pub fn write_u8(&mut self, byte: u8) {
138 if let Some(m_byte) = self.buffer.get_mut(self.position) {
139 self.position += 1;
140 *m_byte = byte;
141 }
142 }
143 /// Check if the byte writer can support
144 /// the following write
145 ///
146 /// # Example
147 /// ```
148 /// use zune_core::bytestream::ZByteWriter;
149 /// let mut data = [0;10];
150 /// let mut stream = ZByteWriter::new(&mut data);
151 /// assert!(stream.has(5));
152 /// assert!(!stream.has(100));
153 /// ```
154 pub const fn has(&self, bytes: usize) -> bool {
155 self.position.saturating_add(bytes) <= self.buffer.len()
156 }
157
158 /// Get length of the underlying buffer.
159 #[inline]
160 pub const fn len(&self) -> usize {
161 self.buffer.len()
162 }
163 /// Return true if the underlying buffer stream is empty
164 #[inline]
165 pub const fn is_empty(&self) -> bool {
166 self.len() == 0
167 }
168
169 /// Return true whether or not we read to the end of the
170 /// buffer and have no more bytes left.
171 ///
172 /// If this is true, all non error variants will silently discard the
173 /// byte and all error variants will return an error on writing a byte
174 /// if any write occurs
175 ///
176 ///
177 #[inline]
178 pub const fn eof(&self) -> bool {
179 self.position >= self.len()
180 }
181
182 /// Rewind the position of the internal cursor back by `by` bytes
183 ///
184 /// The position saturates at zero
185 ///
186 /// # Example
187 /// ```
188 /// use zune_core::bytestream::ZByteWriter;
189 /// let bytes = &mut [1,2,4];
190 /// let mut stream = ZByteWriter::new(bytes);
191 /// stream.write_u16_be(23);
192 /// // now internal cursor is at position 2.
193 /// // lets rewind it
194 /// stream.rewind(usize::MAX);
195 /// assert_eq!(stream.position(),0);
196 /// ```
197 #[inline]
198 pub fn rewind(&mut self, by: usize) {
199 self.position = self.position.saturating_sub(by);
200 }
201 /// Move the internal cursor forward some bytes
202 ///
203 ///
204 /// This saturates at maximum value of usize in your platform.
205 #[inline]
206 pub fn skip(&mut self, by: usize) {
207 self.position = self.position.saturating_add(by);
208 }
209
210 /// Look ahead position bytes and return a reference
211 /// to num_bytes from that position, or an error if the
212 /// peek would be out of bounds.
213 ///
214 /// This doesn't increment the position, bytes would have to be discarded
215 /// at a later point.
216 #[inline]
217 pub fn peek_at(&'a self, position: usize, num_bytes: usize) -> Result<&'a [u8], &'static str> {
218 let start = self.position + position;
219 let end = self.position + position + num_bytes;
220
221 match self.buffer.get(start..end) {
222 Some(bytes) => Ok(bytes),
223 None => Err(ERROR_MSG)
224 }
225 }
226
227 /// Set position for the internal cursor
228 ///
229 /// Further calls to write bytes will proceed from the
230 /// position set
231 pub fn set_position(&mut self, position: usize) {
232 self.position = position;
233 }
234}
235
236macro_rules! write_single_type {
237 ($name:tt,$name2:tt,$name3:tt,$name4:tt,$name5:tt,$name6:tt,$int_type:tt) => {
238 impl<'a> ZByteWriter<'a>
239 {
240 #[inline(always)]
241 fn $name(&mut self, byte: $int_type, mode: Mode) -> Result<(), &'static str>
242 {
243 const SIZE: usize = size_of::<$int_type>();
244
245 match self.buffer.get_mut(self.position..self.position + SIZE)
246 {
247 Some(m_byte) =>
248 {
249 self.position += SIZE;
250 // get bits, depending on mode.
251 // This should be inlined and not visible in
252 // the generated binary since mode is a compile
253 // time constant.
254 let bytes = match mode
255 {
256 Mode::BE => byte.to_be_bytes(),
257 Mode::LE => byte.to_le_bytes()
258 };
259
260 m_byte.copy_from_slice(&bytes);
261
262 Ok(())
263 }
264 None => Err(ERROR_MSG)
265 }
266 }
267 #[inline(always)]
268 fn $name2(&mut self, byte: $int_type, mode: Mode)
269 {
270 const SIZE: usize = size_of::<$int_type>();
271
272 if let Some(m_byte) = self.buffer.get_mut(self.position..self.position + SIZE)
273 {
274 self.position += SIZE;
275 // get bits, depending on mode.
276 // This should be inlined and not visible in
277 // the generated binary since mode is a compile
278 // time constant.
279 let bytes = match mode
280 {
281 Mode::BE => byte.to_be_bytes(),
282 Mode::LE => byte.to_le_bytes()
283 };
284
285 m_byte.copy_from_slice(&bytes);
286 }
287 }
288
289 #[doc=concat!("Write ",stringify!($int_type)," as a big endian integer")]
290 #[doc=concat!("Returning an error if the underlying buffer cannot support a ",stringify!($int_type)," write.")]
291 #[inline]
292 pub fn $name3(&mut self, byte: $int_type) -> Result<(), &'static str>
293 {
294 self.$name(byte, Mode::BE)
295 }
296
297 #[doc=concat!("Write ",stringify!($int_type)," as a little endian integer")]
298 #[doc=concat!("Returning an error if the underlying buffer cannot support a ",stringify!($int_type)," write.")]
299 #[inline]
300 pub fn $name4(&mut self, byte: $int_type) -> Result<(), &'static str>
301 {
302 self.$name(byte, Mode::LE)
303 }
304
305 #[doc=concat!("Write ",stringify!($int_type)," as a big endian integer")]
306 #[doc=concat!("Or don't write anything if the reader cannot support a ",stringify!($int_type)," write.")]
307 #[doc=concat!("\nShould be combined with the [`has`](Self::has) method to ensure a write succeeds")]
308 #[inline]
309 pub fn $name5(&mut self, byte: $int_type)
310 {
311 self.$name2(byte, Mode::BE)
312 }
313 #[doc=concat!("Write ",stringify!($int_type)," as a little endian integer")]
314 #[doc=concat!("Or don't write anything if the reader cannot support a ",stringify!($int_type)," write.")]
315 #[doc=concat!("Should be combined with the [`has`](Self::has) method to ensure a write succeeds")]
316 #[inline]
317 pub fn $name6(&mut self, byte: $int_type)
318 {
319 self.$name2(byte, Mode::LE)
320 }
321 }
322 };
323}
324
325write_single_type!(
326 write_u64_inner_or_die,
327 write_u64_inner_or_none,
328 write_u64_be_err,
329 write_u64_le_err,
330 write_u64_be,
331 write_u64_le,
332 u64
333);
334
335write_single_type!(
336 write_u32_inner_or_die,
337 write_u32_inner_or_none,
338 write_u32_be_err,
339 write_u32_le_err,
340 write_u32_be,
341 write_u32_le,
342 u32
343);
344
345write_single_type!(
346 write_u16_inner_or_die,
347 write_u16_inner_or_none,
348 write_u16_be_err,
349 write_u16_le_err,
350 write_u16_be,
351 write_u16_le,
352 u16
353);
354