BufMut 与 UninitSlice

pub unsafe trait BufMut {
    fn remaining_mut(&self) -> usize;

    unsafe fn advance_mut(&mut self, cnt: usize);

    // ...
    fn chunk_mut(&mut self) -> &mut UninitSlice;
    // ...
}

/// Uninitialized byte slice.
///
/// Returned by `BufMut::chunk_mut()`, the referenced byte slice may be
/// uninitialized. The wrapper provides safe access without introducing
/// undefined behavior.
///
/// The safety invariants of this wrapper are:
///
///  1. Reading from an `UninitSlice` is undefined behavior.
///  2. Writing uninitialized bytes to an `UninitSlice` is undefined behavior.
///
/// The difference between `&mut UninitSlice` and `&mut [MaybeUninit<u8>]` is
/// that it is possible in safe code to write uninitialized bytes to an
/// `&mut [MaybeUninit<u8>]`, which this type prohibits.
#[repr(transparent)]
pub struct UninitSlice([MaybeUninit<u8>]);

use bytes::{BufMut, BytesMut};

fn main() {
  let mut buf = BytesMut::with_capacity(16);
  let first = buf.chunk_mut()[0];
  println!("{first}");
}

error[E0277]: the type `UninitSlice` cannot be indexed by `{integer}`
 --> src/bin/read_uninit.rs:5:31
  |
5 |   let first = buf.chunk_mut()[0];
  |                               ^ `UninitSlice` cannot be indexed by `{integer}`
  |
  = help: the trait `Index<{integer}>` is not implemented for `UninitSlice`
  = help: `UninitSlice` implements trait `Index<Idx>`:
            Index<RangeFull>
            Index<RangeTo<usize>>
            Index<std::ops::Range<usize>>
            Index<std::ops::RangeFrom<usize>>
            Index<std::ops::RangeInclusive<usize>>
            Index<std::ops::RangeToInclusive<usize>>

For more information about this error, try `rustc --explain E0277`.
error: could not compile `exp-bytes-bufmut` (bin "read_uninit") due to 1 previous error

        // SAFETY: ReadBuf and poll_read promise not to set any uninitialized
        // bytes onto `dst`.
        let dst = unsafe { self.read_buf.chunk_mut().as_uninit_slice_mut() };
        let mut buf = ReadBuf::uninit(dst);

    #[inline]
    unsafe fn advance_mut(&mut self, cnt: usize) {
        let len = self.len();
        let remaining = self.capacity() - len;

        if remaining < cnt {
            panic_advance(&TryGetError {
                requested: cnt,
                available: remaining,
            });
        }

        // Addition will not overflow since the sum is at most the capacity.
        self.set_len(len + cnt);
    }

bytes = "=1.12.1"

use std::mem::MaybeUninit;

use bytes::{BufMut, BytesMut};

fn main() {
  let mut buf = BytesMut::with_capacity(16);
  let dst = buf.chunk_mut();
  println!("chunk_mut: {:?}, len {}", dst, dst.len());

  // 只用 UninitSlice 的安全方法写：copy_from_slice 要求长度正好相等
  dst[..5].copy_from_slice(b"hello");
  dst.write_byte(5, b' ');
  // SAFETY: 上面刚初始化了 chunk_mut 的前 6 个字节
  unsafe { buf.advance_mut(6) };
  println!("after advance_mut(6): {:?}, len {}, cap {}", buf, buf.len(), buf.capacity());

  // 把未初始化的尾部交给一个「只认裸指针」的写入者，比如 read(2)
  let dst = buf.chunk_mut();
  let n = fake_read(dst.as_mut_ptr(), dst.len());
  // SAFETY: fake_read 保证写满了前 n 个字节
  unsafe { buf.advance_mut(n) };
  println!("after fake_read: {:?}", buf);

  // &mut [MaybeUninit<u8>] 也是 BufMut：栈上的未初始化数组，不用先清零
  let mut storage = [MaybeUninit::<u8>::uninit(); 8];
  let mut out = &mut storage[..];
  out.put_u32(0xdead_beef);
  out.put_slice(b"ok");
  println!("MaybeUninit slice: remaining_mut {}", out.remaining_mut());
  // SAFETY: 前 6 个字节刚由 put_u32 和 put_slice 写过
  let written = unsafe { std::slice::from_raw_parts(storage.as_ptr().cast::<u8>(), 6) };
  println!("written = {written:x?}");
}

