diff options
Diffstat (limited to 'third_party/rust/object/src/write/util.rs')
-rw-r--r-- | third_party/rust/object/src/write/util.rs | 260 |
1 files changed, 260 insertions, 0 deletions
diff --git a/third_party/rust/object/src/write/util.rs b/third_party/rust/object/src/write/util.rs new file mode 100644 index 0000000000..b05b14d927 --- /dev/null +++ b/third_party/rust/object/src/write/util.rs @@ -0,0 +1,260 @@ +use alloc::vec::Vec; +#[cfg(feature = "std")] +use std::{io, mem}; + +use crate::pod::{bytes_of, bytes_of_slice, Pod}; + +/// Trait for writable buffer. +#[allow(clippy::len_without_is_empty)] +pub trait WritableBuffer { + /// Returns position/offset for data to be written at. + /// + /// Should only be used in debug assertions + fn len(&self) -> usize; + + /// Reserves specified number of bytes in the buffer. + /// + /// This will be called exactly once before writing anything to the buffer, + /// and the given size is the exact total number of bytes that will be written. + fn reserve(&mut self, size: usize) -> Result<(), ()>; + + /// Writes zero bytes at the end of the buffer until the buffer + /// has the specified length. + fn resize(&mut self, new_len: usize); + + /// Writes the specified slice of bytes at the end of the buffer. + fn write_bytes(&mut self, val: &[u8]); + + /// Writes the specified `Pod` type at the end of the buffer. + fn write_pod<T: Pod>(&mut self, val: &T) + where + Self: Sized, + { + self.write_bytes(bytes_of(val)) + } + + /// Writes the specified `Pod` slice at the end of the buffer. + fn write_pod_slice<T: Pod>(&mut self, val: &[T]) + where + Self: Sized, + { + self.write_bytes(bytes_of_slice(val)) + } +} + +impl<'a> dyn WritableBuffer + 'a { + /// Writes the specified `Pod` type at the end of the buffer. + pub fn write<T: Pod>(&mut self, val: &T) { + self.write_bytes(bytes_of(val)) + } + + /// Writes the specified `Pod` slice at the end of the buffer. + pub fn write_slice<T: Pod>(&mut self, val: &[T]) { + self.write_bytes(bytes_of_slice(val)) + } +} + +impl WritableBuffer for Vec<u8> { + #[inline] + fn len(&self) -> usize { + self.len() + } + + #[inline] + fn reserve(&mut self, size: usize) -> Result<(), ()> { + debug_assert!(self.is_empty()); + self.reserve(size); + Ok(()) + } + + #[inline] + fn resize(&mut self, new_len: usize) { + debug_assert!(new_len >= self.len()); + self.resize(new_len, 0); + } + + #[inline] + fn write_bytes(&mut self, val: &[u8]) { + debug_assert!(self.len() + val.len() <= self.capacity()); + self.extend_from_slice(val) + } +} + +/// A [`WritableBuffer`] that streams data to a [`Write`](std::io::Write) implementation. +/// +/// [`Self::result`] must be called to determine if an I/O error occurred during writing. +/// +/// It is advisable to use a buffered writer like [`BufWriter`](std::io::BufWriter) +/// instead of an unbuffered writer like [`File`](std::fs::File). +#[cfg(feature = "std")] +#[derive(Debug)] +pub struct StreamingBuffer<W> { + writer: W, + len: usize, + result: Result<(), io::Error>, +} + +#[cfg(feature = "std")] +impl<W> StreamingBuffer<W> { + /// Create a new `StreamingBuffer` backed by the given writer. + pub fn new(writer: W) -> Self { + StreamingBuffer { + writer, + len: 0, + result: Ok(()), + } + } + + /// Unwraps this [`StreamingBuffer`] giving back the original writer. + pub fn into_inner(self) -> W { + self.writer + } + + /// Returns any error that occurred during writing. + pub fn result(&mut self) -> Result<(), io::Error> { + mem::replace(&mut self.result, Ok(())) + } +} + +#[cfg(feature = "std")] +impl<W: io::Write> WritableBuffer for StreamingBuffer<W> { + #[inline] + fn len(&self) -> usize { + self.len + } + + #[inline] + fn reserve(&mut self, _size: usize) -> Result<(), ()> { + Ok(()) + } + + #[inline] + fn resize(&mut self, new_len: usize) { + debug_assert!(self.len <= new_len); + while self.len < new_len { + let write_amt = (new_len - self.len - 1) % 1024 + 1; + self.write_bytes(&[0; 1024][..write_amt]); + } + } + + #[inline] + fn write_bytes(&mut self, val: &[u8]) { + if self.result.is_ok() { + self.result = self.writer.write_all(val); + } + self.len += val.len(); + } +} + +/// A trait for mutable byte slices. +/// +/// It provides convenience methods for `Pod` types. +pub(crate) trait BytesMut { + fn write_at<T: Pod>(self, offset: usize, val: &T) -> Result<(), ()>; +} + +impl<'a> BytesMut for &'a mut [u8] { + #[inline] + fn write_at<T: Pod>(self, offset: usize, val: &T) -> Result<(), ()> { + let src = bytes_of(val); + let dest = self.get_mut(offset..).ok_or(())?; + let dest = dest.get_mut(..src.len()).ok_or(())?; + dest.copy_from_slice(src); + Ok(()) + } +} + +/// Write an unsigned number using the LEB128 encoding to a buffer. +/// +/// Returns the number of bytes written. +pub(crate) fn write_uleb128(buf: &mut Vec<u8>, mut val: u64) -> usize { + let mut len = 0; + loop { + let mut byte = (val & 0x7f) as u8; + val >>= 7; + let done = val == 0; + if !done { + byte |= 0x80; + } + + buf.push(byte); + len += 1; + + if done { + return len; + } + } +} + +/// Write a signed number using the LEB128 encoding to a buffer. +/// +/// Returns the number of bytes written. +#[allow(dead_code)] +pub(crate) fn write_sleb128<W>(buf: &mut Vec<u8>, mut val: i64) -> usize { + let mut len = 0; + loop { + let mut byte = val as u8; + // Keep the sign bit for testing + val >>= 6; + let done = val == 0 || val == -1; + if done { + byte &= !0x80; + } else { + // Remove the sign bit + val >>= 1; + byte |= 0x80; + } + + buf.push(byte); + len += 1; + + if done { + return len; + } + } +} + +pub(crate) fn align(offset: usize, size: usize) -> usize { + (offset + (size - 1)) & !(size - 1) +} + +#[allow(dead_code)] +pub(crate) fn align_u32(offset: u32, size: u32) -> u32 { + (offset + (size - 1)) & !(size - 1) +} + +#[allow(dead_code)] +pub(crate) fn align_u64(offset: u64, size: u64) -> u64 { + (offset + (size - 1)) & !(size - 1) +} + +pub(crate) fn write_align(buffer: &mut dyn WritableBuffer, size: usize) { + let new_len = align(buffer.len(), size); + buffer.resize(new_len); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bytes_mut() { + let data = vec![0x01, 0x23, 0x45, 0x67]; + + let mut bytes = data.clone(); + bytes.extend_from_slice(bytes_of(&u16::to_be(0x89ab))); + assert_eq!(bytes, [0x01, 0x23, 0x45, 0x67, 0x89, 0xab]); + + let mut bytes = data.clone(); + assert_eq!(bytes.write_at(0, &u16::to_be(0x89ab)), Ok(())); + assert_eq!(bytes, [0x89, 0xab, 0x45, 0x67]); + + let mut bytes = data.clone(); + assert_eq!(bytes.write_at(2, &u16::to_be(0x89ab)), Ok(())); + assert_eq!(bytes, [0x01, 0x23, 0x89, 0xab]); + + assert_eq!(bytes.write_at(3, &u16::to_be(0x89ab)), Err(())); + assert_eq!(bytes.write_at(4, &u16::to_be(0x89ab)), Err(())); + assert_eq!(vec![].write_at(0, &u32::to_be(0x89ab)), Err(())); + } +} |