messagepack_core/
io.rs

1//! Minimal write abstraction used by encoders.
2
3/// Minimal `Write`‑like trait used by encoders to avoid committing to a
4/// specific I/O model.
5pub trait IoWrite {
6    /// Error type produced by the writer.
7    type Error: core::error::Error;
8    /// Write all bytes from `buf`.
9    fn write(&mut self, buf: &[u8]) -> Result<(), Self::Error>;
10}
11
12/// `SliceWriter` Error
13#[derive(Debug, Copy, Clone, PartialOrd, Ord, PartialEq, Eq)]
14pub enum WError {
15    /// buffer is full
16    BufferFull,
17}
18
19impl core::fmt::Display for WError {
20    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
21        match self {
22            WError::BufferFull => write!(f, "Buffer is full"),
23        }
24    }
25}
26
27impl core::error::Error for WError {}
28
29/// Simple writer that writes into a mutable byte slice.
30pub struct SliceWriter<'a> {
31    buf: &'a mut [u8],
32    cursor: usize,
33}
34
35impl<'a> SliceWriter<'a> {
36    /// Create a new writer over the given buffer.
37    pub fn from_slice(buf: &'a mut [u8]) -> Self {
38        Self { buf, cursor: 0 }
39    }
40
41    fn len(&self) -> usize {
42        self.buf.len() - self.cursor
43    }
44}
45
46impl IoWrite for SliceWriter<'_> {
47    type Error = WError;
48
49    fn write(&mut self, buf: &[u8]) -> Result<(), Self::Error> {
50        if self.len() >= buf.len() {
51            let to = &mut self.buf[self.cursor..self.cursor + buf.len()];
52            to.copy_from_slice(buf);
53            self.cursor += buf.len();
54            Ok(())
55        } else {
56            Err(WError::BufferFull)
57        }
58    }
59}
60
61#[cfg(not(any(test, feature = "std")))]
62impl IoWrite for &mut [u8] {
63    type Error = WError;
64
65    fn write(&mut self, buf: &[u8]) -> Result<(), Self::Error> {
66        SliceWriter::from_slice(self).write(buf)
67    }
68}
69
70/// Simple writer that writes into a `&mut Vec<u8>`.
71#[cfg(feature = "alloc")]
72pub struct VecRefWriter<'a> {
73    vec: &'a mut alloc::vec::Vec<u8>,
74}
75
76#[cfg(feature = "alloc")]
77impl<'a> VecRefWriter<'a> {
78    /// Create a new writer
79    pub fn new(vec: &'a mut alloc::vec::Vec<u8>) -> Self {
80        Self { vec }
81    }
82}
83
84#[cfg(feature = "alloc")]
85impl IoWrite for VecRefWriter<'_> {
86    type Error = core::convert::Infallible;
87
88    fn write(&mut self, buf: &[u8]) -> Result<(), Self::Error> {
89        self.vec.extend_from_slice(buf);
90        Ok(())
91    }
92}
93
94#[cfg(all(not(test), feature = "alloc", not(feature = "std")))]
95impl IoWrite for alloc::vec::Vec<u8> {
96    type Error = core::convert::Infallible;
97
98    fn write(&mut self, buf: &[u8]) -> Result<(), Self::Error> {
99        VecRefWriter::new(self).write(buf)
100    }
101}
102
103/// Simple writer that writes into a `Vec<u8>`.
104#[cfg(feature = "alloc")]
105pub struct VecWriter {
106    vec: alloc::vec::Vec<u8>,
107}
108
109#[cfg(feature = "alloc")]
110impl VecWriter {
111    /// Create a new writer
112    pub fn new() -> Self {
113        Self {
114            vec: alloc::vec::Vec::new(),
115        }
116    }
117    /// Get the inner vector
118    pub fn into_vec(self) -> alloc::vec::Vec<u8> {
119        self.vec
120    }
121}
122
123#[cfg(feature = "alloc")]
124impl Default for VecWriter {
125    fn default() -> Self {
126        Self::new()
127    }
128}
129
130#[cfg(feature = "alloc")]
131impl IoWrite for VecWriter {
132    type Error = core::convert::Infallible;
133    fn write(&mut self, buf: &[u8]) -> Result<(), Self::Error> {
134        self.vec.extend_from_slice(buf);
135        Ok(())
136    }
137}
138
139#[cfg(any(test, feature = "std"))]
140impl<W> IoWrite for W
141where
142    W: std::io::Write,
143{
144    type Error = std::io::Error;
145
146    fn write(&mut self, buf: &[u8]) -> Result<(), Self::Error> {
147        self.write_all(buf)
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[test]
156    #[should_panic]
157    fn buffer_full() {
158        let buf: &mut [u8] = &mut [0u8];
159        let mut writer = SliceWriter::from_slice(buf);
160        writer.write(&[1, 2]).unwrap();
161    }
162}