从零写一个最小 Bytes

[dependencies]
bytes = "=1.12.1"

pub struct Bytes {
  ptr: *const u8,
  len: usize,
  // 交给 vtable 解释：静态是空指针，共享是 Shared 的地址。
  // 没有可提升存储，没人会 CAS 它，所以不用 AtomicPtr
  data: *mut (),
  vtable: &'static Vtable,
}

struct Vtable {
  clone: unsafe fn(*mut (), *const u8, usize) -> Bytes,
  drop: unsafe fn(*mut (), *const u8, usize),
}

// SAFETY: 每张表自己保证：静态数据只读；共享块只读，计数是原子的
unsafe impl Send for Bytes {}
unsafe impl Sync for Bytes {}

impl Clone for Bytes {
  fn clone(&self) -> Bytes {
    // SAFETY: data、ptr、len 和 vtable 是同一个构造者给出的一套
    unsafe { (self.vtable.clone)(self.data, self.ptr, self.len) }
  }
}

impl Drop for Bytes {
  fn drop(&mut self) {
    // SAFETY: 同上；之后不再使用 self
    unsafe { (self.vtable.drop)(self.data, self.ptr, self.len) }
  }
}

impl Bytes {
  pub const fn from_static(s: &'static [u8]) -> Bytes {
    Bytes { ptr: s.as_ptr(), len: s.len(), data: ptr::null_mut(), vtable: &STATIC_VTABLE }
  }
}

const STATIC_VTABLE: Vtable = Vtable { clone: static_clone, drop: static_drop };

unsafe fn static_clone(_: *mut (), ptr: *const u8, len: usize) -> Bytes {
  Bytes { ptr, len, data: ptr::null_mut(), vtable: &STATIC_VTABLE }
}

unsafe fn static_drop(_: *mut (), _: *const u8, _: usize) {}

struct Shared {
  #[allow(dead_code)] // 从不读，只靠它的 Drop 释放缓冲
  vec: Vec<u8>,
  ref_cnt: AtomicUsize,
}

// BytesMut 拿 Shared 地址的最低位当 KIND_ARC，所以对齐必须是偶数
const _: () = assert!(align_of::<Shared>() % 2 == 0);

const SHARED_VTABLE: Vtable = Vtable { clone: shared_clone, drop: shared_drop };

unsafe fn increment_shared(shared: *mut Shared) {
  // 能 clone 说明自己持有一份，计数不会在这期间归零，Relaxed 就够
  let old = unsafe { (*shared).ref_cnt.fetch_add(1, Relaxed) };
  if old > isize::MAX as usize {
    std::process::abort();
  }
}

unsafe fn release_shared(shared: *mut Shared) {
  // Release：本句柄对数据的读都排在减一之前
  if unsafe { (*shared).ref_cnt.fetch_sub(1, Release) } != 1 {
    return;
  }
  // 最后一个：Acquire 读和所有 Release 配对，再释放 Vec 和 Shared
  unsafe { (*shared).ref_cnt.load(Acquire) };
  drop(unsafe { Box::from_raw(shared) });
}

impl From<Vec<u8>> for Bytes {
  /// 总是分配一个 Shared。bytes 在 len == cap 时不分配（可提升存储）
  fn from(vec: Vec<u8>) -> Bytes {
    let (ptr, len) = (vec.as_ptr(), vec.len());
    let shared = Box::into_raw(Box::new(Shared { vec, ref_cnt: AtomicUsize::new(1) }));
    Bytes { ptr, len, data: shared.cast(), vtable: &SHARED_VTABLE }
  }
}

  pub fn slice(&self, range: Range<usize>) -> Bytes {
    assert!(range.start <= range.end, "range start must not be greater than end");
    assert!(range.end <= self.len, "range end out of bounds");
    let mut ret = self.clone();
    // SAFETY: range 在 [0, len] 之内
    ret.ptr = unsafe { ret.ptr.add(range.start) };
    ret.len = range.end - range.start;
    ret
  }

const KIND_ARC: usize = 0b0; // data 是 Shared 的地址
const KIND_VEC: usize = 0b1; // data 是 (pos << 1) | 1，独占一个 Vec
const KIND_MASK: usize = 0b1;
const VEC_POS_OFFSET: usize = 1;

pub struct BytesMut {
  ptr: NonNull<u8>,
  len: usize,
  cap: usize,
  data: *mut Shared,
}

  pub fn advance(&mut self, n: usize) {
    assert!(n <= self.len, "advance out of bounds");
    if self.kind() == KIND_VEC {
      // pos 不超过分配大小（≤ isize::MAX），左移一位不会溢出
      let pos = self.vec_pos() + n;
      self.data = ptr::without_provenance_mut((pos << VEC_POS_OFFSET) | KIND_VEC);
    }
    // SAFETY: n <= len <= cap
    self.ptr = unsafe { self.ptr.add(n) };
    self.len -= n;
    self.cap -= n;
  }