// 模拟一个系统调用：往 ptr 写 len 以内的字节，返回写了多少
fn fake_read(ptr: *mut u8, len: usize) -> usize {
  let src = b"world";
  let n = src.len().min(len);
  // SAFETY: 调用方保证 ptr 对 len 个字节可写
  unsafe { std::ptr::copy_nonoverlapping(src.as_ptr(), ptr, n) };
  n
}

chunk_mut: UninitSlice[...], len 16
after advance_mut(6): b"hello ", len 6, cap 16
after fake_read: b"hello world"
MaybeUninit slice: remaining_mut 2
written = [de, ad, be, ef, 6f, 6b]

unsafe impl BufMut for Vec<u8> {
    #[inline]
    fn remaining_mut(&self) -> usize {
        // A vector can never have more than isize::MAX bytes
        isize::MAX as usize - self.len()
    }

    #[inline]
    fn put_slice(&mut self, mut src: &[u8]) {
        if self.remaining_mut() < src.len() {
            panic_advance(&TryGetError {
                requested: src.len(),
                available: self.remaining_mut(),
            });
        }

        while !src.is_empty() {
            let dst = self.chunk_mut();
            let cnt = usize::min(src.len(), dst.len());

            dst[..cnt].copy_from_slice(&src[..cnt]);
            src = &src[cnt..];

            // SAFETY: We just initialized `cnt` bytes in `self`.
            unsafe { self.advance_mut(cnt) };
        }
    }

    fn put<T: Buf>(&mut self, mut src: T)
    where
        Self: Sized,
    {
        if !src.has_remaining() {
            // prevent calling `copy_to_bytes`->`put`->`copy_to_bytes` infintely when src is empty
            return;
        } else if self.capacity() == 0 {
            // When capacity is zero, try reusing allocation of `src`.
            let src_copy = src.copy_to_bytes(src.remaining());
            drop(src);
            match src_copy.try_into_mut() {
                Ok(bytes_mut) => *self = bytes_mut,
                Err(bytes) => self.extend_from_slice(&bytes),
            }
        } else {
            // In case the src isn't contiguous, reserve upfront.
            self.reserve(src.remaining());

            while src.has_remaining() {
                let s = src.chunk();
                let l = s.len();
                self.extend_from_slice(s);
                src.advance(l);
            }
        }
    }

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering};

use bytes::{Buf, BufMut, Bytes, BytesMut};

struct Counting;

static ALLOCS: AtomicUsize = AtomicUsize::new(0);
static REALLOCS: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
  unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
    ALLOCS.fetch_add(1, Ordering::Relaxed);
    unsafe { System.alloc(layout) }
  }
  unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
    unsafe { System.dealloc(ptr, layout) }
  }
  unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
    REALLOCS.fetch_add(1, Ordering::Relaxed);
    unsafe { System.realloc(ptr, layout, new_size) }
  }
}

#[global_allocator]
static GLOBAL: Counting = Counting;

fn counts() -> (usize, usize) {
  (ALLOCS.load(Ordering::Relaxed), REALLOCS.load(Ordering::Relaxed))
}

