zune_core/bytestream/writer.rs
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 */
8use crate::bytestream::{ZByteIoError, ZByteWriterTrait};
9
10mod no_std_writer;
11mod std_writer;
12
13enum Mode {
14 // Big endian
15 BE,
16 // Little Endian
17 LE
18}
19
20/// Encapsulates a simple Byte writer with
21/// support for Endian aware writes
22pub struct ZWriter<T: ZByteWriterTrait> {
23 buffer: T,
24 bytes_written: usize
25}
26
27impl<T: ZByteWriterTrait> ZWriter<T> {
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, ZByteIoError> {
43 let bytes_written = self.buffer.write_bytes(buf)?;
44 self.bytes_written += bytes_written;
45 Ok(bytes_written)
46 }
47 /// Write all bytes from `buf` into the bytestream and return
48 /// and panic if not all bytes were written to the bytestream
49 ///
50 /// # Arguments
51 /// - `buf`: The bytes to be written into the bytestream
52 ///
53 ///# Returns
54 /// - `Ok(())`: Indicates all bytes were written into the bytestream
55 /// - `Err(&static str)`: In case all the bytes could not be written
56 /// to the stream
57 pub fn write_all(&mut self, buf: &[u8]) -> Result<(), ZByteIoError> {
58 self.buffer.write_all_bytes(buf)?;
59 self.bytes_written += buf.len();
60 Ok(())
61 }
62 /// Create a new bytestream writer
63 /// Bytes are written from the start to the end and not assumptions
64 /// are made of the nature of the underlying stream
65 ///
66 /// # Arguments
67 pub fn new(data: T) -> ZWriter<T> {
68 ZWriter {
69 buffer: data,
70 bytes_written: 0
71 }
72 }
73
74 /// Write a single byte into the bytestream or error out
75 /// if there is not enough space
76 ///
77 /// # Example
78 /// ```
79 /// use zune_core::bytestream::ZWriter;
80 /// let mut buf = [0;10];
81 /// let mut stream = ZWriter::new(&mut buf[..]);
82 /// assert!(stream.write_u8_err(34).is_ok());
83 /// ```
84 /// No space
85 /// ```
86 /// use zune_core::bytestream::ZWriter;
87 /// let mut no_space = [];
88 /// let mut stream = ZWriter::new(&mut no_space[..]);
89 /// assert!(stream.write_u8_err(32).is_err());
90 /// ```
91 ///
92 #[inline]
93 pub fn write_u8_err(&mut self, byte: u8) -> Result<(), ZByteIoError> {
94 self.write_const_bytes(&[byte])
95 }
96 /// Write a fixed compile time known number of bytes to the sink
97 ///
98 /// This is provided since some implementations can optimize such writes by eliminating
99 /// some redundant code.
100 #[inline]
101 pub fn write_const_bytes<const N: usize>(
102 &mut self, byte: &[u8; N]
103 ) -> Result<(), ZByteIoError> {
104 self.buffer.write_const_bytes(byte)?;
105 self.bytes_written += N;
106 Ok(())
107 }
108
109 /// Write a single byte in the stream or don't write
110 /// anything if the buffer is full and cannot support the byte read
111 ///
112 #[inline]
113 pub fn write_u8(&mut self, byte: u8) {
114 let _ = self.write_const_bytes(&[byte]);
115 }
116 /// Return the number of bytes written by this encoder
117 ///
118 /// The encoder keeps information of how many bytes were written and this method
119 /// returns that value.
120 ///
121 /// # Returns
122 /// Number of bytes written
123 pub fn bytes_written(&self) -> usize {
124 self.bytes_written
125 }
126
127 /// Reserve some additional space to write.
128 ///
129 /// Some sinks like `Vec<u8>` allow reallocation and to prevent too much reallocation
130 /// one can use this to reserve additional space to encode
131 ///
132 /// # Example
133 ///
134 /// ```
135 /// use zune_core::bytestream::ZWriter;
136 /// let space_needed = 10; // Assume the image will fit into 10 bytes
137 /// let mut output = Vec::new();
138 /// let mut sink = ZWriter::new(&mut output);
139 /// // now reserve some space
140 ///sink.reserve(space_needed).unwrap();
141 /// // at this point, we can assume that ZWriter allocated space for output
142 /// ```
143 pub fn reserve(&mut self, additional: usize) -> Result<(), ZByteIoError> {
144 self.buffer.reserve_capacity(additional)
145 }
146 /// Consume the writer and return the inner sink
147 /// we were writing to.
148 ///
149 /// After this, the writer can no longer be used
150 pub fn inner(self) -> T {
151 self.buffer
152 }
153 /// Return an immutable reference to the inner sink
154 pub fn inner_ref(&self) -> &T {
155 &self.buffer
156 }
157 /// Return a mutable reference to the inner sink
158 pub fn inner_mut(&mut self) -> &mut T {
159 &mut self.buffer
160 }
161}
162
163macro_rules! write_single_type {
164 ($name:tt,$name2:tt,$name3:tt,$name4:tt,$name5:tt,$name6:tt,$int_type:tt) => {
165 impl<T:ZByteWriterTrait> ZWriter<T>
166 {
167 #[inline(always)]
168 fn $name(&mut self, byte: $int_type, mode: Mode) -> Result<(), ZByteIoError>
169 {
170
171 // get bits, depending on mode.
172 // This should be inlined and not visible in
173 // the generated binary since mode is a compile
174 // time constant.
175 let bytes = match mode
176 {
177 Mode::BE => byte.to_be_bytes(),
178 Mode::LE => byte.to_le_bytes()
179 };
180 self.write_const_bytes(&bytes)
181 }
182 #[inline(always)]
183 fn $name2(&mut self, byte: $int_type, mode: Mode)
184 {
185
186 // get bits, depending on mode.
187 // This should be inlined and not visible in
188 // the generated binary since mode is a compile
189 // time constant.
190 let bytes = match mode
191 {
192 Mode::BE => byte.to_be_bytes(),
193 Mode::LE => byte.to_le_bytes()
194 };
195 let _ = self.write_const_bytes(&bytes);
196
197
198 }
199
200 #[doc=concat!("Write ",stringify!($int_type)," as a big endian integer")]
201 #[doc=concat!("Returning an error if the underlying buffer cannot support a ",stringify!($int_type)," write.")]
202 #[inline]
203 pub fn $name3(&mut self, byte: $int_type) -> Result<(), ZByteIoError>
204 {
205 self.$name(byte, Mode::BE)
206 }
207
208 #[doc=concat!("Write ",stringify!($int_type)," as a little endian integer")]
209 #[doc=concat!("Returning an error if the underlying buffer cannot support a ",stringify!($int_type)," write.")]
210 #[inline]
211 pub fn $name4(&mut self, byte: $int_type) -> Result<(), ZByteIoError>
212 {
213 self.$name(byte, Mode::LE)
214 }
215
216 #[doc=concat!("Write ",stringify!($int_type)," as a big endian integer")]
217 #[doc=concat!("Or don't write anything if the reader cannot support a ",stringify!($int_type)," write.")]
218 #[inline]
219 pub fn $name5(&mut self, byte: $int_type)
220 {
221 self.$name2(byte, Mode::BE)
222 }
223 #[doc=concat!("Write ",stringify!($int_type)," as a little endian integer")]
224 #[doc=concat!("Or don't write anything if the reader cannot support a ",stringify!($int_type)," write.")]
225 #[inline]
226 pub fn $name6(&mut self, byte: $int_type)
227 {
228 self.$name2(byte, Mode::LE)
229 }
230 }
231 };
232}
233
234write_single_type!(
235 write_u64_inner_or_die,
236 write_u64_inner_or_none,
237 write_u64_be_err,
238 write_u64_le_err,
239 write_u64_be,
240 write_u64_le,
241 u64
242);
243
244write_single_type!(
245 write_u32_inner_or_die,
246 write_u32_inner_or_none,
247 write_u32_be_err,
248 write_u32_le_err,
249 write_u32_be,
250 write_u32_le,
251 u32
252);
253
254write_single_type!(
255 write_u16_inner_or_die,
256 write_u16_inner_or_none,
257 write_u16_be_err,
258 write_u16_le_err,
259 write_u16_be,
260 write_u16_le,
261 u16
262);