Skip to main content

wincode/io/
std_write.rs

1use {
2    crate::io::{WriteResult, Writer, slice::SliceMutUnchecked, write_size_limit},
3    std::io::{BufWriter, Cursor, Write},
4};
5
6/// [`Writer`] adapter over any [`std::io::Write`] sink.
7///
8/// Wraps any `W: std::io::Write` and exposes it as a wincode [`Writer`], allowing
9/// serialization into files, network streams, or other I/O sinks.
10///
11/// # Examples
12///
13/// Serialize a tuple into a `Vec<u8>` via `WriteAdapter`:
14///
15/// ```
16/// use wincode::{Serialize, io::{Writer, std_write::WriteAdapter}};
17///
18/// let tuple = (42u32, true, 1234567890i64);
19/// let mut buf = Vec::new();
20/// let mut writer = WriteAdapter::new(&mut buf);
21/// <(u32, bool, i64)>::serialize_into(&mut writer, &tuple).unwrap();
22/// writer.finish().unwrap();
23/// assert_eq!(buf, wincode::serialize(&tuple).unwrap());
24/// ```
25#[derive(Debug)]
26pub struct WriteAdapter<W: ?Sized>(W);
27
28impl<W: Write> WriteAdapter<W> {
29    pub fn new(writer: W) -> Self {
30        Self(writer)
31    }
32}
33
34impl<W: Write + ?Sized> Writer for WriteAdapter<W> {
35    fn write(&mut self, src: &[u8]) -> WriteResult<()> {
36        Ok(self.0.write_all(src)?)
37    }
38
39    fn finish(&mut self) -> WriteResult<()> {
40        Ok(self.0.flush()?)
41    }
42}
43
44impl<W: Write + ?Sized> Writer for BufWriter<W> {
45    fn write(&mut self, src: &[u8]) -> WriteResult<()> {
46        Ok(self.write_all(src)?)
47    }
48
49    fn finish(&mut self) -> WriteResult<()> {
50        Ok(self.flush()?)
51    }
52}
53
54#[inline]
55fn cursor_slice_as_trusted_for(
56    cursor: &mut Cursor<impl AsMut<[u8]>>,
57    n_bytes: usize,
58) -> WriteResult<impl Writer> {
59    let Ok(pos) = usize::try_from(cursor.position()) else {
60        return Err(write_size_limit(usize::MAX));
61    };
62
63    let inner = cursor.get_mut().as_mut();
64    let next_pos = pos.saturating_add(n_bytes);
65    if next_pos > inner.len() {
66        return Err(write_size_limit(n_bytes));
67    }
68
69    cursor.set_position(next_pos as u64);
70    let slice = &mut cursor.get_mut().as_mut()[pos..next_pos];
71    // SAFETY: by calling `as_trusted_for`, caller guarantees they
72    // will fully initialize `n_bytes` of memory and will not write
73    // beyond the bounds of the slice.
74    Ok(unsafe { SliceMutUnchecked::new(slice) })
75}
76
77impl Writer for Cursor<&mut [u8]> {
78    #[inline]
79    fn write(&mut self, src: &[u8]) -> WriteResult<()> {
80        Ok(self.write_all(src)?)
81    }
82
83    #[inline]
84    fn finish(&mut self) -> WriteResult<()> {
85        Ok(self.flush()?)
86    }
87
88    #[inline]
89    unsafe fn as_trusted_for(&mut self, n_bytes: usize) -> WriteResult<impl Writer> {
90        cursor_slice_as_trusted_for(self, n_bytes)
91    }
92}
93
94impl<const N: usize> Writer for Cursor<[u8; N]> {
95    #[inline]
96    fn write(&mut self, src: &[u8]) -> WriteResult<()> {
97        Ok(self.write_all(src)?)
98    }
99
100    #[inline]
101    fn finish(&mut self) -> WriteResult<()> {
102        Ok(self.flush()?)
103    }
104
105    #[inline]
106    unsafe fn as_trusted_for(&mut self, n_bytes: usize) -> WriteResult<impl Writer> {
107        cursor_slice_as_trusted_for(self, n_bytes)
108    }
109}
110
111impl Writer for Cursor<Box<[u8]>> {
112    #[inline]
113    fn write(&mut self, src: &[u8]) -> WriteResult<()> {
114        Ok(self.write_all(src)?)
115    }
116
117    #[inline]
118    fn finish(&mut self) -> WriteResult<()> {
119        Ok(self.flush()?)
120    }
121
122    #[inline]
123    unsafe fn as_trusted_for(&mut self, n_bytes: usize) -> WriteResult<impl Writer> {
124        cursor_slice_as_trusted_for(self, n_bytes)
125    }
126}
127
128#[inline]
129fn cursor_vec_as_trusted_for(
130    cursor: &mut Cursor<impl AsMut<Vec<u8>>>,
131    n_bytes: usize,
132) -> WriteResult<impl Writer> {
133    let Ok(pos) = usize::try_from(cursor.position()) else {
134        return Err(write_size_limit(usize::MAX));
135    };
136
137    let vec = cursor.get_mut().as_mut();
138    crate::io::cursor::vec::prepare_write(vec, pos, n_bytes)?;
139
140    // SAFETY: by calling `as_trusted_for`, caller guarantees they
141    // will fully initialize `n_bytes` of memory and will not write
142    // beyond the bounds of the slice.
143    Ok(unsafe { CursorVecUnchecked::new(cursor) })
144}
145
146struct CursorVecUnchecked<'a, T> {
147    inner: &'a mut Cursor<T>,
148}
149
150impl<'a, T> CursorVecUnchecked<'a, T> {
151    /// # Safety
152    ///
153    /// The caller must ensure that `inner.position()` fits within the backing
154    /// vector's capacity, that any gap before that position has already been
155    /// initialized, and that all writes through this writer stay within the
156    /// trusted window reserved by `as_trusted_for`.
157    const unsafe fn new(inner: &'a mut Cursor<T>) -> Self {
158        Self { inner }
159    }
160}
161
162impl<'a, T> Writer for CursorVecUnchecked<'a, T>
163where
164    T: AsMut<Vec<u8>>,
165{
166    #[inline]
167    fn write(&mut self, src: &[u8]) -> WriteResult<()> {
168        let cur_pos = self.inner.position();
169        let inner = self.inner.get_mut().as_mut();
170        let cur_len = inner.len();
171        // `cursor_vec_as_trusted_for` checked that the trusted window end fits in
172        // `usize`, and the trusted contract requires `next_pos` to stay within that
173        // window.
174        #[expect(clippy::arithmetic_side_effects)]
175        let next_pos = cur_pos + src.len() as u64;
176
177        // SAFETY:
178        // - `cursor_vec_as_trusted_for` ensured sufficient capacity for the trusted window before
179        //   constructing this writer.
180        // - The trusted-writer contract requires all writes through this writer to stay
181        //   within that reserved window.
182        // - Given Rust's aliasing rules, we can assume that `src` does not overlap with
183        //   the internal buffer.
184        unsafe {
185            core::ptr::copy_nonoverlapping(
186                src.as_ptr(),
187                // `cursor_vec_as_trusted_for` checked that the initial cursor
188                // position plus the trusted window fits in `usize`. The trusted-writer
189                // contract requires all writes through this writer to stay within that
190                // reserved window, so `cur_pos` also fits in `usize`.
191                inner.as_mut_ptr().add(cur_pos as usize),
192                src.len(),
193            );
194        }
195
196        if next_pos > cur_len as u64 {
197            // SAFETY: any gap before the trusted window was initialized before
198            // constructing this writer, and the copy above initialized `cur_pos..next_pos`.
199            unsafe {
200                // `cursor_vec_as_trusted_for` checked that the trusted window end fits in
201                // `usize`, and the trusted contract requires `next_pos` to stay within that
202                // window.
203                inner.set_len(next_pos as usize)
204            }
205        }
206
207        self.inner.set_position(next_pos);
208
209        Ok(())
210    }
211}
212
213impl Writer for Cursor<Vec<u8>> {
214    #[inline]
215    fn write(&mut self, src: &[u8]) -> WriteResult<()> {
216        Ok(self.write_all(src)?)
217    }
218
219    #[inline]
220    fn finish(&mut self) -> WriteResult<()> {
221        Ok(self.flush()?)
222    }
223
224    #[inline]
225    unsafe fn as_trusted_for(&mut self, n_bytes: usize) -> WriteResult<impl Writer> {
226        cursor_vec_as_trusted_for(self, n_bytes)
227    }
228}
229
230impl Writer for Cursor<&mut Vec<u8>> {
231    #[inline]
232    fn write(&mut self, src: &[u8]) -> WriteResult<()> {
233        Ok(self.write_all(src)?)
234    }
235
236    #[inline]
237    fn finish(&mut self) -> WriteResult<()> {
238        Ok(self.flush()?)
239    }
240
241    #[inline]
242    unsafe fn as_trusted_for(&mut self, n_bytes: usize) -> WriteResult<impl Writer> {
243        cursor_vec_as_trusted_for(self, n_bytes)
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use {
250        super::*,
251        crate::{
252            io::WriteError,
253            serde::{Serialize, serialize, serialized_size},
254        },
255    };
256
257    const MAGIC: u64 = 0xdeadbeef_cafebabe;
258    const DATA: &[(u32, bool, &u64)] = &[
259        (1u32, false, &MAGIC),
260        (2u32, true, &MAGIC),
261        (3u32, false, &MAGIC),
262    ];
263
264    fn assert_serializes_data(mut writer: impl Writer) {
265        <[(u32, bool, &u64)]>::serialize_into(writer.by_ref(), DATA).unwrap();
266        writer.finish().unwrap();
267    }
268
269    #[test]
270    fn write_adapter_serialize_tuples() {
271        let mut buf = Vec::new();
272        assert_serializes_data(WriteAdapter::new(&mut buf));
273        assert_eq!(buf, serialize(DATA).unwrap());
274    }
275
276    #[test]
277    fn buf_writer_serialize_tuples() {
278        let mut buf = Vec::new();
279        assert_serializes_data(BufWriter::new(&mut buf));
280        assert_eq!(buf, serialize(DATA).unwrap());
281    }
282
283    #[test]
284    fn cursor_vec_writer_serialize_tuples() {
285        let mut buf = Cursor::new(Vec::new());
286        assert_serializes_data(&mut buf);
287        assert_eq!(buf.into_inner(), serialize(DATA).unwrap());
288    }
289
290    #[test]
291    fn cursor_slice_writer_serialize_tuples() {
292        let size = serialized_size(DATA).unwrap() as usize;
293        let mut buf = Cursor::new(vec![0; size].into_boxed_slice());
294        assert_serializes_data(&mut buf);
295        assert_eq!(buf.into_inner().as_ref(), serialize(DATA).unwrap());
296    }
297
298    fn write_trusted(writer: &mut impl Writer, bytes: &[u8]) {
299        let mut trusted = unsafe { writer.as_trusted_for(bytes.len()) }.unwrap();
300        trusted.write(bytes).unwrap();
301        trusted.finish().unwrap();
302    }
303
304    macro_rules! with_vec_cursors {
305        ($inner:expr, |$reader: ident| $body: block) => {{
306            {
307                let mut $reader = Cursor::new($inner.clone());
308                $body
309            }
310            {
311                let mut $reader = Cursor::new(&mut $inner);
312                $body
313            }
314        }};
315    }
316
317    #[test]
318    fn cursor_vec_trusted_append_with_spare_capacity() {
319        let mut inner = Vec::with_capacity(8);
320        with_vec_cursors!(inner, |cursor| {
321            cursor.write_all(&[1, 2, 3]).unwrap();
322
323            write_trusted(&mut cursor, &[4, 5]);
324            cursor.finish().unwrap();
325
326            assert_eq!(&*cursor.into_inner(), &vec![1, 2, 3, 4, 5]);
327        });
328    }
329
330    #[test]
331    fn cursor_vec_trusted_overwrite_then_extend() {
332        let mut inner = vec![1, 2, 3, 4];
333        with_vec_cursors!(inner, |cursor| {
334            cursor.set_position(2);
335
336            write_trusted(&mut cursor, &[9, 8, 7, 6]);
337            cursor.finish().unwrap();
338
339            assert_eq!(&cursor.into_inner()[..], &vec![1, 2, 9, 8, 7, 6]);
340        });
341    }
342
343    #[test]
344    fn cursor_vec_trusted_overwrite_preserves_tail() {
345        let mut inner = vec![1, 2, 3, 4, 5, 6];
346        with_vec_cursors!(inner, |cursor| {
347            cursor.set_position(2);
348
349            write_trusted(&mut cursor, &[9, 8]);
350            cursor.finish().unwrap();
351
352            assert_eq!(&cursor.into_inner()[..], &vec![1, 2, 9, 8, 5, 6]);
353        });
354    }
355
356    #[test]
357    fn cursor_vec_trusted_zero_fills_gap() {
358        let mut inner = Vec::with_capacity(16);
359        inner.extend_from_slice(&[1, 2, 3]);
360        with_vec_cursors!(inner, |cursor| {
361            cursor.set_position(6);
362
363            write_trusted(&mut cursor, &[9, 10]);
364            cursor.finish().unwrap();
365
366            assert_eq!(&cursor.into_inner()[..], &vec![1, 2, 3, 0, 0, 0, 9, 10]);
367        });
368    }
369
370    #[test]
371    fn cursor_vec_trusted_does_not_extend_len_before_write() {
372        let mut inner = Vec::with_capacity(16);
373        with_vec_cursors!(inner, |cursor| {
374            {
375                let _trusted = unsafe { cursor.as_trusted_for(8) }.unwrap();
376            }
377
378            cursor.finish().unwrap();
379
380            assert_eq!(cursor.into_inner().len(), 0);
381        });
382    }
383
384    macro_rules! with_slice_cursors {
385        ($inner:expr, |$reader: ident| $body: block) => {{
386            {
387                let mut $reader = Cursor::new($inner);
388                $body
389            }
390            {
391                let mut inner = $inner;
392                let mut $reader = Cursor::new(&mut inner[..]);
393                $body
394            }
395            {
396                let mut $reader = Cursor::new(Box::from($inner));
397                $body
398            }
399        }};
400    }
401
402    #[test]
403    fn cursor_mut_slice_trusted_writes_in_bounds() {
404        with_slice_cursors!([1, 2, 3, 4, 5], |cursor| {
405            let pos = {
406                cursor.set_position(1);
407
408                write_trusted(&mut cursor, &[9, 8, 7]);
409                cursor.finish().unwrap();
410                cursor.position()
411            };
412            assert_eq!(&cursor.get_ref()[..], &[1, 9, 8, 7, 5]);
413            assert_eq!(pos, 4);
414        });
415    }
416
417    #[test]
418    fn cursor_slice_trusted_out_of_bounds_errors() {
419        with_slice_cursors!([1, 2, 3], |cursor| {
420            cursor.set_position(2);
421            let result = unsafe { cursor.as_trusted_for(2) };
422            assert!(matches!(result, Err(WriteError::WriteSizeLimit(2))));
423        });
424    }
425}