unsafe fn rebuild_vec(ptr: *mut u8, len: usize, cap: usize, off: usize) -> Vec<u8> {
  unsafe { Vec::from_raw_parts(ptr.sub(off), off + len, off + cap) }
}

  pub fn split_to(&mut self, at: usize) -> BytesMut {
    assert!(at <= self.len, "split_to out of bounds");
    if self.kind() == KIND_ARC {
      // SAFETY: KIND_ARC 时 data 是活着的 Shared
      unsafe { increment_shared(self.data) };
    } else {
      self.promote_to_shared(2); // 一份给 self，一份给返回值
    }
    let other = BytesMut { ptr: self.ptr, len: at, cap: at, data: self.data };
    // SAFETY: at <= len <= cap；两个句柄从此看不重叠的两段
    self.ptr = unsafe { self.ptr.add(at) };
    self.len -= at;
    self.cap -= at;
    other
  }

  fn promote_to_shared(&mut self, ref_cnt: usize) {
    // 先分配再拼回 Vec：Box 分配失败展开时，缓冲仍只归 self 所有（bytes 1.12.1 的 #837）
    let uninit = Box::<Shared>::new_uninit();
    // SAFETY: KIND_VEC 时 ptr - pos 是 Vec 的起点，len、cap 加回 pos 就是原来的 Vec
    let vec = unsafe { rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, self.vec_pos()) };
    let shared = Box::write(uninit, Shared { vec, ref_cnt: AtomicUsize::new(ref_cnt) });
    self.data = Box::into_raw(shared); // 偶数对齐，最低位天然是 KIND_ARC
  }

  pub fn freeze(self) -> Bytes {
    let me = ManuallyDrop::new(self);
    if me.kind() == KIND_ARC {
      // 计数原样交给 Bytes：同一个 Shared，换一张 vtable
      return Bytes { ptr: me.ptr.as_ptr(), len: me.len, data: me.data.cast(), vtable: &SHARED_VTABLE };
    }
    let off = me.vec_pos();
    // SAFETY: 同 promote_to_shared；me 不会再 drop
    let vec = unsafe { rebuild_vec(me.ptr.as_ptr(), me.len, me.cap, off) };
    let mut b = Bytes::from(vec); // 分配一个 Shared
    // SAFETY: off <= vec.len()，跳过前面已经 advance 掉的部分
    b.ptr = unsafe { b.ptr.add(off) };
    b.len -= off;
    b
  }

type Count = (usize, usize, usize); // (分配次数, 分配字节, 释放次数)

// 跑 f，返回 (结果, 这一步的计数)；结果在计数之后才由调用方 drop
fn step<R>(f: impl FnOnce() -> R) -> (R, Count) {
  let (a, b, d) = (ALLOCS.load(Relaxed), BYTES.load(Relaxed), FREES.load(Relaxed));
  let r = f();
  (r, (ALLOCS.load(Relaxed) - a, BYTES.load(Relaxed) - b, FREES.load(Relaxed) - d))
}

// 同一段代码，$b / $bm 换成 mini 或 bytes 的类型；返回每一步的计数和最后的内容
macro_rules! scenario {
  ($b:ty, $bm:ty) => {{
    // ...
    let v = vec![7u8; 1024];
    let (b, r) = step(|| <$b>::from(v));
    rows.push(("from(Vec len==cap)", r));
    let (c1, r) = step(|| b.clone());
    rows.push(("  .clone() #1", r));
    // ...
    let (mut buf, r) = step(|| <$bm>::with_capacity(8192));
    rows.push(("BytesMut::with_capacity", r));
    let req = b"POST /upload HTTP/1.1\r\nHost: a\r\nContent-Length: 11\r\n\r\nhello world";
    let (_, r) = step(|| buf.put_slice(&req[..]));
    rows.push(("  put_slice(65)", r));
    let (head, r) = step(|| buf.split_to(54));
    rows.push(("  split_to(54)", r));
    let (head, r) = step(|| head.freeze());
    rows.push(("  head.freeze()", r));
    // ...
  }};
}

------------------------------------------------------------------
step                                    mini        bytes 1.12.1
from_static                          0(0B)/0             0(0B)/0
  .clone()                           0(0B)/0             0(0B)/0
  drop both                          0(0B)/0             0(0B)/0
from(Vec len==cap)                  1(32B)/0             0(0B)/0  *
  .clone() #1                        0(0B)/0            1(24B)/0  *
  .clone() #2                        0(0B)/0             0(0B)/0
  .slice(10..20)                     0(0B)/0             0(0B)/0
  drop all four                      0(0B)/2             0(0B)/2
from(Vec len<cap)                   1(32B)/0            1(24B)/0  *
  drop                               0(0B)/2             0(0B)/2
BytesMut::with_capacity           1(8192B)/0          1(8192B)/0
  put_slice(65)                      0(0B)/0             0(0B)/0
  split_to(54)                      1(32B)/0            1(40B)/0  *
  head.freeze()                      0(0B)/0             0(0B)/0
  head.slice(5..12)                  0(0B)/0             0(0B)/0
  split_to(11).freeze()              0(0B)/0             0(0B)/0
  drop uri, head, body               0(0B)/0             0(0B)/0
  drop buf                           0(0B)/2             0(0B)/2
KIND_VEC advance(4)                  0(0B)/0             0(0B)/0
  .freeze() (len<cap)               1(32B)/0            1(24B)/0  *
  drop                               0(0B)/2             0(0B)/2
contents equal: true
live bytes after scenario: mini=0 bytes=0
size_of Bytes: mini=32 bytes=32; BytesMut: mini=32 bytes=32

(stdout warmed up)
rounds=1000 threads=8 warm-up live=0B, later rounds with different live: 0

mini  with_capacity(0): allocs=0 (0B)
mini  split_to(0):      allocs=1 (32B)
bytes with_capacity(0): allocs=0 (0B)
bytes split_to(0):      allocs=1 (40B)