use std::{borrow, fmt, io, ops::DerefMut, ptr};
use crate::{Buf, BufMut, Bytes, buf::UninitSlice, stvec::StorageVec};
pub struct BytesMut {
pub(crate) storage: StorageVec,
}
impl BytesMut {
#[inline]
#[must_use]
pub fn with_capacity(capacity: usize) -> BytesMut {
BytesMut {
storage: StorageVec::with_capacity(capacity),
}
}
#[inline]
#[must_use]
pub fn copy_from_slice<T: AsRef<[u8]>>(src: T) -> Self {
let slice = src.as_ref();
BytesMut {
storage: StorageVec::from_slice(slice.len(), slice),
}
}
#[inline]
#[must_use]
pub fn new() -> BytesMut {
BytesMut {
storage: StorageVec::with_capacity(crate::storage::MIN_CAPACITY),
}
}
#[inline]
pub fn len(&self) -> usize {
self.storage.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.storage.len() == 0
}
#[inline]
pub fn capacity(&self) -> usize {
self.storage.capacity()
}
#[inline]
pub fn is_unique(&self) -> bool {
self.storage.is_unique()
}
#[inline]
#[must_use]
pub fn freeze(self) -> Bytes {
Bytes {
storage: self.storage.freeze(),
}
}
#[inline]
#[must_use]
pub fn take(&mut self) -> Bytes {
Bytes {
storage: self.storage.split_to(self.len()),
}
}
#[inline]
#[must_use]
pub fn split_to(&mut self, at: usize) -> Bytes {
self.split_to_checked(at)
.expect("at value must be <= self.len()`")
}
#[inline]
pub fn advance_to(&mut self, cnt: usize) {
unsafe {
self.storage.set_start(cnt);
}
}
#[inline]
#[must_use]
pub fn split_to_checked(&mut self, at: usize) -> Option<Bytes> {
if at <= self.len() {
Some(Bytes {
storage: self.storage.split_to(at),
})
} else {
None
}
}
#[inline]
pub fn truncate(&mut self, len: usize) {
self.storage.truncate(len);
}
#[inline]
pub fn clear(&mut self) {
self.truncate(0);
}
#[inline]
pub fn resize(&mut self, new_len: usize, value: u8) {
self.storage.resize(new_len, value);
}
#[inline]
pub unsafe fn set_len(&mut self, len: usize) {
self.storage.set_len(len);
}
#[inline]
pub fn reserve(&mut self, additional: usize) {
self.storage.reserve(additional);
}
#[inline]
pub fn reserve_exact(&mut self, additional: usize) {
self.storage.reserve_exact(additional);
}
#[inline]
pub fn reserve_capacity(&mut self, cap: usize) {
self.storage.reserve_capacity(cap);
}
#[inline]
pub fn extend_from_slice(&mut self, extend: &[u8]) {
self.put_slice(extend);
}
#[inline]
pub fn iter(&'_ self) -> std::slice::Iter<'_, u8> {
self.chunk().iter()
}
}
impl_buf!(BytesMut {});
impl_slice_traits!(BytesMut);
impl_partial_eq!(BytesMut);
impl BufMut for BytesMut {
#[inline]
fn remaining_mut(&self) -> usize {
self.storage.remaining()
}
#[inline]
unsafe fn advance_mut(&mut self, cnt: usize) {
self.storage.set_len(self.len() + cnt);
}
#[inline]
fn chunk_mut(&mut self) -> &mut UninitSlice {
self.storage.spare_mut()
}
#[inline]
fn put<T: Buf>(&mut self, mut src: T)
where
Self: Sized,
{
self.reserve(src.remaining());
while src.has_remaining() {
let chunk = src.chunk();
let len = chunk.len();
self.put_slice(chunk);
src.advance(len);
}
}
#[inline]
fn put_slice(&mut self, src: &[u8]) {
self.reserve(src.len());
self.storage.put_slice_partial(src);
}
#[inline]
fn put_u8(&mut self, n: u8) {
self.reserve(1);
self.storage.put_u8(n);
}
#[inline]
fn put_i8(&mut self, n: i8) {
self.put_u8(n as u8);
}
}
unsafe impl bytes::buf::BufMut for BytesMut {
#[inline]
fn remaining_mut(&self) -> usize {
usize::MAX - self.len()
}
#[inline]
unsafe fn advance_mut(&mut self, cnt: usize) {
let remaining = BufMut::remaining_mut(self);
assert!(
cnt <= remaining,
"cannot advance past `remaining_mut`: {cnt:?} <= {remaining:?}"
);
BufMut::advance_mut(self, cnt);
}
#[inline]
fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
if BufMut::remaining_mut(self) == 0 {
self.reserve(64);
}
unsafe {
let ptr = self.storage.as_ptr();
bytes::buf::UninitSlice::from_raw_parts_mut(
ptr.add(self.len()),
BufMut::remaining_mut(self),
)
}
}
#[inline]
fn put<T: bytes::buf::Buf>(&mut self, mut src: T)
where
Self: Sized,
{
self.reserve(src.remaining());
while src.has_remaining() {
let chunk = src.chunk();
let len = chunk.len();
BufMut::put_slice(self, chunk);
src.advance(len);
}
}
#[inline]
fn put_slice(&mut self, src: &[u8]) {
BufMut::put_slice(self, src);
}
#[inline]
fn put_bytes(&mut self, val: u8, cnt: usize) {
self.reserve(cnt);
unsafe {
ptr::write_bytes(self.storage.as_ptr().add(self.len()), val, cnt);
BufMut::advance_mut(self, cnt);
}
}
#[inline]
fn put_u8(&mut self, n: u8) {
BufMut::put_u8(self, n);
}
#[inline]
fn put_i8(&mut self, n: i8) {
BufMut::put_i8(self, n);
}
}
impl AsMut<[u8]> for BytesMut {
#[inline]
fn as_mut(&mut self) -> &mut [u8] {
self.storage.as_mut()
}
}
impl DerefMut for BytesMut {
#[inline]
fn deref_mut(&mut self) -> &mut [u8] {
self.storage.as_mut()
}
}
impl Eq for BytesMut {}
impl PartialEq for BytesMut {
#[inline]
fn eq(&self, other: &BytesMut) -> bool {
self.storage.as_ref() == other.storage.as_ref()
}
}
impl borrow::BorrowMut<[u8]> for BytesMut {
#[inline]
fn borrow_mut(&mut self) -> &mut [u8] {
self.as_mut()
}
}
impl PartialEq<Bytes> for BytesMut {
fn eq(&self, other: &Bytes) -> bool {
other[..] == self[..]
}
}
impl PartialEq<BytesMut> for Bytes {
fn eq(&self, other: &BytesMut) -> bool {
*other == *self
}
}
impl_read!(BytesMut);
impl io::Write for BytesMut {
fn write(&mut self, src: &[u8]) -> Result<usize, io::Error> {
self.extend_from_slice(src);
Ok(src.len())
}
fn flush(&mut self) -> Result<(), io::Error> {
Ok(())
}
}
impl fmt::Write for BytesMut {
#[inline]
fn write_str(&mut self, s: &str) -> fmt::Result {
self.extend_from_slice(s.as_bytes());
Ok(())
}
}
impl Clone for BytesMut {
#[inline]
fn clone(&self) -> BytesMut {
BytesMut::from(&self[..])
}
}
impl FromIterator<u8> for BytesMut {
fn from_iter<T: IntoIterator<Item = u8>>(into_iter: T) -> Self {
let iter = into_iter.into_iter();
let (min, maybe_max) = iter.size_hint();
let mut out = BytesMut::with_capacity(maybe_max.unwrap_or(min));
out.extend(iter);
out
}
}
impl<'a> FromIterator<&'a u8> for BytesMut {
fn from_iter<T: IntoIterator<Item = &'a u8>>(into_iter: T) -> Self {
into_iter.into_iter().copied().collect::<BytesMut>()
}
}
impl Extend<u8> for BytesMut {
fn extend<T>(&mut self, iter: T)
where
T: IntoIterator<Item = u8>,
{
let iter = iter.into_iter();
self.reserve(iter.size_hint().0);
for b in iter {
self.put_u8(b);
}
}
}
impl<'a> Extend<&'a u8> for BytesMut {
fn extend<T>(&mut self, iter: T)
where
T: IntoIterator<Item = &'a u8>,
{
self.extend(iter.into_iter().copied());
}
}
impl From<BytesMut> for Bytes {
#[inline]
fn from(b: BytesMut) -> Self {
b.freeze()
}
}
impl<'a> From<&'a [u8]> for BytesMut {
#[inline]
fn from(src: &'a [u8]) -> BytesMut {
BytesMut::copy_from_slice(src)
}
}
impl<const N: usize> From<[u8; N]> for BytesMut {
#[inline]
fn from(src: [u8; N]) -> BytesMut {
BytesMut::copy_from_slice(src)
}
}
impl<'a, const N: usize> From<&'a [u8; N]> for BytesMut {
#[inline]
fn from(src: &'a [u8; N]) -> BytesMut {
BytesMut::copy_from_slice(src)
}
}
impl<'a> From<&'a str> for BytesMut {
#[inline]
fn from(src: &'a str) -> BytesMut {
BytesMut::from(src.as_bytes())
}
}
impl From<Bytes> for BytesMut {
#[inline]
fn from(src: Bytes) -> BytesMut {
match src.storage.try_into_vec() {
Ok(storage) => BytesMut { storage },
Err(storage) => BytesMut::copy_from_slice(storage.as_ref()),
}
}
}
impl From<&Bytes> for BytesMut {
#[inline]
fn from(src: &Bytes) -> BytesMut {
BytesMut::copy_from_slice(&src[..])
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn growth_is_amortized() {
let mut buf = BytesMut::with_capacity(0);
let mut cap = buf.capacity();
let mut reallocs = 0;
for _ in 0..10_000 {
buf.put_slice(b"abcdefgh");
if buf.capacity() != cap {
reallocs += 1;
cap = buf.capacity();
}
}
assert_eq!(buf.len(), 80_000);
assert!(reallocs <= 16, "reallocs: {reallocs}");
let mut buf = BytesMut::with_capacity(0);
for i in 0..1000 {
std::fmt::Write::write_fmt(&mut buf, format_args!("{i:08}")).unwrap();
}
assert_eq!(buf.len(), 8000);
assert!(buf.capacity() < 16_000);
}
#[test]
fn growth_of_little_data_is_exact() {
let mut buf = BytesMut::copy_from_slice(b"hello");
buf.reserve(64 * 1024);
assert_eq!(buf.capacity(), 5 + 64 * 1024);
let mut buf = BytesMut::with_capacity(1024);
buf.extend_from_slice(&[1; 1024]);
let head = buf.split_to(1000);
buf.reserve(4096);
assert_eq!(buf.capacity(), 24 + 4096);
assert_eq!(&head[..], &[1; 1000][..]);
}
#[test]
fn from_unique_bytes_reuses_buffer() {
let mut buf = BytesMut::with_capacity(256);
buf.extend_from_slice(&[1; 64]);
let b = buf.freeze();
let ptr = b.as_ptr();
let mut m = BytesMut::from(b);
assert_eq!(m.as_ptr(), ptr);
assert_eq!(&m[..], &[1; 64][..]);
assert_eq!(m.capacity(), 256);
m.extend_from_slice(&[2; 192]);
assert_eq!(m.as_ptr(), ptr);
assert_eq!(&m[64..], &[2; 192][..]);
}
#[test]
fn from_unique_bytes_subview() {
let mut buf = BytesMut::with_capacity(256);
buf.extend_from_slice(&[1; 128]);
let mut b = buf.freeze();
let head = b.split_to(32);
drop(head);
b.truncate(64);
let ptr = b.as_ptr();
let mut m = BytesMut::from(b);
assert_eq!(m.as_ptr(), ptr);
assert_eq!(m.len(), 64);
assert_eq!(m.capacity(), 256 - 32);
m.extend_from_slice(&[3; 160]);
assert_eq!(m.as_ptr(), ptr);
assert_eq!(&m[..64], &[1; 64][..]);
assert_eq!(&m[64..], &[3; 160][..]);
}
#[test]
fn from_shared_bytes_copies() {
let b = BytesMut::copy_from_slice([1; 64]).freeze();
let b2 = b.clone();
let mut m = BytesMut::from(b);
assert_ne!(m.as_ptr(), b2.as_ptr());
m[0] = 2;
assert_eq!(&b2[..], &[1; 64][..]);
let mut buf = BytesMut::with_capacity(256);
buf.extend_from_slice(&[1; 64]);
let b = buf.take();
let m = BytesMut::from(b);
assert_ne!(m.as_ptr(), buf.as_ptr());
buf.extend_from_slice(&[2; 64]);
assert_eq!(&m[..], &[1; 64][..]);
}
#[test]
fn from_bytes_synchronizes_with_release() {
let b = BytesMut::copy_from_slice([1; 64]).freeze();
let other = b.clone();
let handle = std::thread::spawn(move || {
let val = other[0];
drop(other);
val
});
let ptr = b.as_ptr();
let mut storage = b.storage;
let mut m = loop {
match storage.try_into_vec() {
Ok(storage) => break BytesMut { storage },
Err(st) => {
storage = st;
std::thread::yield_now();
}
}
};
assert_eq!(m.as_ptr(), ptr);
m[0] = 2;
assert_eq!(handle.join().unwrap(), 1);
}
#[test]
fn from_inline_and_static_bytes() {
let m = BytesMut::from(Bytes::copy_from_slice(b"inline"));
assert_eq!(&m[..], b"inline");
let m = BytesMut::from(Bytes::from_static(&[1; 64]));
assert_eq!(&m[..], &[1; 64][..]);
}
#[test]
fn bvec_read() {
use std::io::Read;
let mut b = BytesMut::copy_from_slice(b"123");
let mut buf = [0; 10];
assert_eq!(b.read(&mut buf).unwrap(), 3);
assert_eq!(b.len(), 0);
assert_eq!(buf, [49, 50, 51, 0, 0, 0, 0, 0, 0, 0]);
}
#[test]
fn from_bytes_ref() {
let b = Bytes::from_static(b"hello");
let mut m = BytesMut::from(&b);
m.extend_from_slice(b"!");
assert_eq!(m, "hello!");
assert_eq!(b, "hello");
}
}