Buf

bytes = "=1.12.1"

pub trait Buf {
    fn remaining(&self) -> usize;

    fn chunk(&self) -> &[u8];

    fn advance(&mut self, cnt: usize);

    // ...
}

    fn try_copy_to_slice(&mut self, mut dst: &mut [u8]) -> Result<(), TryGetError> {
        if self.remaining() < dst.len() {
            return Err(TryGetError {
                requested: dst.len(),
                available: self.remaining(),
            });
        }

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

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

            self.advance(cnt);
        }
        Ok(())
    }

    fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
        if dst.is_empty() {
            return 0;
        }

        if self.has_remaining() {
            dst[0] = IoSlice::new(self.chunk());
            1
        } else {
            0
        }
    }

use std::io::IoSlice;

use bytes::Buf;

// 合法但「吝啬」的 Buf：chunk 每次只给 1 个字节，chunks_vectored 用默认实现
struct Trickle<'a>(&'a [u8]);

impl Buf for Trickle<'_> {
  fn remaining(&self) -> usize {
    self.0.len()
  }
  fn chunk(&self) -> &[u8] {
    &self.0[..self.0.len().min(1)]
  }
  fn advance(&mut self, n: usize) {
    self.0 = &self.0[n..];
  }
}

fn show(label: &str, buf: &impl Buf) {
  let mut io = [IoSlice::new(&[]); 16];
  let n = buf.chunks_vectored(&mut io);
  let parts: Vec<String> = io[..n]
    .iter()
    .map(|s| String::from_utf8_lossy(s).into_owned())
    .collect();
  println!("{label:<16} remaining={:<2} n={n} slices={parts:?}", buf.remaining());
}

fn main() {
  show("Trickle", &Trickle(b"hello"));
  show("&[u8] chain", &(&b"hello"[..]).chain(&b"world"[..]));
  show("Trickle chain", &Trickle(b"hello").chain(&b"world"[..]));
}

Trickle          remaining=5  n=1 slices=["h"]
&[u8] chain      remaining=10 n=2 slices=["hello", "world"]
Trickle chain    remaining=10 n=2 slices=["h", "world"]

    fn get_u32(&mut self) -> u32 {
        buf_get_impl!(self, u32::from_be_bytes);
    }

macro_rules! buf_try_get_impl {
    ($this:ident, $typ:tt::$conv:tt) => {{
        const SIZE: usize = core::mem::size_of::<$typ>();

        if $this.remaining() < SIZE {
            return Err(TryGetError {
                requested: SIZE,
                available: $this.remaining(),
            });
        }

        // try to convert directly from the bytes
        // this Option<ret> trick is to avoid keeping a borrow on self
        // when advance() is called (mut borrow) and to call bytes() only once
        let ret = $this
            .chunk()
            .get(..SIZE)
            .map(|src| unsafe { $typ::$conv(*(src as *const _ as *const [_; SIZE])) });

        if let Some(ret) = ret {
            // if the direct conversion was possible, advance and return
            $this.advance(SIZE);
            return Ok(ret);
        } else {
            // if not we copy the bytes in a temp buffer then convert
            let mut buf = [0; SIZE];
            $this.copy_to_slice(&mut buf); // (do the advance)
            return Ok($typ::$conv(buf));
        }
    }};
    // ...
}

macro_rules! buf_get_impl {
    ($this:ident, $typ:tt::$conv:tt) => {{
        return (|| buf_try_get_impl!($this, $typ::$conv))()
            .unwrap_or_else(|error| panic_advance(&error));
    }};
    // ...
}

use std::collections::VecDeque;

use bytes::{Buf, Bytes};

fn main() {
  let data = [0x0A, 0x0B, 0x0C, 0x0D, 0xEE];

  // 1. 连续：一个 &[u8]
  let mut flat = &data[..];
  println!("flat   chunk={:?} get_u32={:#010x} remaining={}", flat.chunk().len(), flat.get_u32(), flat.remaining());

  // 2. 跨块：前 2 字节在 a，后 3 字节在 b
  let mut chain = (&data[..2]).chain(&data[2..]);
  println!("chain  chunk={:?} get_u32={:#010x} remaining={}", chain.chunk().len(), chain.get_u32(), chain.remaining());

  // 3. VecDeque：环形缓冲绕回之后有两段
  let mut dq: VecDeque<u8> = VecDeque::with_capacity(8);
  dq.extend([0u8; 6]);
  dq.drain(..5); // 留 1 个在第 5 格
  dq.extend(data); // 写进第 6、7 格，再绕回第 0..3 格
  dq.pop_front(); // 去掉留下的那个
  let (s1, s2) = dq.as_slices();
  println!("deque  slices=({}, {}) get_u32={:#010x} remaining={}", s1.len(), s2.len(), dq.get_u32(), dq.remaining());

  // 4. try_get_u32：够与不够
  let mut short = Bytes::from_static(&[1, 2, 3]);
  println!("try_get_u32 on 3 bytes: {:?}, remaining={}", short.try_get_u32(), short.remaining());
  let err = short.try_get_u32().unwrap_err();
  println!("Display: {err}");
  let mut ok = Bytes::from_static(&[1, 2, 3, 4]);
  println!("try_get_u32 on 4 bytes: {:?}, remaining={}", ok.try_get_u32(), ok.remaining());

  // 5. get_int 与 try_get_int：3 字节的 0xFFFFFF
  let ff = [0xFF, 0xFF, 0xFF];
  println!("get_int(3)      = {}", (&ff[..]).get_int(3));
  println!("try_get_int(3)  = {:?}", (&ff[..]).try_get_int(3));
  println!("get_uint(3)     = {}", (&ff[..]).get_uint(3));

  // 6. nbytes == 0
  let mut b = &b"hello"[..];
  println!("get_int(0) = {}, remaining={}", b.get_int(0), b.remaining());
}

