Skip to main content

bincode_next/enc/
write.rs

1//! This module contains writer-based structs and traits.
2//!
3//! Because `std::io::Write` is only limited to `std` and not `core`, we provide our own [Writer].
4#![allow(unsafe_code)]
5
6use crate::error::EncodeError;
7
8#[cfg(feature = "alloc")]
9impl Writer for crate::alloc::vec::Vec<u8> {
10    #[inline(always)]
11    fn write(
12        &mut self,
13        bytes: &[u8],
14    ) -> Result<(), EncodeError> {
15        self.extend_from_slice(bytes);
16        Ok(())
17    }
18
19    #[inline(always)]
20    fn write_u8(
21        &mut self,
22        value: u8,
23    ) -> Result<(), EncodeError> {
24        self.push(value);
25        Ok(())
26    }
27}
28
29/// Trait that indicates that a struct can be used as a destination to encode data too. This is used by [Encode]
30///
31/// [Encode]: ../trait.Encode.html
32pub trait Writer {
33    /// Write `bytes` to the underlying writer. Exactly `bytes.len()` bytes must be written, or else an error should be returned.
34    ///
35    /// # Errors
36    ///
37    /// Returns `EncodeError::UnexpectedEnd` if the writer does not have enough space.
38    fn write(
39        &mut self,
40        bytes: &[u8],
41    ) -> Result<(), EncodeError>;
42
43    /// Write a single byte to the underlying writer.
44    ///
45    /// # Errors
46    ///
47    /// Returns `EncodeError::UnexpectedEnd` if the writer does not have enough space.
48    #[inline(always)]
49    fn write_u8(
50        &mut self,
51        value: u8,
52    ) -> Result<(), EncodeError> {
53        self.write(&[value])
54    }
55
56    /// Write a `u16` to the underlying writer.
57    ///
58    /// # Errors
59    ///
60    /// Returns `EncodeError::UnexpectedEnd` if the writer does not have enough space.
61    #[inline(always)]
62    fn write_u16(
63        &mut self,
64        value: u16,
65    ) -> Result<(), EncodeError> {
66        self.write(&value.to_ne_bytes())
67    }
68
69    /// Write a `u32` to the underlying writer.
70    ///
71    /// # Errors
72    ///
73    /// Returns `EncodeError::UnexpectedEnd` if the writer does not have enough space.
74    #[inline(always)]
75    fn write_u32(
76        &mut self,
77        value: u32,
78    ) -> Result<(), EncodeError> {
79        self.write(&value.to_ne_bytes())
80    }
81
82    /// Write a `u64` to the underlying writer.
83    ///
84    /// # Errors
85    ///
86    /// Returns `EncodeError::UnexpectedEnd` if the writer does not have enough space.
87    #[inline(always)]
88    fn write_u64(
89        &mut self,
90        value: u64,
91    ) -> Result<(), EncodeError> {
92        self.write(&value.to_ne_bytes())
93    }
94
95    /// Write a `u128` to the underlying writer.
96    ///
97    /// # Errors
98    ///
99    /// Returns `EncodeError::UnexpectedEnd` if the writer does not have enough space.
100    #[inline(always)]
101    fn write_u128(
102        &mut self,
103        value: u128,
104    ) -> Result<(), EncodeError> {
105        self.write(&value.to_ne_bytes())
106    }
107}
108
109impl<T: Writer> Writer for &mut T {
110    #[inline(always)]
111    fn write(
112        &mut self,
113        bytes: &[u8],
114    ) -> Result<(), EncodeError> {
115        (**self).write(bytes)
116    }
117
118    #[inline(always)]
119    fn write_u8(
120        &mut self,
121        value: u8,
122    ) -> Result<(), EncodeError> {
123        (**self).write_u8(value)
124    }
125
126    #[inline(always)]
127    fn write_u16(
128        &mut self,
129        value: u16,
130    ) -> Result<(), EncodeError> {
131        (**self).write_u16(value)
132    }
133
134    #[inline(always)]
135    fn write_u32(
136        &mut self,
137        value: u32,
138    ) -> Result<(), EncodeError> {
139        (**self).write_u32(value)
140    }
141
142    #[inline(always)]
143    fn write_u64(
144        &mut self,
145        value: u64,
146    ) -> Result<(), EncodeError> {
147        (**self).write_u64(value)
148    }
149
150    #[inline(always)]
151    fn write_u128(
152        &mut self,
153        value: u128,
154    ) -> Result<(), EncodeError> {
155        (**self).write_u128(value)
156    }
157}
158
159/// A helper struct that implements `Writer` for a `&[u8]` slice.
160///
161/// ```
162/// use bincode_next::enc::write::SliceWriter;
163/// use bincode_next::enc::write::Writer;
164///
165/// let destination = &mut [0u8; 100];
166/// let mut writer = SliceWriter::new(destination);
167/// writer.write(&[1, 2, 3, 4, 5]).unwrap();
168///
169/// assert_eq!(writer.bytes_written(), 5);
170/// assert_eq!(destination[0..6], [1, 2, 3, 4, 5, 0]);
171/// ```
172pub struct SliceWriter<'storage> {
173    slice: &'storage mut [u8],
174    original_length: usize,
175}
176
177impl<'storage> SliceWriter<'storage> {
178    /// Create a new instance of `SliceWriter` with the given byte array.
179    pub const fn new(bytes: &'storage mut [u8]) -> Self {
180        let original = bytes.len();
181        Self {
182            slice: bytes,
183            original_length: original,
184        }
185    }
186
187    /// Return the amount of bytes written so far.
188    #[must_use]
189    pub const fn bytes_written(&self) -> usize {
190        self.original_length - self.slice.len()
191    }
192}
193
194impl Writer for SliceWriter<'_> {
195    #[inline(always)]
196    fn write(
197        &mut self,
198        bytes: &[u8],
199    ) -> Result<(), EncodeError> {
200        if bytes.len() > self.slice.len() {
201            return crate::error::cold_encode_error_unexpected_end();
202        }
203        // SAFETY: `bytes.len() <= self.slice.len()` is checked above.
204        // `copy_nonoverlapping` is safe since the pointers won't overlap and bounds are exact.
205        // Advancing the pointer by `bytes.len()` is safe for the same reason.
206        unsafe {
207            core::ptr::copy_nonoverlapping(bytes.as_ptr(), self.slice.as_mut_ptr(), bytes.len());
208            let ptr = self.slice.as_mut_ptr().add(bytes.len());
209            let len = self.slice.len() - bytes.len();
210            self.slice = core::slice::from_raw_parts_mut(ptr, len);
211        }
212
213        Ok(())
214    }
215
216    #[inline(always)]
217    fn write_u8(
218        &mut self,
219        value: u8,
220    ) -> Result<(), EncodeError> {
221        if self.slice.is_empty() {
222            return crate::error::cold_encode_error_unexpected_end();
223        }
224        // SAFETY: `!self.slice.is_empty()` is checked above.
225        unsafe {
226            *self.slice.as_mut_ptr() = value;
227            let ptr = self.slice.as_mut_ptr().add(1);
228            let len = self.slice.len() - 1;
229            self.slice = core::slice::from_raw_parts_mut(ptr, len);
230        }
231
232        Ok(())
233    }
234
235    #[inline(always)]
236    fn write_u16(
237        &mut self,
238        value: u16,
239    ) -> Result<(), EncodeError> {
240        if self.slice.len() < 2 {
241            return crate::error::cold_encode_error_unexpected_end();
242        }
243        // SAFETY: `self.slice.len() < 2` is checked above.
244        unsafe {
245            core::ptr::write_unaligned(self.slice.as_mut_ptr().cast::<u16>(), value);
246            let ptr = self.slice.as_mut_ptr().add(2);
247            let len = self.slice.len() - 2;
248            self.slice = core::slice::from_raw_parts_mut(ptr, len);
249        }
250
251        Ok(())
252    }
253
254    #[inline(always)]
255    fn write_u32(
256        &mut self,
257        value: u32,
258    ) -> Result<(), EncodeError> {
259        if self.slice.len() < 4 {
260            return crate::error::cold_encode_error_unexpected_end();
261        }
262        // SAFETY: `self.slice.len() < 4` is checked above.
263        unsafe {
264            core::ptr::write_unaligned(self.slice.as_mut_ptr().cast::<u32>(), value);
265            let ptr = self.slice.as_mut_ptr().add(4);
266            let len = self.slice.len() - 4;
267            self.slice = core::slice::from_raw_parts_mut(ptr, len);
268        }
269
270        Ok(())
271    }
272
273    #[inline(always)]
274    fn write_u64(
275        &mut self,
276        value: u64,
277    ) -> Result<(), EncodeError> {
278        if self.slice.len() < 8 {
279            return crate::error::cold_encode_error_unexpected_end();
280        }
281        // SAFETY: `self.slice.len() < 8` is checked above.
282        unsafe {
283            core::ptr::write_unaligned(self.slice.as_mut_ptr().cast::<u64>(), value);
284            let ptr = self.slice.as_mut_ptr().add(8);
285            let len = self.slice.len() - 8;
286            self.slice = core::slice::from_raw_parts_mut(ptr, len);
287        }
288
289        Ok(())
290    }
291
292    #[inline(always)]
293    fn write_u128(
294        &mut self,
295        value: u128,
296    ) -> Result<(), EncodeError> {
297        if self.slice.len() < 16 {
298            return crate::error::cold_encode_error_unexpected_end();
299        }
300        // SAFETY: `self.slice.len() < 16` is checked above.
301        unsafe {
302            core::ptr::write_unaligned(self.slice.as_mut_ptr().cast::<u128>(), value);
303            let ptr = self.slice.as_mut_ptr().add(16);
304            let len = self.slice.len() - 16;
305            self.slice = core::slice::from_raw_parts_mut(ptr, len);
306        }
307
308        Ok(())
309    }
310}
311
312/// A writer that counts how many bytes were written. This is useful for e.g. pre-allocating buffers before writing to them.
313#[derive(Default)]
314pub struct SizeWriter {
315    /// the amount of bytes that were written so far
316    pub bytes_written: usize,
317}
318impl Writer for SizeWriter {
319    #[inline(always)]
320    fn write(
321        &mut self,
322        bytes: &[u8],
323    ) -> Result<(), EncodeError> {
324        self.bytes_written += bytes.len();
325
326        Ok(())
327    }
328
329    #[inline(always)]
330    fn write_u8(
331        &mut self,
332        _: u8,
333    ) -> Result<(), EncodeError> {
334        self.bytes_written += 1;
335
336        Ok(())
337    }
338
339    #[inline(always)]
340    fn write_u16(
341        &mut self,
342        _: u16,
343    ) -> Result<(), EncodeError> {
344        self.bytes_written += 2;
345
346        Ok(())
347    }
348
349    #[inline(always)]
350    fn write_u32(
351        &mut self,
352        _: u32,
353    ) -> Result<(), EncodeError> {
354        self.bytes_written += 4;
355
356        Ok(())
357    }
358
359    #[inline(always)]
360    fn write_u64(
361        &mut self,
362        _: u64,
363    ) -> Result<(), EncodeError> {
364        self.bytes_written += 8;
365
366        Ok(())
367    }
368
369    #[inline(always)]
370    fn write_u128(
371        &mut self,
372        _: u128,
373    ) -> Result<(), EncodeError> {
374        self.bytes_written += 16;
375
376        Ok(())
377    }
378}