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(val: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(val: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///
87/// It is advisable to use a buffered writer like [`BufWriter`](std::io::BufWriter)
88/// instead of an unbuffered writer like [`File`](std::fs::File).
89#[cfg(feature = "std")]
90#[derive(Debug)]
91pub struct StreamingBuffer<W> {
92 writer: W,
93 len: usize,
94 result: Result<(), io::Error>,
95}
96
97#[cfg(feature = "std")]
98impl<W> StreamingBuffer<W> {
99 /// Create a new `StreamingBuffer` backed by the given writer.
100 pub fn new(writer: W) -> Self {
101 StreamingBuffer {
102 writer,
103 len: 0,
104 result: Ok(()),
105 }
106 }
107
108 /// Unwraps this [`StreamingBuffer`] giving back the original writer.
109 pub fn into_inner(self) -> W {
110 self.writer
111 }
112
113 /// Returns any error that occurred during writing.
114 pub fn result(&mut self) -> Result<(), io::Error> {
115 mem::replace(&mut self.result, src:Ok(()))
116 }
117}
118
119#[cfg(feature = "std")]
120impl<W: io::Write> WritableBuffer for StreamingBuffer<W> {
121 #[inline]
122 fn len(&self) -> usize {
123 self.len
124 }
125
126 #[inline]
127 fn reserve(&mut self, _size: usize) -> Result<(), ()> {
128 Ok(())
129 }
130
131 #[inline]
132 fn resize(&mut self, new_len: usize) {
133 debug_assert!(self.len <= new_len);
134 while self.len < new_len {
135 let write_amt = (new_len - self.len - 1) % 1024 + 1;
136 self.write_bytes(&[0; 1024][..write_amt]);
137 }
138 }
139
140 #[inline]
141 fn write_bytes(&mut self, val: &[u8]) {
142 if self.result.is_ok() {
143 self.result = self.writer.write_all(val);
144 }
145 self.len += val.len();
146 }
147}
148
149/// A trait for mutable byte slices.
150///
151/// It provides convenience methods for `Pod` types.
152pub(crate) trait BytesMut {
153 fn write_at<T: Pod>(self, offset: usize, val: &T) -> Result<(), ()>;
154}
155
156impl<'a> BytesMut for &'a mut [u8] {
157 #[inline]
158 fn write_at<T: Pod>(self, offset: usize, val: &T) -> Result<(), ()> {
159 let src: &[u8] = bytes_of(val);
160 let dest: &mut [u8] = self.get_mut(offset..).ok_or(())?;
161 let dest: &mut [u8] = dest.get_mut(..src.len()).ok_or(())?;
162 dest.copy_from_slice(src);
163 Ok(())
164 }
165}
166
167/// Write an unsigned number using the LEB128 encoding to a buffer.
168///
169/// Returns the number of bytes written.
170pub(crate) fn write_uleb128(buf: &mut Vec<u8>, mut val: u64) -> usize {
171 let mut len: usize = 0;
172 loop {
173 let mut byte: u8 = (val & 0x7f) as u8;
174 val >>= 7;
175 let done: bool = val == 0;
176 if !done {
177 byte |= 0x80;
178 }
179
180 buf.push(byte);
181 len += 1;
182
183 if done {
184 return len;
185 }
186 }
187}
188
189/// Write a signed number using the LEB128 encoding to a buffer.
190///
191/// Returns the number of bytes written.
192#[allow(dead_code)]
193pub(crate) fn write_sleb128(buf: &mut Vec<u8>, mut val: i64) -> usize {
194 let mut len: usize = 0;
195 loop {
196 let mut byte: u8 = val as u8;
197 // Keep the sign bit for testing
198 val >>= 6;
199 let done: bool = val == 0 || val == -1;
200 if done {
201 byte &= !0x80;
202 } else {
203 // Remove the sign bit
204 val >>= 1;
205 byte |= 0x80;
206 }
207
208 buf.push(byte);
209 len += 1;
210
211 if done {
212 return len;
213 }
214 }
215}
216
217pub(crate) fn align(offset: usize, size: usize) -> usize {
218 (offset + (size - 1)) & !(size - 1)
219}
220
221#[allow(dead_code)]
222pub(crate) fn align_u32(offset: u32, size: u32) -> u32 {
223 (offset + (size - 1)) & !(size - 1)
224}
225
226#[allow(dead_code)]
227pub(crate) fn align_u64(offset: u64, size: u64) -> u64 {
228 (offset + (size - 1)) & !(size - 1)
229}
230
231pub(crate) fn write_align(buffer: &mut dyn WritableBuffer, size: usize) {
232 let new_len: usize = align(offset:buffer.len(), size);
233 buffer.resize(new_len);
234}
235
236#[cfg(test)]
237mod tests {
238 use super::*;
239
240 #[test]
241 fn bytes_mut() {
242 let data = vec![0x01, 0x23, 0x45, 0x67];
243
244 let mut bytes = data.clone();
245 bytes.extend_from_slice(bytes_of(&u16::to_be(0x89ab)));
246 assert_eq!(bytes, [0x01, 0x23, 0x45, 0x67, 0x89, 0xab]);
247
248 let mut bytes = data.clone();
249 assert_eq!(bytes.write_at(0, &u16::to_be(0x89ab)), Ok(()));
250 assert_eq!(bytes, [0x89, 0xab, 0x45, 0x67]);
251
252 let mut bytes = data.clone();
253 assert_eq!(bytes.write_at(2, &u16::to_be(0x89ab)), Ok(()));
254 assert_eq!(bytes, [0x01, 0x23, 0x89, 0xab]);
255
256 assert_eq!(bytes.write_at(3, &u16::to_be(0x89ab)), Err(()));
257 assert_eq!(bytes.write_at(4, &u16::to_be(0x89ab)), Err(()));
258 assert_eq!(vec![].write_at(0, &u32::to_be(0x89ab)), Err(()));
259 }
260}
261