flat   chunk=5 get_u32=0x0a0b0c0d remaining=1
chain  chunk=2 get_u32=0x0a0b0c0d remaining=1
deque  slices=(2, 3) get_u32=0x0a0b0c0d remaining=1
try_get_u32 on 3 bytes: Err(TryGetError { requested: 4, available: 3 }), remaining=3
Display: Not enough bytes remaining in buffer to read value (requested 4 but only 3 available)
try_get_u32 on 4 bytes: Ok(16909060), remaining=0
get_int(3)      = -1
try_get_int(3)  = Ok(16777215)
get_uint(3)     = 16777215
get_int(0) = 0, remaining=5

[dev-dependencies]
criterion = "0.5"

use std::hint::black_box;

use bytes::Buf;
use criterion::{Criterion, criterion_group, criterion_main};

fn bench(c: &mut Criterion) {
  let data = [0x0Au8, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11];

  // 连续的 &[u8]：快路径
  c.bench_function("slice", |b| {
    b.iter(|| {
      let mut buf = black_box(&data[..]);
      buf.get_u32()
    })
  });
  // Chain，4 字节都在第一段：快路径
  c.bench_function("chain_in_a", |b| {
    b.iter(|| {
      let mut buf = black_box(&data[..4]).chain(black_box(&data[4..]));
      buf.get_u32()
    })
  });
  // Chain，2 + 2 跨段：慢路径
  c.bench_function("chain_split", |b| {
    b.iter(|| {
      let mut buf = black_box(&data[..2]).chain(black_box(&data[2..]));
      buf.get_u32()
    })
  });
}

criterion_group!(benches, bench);
criterion_main!(benches);

slice                   time:   [499.59 ps 503.92 ps 508.39 ps]
chain_in_a              time:   [2.3786 ns 2.4046 ns 2.4313 ns]
chain_split             time:   [6.3531 ns 6.3654 ns 6.3781 ns]
slice                   time:   [514.68 ps 519.39 ps 523.69 ps]
chain_in_a              time:   [2.3307 ns 2.3482 ns 2.3689 ns]
chain_split             time:   [6.3170 ns 6.3243 ns 6.3320 ns]
slice                   time:   [496.23 ps 498.66 ps 500.98 ps]
chain_in_a              time:   [2.3384 ns 2.3463 ns 2.3542 ns]
chain_split             time:   [6.3247 ns 6.3312 ns 6.3383 ns]

/// Panic with a nice error message.
#[cold]
fn panic_advance(error_info: &TryGetError) -> ! {
    panic!(
        "advance out of bounds: the len is {} but advancing by {}",
        error_info.available, error_info.requested
    );
}

use std::collections::VecDeque;
use std::panic;

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

fn try_it(label: &str, f: impl FnOnce() + panic::UnwindSafe) {
  match panic::catch_unwind(f) {
    Ok(()) => println!("{label:<28} ok"),
    Err(e) => {
      let msg = e
        .downcast_ref::<String>()
        .cloned()
        .or_else(|| e.downcast_ref::<&str>().map(|s| s.to_string()))
        .unwrap_or_default();
      println!("{label:<28} panic: {msg}");
    }
  }
}

fn main() {
  panic::set_hook(Box::new(|_| {})); // 不让默认 hook 打印，下面自己打印消息

  try_it("&[u8] get_u32 (3B)", || {
    (&[1u8, 2, 3][..]).get_u32();
  });
  try_it("Bytes get_u32 (3B)", || {
    Bytes::from_static(&[1, 2, 3]).get_u32();
  });
  try_it("Chain get_u32 (1+2B)", || {
    (&[1u8][..]).chain(&[2u8, 3][..]).get_u32();
  });
  try_it("&[u8] get_u8 (0B)", || {
    (&[][..]).get_u8();
  });
  try_it("&[u8] advance(5) (3B)", || {
    (&[1u8, 2, 3][..]).advance(5);
  });
  try_it("Bytes advance(5) (3B)", || {
    Bytes::from_static(&[1, 2, 3]).advance(5);
  });
  try_it("BytesMut advance(5) (3B)", || {
    BytesMut::from(&[1u8, 2, 3][..]).advance(5);
  });
  try_it("VecDeque advance(5) (3B)", || {
    VecDeque::from(vec![1u8, 2, 3]).advance(5);
  });
  try_it("&[u8] copy_to_bytes(5) (3B)", || {
    (&[1u8, 2, 3][..]).copy_to_bytes(5);
  });
  try_it("Bytes copy_to_bytes(5) (3B)", || {
    Bytes::from_static(&[1, 2, 3]).copy_to_bytes(5);
  });
  try_it("&[u8] get_uint(9) (16B)", || {
    (&[0u8; 16][..]).get_uint(9);
  });
}

