use std::alloc::{Layout, alloc, dealloc, handle_alloc_error, realloc};
use std::borrow::Cow;
use std::fmt;
use std::marker::PhantomData;
use std::mem::{ManuallyDrop, MaybeUninit};
use std::ptr::{self, NonNull};
#[repr(C)]
pub struct BytesWriter<'a> {
data: *mut u8,
len: usize,
capacity: usize,
backing: Backing<'a>,
}
impl Default for BytesWriter<'_> {
fn default() -> Self {
Self::new()
}
}
impl<'a> From<&'a mut [MaybeUninit<u8>]> for BytesWriter<'a> {
fn from(value: &'a mut [MaybeUninit<u8>]) -> Self {
BytesWriter {
capacity: value.len(),
data: value.as_mut_ptr().cast(),
len: 0,
backing: Backing::Borrowed {
marker: PhantomData,
},
}
}
}
impl<'a> From<&'a mut Vec<u8>> for BytesWriter<'a> {
fn from(value: &'a mut Vec<u8>) -> Self {
BytesWriter {
data: value.as_mut_ptr(),
len: value.len(),
capacity: value.capacity(),
backing: Backing::Vec {
bytes: NonNull::from(value),
offset: 0,
marker: PhantomData,
},
}
}
}
impl Drop for BytesWriter<'_> {
fn drop(&mut self) {
match self.backing {
Backing::Owned | Backing::Write { .. } => {
unsafe {
if self.capacity != 0 {
let layout = Layout::from_size_align_unchecked(self.capacity, 1);
dealloc(self.data, layout);
}
}
}
Backing::Vec {
mut bytes, offset, ..
} => {
unsafe {
let vec = bytes.as_mut();
vec.set_len(offset + self.len);
}
}
Backing::Borrowed { .. } => (),
}
}
}
pub(crate) enum Backing<'a> {
Owned,
Vec {
bytes: NonNull<Vec<u8>>,
offset: usize,
marker: PhantomData<&'a ()>,
},
Borrowed {
marker: PhantomData<&'a ()>,
},
Write {
written: usize,
error: Option<std::io::Error>,
writer: &'a mut (dyn std::io::Write + Send),
},
}
impl<'a> BytesWriter<'a> {
pub fn new_writer(writer: &'a mut (dyn std::io::Write + Send)) -> BytesWriter<'a> {
BytesWriter {
data: safe_alloc(4096),
len: 0,
capacity: 4096,
backing: Backing::Write {
written: 0,
error: None,
writer,
},
}
}
pub(crate) fn from_vec_suffix(value: &'a mut Vec<u8>) -> Self {
let offset = value.len();
BytesWriter {
data: unsafe { value.as_mut_ptr().add(offset) },
len: 0,
capacity: value.capacity() - offset,
backing: Backing::Vec {
bytes: NonNull::from(value),
offset,
marker: PhantomData,
},
}
}
pub fn last_2(&mut self) -> Option<&mut [u8; 2]> {
if self.len < 2 {
return None;
}
Some(unsafe { &mut *self.data.add(self.len - 2).cast() })
}
pub fn last(&mut self) -> Option<&mut u8> {
if self.len == 0 {
return None;
}
Some(unsafe { &mut *self.data.add(self.len - 1) })
}
pub fn clear(&mut self) {
self.len = 0;
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub unsafe fn set_len(&mut self, new_len: usize) {
self.len = new_len;
}
pub fn saturting_pop(&mut self) {
self.len = self.len.saturating_sub(1);
}
#[inline]
pub fn new() -> BytesWriter<'a> {
BytesWriter {
data: NonNull::<u8>::dangling().as_ptr(),
len: 0,
capacity: 0,
backing: Backing::Owned,
}
}
pub fn with_capacity(capacity: usize) -> BytesWriter<'a> {
let data = safe_alloc(capacity);
BytesWriter {
data,
len: 0,
capacity,
backing: Backing::Owned,
}
}
pub fn buffer_slice(&self) -> &[u8] {
unsafe { std::slice::from_raw_parts(self.data, self.len) }
}
pub fn into_vec(self) -> Vec<u8> {
let this = ManuallyDrop::new(self);
match this.backing {
Backing::Vec {
mut bytes, offset, ..
} => {
unsafe {
let vec = bytes.as_mut();
vec.set_len(offset + this.len);
std::mem::take(vec)
}
}
Backing::Borrowed { .. } => this.buffer_slice().into(),
Backing::Write { .. } | Backing::Owned => {
unsafe { Vec::from_raw_parts(this.data, this.len, this.capacity) }
}
}
}
pub fn owned_into_vec(self) -> Vec<u8> {
let mut this = ManuallyDrop::new(self);
if let Backing::Owned = this.backing {
unsafe { Vec::from_raw_parts(this.data, this.len, this.capacity) }
} else {
unsafe { ManuallyDrop::drop(&mut this) };
panic!("Expected write buffer to backed by an owned allocation");
}
}
pub fn into_backed_with_extended_slice(self) -> &'a [u8] {
let mut this = ManuallyDrop::new(self);
let len = this.len;
if let Backing::Vec { bytes, offset, .. } = &mut this.backing {
let (data, start, suffix_len) = unsafe {
let bytes = bytes.as_mut();
let oldlen = bytes.len();
let end = *offset + len;
bytes.set_len(end);
let start = oldlen.min(end);
(bytes.as_mut_ptr(), start, end - start)
};
return unsafe { std::slice::from_raw_parts(data.add(start), suffix_len) };
}
unsafe { ManuallyDrop::drop(&mut this) };
panic!("Expected write buffer to backed by a Vec<u8>");
}
pub unsafe fn into_cow_utf8_unchecked(self) -> Cow<'a, str> {
let mut this = ManuallyDrop::new(self);
let data = this.data;
let len = this.len;
let capacity = this.capacity;
match &this.backing {
Backing::Owned => Cow::Owned(unsafe {
String::from_utf8_unchecked(Vec::from_raw_parts(data, len, capacity))
}),
Backing::Borrowed { .. } => Cow::Borrowed(unsafe {
std::str::from_utf8_unchecked(std::slice::from_raw_parts(data, len))
}),
_ => {
unsafe { ManuallyDrop::drop(&mut this) };
panic!("Expected Borrowed or owneded Instance");
}
}
}
pub fn into_cow(self) -> Cow<'a, [u8]> {
let mut this = ManuallyDrop::new(self);
let data = this.data;
let len = this.len;
let capacity = this.capacity;
match &this.backing {
Backing::Owned => Cow::Owned(unsafe { Vec::from_raw_parts(data, len, capacity) }),
Backing::Borrowed { .. } => {
Cow::Borrowed(unsafe { std::slice::from_raw_parts(data, len) })
}
_ => {
unsafe { ManuallyDrop::drop(&mut this) };
panic!("Expected Borrowed or owneded Instance");
}
}
}
pub fn into_write_finish(mut self) -> Result<usize, std::io::Error> {
self.flush();
match &mut self.backing {
Backing::Write { written, error, .. } => {
if let Some(error) = error.take() {
return Err(error);
}
Ok(*written)
}
_ => {
panic!("Expected Write Instance");
}
}
}
#[inline]
pub fn as_mut_ptr(&self) -> *mut u8 {
self.data
}
#[inline]
pub fn len(&self) -> usize {
self.len
}
#[inline]
pub fn capacity(&self) -> usize {
self.capacity
}
#[inline]
pub unsafe fn advance(&mut self, additional: usize) {
self.len += additional;
}
#[inline]
pub(crate) unsafe fn reserve_small(&mut self, size: usize) {
debug_assert!(size <= isize::MAX as usize);
if self.len + size <= self.capacity {
return;
}
self.reserve_internal(size, true);
}
pub unsafe fn push_as_bytes<T>(&mut self, data: &T) {
self.push_bytes(unsafe {
::std::slice::from_raw_parts(data as *const T as *const u8, ::std::mem::size_of::<T>())
});
}
#[inline]
pub fn push_bytes(&mut self, data: &[u8]) {
let size = data.len();
unsafe {
self.reserve_small(size);
let p = self.data.add(self.len);
std::ptr::copy_nonoverlapping(data.as_ptr(), p, size);
self.len += size;
}
debug_assert!(self.len <= self.capacity);
}
#[inline]
pub fn push(&mut self, byte: u8) {
let len = self.len;
if len == self.capacity {
self.reserve_internal(1, false);
}
unsafe {
*self.data.add(len) = byte;
self.len = len + 1;
}
}
#[inline]
pub fn push_char(&mut self, data: char) {
unsafe {
self.reserve_small(4);
let mut buffer = [0u8; 4];
let result = data.encode_utf8(&mut buffer);
let result = result.as_bytes();
std::ptr::copy_nonoverlapping(result.as_ptr(), self.data.add(self.len), result.len());
self.len += result.len();
}
}
fn flush(&mut self) {
if let Backing::Write {
written,
error,
writer,
} = &mut self.backing
{
use std::io::Write;
let buffered = unsafe { std::slice::from_raw_parts(self.data, self.len) };
if let Err(err) = writer.write_all(buffered) {
*error = Some(err);
}
*written += self.len;
self.len = 0;
}
}
#[cold]
fn reserve_internal(&mut self, size: usize, can_write: bool) {
debug_assert!(size <= isize::MAX as usize);
if let Backing::Write {
written,
error,
writer,
} = &mut self.backing
{
if can_write {
use std::io::Write;
let buffered = unsafe { std::slice::from_raw_parts(self.data, self.len) };
if let Err(err) = writer.write_all(buffered) {
*error = Some(err);
return;
}
*written += self.len;
self.len = 0;
if self.capacity >= size {
return;
}
}
}
let new_capacity = std::cmp::max(self.capacity * 2, self.capacity + size);
debug_assert!(new_capacity > self.capacity);
if let Backing::Borrowed { .. } = &self.backing {
let new_data = safe_alloc(new_capacity);
unsafe {
ptr::copy_nonoverlapping(self.data, new_data, self.len);
}
self.backing = Backing::Owned;
self.data = new_data;
} else if let Backing::Vec { bytes, offset, .. } = &mut self.backing {
let old_capacity = *offset + self.capacity;
let new_total_capacity = *offset + new_capacity;
let data = unsafe {
let vec = bytes.as_mut();
let original_len = vec.len();
let data = safe_realloc(vec.as_mut_ptr(), old_capacity, new_total_capacity);
bytes.write(Vec::from_raw_parts(data, original_len, new_total_capacity));
data
};
self.data = unsafe { data.add(*offset) };
self.capacity = new_capacity;
return;
} else {
self.data = unsafe { safe_realloc(self.data, self.capacity, new_capacity) };
}
self.capacity = new_capacity;
debug_assert!(!self.data.is_null());
debug_assert!(self.len <= self.capacity);
}
}
#[inline(never)]
fn safe_alloc(capacity: usize) -> *mut u8 {
assert!(capacity > 0);
assert!(capacity <= isize::MAX as usize, "capacity is too large");
unsafe {
let layout = Layout::from_size_align_unchecked(capacity, 1);
let data = alloc(layout);
if data.is_null() {
handle_alloc_error(layout);
}
data
}
}
#[cold]
#[inline(never)]
unsafe fn safe_realloc(ptr: *mut u8, capacity: usize, new_capacity: usize) -> *mut u8 {
assert!(new_capacity > 0);
assert!(new_capacity <= isize::MAX as usize, "capacity is too large");
let data = if capacity == 0 {
unsafe {
let new_layout = Layout::from_size_align_unchecked(new_capacity, 1);
alloc(new_layout)
}
} else {
unsafe {
let old_layout = Layout::from_size_align_unchecked(capacity, 1);
realloc(ptr, old_layout, new_capacity)
}
};
if data.is_null() {
unsafe {
handle_alloc_error(Layout::from_size_align_unchecked(new_capacity, 1));
}
}
data
}
impl fmt::Write for BytesWriter<'_> {
#[inline]
fn write_str(&mut self, s: &str) -> fmt::Result {
BytesWriter::push_bytes(self, s.as_bytes());
Ok(())
}
}
unsafe impl Send for BytesWriter<'_> {}
unsafe impl Sync for BytesWriter<'_> {}
pub trait IntoByteWriter<'a> {
type Output;
fn into_byte_writer(self) -> BytesWriter<'a>;
fn finish_writing(buffer: BytesWriter<'a>) -> Self::Output;
}
impl<'a> IntoByteWriter<'a> for &'a mut Vec<u8> {
type Output = &'a [u8];
fn into_byte_writer(self) -> BytesWriter<'a> {
BytesWriter::from(self)
}
fn finish_writing(buffer: BytesWriter<'a>) -> &'a [u8] {
buffer.into_backed_with_extended_slice()
}
}
pub type DynWrite<'a> = &'a mut (dyn std::io::Write + Send);
impl<'a> IntoByteWriter<'a> for DynWrite<'a> {
type Output = Result<usize, std::io::Error>;
fn into_byte_writer(self) -> BytesWriter<'a> {
BytesWriter::new_writer(self)
}
fn finish_writing(buffer: BytesWriter<'a>) -> Result<usize, std::io::Error> {
buffer.into_write_finish()
}
}
impl<'a> IntoByteWriter<'a> for &'a mut [MaybeUninit<u8>] {
type Output = Cow<'a, [u8]>;
fn into_byte_writer(self) -> BytesWriter<'a> {
BytesWriter::from(self)
}
fn finish_writing(buffer: BytesWriter<'a>) -> Cow<'a, [u8]> {
buffer.into_cow()
}
}