// 闭包返回 (说明, 数值)，在计数之后才格式化，println 的分配不算进去
fn report(name: &str, f: impl FnOnce() -> (&'static str, usize)) {
  let (a0, r0) = counts();
  let (label, value) = f();
  let (a1, r1) = counts();
  println!("{name:<42} alloc +{} realloc +{}  {label} {value}", a1 - a0, r1 - r0);
}

fn main() {
  println!("start");
  let data = vec![7u8; 4096];

  // 源：Vec 转来的 Bytes（len == cap，可提升存储），唯一拥有
  let src = Bytes::from(data.clone());
  let p = src.as_ptr();
  report("BytesMut::new().put(Bytes from Vec)", || {
    let mut dst = BytesMut::new();
    dst.put(src);
    ("same buffer:", (dst.as_ptr() == p) as usize)
  });

  let src = Bytes::from(data.clone());
  report("BytesMut::new().put_slice(&bytes)", || {
    let mut dst = BytesMut::new();
    dst.put_slice(&src);
    ("len", dst.len())
  });
  drop(src);

  let src = Bytes::from(data.clone());
  let keep = src.clone();
  report("BytesMut::new().put(Bytes, 有别的 clone)", || {
    let mut dst = BytesMut::new();
    dst.put(src);
    ("same buffer:", (dst.as_ptr() == keep.as_ptr()) as usize)
  });
  drop(keep);

  let src = Bytes::from_static(&[7u8; 4096]);
  report("BytesMut::new().put(Bytes::from_static)", || {
    let mut dst = BytesMut::new();
    dst.put(src);
    ("len", dst.len())
  });

  let src = Bytes::from(data.clone());
  let p = src.as_ptr();
  report("BytesMut::with_capacity(1).put(Bytes)", || {
    let mut dst = BytesMut::with_capacity(1);
    dst.put(src);
    ("same buffer:", (dst.as_ptr() == p) as usize)
  });

  let src = Bytes::from(data.clone());
  report("Vec::new().put(Bytes)", || {
    let mut dst: Vec<u8> = Vec::new();
    dst.put(src);
    ("cap", dst.capacity())
  });

  report("BytesMut::new().put(&[u8])", || {
    let mut dst = BytesMut::new();
    dst.put(&data[..]);
    ("len", dst.len())
  });

  // #794：不连续的源，put 先 reserve 总长度
  let (a, b, c) = (vec![1u8; 1000], vec![2u8; 1000], vec![3u8; 1000]);
  report("with_capacity(1).put(a.chain(b).chain(c))", || {
    let mut dst = BytesMut::with_capacity(1);
    dst.put((&a[..]).chain(&b[..]).chain(&c[..]));
    ("cap", dst.capacity())
  });
  report("with_capacity(1) + 3 x put_slice", || {
    let mut dst = BytesMut::with_capacity(1);
    dst.put_slice(&a);
    dst.put_slice(&b);
    dst.put_slice(&c);
    ("cap", dst.capacity())
  });

  // #803：空源不再递归
  report("BytesMut::new().put(&[] as &[u8])", || {
    let mut dst = BytesMut::new();
    dst.put(&[] as &[u8]);
    ("cap", dst.capacity())
  });
}

start
BytesMut::new().put(Bytes from Vec)        alloc +0 realloc +0  same buffer: 1
BytesMut::new().put_slice(&bytes)          alloc +1 realloc +0  len 4096
BytesMut::new().put(Bytes, 有别的 clone)      alloc +1 realloc +0  same buffer: 0
BytesMut::new().put(Bytes::from_static)    alloc +1 realloc +0  len 4096
BytesMut::with_capacity(1).put(Bytes)      alloc +1 realloc +1  same buffer: 0
Vec::new().put(Bytes)                      alloc +1 realloc +0  cap 4096
BytesMut::new().put(&[u8])                 alloc +1 realloc +0  len 4096
with_capacity(1).put(a.chain(b).chain(c))  alloc +1 realloc +1  cap 3000
with_capacity(1) + 3 x put_slice           alloc +1 realloc +3  cap 4000
BytesMut::new().put(&[] as &[u8])          alloc +0 realloc +0  cap 0

    #[inline]
    fn chunk_mut(&mut self) -> &mut UninitSlice {
        if self.capacity() == self.len() {
            self.reserve(64); // Grow the vec
        }

        let cap = self.capacity();
        let len = self.len();

        let ptr = self.as_mut_ptr();
        // SAFETY: Since `ptr` is valid for `cap` bytes, `ptr.add(len)` must be
        // valid for `cap - len` bytes. The subtraction will not underflow since
        // `len <= cap`.
        unsafe { UninitSlice::from_raw_parts_mut(ptr.add(len), cap - len) }
    }

    #[inline]
    fn chunk_mut(&mut self) -> &mut UninitSlice {
        if self.capacity() == self.len() {
            self.reserve(64);
        }
        self.spare_capacity_mut().into()
    }

    #[inline]
    unsafe fn advance_mut(&mut self, cnt: usize) {
        if self.len() < cnt {
            panic_advance(&TryGetError {
                requested: cnt,
                available: self.len(),
            });
        }

        // Lifetime dance taken from `impl Write for &mut [u8]`.
        let (_, b) = core::mem::take(self).split_at_mut(cnt);
        *self = b;
    }

Vec<u8>   put_u8 x300: [0, 8, 16, 32, 64, 128, 256, 512]
BytesMut  put_u8 x300: [0, 8, 16, 32, 64, 128, 256, 512]
Vec<u8>   chunk_mut x300: [0, 64, 128, 256, 512]
BytesMut  chunk_mut x300: [0, 64, 128, 256, 512]
remaining_mut: Vec 9223372036854775807 / BytesMut 9223372036854775807 / isize::MAX 9223372036854775807
&mut [u8; 8] after put_u16: remaining_mut 6
arr = [171, 205, 0, 0, 0, 0, 0, 0]