object/write/
util.rs

1use alloc::vec::Vec;
2#[cfg(feature = "std")]
3use std::{io, mem};
4
5use crate::pod::{bytes_of, bytes_of_slice, Pod};
6
7/// Trait for writable buffer.
8#[allow(clippy::len_without_is_empty)]
9pub trait WritableBuffer {
10    /// Returns position/offset for data to be written at.
11    ///
12    /// Should only be used in debug assertions
13    fn len(&self) -> usize;
14
15    /// Reserves specified number of bytes in the buffer.
16    ///
17    /// This will be called exactly once before writing anything to the buffer,
18    /// and the given size is the exact total number of bytes that will be written.
19    fn reserve(&mut self, size: usize) -> Result<(), ()>;
20
21    /// Writes zero bytes at the end of the buffer until the buffer
22    /// has the specified length.
23    fn resize(&mut self, new_len: usize);
24
25    /// Writes the specified slice of bytes at the end of the buffer.
26    fn write_bytes(&mut self, val: &[u8]);
27
28    /// Writes the specified `Pod` type at the end of the buffer.
29    fn write_pod<T: Pod>(&mut self, val: &T)
30    where
31        Self: Sized,
32    {
33        self.write_bytes(bytes_of(val))
34    }
35
36    /// Writes the specified `Pod` slice at the end of the buffer.
37    fn write_pod_slice<T: Pod>(&mut self, val: &[T])
38    where
39        Self: Sized,
40    {
41        self.write_bytes(bytes_of_slice(val))
42    }
43}
44
45impl<'a> dyn WritableBuffer + 'a {
46    /// Writes the specified `Pod` type at the end of the buffer.
47    pub fn write<T: Pod>(&mut self, val: &T) {
48        self.write_bytes(bytes_of(val))
49    }
50
51    /// Writes the specified `Pod` slice at the end of the buffer.
52    pub fn write_slice<T: Pod>(&mut self, val: &[T]) {
53        self.write_bytes(bytes_of_slice(val))
54    }
55}
56
57impl WritableBuffer for Vec<u8> {
58    #[inline]
59    fn len(&self) -> usize {
60        self.len()
61    }
62
63    #[inline]
64    fn reserve(&mut self, size: usize) -> Result<(), ()> {
65        debug_assert!(self.is_empty());
66        self.reserve(size);
67        Ok(())
68    }
69
70    #[inline]
71    fn resize(&mut self, new_len: usize) {
72        debug_assert!(new_len >= self.len());
73        self.resize(new_len, 0);
74    }
75
76    #[inline]
77    fn write_bytes(&mut self, val: &[u8]) {
78        debug_assert!(self.len() + val.len() <= self.capacity());
79        self.extend_from_slice(val)
80    }
81}
82
83/// A [`WritableBuffer`] that streams data to a [`Write`](std::io::Write) implementation.
84///
85/// [`Self::result`] must be called to determine if an I/O error occurred during writing.
86/// Alternatively, [`Self::flush`] will both check for errors and flush.
87///
88/// It is advisable to use a buffered writer like [`BufWriter`](std::io::BufWriter)
89/// instead of an unbuffered writer like [`File`](std::fs::File).
90#[cfg(feature = "std")]
91#[derive(Debug)]
92pub struct StreamingBuffer<W> {
93    writer: W,
94    len: usize,
95    result: Result<(), io::Error>,
96}
97
98#[cfg(feature = "std")]
99impl<W> StreamingBuffer<W> {
100    /// Create a new `StreamingBuffer` backed by the given writer.
101    pub fn new(writer: W) -> Self {
102        StreamingBuffer {
103            writer,
104            len: 0,
105            result: Ok(()),
106        }
107    }
108
109    /// Unwraps this [`StreamingBuffer`] giving back the original writer.
110    pub fn into_inner(self) -> W {
111        self.writer
112    }
113
114    /// Returns any error that occurred during writing.
115    pub fn result(&mut self) -> Result<(), io::Error> {
116        mem::replace(&mut self.result, Ok(()))
117    }
118}
119
120#[cfg(feature = "std")]
121impl<W: io::Write> StreamingBuffer<W> {
122    /// Flushes after first checking if any error previously occurred during writing.
123    pub fn flush(&mut self) -> Result<(), io::Error> {
124        self.result()?;
125        self.writer.flush()
126    }
127}
128
129#[cfg(feature = "std")]
130impl<W: io::Write> WritableBuffer for StreamingBuffer<W> {
131    #[inline]
132    fn len(&self) -> usize {
133        self.len
134    }
135
136    #[inline]
137    fn reserve(&mut self, _size: usize) -> Result<(), ()> {
138        Ok(())
139    }
140
141    #[inline]
142    fn resize(&mut self, new_len: usize) {
143        debug_assert!(self.len <= new_len);
144        while self.len < new_len {
145            let write_amt = (new_len - self.len - 1) % 1024 + 1;
146            self.write_bytes(&[0; 1024][..write_amt]);
147        }
148    }
149
150    #[inline]
151    fn write_bytes(&mut self, val: &[u8]) {
152        if self.result.is_ok() {
153            self.result = self.writer.write_all(val);
154        }
155        self.len += val.len();
156    }
157}
158
159/// A trait for mutable byte slices.
160///
161/// It provides convenience methods for `Pod` types.
162pub(crate) trait BytesMut {
163    fn write_at<T: Pod>(self, offset: usize, val: &T) -> Result<(), ()>;
164}
165
166impl<'a> BytesMut for &'a mut [u8] {
167    #[inline]
168    fn write_at<T: Pod>(self, offset: usize, val: &T) -> Result<(), ()> {
169        let src = bytes_of(val);
170        let dest = self.get_mut(offset..).ok_or(())?;
171        let dest = dest.get_mut(..src.len()).ok_or(())?;
172        dest.copy_from_slice(src);
173        Ok(())
174    }
175}
176
177/// Write an unsigned number using the LEB128 encoding to a buffer.
178///
179/// Returns the number of bytes written.
180#[allow(dead_code)]
181pub(crate) fn write_uleb128(buf: &mut Vec<u8>, mut val: u64) -> usize {
182    let mut len = 0;
183    loop {
184        let mut byte = (val & 0x7f) as u8;
185        val >>= 7;
186        let done = val == 0;
187        if !done {
188            byte |= 0x80;
189        }
190
191        buf.push(byte);
192        len += 1;
193
194        if done {
195            return len;
196        }
197    }
198}
199
200/// Write a signed number using the LEB128 encoding to a buffer.
201///
202/// Returns the number of bytes written.
203#[allow(dead_code)]
204pub(crate) fn write_sleb128(buf: &mut Vec<u8>, mut val: i64) -> usize {
205    let mut len = 0;
206    loop {
207        let mut byte = val as u8;
208        // Keep the sign bit for testing
209        val >>= 6;
210        let done = val == 0 || val == -1;
211        if done {
212            byte &= !0x80;
213        } else {
214            // Remove the sign bit
215            val >>= 1;
216            byte |= 0x80;
217        }
218
219        buf.push(byte);
220        len += 1;
221
222        if done {
223            return len;
224        }
225    }
226}
227
228pub(crate) fn align(offset: usize, size: usize) -> usize {
229    (offset + (size - 1)) & !(size - 1)
230}
231
232#[allow(dead_code)]
233pub(crate) fn align_u32(offset: u32, size: u32) -> u32 {
234    (offset + (size - 1)) & !(size - 1)
235}
236
237#[allow(dead_code)]
238pub(crate) fn align_u64(offset: u64, size: u64) -> u64 {
239    (offset + (size - 1)) & !(size - 1)
240}
241
242pub(crate) fn write_align(buffer: &mut dyn WritableBuffer, size: usize) {
243    let new_len = align(buffer.len(), size);
244    buffer.resize(new_len);
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    #[test]
252    fn bytes_mut() {
253        let data = vec![0x01, 0x23, 0x45, 0x67];
254
255        let mut bytes = data.clone();
256        bytes.extend_from_slice(bytes_of(&u16::to_be(0x89ab)));
257        assert_eq!(bytes, [0x01, 0x23, 0x45, 0x67, 0x89, 0xab]);
258
259        let mut bytes = data.clone();
260        assert_eq!(bytes.write_at(0, &u16::to_be(0x89ab)), Ok(()));
261        assert_eq!(bytes, [0x89, 0xab, 0x45, 0x67]);
262
263        let mut bytes = data.clone();
264        assert_eq!(bytes.write_at(2, &u16::to_be(0x89ab)), Ok(()));
265        assert_eq!(bytes, [0x01, 0x23, 0x89, 0xab]);
266
267        assert_eq!(bytes.write_at(3, &u16::to_be(0x89ab)), Err(()));
268        assert_eq!(bytes.write_at(4, &u16::to_be(0x89ab)), Err(()));
269        assert_eq!([].write_at(0, &u32::to_be(0x89ab)), Err(()));
270    }
271}
object/write/util.rs

object/write/
util.rs