&[u8] get_u32 (3B)           panic: advance out of bounds: the len is 3 but advancing by 4
Bytes get_u32 (3B)           panic: advance out of bounds: the len is 3 but advancing by 4
Chain get_u32 (1+2B)         panic: advance out of bounds: the len is 3 but advancing by 4
&[u8] get_u8 (0B)            panic: advance out of bounds: the len is 0 but advancing by 1
&[u8] advance(5) (3B)        panic: advance out of bounds: the len is 3 but advancing by 5
Bytes advance(5) (3B)        panic: cannot advance past `remaining`: 5 <= 3
BytesMut advance(5) (3B)     panic: cannot advance past `remaining`: 5 <= 3
VecDeque advance(5) (3B)     panic: range end index 5 out of range for slice of length 3
&[u8] copy_to_bytes(5) (3B)  panic: advance out of bounds: the len is 3 but advancing by 5
Bytes copy_to_bytes(5) (3B)  panic: split_to out of bounds: 5 <= 3
&[u8] get_uint(9) (16B)      panic: size too large: the integer type can fit 8 bytes, but nbytes is 9

    fn try_get_u32(&mut self) -> Result<u32, TryGetError> {
        buf_try_get_impl!(self, u32::from_be_bytes)
    }

    fn get_int(&mut self, nbytes: usize) -> i64 {
        sign_extend(self.get_uint(nbytes), nbytes)
    }

// https://en.wikipedia.org/wiki/Sign_extension
fn sign_extend(val: u64, nbytes: usize) -> i64 {
    if nbytes == 0 {
        // avoid `val << 64` panic
        0
    } else {
        let shift = (8 - nbytes) * 8;
        (val << shift) as i64 >> shift
    }
}

    fn try_get_int(&mut self, nbytes: usize) -> Result<i64, TryGetError> {
        buf_try_get_impl!(be => self, i64, nbytes);
    }

    fn copy_to_bytes(&mut self, len: usize) -> crate::Bytes {
        use super::BufMut;

        if self.remaining() < len {
            panic_advance(&TryGetError {
                requested: len,
                available: self.remaining(),
            });
        }

        let mut ret = crate::BytesMut::with_capacity(len);
        ret.put(self.take(len));
        ret.freeze()
    }

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

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

struct Counting;
static ALLOCS: AtomicUsize = AtomicUsize::new(0);
static BYTES: AtomicUsize = AtomicUsize::new(0);

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

#[global_allocator]
static GLOBAL: Counting = Counting;

// 对同一个 Buf 连做三次 copy_to_bytes(100)，分别记分配次数和字节数
fn run<B: Buf>(label: &str, mut buf: B) {
  let mut cols = Vec::with_capacity(3);
  let mut keep = Vec::with_capacity(3);
  for _ in 0..3 {
    let (a0, b0) = (ALLOCS.load(Relaxed), BYTES.load(Relaxed));
    let piece = buf.copy_to_bytes(100);
    let (a1, b1) = (ALLOCS.load(Relaxed), BYTES.load(Relaxed));
    cols.push(format!("{}次/{}B", a1 - a0, b1 - b0));
    keep.push(piece);
  }
  println!("{label:<20} {}", cols.join("  "));
}

fn main() {
  let raw = vec![7u8; 1000];
  let vec_for_bytes = raw.clone(); // len == cap
  let mut bm = BytesMut::with_capacity(1000);
  bm.extend_from_slice(&raw);
  let deque: VecDeque<u8> = raw.iter().copied().collect();

  run("&[u8]", &raw[..]);
  run("Bytes::from_static", Bytes::from_static(&[7u8; 1000]));
  run("Bytes::from(Vec)", Bytes::from(vec_for_bytes));
  run("BytesMut", bm);
  run("VecDeque<u8>", deque);
}

&[u8]                1次/100B  1次/100B  1次/100B
Bytes::from_static   0次/0B  0次/0B  0次/0B
Bytes::from(Vec)     1次/24B  0次/0B  0次/0B
BytesMut             1次/40B  0次/0B  0次/0B
VecDeque<u8>         1次/100B  1次/100B  1次/100B

    fn advance(&mut self, cnt: usize) {
        if self.len() < cnt {
            panic_advance(&TryGetError {
                requested: cnt,
                available: self.len(),
            });
        }

        *self = &self[cnt..];
    }

impl<T: Buf + ?Sized> Buf for &mut T {
    deref_forward_buf!();
}

impl<T: Buf + ?Sized> Buf for Box<T> {
    deref_forward_buf!();
}

Box<dyn Buf>     allocs=0
Wrapper          allocs=1