//! Immutable and mutable collections of I/O buffers. //! //! [`IoBufs`] and [`IoBufsMut`] present multiple buffers as one logical byte //! stream. [`Builder`] assembles an [`IoBufs`] from encoded bytes and existing //! byte segments, and [`EncodeExt`] encodes values into [`IoBufs`] through a //! [`Builder`]. use super::{ buf::{IoBuf, IoBufMut}, panic_advance, }; use bytes::{Buf as _, BufMut, Bytes, BytesMut}; use commonware_codec::{Buf, BufsMut, EncodeSize, Input, OwnedBytes, Write}; use std::{collections::VecDeque, io::IoSlice}; /// Container for one or more immutable buffers. #[derive(Clone, Debug)] pub struct IoBufs { inner: IoBufsInner, } /// Internal immutable representation. /// /// - Representation is canonical and minimal for readable data: /// - `Single` is the only representation for empty data and one-chunk data. /// - `Chunked` is used only when four or more readable chunks remain. /// - `Pair`, `Triple`, and `Chunked` never store empty chunks. #[derive(Debug)] enum IoBufsInner { /// Single buffer (fast path). Single(IoBuf), /// Two buffers (fast path). Pair([IoBuf; 2]), /// Three buffers (fast path). Triple([IoBuf; 3]), /// Four or more buffers. Chunked(VecDeque), } // Clones `Pair` and `Triple` chunks one at a time so each stays in registers // until it is stored into the result. The derived clone of a chunk array builds // a temporary array, to drop the chunks already cloned if a later clone // unwinds, and then copies it into place. That copy reloads the array with // loads wider than the stores that filled it and stalls store forwarding. impl Clone for IoBufsInner { #[inline] fn clone(&self) -> Self { match self { Self::Single(buf) => Self::Single(buf.clone()), Self::Pair([a, b]) => Self::Pair([a.clone(), b.clone()]), Self::Triple([a, b, c]) => Self::Triple([a.clone(), b.clone(), c.clone()]), Self::Chunked(bufs) => Self::Chunked(bufs.clone()), } } } impl Default for IoBufs { fn default() -> Self { Self { inner: IoBufsInner::Single(IoBuf::default()), } } } impl IoBufs { /// Build canonical immutable chunk storage from readable chunks. /// /// Empty chunks are removed before representation selection. fn from_chunks_iter(chunks: impl IntoIterator) -> Self { let mut iter = chunks.into_iter().filter(|buf| !buf.is_empty()); let first = match iter.next() { Some(first) => first, None => return Self::default(), }; let second = match iter.next() { Some(second) => second, None => { return Self { inner: IoBufsInner::Single(first), }; } }; let third = match iter.next() { Some(third) => third, None => { return Self { inner: IoBufsInner::Pair([first, second]), }; } }; let fourth = match iter.next() { Some(fourth) => fourth, None => { return Self { inner: IoBufsInner::Triple([first, second, third]), }; } }; let mut bufs = VecDeque::with_capacity(4); bufs.push_back(first); bufs.push_back(second); bufs.push_back(third); bufs.push_back(fourth); bufs.extend(iter); Self { inner: IoBufsInner::Chunked(bufs), } } /// Re-establish canonical immutable representation invariants. fn canonicalize(&mut self) { let inner = std::mem::replace(&mut self.inner, IoBufsInner::Single(IoBuf::default())); self.inner = match inner { IoBufsInner::Single(buf) => { if buf.is_empty() { IoBufsInner::Single(IoBuf::default()) } else { IoBufsInner::Single(buf) } } IoBufsInner::Pair([a, b]) => Self::from_chunks_iter([a, b]).inner, IoBufsInner::Triple([a, b, c]) => Self::from_chunks_iter([a, b, c]).inner, IoBufsInner::Chunked(bufs) => Self::from_chunks_iter(bufs).inner, }; } /// Returns a reference to the single contiguous buffer, if present. /// /// Returns `Some` only when all remaining data is in one contiguous buffer. pub const fn as_single(&self) -> Option<&IoBuf> { match &self.inner { IoBufsInner::Single(buf) => Some(buf), _ => None, } } /// Consume this container and return the single buffer if present. /// /// Returns `Ok(IoBuf)` only when all remaining data is already contained in /// a single chunk. Returns `Err(Self)` with the original container /// otherwise. pub fn try_into_single(self) -> Result { match self.inner { IoBufsInner::Single(buf) => Ok(buf), inner => Err(Self { inner }), } } /// Number of bytes remaining across all buffers. #[inline] pub fn len(&self) -> usize { self.remaining() } /// Number of non-empty readable chunks. #[inline] pub fn chunk_count(&self) -> usize { // This assumes canonical form. match &self.inner { IoBufsInner::Single(buf) => { if buf.is_empty() { 0 } else { 1 } } IoBufsInner::Pair(_) => 2, IoBufsInner::Triple(_) => 3, IoBufsInner::Chunked(bufs) => bufs.len(), } } /// Returns the readable chunk at the given index without advancing the buffers. /// /// Returns `None` when `index` is greater than or equal to [`Self::chunk_count`]. pub fn chunk_at(&self, index: usize) -> Option<&[u8]> { match &self.inner { IoBufsInner::Single(buf) => (index == 0 && !buf.is_empty()).then(|| buf.as_ref()), IoBufsInner::Pair(bufs) => bufs.get(index).map(AsRef::as_ref), IoBufsInner::Triple(bufs) => bufs.get(index).map(AsRef::as_ref), IoBufsInner::Chunked(bufs) => bufs.get(index).map(AsRef::as_ref), } } /// Iterates over the readable buffer owners without coalescing them. pub(crate) fn iter(&self) -> impl ExactSizeIterator + Clone { (0..self.chunk_count()).map(|index| match &self.inner { IoBufsInner::Single(buf) => { assert_eq!(index, 0); buf } IoBufsInner::Pair(bufs) => &bufs[index], IoBufsInner::Triple(bufs) => &bufs[index], IoBufsInner::Chunked(bufs) => &bufs[index], }) } /// Whether all buffers are empty. #[inline] pub const fn is_empty(&self) -> bool { // Only Single can be empty. All other variants contain readable bytes. match &self.inner { IoBufsInner::Single(buf) => buf.is_empty(), IoBufsInner::Pair(_) | IoBufsInner::Triple(_) | IoBufsInner::Chunked(_) => false, } } /// Whether this contains a single contiguous buffer. /// /// When true, `chunk()` returns all remaining bytes. #[inline] pub const fn is_single(&self) -> bool { matches!(self.inner, IoBufsInner::Single(_)) } /// Visit each readable chunk in order without coalescing. #[inline] pub fn for_each_chunk(&self, mut f: impl FnMut(&[u8])) { match &self.inner { IoBufsInner::Single(buf) => { let chunk = buf.as_ref(); if !chunk.is_empty() { f(chunk); } } IoBufsInner::Pair(pair) => { for buf in pair { let chunk = buf.as_ref(); if !chunk.is_empty() { f(chunk); } } } IoBufsInner::Triple(triple) => { for buf in triple { let chunk = buf.as_ref(); if !chunk.is_empty() { f(chunk); } } } IoBufsInner::Chunked(bufs) => { for buf in bufs { let chunk = buf.as_ref(); if !chunk.is_empty() { f(chunk); } } } } } /// Prepend a buffer to the front. /// /// Empty input buffers are ignored. pub fn prepend(&mut self, buf: IoBuf) { if buf.is_empty() { return; } let inner = std::mem::replace(&mut self.inner, IoBufsInner::Single(IoBuf::default())); self.inner = match inner { IoBufsInner::Single(existing) if existing.is_empty() => IoBufsInner::Single(buf), IoBufsInner::Single(existing) => IoBufsInner::Pair([buf, existing]), IoBufsInner::Pair([a, b]) => IoBufsInner::Triple([buf, a, b]), IoBufsInner::Triple([a, b, c]) => { let mut bufs = VecDeque::with_capacity(4); bufs.push_back(buf); bufs.push_back(a); bufs.push_back(b); bufs.push_back(c); IoBufsInner::Chunked(bufs) } IoBufsInner::Chunked(mut bufs) => { bufs.push_front(buf); IoBufsInner::Chunked(bufs) } }; } /// Append a buffer to the back. /// /// Empty input buffers are ignored. pub fn append(&mut self, buf: IoBuf) { if buf.is_empty() { return; } let inner = std::mem::replace(&mut self.inner, IoBufsInner::Single(IoBuf::default())); self.inner = match inner { IoBufsInner::Single(existing) if existing.is_empty() => IoBufsInner::Single(buf), IoBufsInner::Single(existing) => IoBufsInner::Pair([existing, buf]), IoBufsInner::Pair([a, b]) => IoBufsInner::Triple([a, b, buf]), IoBufsInner::Triple([a, b, c]) => { let mut bufs = VecDeque::with_capacity(4); bufs.push_back(a); bufs.push_back(b); bufs.push_back(c); bufs.push_back(buf); IoBufsInner::Chunked(bufs) } IoBufsInner::Chunked(mut bufs) => { bufs.push_back(buf); IoBufsInner::Chunked(bufs) } }; } /// Splits the buffer(s) into two at the given index. /// /// Afterwards `self` contains bytes `[at, len)`, and the returned /// [`IoBufs`] contains bytes `[0, at)`. /// /// Whole chunks are moved without copying. If the split point lands inside /// a chunk, the chunk is split zero-copy via [`IoBuf::split_to`]. /// /// # Panics /// /// Panics if `at > len`. pub fn split_to(&mut self, at: usize) -> Self { if at == 0 { return Self::default(); } let remaining = self.remaining(); assert!( at <= remaining, "split_to out of bounds: {:?} <= {:?}", at, remaining, ); if at == remaining { return std::mem::take(self); } // A prefix within the front chunk of a chunked representation splits // that chunk in place, without moving the deque out and back. The // smaller representations handle that case in their own arms below. if matches!(self.inner, IoBufsInner::Chunked(_)) && let Some(prefix) = self.split_front(at) { return Self::from(prefix); } let inner = std::mem::replace(&mut self.inner, IoBufsInner::Single(IoBuf::default())); match inner { IoBufsInner::Single(mut buf) => { // Delegate directly and keep remainder as single let prefix = buf.split_to(at); self.inner = IoBufsInner::Single(buf); Self::from(prefix) } IoBufsInner::Pair([mut a, mut b]) => { let a_len = a.remaining(); if at < a_len { // Split stays entirely in chunk `a`. let prefix = a.split_to(at); self.inner = IoBufsInner::Pair([a, b]); return Self::from(prefix); } if at == a_len { // Exact chunk boundary: move `a` out, keep `b`. self.inner = IoBufsInner::Single(b); return Self::from(a); } // Split crosses from `a` into `b`. let b_prefix_len = at - a_len; let b_prefix = b.split_to(b_prefix_len); self.inner = IoBufsInner::Single(b); Self { inner: IoBufsInner::Pair([a, b_prefix]), } } IoBufsInner::Triple([mut a, mut b, mut c]) => { let a_len = a.remaining(); if at < a_len { // Split stays entirely in chunk `a`. let prefix = a.split_to(at); self.inner = IoBufsInner::Triple([a, b, c]); return Self::from(prefix); } if at == a_len { // Exact boundary after `a`. self.inner = IoBufsInner::Pair([b, c]); return Self::from(a); } let mut remaining = at - a_len; let b_len = b.remaining(); if remaining < b_len { // Split lands inside `b`. let b_prefix = b.split_to(remaining); self.inner = IoBufsInner::Pair([b, c]); return Self { inner: IoBufsInner::Pair([a, b_prefix]), }; } if remaining == b_len { // Exact boundary after `b`. self.inner = IoBufsInner::Single(c); return Self { inner: IoBufsInner::Pair([a, b]), }; } // Split reaches into `c`. remaining -= b_len; let c_prefix = c.split_to(remaining); self.inner = IoBufsInner::Single(c); Self { inner: IoBufsInner::Triple([a, b, c_prefix]), } } IoBufsInner::Chunked(mut bufs) => { // Count the chunks the prefix takes whole. The split lands // inside the next one when bytes remain. let mut remaining = at; let mut whole = 0; for buf in &bufs { let avail = buf.remaining(); if remaining == 0 || remaining < avail { break; } remaining -= avail; whole += 1; } // Split in place the chunk that the prefix ends inside. The // bounds check above guarantees that chunk exists. let piece = (remaining > 0).then(|| bufs[whole].split_to(remaining)); // Move out the side with fewer chunks, counting the piece // toward the prefix, and keep the deque for the other, so the // copying and any new deque scale with the shorter side. A // prefix of fewer than four chunks from a large container // needs no new deque. let prefix_chunks = whole + usize::from(piece.is_some()); let suffix_chunks = bufs.len() - whole; if prefix_chunks <= suffix_chunks { let prefix = Self::from_exact_chunks(bufs.drain(..whole).chain(piece), prefix_chunks); self.inner = Self::from_chunk_deque(bufs).inner; prefix } else { self.inner = Self::from_exact_chunks(bufs.drain(whole..), suffix_chunks).inner; bufs.extend(piece); Self::from_chunk_deque(bufs) } } } } /// Builds canonical storage from a deque of non-empty chunks and keeps the /// deque when it holds four or more. fn from_chunk_deque(bufs: VecDeque) -> Self { if bufs.len() >= 4 { Self { inner: IoBufsInner::Chunked(bufs), } } else { Self::from_chunks_iter(bufs) } } /// Builds canonical storage from `len` non-empty chunks and allocates a /// deque only when there are four or more. fn from_exact_chunks(chunks: impl Iterator, len: usize) -> Self { if len >= 4 { Self { inner: IoBufsInner::Chunked(chunks.collect()), } } else { Self::from_chunks_iter(chunks) } } /// Splits off the first `len` bytes as a view of the front chunk when that chunk holds them /// all, and returns `None` when they span chunks. /// /// A prefix that leaves part of the front chunk returns right after its split. One that /// consumes the whole front chunk of a multi-chunk representation restores the canonical form /// through the outlined [`Self::remove_front`]. #[inline(always)] fn split_front(&mut self, len: usize) -> Option { let front = match &mut self.inner { IoBufsInner::Single(buf) => return (buf.len() >= len).then(|| buf.split_to(len)), IoBufsInner::Pair([front, ..]) | IoBufsInner::Triple([front, ..]) => front, IoBufsInner::Chunked(bufs) => bufs.front_mut()?, }; if front.len() > len { return Some(front.split_to(len)); } if front.len() < len { return None; } let prefix = std::mem::take(front); self.remove_front(); Some(prefix) } /// Removes the empty front chunk of a multi-chunk representation and /// restores the canonical form. fn remove_front(&mut self) { match &mut self.inner { IoBufsInner::Single(_) => {} IoBufsInner::Pair(_) | IoBufsInner::Triple(_) => self.canonicalize(), IoBufsInner::Chunked(bufs) => { bufs.pop_front(); if bufs.len() <= 3 { self.canonicalize(); } } } } /// Takes a prefix that spans chunks, copying it once. // Kept out of line so the inlined front-chunk fast path stays small. #[inline(never)] fn take_spanning(&mut self, len: usize) -> T { // A prefix that spans chunks is copied once, into one `IoBuf` when `T` claims `IoBuf` // views and into `Bytes` otherwise. Claiming an empty view tells which case applies. if T::claim(IoBuf::default()).is_ok() { T::claim(self.split_to(len).coalesce()) .unwrap_or_else(|view| T::from_bytes(view.into())) } else { T::from_bytes(self.copy_to_bytes(len)) } } /// Coalesce all remaining bytes into a single contiguous [`IoBuf`]. /// /// Zero-copy if only one buffer. Copies into one native heap allocation /// if multiple buffers, so the result supports zero-copy /// [`IoBuf::try_into_mut`]. #[inline] pub fn coalesce(self) -> IoBuf { match self.inner { IoBufsInner::Single(buf) => buf, inner => { let bufs = Self { inner }; let mut out = IoBufMut::with_capacity(bufs.remaining()); bufs.for_each_chunk(|chunk| out.put_slice(chunk)); out.freeze() } } } } impl Buf for IoBufs { /// Returns the prefix as a view of the front chunk when that chunk holds it all, and copies it /// once when it spans chunks. A `T` that claims [`IoBuf`] views receives the [`IoBuf`], and /// any other `T` is built from its bytes. // Inlined with `split_front` so a front-chunk view is built in registers at the call site. A // view returned through a stack slot can be reloaded with loads wider than the stores that // filled it, and that stalls store forwarding. #[inline(always)] fn take_owned(&mut self, len: usize) -> T { if let Some(view) = self.split_front(len) { return T::claim(view).unwrap_or_else(|view| T::from_bytes(view.into())); } self.take_spanning(len) } } impl bytes::Buf for IoBufs { #[inline] fn has_remaining(&self) -> bool { !self.is_empty() } #[inline] fn remaining(&self) -> usize { match &self.inner { IoBufsInner::Single(buf) => buf.remaining(), IoBufsInner::Pair([a, b]) => a.remaining().saturating_add(b.remaining()), IoBufsInner::Triple([a, b, c]) => a .remaining() .saturating_add(b.remaining()) .saturating_add(c.remaining()), IoBufsInner::Chunked(bufs) => bufs .iter() .map(|b| b.remaining()) .fold(0, usize::saturating_add), } } #[inline] fn chunk(&self) -> &[u8] { match &self.inner { IoBufsInner::Single(buf) => buf.chunk(), IoBufsInner::Pair([a, b]) => { if a.remaining() > 0 { a.chunk() } else if b.remaining() > 0 { b.chunk() } else { &[] } } IoBufsInner::Triple([a, b, c]) => { if a.remaining() > 0 { a.chunk() } else if b.remaining() > 0 { b.chunk() } else if c.remaining() > 0 { c.chunk() } else { &[] } } IoBufsInner::Chunked(bufs) => { for buf in bufs.iter() { if buf.remaining() > 0 { return buf.chunk(); } } &[] } } } #[inline] fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize { if dst.is_empty() { return 0; } match &self.inner { IoBufsInner::Single(buf) => { let chunk = buf.chunk(); if !chunk.is_empty() { dst[0] = IoSlice::new(chunk); return 1; } 0 } IoBufsInner::Pair([a, b]) => fill_vectored_from_chunks(dst, [a.chunk(), b.chunk()]), IoBufsInner::Triple([a, b, c]) => { fill_vectored_from_chunks(dst, [a.chunk(), b.chunk(), c.chunk()]) } IoBufsInner::Chunked(bufs) => { fill_vectored_from_chunks(dst, bufs.iter().map(|buf| buf.chunk())) } } } #[inline] fn advance(&mut self, cnt: usize) { let should_canonicalize = match &mut self.inner { IoBufsInner::Single(buf) => { buf.advance(cnt); false } IoBufsInner::Pair(pair) => advance_small_chunks(pair.as_mut_slice(), cnt), IoBufsInner::Triple(triple) => advance_small_chunks(triple.as_mut_slice(), cnt), IoBufsInner::Chunked(bufs) => { advance_chunked_front(bufs, cnt); bufs.len() <= 3 } }; if should_canonicalize { self.canonicalize(); } } #[inline] fn copy_to_bytes(&mut self, len: usize) -> Bytes { let (result, needs_canonicalize) = match &mut self.inner { IoBufsInner::Single(buf) => return buf.copy_to_bytes(len), IoBufsInner::Pair(pair) => { copy_to_bytes_small_chunks(pair, len, "IoBufs::copy_to_bytes: not enough data") } IoBufsInner::Triple(triple) => { copy_to_bytes_small_chunks(triple, len, "IoBufs::copy_to_bytes: not enough data") } IoBufsInner::Chunked(bufs) => { copy_to_bytes_chunked(bufs, len, "IoBufs::copy_to_bytes: not enough data") } }; if needs_canonicalize { self.canonicalize(); } result } } /// Zero-copy: wraps the buffer as the single chunk. impl From for IoBufs { fn from(buf: IoBuf) -> Self { Self { inner: IoBufsInner::Single(buf), } } } /// Zero-copy: freezes the buffer and wraps it as the single chunk. impl From for IoBufs { fn from(buf: IoBufMut) -> Self { Self { inner: IoBufsInner::Single(buf.freeze()), } } } /// Zero-copy via `From for IoBuf`. impl From for IoBufs { fn from(bytes: Bytes) -> Self { Self::from(IoBuf::from(bytes)) } } /// Zero-copy via `From for IoBuf`. impl From for IoBufs { fn from(bytes: BytesMut) -> Self { Self::from(IoBuf::from(bytes)) } } /// Zero-copy via `From> for IoBuf`. impl From> for IoBufs { fn from(vec: Vec) -> Self { Self::from(IoBuf::from(vec)) } } /// Zero-copy: collects the chunks, dropping empty ones. impl From> for IoBufs { fn from(bufs: Vec) -> Self { Self::from_iter(bufs) } } /// Zero-copy: collects the chunks, dropping empty ones. impl From<[IoBuf; N]> for IoBufs { fn from(bufs: [IoBuf; N]) -> Self { Self::from_iter(bufs) } } /// Zero-copy: collects the chunks, dropping empty ones. impl FromIterator for IoBufs { fn from_iter>(bufs: T) -> Self { Self::from_chunks_iter(bufs) } } /// Zero-copy: creates a static view with no owner. impl From<&'static [u8; N]> for IoBufs { fn from(array: &'static [u8; N]) -> Self { Self::from(IoBuf::from(array)) } } /// Zero-copy: creates a static view with no owner. impl From<&'static [u8]> for IoBufs { fn from(slice: &'static [u8]) -> Self { Self::from(IoBuf::from(slice)) } } /// Container for one or more mutable buffers. /// /// # Capacity retention /// /// The intended usage is fill-then-read: write into the container (through /// [`BufMut`], [`Self::copy_from_slice`], or /// [`Blob::read_at_buf`](crate::Blob::read_at_buf)), then consume it through /// [`bytes::Buf`]. Caller-reserved write capacity generally survives read /// operations, with three exceptions: /// - The deque-backed read paths (four or more chunks) skip past a chunk /// with no readable bytes by popping it, so a never-filled chunk ordered /// before readable data loses its capacity when a read crosses it. /// - The same paths pop any chunk whose readable bytes a read fully /// consumes, releasing its writable tail. The two and three-chunk shapes /// advance in place and retain such tails, so `remaining_mut()` after an /// identical read can differ by shape. /// - A nonzero `copy_to_bytes` that exactly drains the front chunk's /// readable bytes consumes that chunk's whole handle on every shape (see /// [`IoBufMut`]'s `copy_to_bytes` doc), so spare capacity behind an /// exactly-drained chunk is released rather than retained. /// /// Interleaving writes after reads is outside this contract. A read can /// leave a drained chunk's surviving tail ordered before later readable /// bytes, and a subsequent [`BufMut`] write fills that tail first. #[derive(Debug)] pub struct IoBufsMut { inner: IoBufsMutInner, } /// Internal mutable representation. /// /// Construction and canonicalization keep every chunk that still owns /// storage (`capacity() > 0`), readable or not, so caller-reserved write /// capacity generally survives read operations. The reader-facing rules and /// the accepted exceptions are documented on [`IoBufsMut`] under "Capacity /// retention". Only fully-drained chunks (capacity consumed by `advance`) /// and empty defaults are removed as the shape collapses. #[derive(Debug)] enum IoBufsMutInner { /// Single buffer (common case, no allocation). Single(IoBufMut), /// Two buffers (fast path, no VecDeque allocation). Pair([IoBufMut; 2]), /// Three buffers (fast path, no VecDeque allocation). Triple([IoBufMut; 3]), /// Four or more buffers. Chunked(VecDeque), } impl Default for IoBufsMut { fn default() -> Self { Self { inner: IoBufsMutInner::Single(IoBufMut::default()), } } } impl IoBufsMut { /// Build mutable chunk storage from already-filtered chunks. /// /// This helper intentionally does not filter. Callers route through /// [`Self::from_writable_chunks_iter`] so storage-owning chunks are kept. fn from_chunks_iter(chunks: impl IntoIterator) -> Self { let mut iter = chunks.into_iter(); let first = match iter.next() { Some(first) => first, None => return Self::default(), }; let second = match iter.next() { Some(second) => second, None => { return Self { inner: IoBufsMutInner::Single(first), }; } }; let third = match iter.next() { Some(third) => third, None => { return Self { inner: IoBufsMutInner::Pair([first, second]), }; } }; let fourth = match iter.next() { Some(fourth) => fourth, None => { return Self { inner: IoBufsMutInner::Triple([first, second, third]), }; } }; let mut bufs = VecDeque::with_capacity(4); bufs.push_back(first); bufs.push_back(second); bufs.push_back(third); bufs.push_back(fourth); bufs.extend(iter); Self { inner: IoBufsMutInner::Chunked(bufs), } } /// Build canonical mutable chunk storage from writable chunks. /// /// Keeps chunks that still own storage, readable or not (`capacity()` /// covers both readable bytes and the writable tail), so a never-filled /// chunk's reserved capacity is not discarded. Fully-drained chunks /// (capacity consumed by `advance`) and empty defaults are removed. fn from_writable_chunks_iter(chunks: impl IntoIterator) -> Self { Self::from_chunks_iter(chunks.into_iter().filter(|buf| buf.capacity() > 0)) } /// Re-establish canonical mutable representation invariants. /// /// Uses the same storage-keeping filter as construction, so chunks that /// were drained of readable bytes but still own writable capacity survive. /// Reads then leave `remaining_mut()` unchanged, apart from the exceptions /// listed under "Capacity retention" on [`IoBufsMut`]. fn canonicalize(&mut self) { let inner = std::mem::replace(&mut self.inner, IoBufsMutInner::Single(IoBufMut::default())); self.inner = match inner { IoBufsMutInner::Single(buf) => IoBufsMutInner::Single(buf), IoBufsMutInner::Pair([a, b]) => Self::from_writable_chunks_iter([a, b]).inner, IoBufsMutInner::Triple([a, b, c]) => Self::from_writable_chunks_iter([a, b, c]).inner, IoBufsMutInner::Chunked(bufs) => Self::from_writable_chunks_iter(bufs).inner, }; } #[inline] fn for_each_chunk_mut(&mut self, mut f: impl FnMut(&mut IoBufMut)) { match &mut self.inner { IoBufsMutInner::Single(buf) => f(buf), IoBufsMutInner::Pair(pair) => { for buf in pair.iter_mut() { f(buf); } } IoBufsMutInner::Triple(triple) => { for buf in triple.iter_mut() { f(buf); } } IoBufsMutInner::Chunked(bufs) => { for buf in bufs.iter_mut() { f(buf); } } } } /// Returns a reference to the single contiguous buffer, if present. /// /// Returns `Some` only when this is currently represented as one chunk. pub const fn as_single(&self) -> Option<&IoBufMut> { match &self.inner { IoBufsMutInner::Single(buf) => Some(buf), _ => None, } } /// Returns a mutable reference to the single contiguous buffer, if present. /// /// Returns `Some` only when this is currently represented as one chunk. pub const fn as_single_mut(&mut self) -> Option<&mut IoBufMut> { match &mut self.inner { IoBufsMutInner::Single(buf) => Some(buf), _ => None, } } /// Consume this container and return the single buffer if present. /// /// Returns `Ok(IoBufMut)` only when readable data is represented as one /// chunk. Returns `Err(Self)` with the original container otherwise. pub fn try_into_single(self) -> Result { match self.inner { IoBufsMutInner::Single(buf) => Ok(buf), inner => Err(Self { inner }), } } /// Number of bytes remaining across all buffers. #[inline] pub fn len(&self) -> usize { self.remaining() } /// Whether all buffers are empty. #[inline] pub fn is_empty(&self) -> bool { match &self.inner { IoBufsMutInner::Single(buf) => buf.is_empty(), IoBufsMutInner::Pair([a, b]) => a.is_empty() && b.is_empty(), IoBufsMutInner::Triple([a, b, c]) => a.is_empty() && b.is_empty() && c.is_empty(), IoBufsMutInner::Chunked(bufs) => bufs.iter().all(IoBufMut::is_empty), } } /// Whether this contains a single contiguous buffer. /// /// When true, `chunk()` returns all remaining bytes. #[inline] pub const fn is_single(&self) -> bool { matches!(self.inner, IoBufsMutInner::Single(_)) } /// Freeze into immutable [`IoBufs`]. /// /// Each chunk freezes through [`IoBufMut::freeze`], so a chunk whose spare /// capacity is too large to record in its handle allocates its owner record. pub fn freeze(self) -> IoBufs { match self.inner { IoBufsMutInner::Single(buf) => IoBufs::from(buf.freeze()), IoBufsMutInner::Pair([a, b]) => IoBufs::from_chunks_iter([a.freeze(), b.freeze()]), IoBufsMutInner::Triple([a, b, c]) => { IoBufs::from_chunks_iter([a.freeze(), b.freeze(), c.freeze()]) } IoBufsMutInner::Chunked(bufs) => { IoBufs::from_chunks_iter(bufs.into_iter().map(IoBufMut::freeze)) } } } /// Coalesce all buffers into a single contiguous [`IoBufMut`]. /// /// Zero-copy if only one buffer. Otherwise copies into one allocation of /// exactly the readable length. pub fn coalesce(self) -> IoBufMut { match self.inner { IoBufsMutInner::Single(buf) => buf, inner => { let mut bufs = Self { inner }; let mut out = IoBufMut::with_capacity(bufs.len()); bufs.for_each_chunk_mut(|buf| out.put_slice(buf.as_ref())); out } } } /// Returns the total capacity across all buffers. pub fn capacity(&self) -> usize { match &self.inner { IoBufsMutInner::Single(buf) => buf.capacity(), IoBufsMutInner::Pair([a, b]) => a.capacity().saturating_add(b.capacity()), IoBufsMutInner::Triple([a, b, c]) => a .capacity() .saturating_add(b.capacity()) .saturating_add(c.capacity()), IoBufsMutInner::Chunked(bufs) => bufs .iter() .map(|b| b.capacity()) .fold(0, usize::saturating_add), } } /// Sets the length of the buffer(s) to `len`, distributing across chunks /// while preserving the current chunk layout. /// /// This is useful for APIs that must fill caller-provided buffer structure /// in place (for example [`Blob::read_at_buf`](crate::Blob::read_at_buf)). /// /// # Safety /// /// Caller must initialize all `len` bytes before the buffer is read. /// /// # Panics /// /// Panics if `len` exceeds total capacity. pub(crate) unsafe fn set_len(&mut self, len: usize) { let capacity = self.capacity(); assert!( len <= capacity, "set_len({len}) exceeds capacity({capacity})" ); let mut remaining = len; self.for_each_chunk_mut(|buf| { let cap = buf.capacity(); let to_set = remaining.min(cap); // SAFETY: forwarded from this method's contract. The caller // initializes all `len` bytes before any read. unsafe { buf.set_len(to_set) }; remaining -= to_set; }); } /// Copy data from a slice into the buffers. /// /// # Panics /// /// Panics if the slice length doesn't match the total buffer length. pub fn copy_from_slice(&mut self, src: &[u8]) { assert_eq!( src.len(), self.len(), "source slice length must match buffer length" ); let mut offset = 0; self.for_each_chunk_mut(|buf| { let len = buf.len(); buf.as_mut().copy_from_slice(&src[offset..offset + len]); offset += len; }); } } impl Input for IoBufsMut { type Buf = IoBufs; } /// Freezes every chunk (see [`IoBufsMut::freeze`]). impl From for IoBufs { fn from(bufs: IoBufsMut) -> Self { bufs.freeze() } } impl bytes::Buf for IoBufsMut { #[inline] fn has_remaining(&self) -> bool { !self.is_empty() } #[inline] fn remaining(&self) -> usize { match &self.inner { IoBufsMutInner::Single(buf) => buf.remaining(), IoBufsMutInner::Pair([a, b]) => a.remaining().saturating_add(b.remaining()), IoBufsMutInner::Triple([a, b, c]) => a .remaining() .saturating_add(b.remaining()) .saturating_add(c.remaining()), IoBufsMutInner::Chunked(bufs) => bufs .iter() .map(|b| b.remaining()) .fold(0, usize::saturating_add), } } #[inline] fn chunk(&self) -> &[u8] { match &self.inner { IoBufsMutInner::Single(buf) => buf.chunk(), IoBufsMutInner::Pair([a, b]) => { if a.remaining() > 0 { a.chunk() } else if b.remaining() > 0 { b.chunk() } else { &[] } } IoBufsMutInner::Triple([a, b, c]) => { if a.remaining() > 0 { a.chunk() } else if b.remaining() > 0 { b.chunk() } else if c.remaining() > 0 { c.chunk() } else { &[] } } IoBufsMutInner::Chunked(bufs) => { for buf in bufs.iter() { if buf.remaining() > 0 { return buf.chunk(); } } &[] } } } #[inline] fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize { if dst.is_empty() { return 0; } match &self.inner { IoBufsMutInner::Single(buf) => { let chunk = buf.chunk(); if !chunk.is_empty() { dst[0] = IoSlice::new(chunk); return 1; } 0 } IoBufsMutInner::Pair([a, b]) => fill_vectored_from_chunks(dst, [a.chunk(), b.chunk()]), IoBufsMutInner::Triple([a, b, c]) => { fill_vectored_from_chunks(dst, [a.chunk(), b.chunk(), c.chunk()]) } IoBufsMutInner::Chunked(bufs) => { fill_vectored_from_chunks(dst, bufs.iter().map(|buf| buf.chunk())) } } } #[inline] fn advance(&mut self, cnt: usize) { let should_canonicalize = match &mut self.inner { IoBufsMutInner::Single(buf) => { buf.advance(cnt); false } IoBufsMutInner::Pair(pair) => advance_small_chunks(pair.as_mut_slice(), cnt), IoBufsMutInner::Triple(triple) => advance_small_chunks(triple.as_mut_slice(), cnt), IoBufsMutInner::Chunked(bufs) => { advance_chunked_front(bufs, cnt); bufs.len() <= 3 } }; if should_canonicalize { self.canonicalize(); } } #[inline] fn copy_to_bytes(&mut self, len: usize) -> Bytes { // Zero-length drains must not disturb chunk state. The deque-backed // path pops chunks with no readable bytes, and popping a never-filled // chunk would discard its reserved capacity. if len == 0 { return Bytes::new(); } let (result, needs_canonicalize) = match &mut self.inner { IoBufsMutInner::Single(buf) => return buf.copy_to_bytes(len), IoBufsMutInner::Pair(pair) => { copy_to_bytes_small_chunks(pair, len, "IoBufsMut::copy_to_bytes: not enough data") } IoBufsMutInner::Triple(triple) => { copy_to_bytes_small_chunks(triple, len, "IoBufsMut::copy_to_bytes: not enough data") } IoBufsMutInner::Chunked(bufs) => { copy_to_bytes_chunked(bufs, len, "IoBufsMut::copy_to_bytes: not enough data") } }; if needs_canonicalize { self.canonicalize(); } result } } // SAFETY: Delegates to IoBufMut which implements BufMut safely. unsafe impl BufMut for IoBufsMut { #[inline] fn remaining_mut(&self) -> usize { match &self.inner { IoBufsMutInner::Single(buf) => buf.remaining_mut(), IoBufsMutInner::Pair([a, b]) => a.remaining_mut().saturating_add(b.remaining_mut()), IoBufsMutInner::Triple([a, b, c]) => a .remaining_mut() .saturating_add(b.remaining_mut()) .saturating_add(c.remaining_mut()), IoBufsMutInner::Chunked(bufs) => bufs .iter() .map(|b| b.remaining_mut()) .fold(0, usize::saturating_add), } } #[inline] unsafe fn advance_mut(&mut self, cnt: usize) { // SAFETY: The caller guarantees `cnt <= self.remaining_mut()`, so each // delegated advance remains within the writable chunks. unsafe { // On failure, every writable byte was consumed before the chunks ran // out, so the advanced amount (`cnt - remaining`) is exactly what was // available. let mut remaining = cnt; let advanced = match &mut self.inner { IoBufsMutInner::Single(buf) => { buf.advance_mut(cnt); return; } IoBufsMutInner::Pair(pair) => advance_mut_in_chunks(pair, &mut remaining), IoBufsMutInner::Triple(triple) => advance_mut_in_chunks(triple, &mut remaining), IoBufsMutInner::Chunked(bufs) => { let (first, second) = bufs.as_mut_slices(); advance_mut_in_chunks(first, &mut remaining) || advance_mut_in_chunks(second, &mut remaining) } }; if !advanced { panic_advance(cnt, cnt - remaining); } } } #[inline] fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice { match &mut self.inner { IoBufsMutInner::Single(buf) => buf.chunk_mut(), IoBufsMutInner::Pair(pair) => { if pair[0].remaining_mut() > 0 { pair[0].chunk_mut() } else if pair[1].remaining_mut() > 0 { pair[1].chunk_mut() } else { bytes::buf::UninitSlice::new(&mut []) } } IoBufsMutInner::Triple(triple) => { if triple[0].remaining_mut() > 0 { triple[0].chunk_mut() } else if triple[1].remaining_mut() > 0 { triple[1].chunk_mut() } else if triple[2].remaining_mut() > 0 { triple[2].chunk_mut() } else { bytes::buf::UninitSlice::new(&mut []) } } IoBufsMutInner::Chunked(bufs) => { for buf in bufs.iter_mut() { if buf.remaining_mut() > 0 { return buf.chunk_mut(); } } bytes::buf::UninitSlice::new(&mut []) } } } } /// Zero-copy: wraps the buffer as the single chunk. impl From for IoBufsMut { fn from(buf: IoBufMut) -> Self { Self { inner: IoBufsMutInner::Single(buf), } } } /// Convert a [`Vec`] into a single-buffer [`IoBufsMut`]. /// /// Adopts the vec's allocation and capacity via `From> for IoBufMut`. impl From> for IoBufsMut { fn from(vec: Vec) -> Self { Self { inner: IoBufsMutInner::Single(IoBufMut::from(vec)), } } } /// Reuses the allocation via `From for IoBufMut` when no other /// handle shares it and copies otherwise. Either way at least the reserved /// capacity is kept. impl From for IoBufsMut { fn from(bytes: BytesMut) -> Self { Self { inner: IoBufsMutInner::Single(IoBufMut::from(bytes)), } } } /// Zero-copy: collects the chunks, dropping zero-capacity ones. impl From> for IoBufsMut { fn from(bufs: Vec) -> Self { Self::from_iter(bufs) } } /// Zero-copy: collects the chunks, dropping zero-capacity ones. impl From<[IoBufMut; N]> for IoBufsMut { fn from(bufs: [IoBufMut; N]) -> Self { Self::from_iter(bufs) } } /// Zero-copy: collects the chunks, dropping zero-capacity ones. impl FromIterator for IoBufsMut { fn from_iter>(bufs: T) -> Self { Self::from_writable_chunks_iter(bufs) } } /// Copies via `From<[u8; N]> for IoBufMut`. impl From<[u8; N]> for IoBufsMut { fn from(array: [u8; N]) -> Self { Self { inner: IoBufsMutInner::Single(IoBufMut::from(array)), } } } /// Drain `len` readable bytes from a small fixed chunk array (`Pair`/`Triple`). /// /// Returns drained bytes plus whether the caller should canonicalize afterward. #[inline] fn copy_to_bytes_small_chunks( chunks: &mut [B; N], len: usize, not_enough_data_msg: &str, ) -> (Bytes, bool) { let total = chunks .iter() .map(|buf| buf.remaining()) .fold(0, usize::saturating_add); assert!(total >= len, "{not_enough_data_msg}"); if chunks[0].remaining() >= len { let bytes = chunks[0].copy_to_bytes(len); return (bytes, chunks[0].remaining() == 0); } let mut out = BytesMut::with_capacity(len); let mut remaining = len; for buf in chunks.iter_mut() { if remaining == 0 { break; } let to_copy = remaining.min(buf.remaining()); out.extend_from_slice(&buf.chunk()[..to_copy]); buf.advance(to_copy); remaining -= to_copy; } // Slow path always consumes past chunk 0, so canonicalization is required. (out.freeze(), true) } /// Drain `len` readable bytes from a deque-backed chunk representation. /// /// Returns drained bytes plus whether the caller should canonicalize afterward. #[inline] fn copy_to_bytes_chunked( bufs: &mut VecDeque, len: usize, not_enough_data_msg: &str, ) -> (Bytes, bool) { while bufs.front().is_some_and(|buf| buf.remaining() == 0) { bufs.pop_front(); } if bufs.front().is_none() { assert_eq!(len, 0, "{not_enough_data_msg}"); // The deque is empty now, so the container must collapse back to the // canonical Single representation. return (Bytes::new(), true); } if bufs.front().is_some_and(|front| front.remaining() >= len) { let front = bufs.front_mut().expect("front checked above"); let bytes = front.copy_to_bytes(len); if front.remaining() == 0 { bufs.pop_front(); } return (bytes, bufs.len() <= 3); } let total = bufs .iter() .map(|buf| buf.remaining()) .fold(0, usize::saturating_add); assert!(total >= len, "{not_enough_data_msg}"); let mut out = BytesMut::with_capacity(len); let mut remaining = len; while remaining > 0 { let front = bufs .front_mut() .expect("remaining > 0 implies non-empty bufs"); let to_copy = remaining.min(front.remaining()); out.extend_from_slice(&front.chunk()[..to_copy]); front.advance(to_copy); if front.remaining() == 0 { bufs.pop_front(); } remaining -= to_copy; } (out.freeze(), bufs.len() <= 3) } /// Advance across a [`VecDeque`] of chunks by consuming from the front. #[inline] fn advance_chunked_front(bufs: &mut VecDeque, mut cnt: usize) { while cnt > 0 { let front = bufs.front_mut().expect("cannot advance past end of buffer"); let avail = front.remaining(); if avail == 0 { bufs.pop_front(); continue; } if cnt < avail { front.advance(cnt); break; } front.advance(avail); bufs.pop_front(); cnt -= avail; } } /// Advance across a small fixed set of chunks (`Pair`/`Triple`). /// /// Returns `true` when one or more chunks became (or were) empty, so callers /// can canonicalize once after the operation. #[inline] fn advance_small_chunks(chunks: &mut [B], mut cnt: usize) -> bool { let mut idx = 0; let mut needs_canonicalize = false; while cnt > 0 { let chunk = chunks .get_mut(idx) .expect("cannot advance past end of buffer"); let avail = chunk.remaining(); if avail == 0 { idx += 1; needs_canonicalize = true; continue; } if cnt < avail { chunk.advance(cnt); return needs_canonicalize; } chunk.advance(avail); cnt -= avail; idx += 1; needs_canonicalize = true; } needs_canonicalize } /// Advance writable cursors across `chunks` by up to `*remaining` bytes. /// /// Returns `true` when the full request has been satisfied. /// /// # Safety /// /// Forwards to [`BufMut::advance_mut`], so callers must ensure the advanced /// region has been initialized according to [`BufMut`]'s contract. #[inline] unsafe fn advance_mut_in_chunks(chunks: &mut [B], remaining: &mut usize) -> bool { if *remaining == 0 { return true; } for buf in chunks.iter_mut() { let avail = buf.chunk_mut().len(); if avail == 0 { continue; } if *remaining <= avail { // SAFETY: Upheld by this function's safety contract. unsafe { buf.advance_mut(*remaining) }; *remaining = 0; return true; } // SAFETY: Upheld by this function's safety contract. unsafe { buf.advance_mut(avail) }; *remaining -= avail; } false } /// Fill `dst` with `IoSlice`s built from `chunks`. /// /// Empty chunks are skipped. At most `dst.len()` slices are written. /// Returns the number of slices written. #[inline] fn fill_vectored_from_chunks<'a, I>(dst: &mut [IoSlice<'a>], chunks: I) -> usize where I: IntoIterator, { let mut written = 0; for chunk in chunks .into_iter() .filter(|chunk| !chunk.is_empty()) .take(dst.len()) { dst[written] = IoSlice::new(chunk); written += 1; } written } /// Assembles [`IoBufs`] from a mix of inline writes and zero-copy pieces. /// /// All inline writes go into a single buffer, and [`BufsMut::push`] records /// the inline offset at which each pushed piece belongs. [`Builder::finish`] /// freezes the buffer once, splits it at those offsets with /// [`IoBuf::split_to`], and interleaves the pushed chunks with the inline /// segments. /// /// A pushed [`IoBuf`] becomes a chunk as is. Any other pushed value converts /// into [`Bytes`], and the chunk owns that `Bytes` without copying. /// /// The inline buffer has a fixed capacity set at construction and will not /// grow. Callers must ensure the capacity accounts for all inline /// (non-pushed) bytes that will be written. Exceeding it will panic. /// /// ```text /// builder.put_u16(99); // inline /// builder.push(shard_payload.clone()); // zero-copy: shares the Bytes /// builder.put_u32(checksum); // inline /// let output = builder.finish(); /// /// // inline buffer [ 99 | checksum ] one allocation, frozen once /// // pushed Bytes [ ----- 1 MB shard ----- ] the caller's, never copied /// // /// // output [ 99 ][ ----- 1 MB shard ----- ][ checksum ] /// // inline pushed Bytes inline /// ``` pub struct Builder { // Single working buffer for all inline writes. buf: IoBufMut, // Each entry is (offset_in_buf, pushed_chunk) recording where a push // interrupts the inline byte stream. pushes: Vec<(usize, IoBuf)>, } impl Builder { /// Creates a builder whose inline buffer holds exactly `capacity` bytes. /// /// A zero capacity allocates nothing, so the builder accepts only pushed /// chunks. /// /// # Panics /// /// Panics if `capacity` exceeds `isize::MAX`. Writing more than `capacity` /// inline bytes panics at the write. pub fn new(capacity: usize) -> Self { Self { buf: IoBufMut::with_capacity(capacity), pushes: Vec::new(), } } /// Freezes the inline buffer and assembles [`IoBufs`] by splitting it at /// the recorded push boundaries. pub fn finish(self) -> IoBufs { if self.pushes.is_empty() { return IoBufs::from(self.buf.freeze()); } // Segments split off the owned frozen buffer. A first split of a // descriptor stores the shared record with a plain write, and the // remainder moves into the result without taking another reference. let mut frozen = self.buf.freeze(); let mut result = IoBufs::default(); let mut pos = 0; for (offset, pushed) in self.pushes { if offset > pos { result.append(frozen.split_to(offset - pos)); } result.append(pushed); pos = offset; } result.append(frozen); result } } // SAFETY: All methods delegate directly to `self.buf`, an `IoBufMut` with a // sound `BufMut` implementation. The inline buffer has fixed capacity. Writes // that exceed it panic in bytes' `BufMut` trait defaults (which check // `remaining_mut`) or, for a direct `advance_mut`, in `IoBufMut::advance_mut`. unsafe impl BufMut for Builder { #[inline] fn remaining_mut(&self) -> usize { self.buf.remaining_mut() } #[inline] unsafe fn advance_mut(&mut self, cnt: usize) { // SAFETY: The caller guarantees that `cnt` is within the writable chunk. unsafe { self.buf.advance_mut(cnt); } } #[inline] fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice { self.buf.chunk_mut() } } impl BufsMut for Builder { fn push(&mut self, chunk: impl Into + 'static) { let chunk = IoBuf::claim(chunk).unwrap_or_else(|chunk| IoBuf::from(chunk.into())); if !chunk.is_empty() { self.pushes.push((self.buf.len(), chunk)); } } } /// Extension trait for encoding values that implement [`Write`] and /// [`EncodeSize`] into [`IoBufs`]. pub trait EncodeExt: EncodeSize + Write { /// Encode this value into [`IoBufs`] through a [`Builder`]. /// /// Override [`Write::write_bufs`] to avoid copying large byte fields. /// /// # Panics /// /// Panics if [`Write::write_bufs`] writes more inline bytes than /// [`EncodeSize::encode_inline_size`], or if [`EncodeSize::encode_size`] /// does not match the total bytes written. fn encode_iobufs(&self) -> IoBufs { let len = self.encode_size(); let mut builder = Builder::new(self.encode_inline_size()); self.write_bufs(&mut builder); let bufs = builder.finish(); assert_eq!( bufs.remaining(), len, "write_bufs() must write exactly encode_size() bytes" ); bufs } } impl EncodeExt for T {} #[cfg(all(test, not(feature = "loom")))] #[commonware_macros::test_group("miri")] mod tests { use super::*; use bytes::{Bytes, BytesMut}; use commonware_codec::{Decode, Encode, types::lazy::Lazy}; use commonware_utils::range::NonEmptyRange; use std::{ alloc::{GlobalAlloc, Layout, System}, cell::Cell, collections::{BTreeMap, HashMap}, sync::Arc, }; /// Counts this thread's allocations during [`allocation_calls_during`] and delegates every /// call to the system allocator. struct TrackingAllocator; std::thread_local! { /// This thread's allocation count, or `None` while tracking is off. static ALLOCATION_CALLS: Cell> = const { Cell::new(None) }; } #[global_allocator] static ALLOCATOR: TrackingAllocator = TrackingAllocator; /// Counts an allocation when tracking is enabled on this thread. fn record_allocation() { ALLOCATION_CALLS.with(|calls| { if let Some(count) = calls.get() { calls.set(Some(count + 1)); } }); } // SAFETY: Every operation delegates to the system allocator with unchanged // arguments. Recording an allocation only updates a thread-local cell. unsafe impl GlobalAlloc for TrackingAllocator { unsafe fn alloc(&self, layout: Layout) -> *mut u8 { record_allocation(); // SAFETY: The caller supplies the layout required by GlobalAlloc. unsafe { System.alloc(layout) } } unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { record_allocation(); // SAFETY: The caller supplies the layout required by GlobalAlloc. unsafe { System.alloc_zeroed(layout) } } unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { // SAFETY: This wrapper delegates every allocation to System, so the // pointer came from System with this layout. unsafe { System.dealloc(ptr, layout) } } unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { record_allocation(); // SAFETY: The pointer and layout came from System through this wrapper, // and the caller supplies the new size required by GlobalAlloc. unsafe { System.realloc(ptr, layout, new_size) } } } /// Runs `f` and returns its result with the number of allocations it made on this thread. fn allocation_calls_during(f: impl FnOnce() -> T) -> (T, usize) { ALLOCATION_CALLS.with(|calls| assert!(calls.replace(Some(0)).is_none())); let result = f(); let calls = ALLOCATION_CALLS.with(|calls| { calls .replace(None) .expect("allocation tracking must stay enabled") }); (result, calls) } /// A prefix held by the front chunk is handed out as a view of that chunk, in every /// representation, and consuming the chunk keeps the container canonical. Only the first /// partial view of an unshared chunk allocates, for the chunk's owner record. A prefix that /// spans chunks is copied once into the requested type. #[test] fn test_take_owned_front_chunk() { let chunk = |byte: u8| IoBuf::copy_from_slice(&[byte; 8]); let mut bufs = IoBufs::default(); for byte in 1..=5 { bufs.append(chunk(byte)); } assert!(matches!(bufs.inner, IoBufsInner::Chunked(_))); // Chunked: a partial take leaves the chunk in place, a take of the remainder removes it. // The partial take shares the unshared chunk, which allocates its owner record. let front_ptr = bufs.chunk().as_ptr(); let (view, calls) = allocation_calls_during(|| bufs.take_owned::(3)); assert_eq!(calls, 1); assert_eq!(view, [1u8; 3]); assert_eq!(view.as_ptr(), front_ptr); let (view, calls) = allocation_calls_during(|| bufs.take_owned::(5)); assert_eq!(calls, 0); assert_eq!(view, [1u8; 5]); assert!(matches!(bufs.inner, IoBufsInner::Chunked(ref chunks) if chunks.len() == 4)); // Consuming the next chunk collapses to the small representation. let (view, calls) = allocation_calls_during(|| bufs.take_owned::(8)); assert_eq!(calls, 0); assert_eq!(view, [2u8; 8]); assert!(matches!(bufs.inner, IoBufsInner::Triple(_))); // Triple and Pair front takes are views too, and a `Bytes` field gets the view's bytes. let (view, calls) = allocation_calls_during(|| bufs.take_owned::(8)); assert_eq!(calls, 0); assert_eq!(view, [3u8; 8]); assert!(matches!(bufs.inner, IoBufsInner::Pair(_))); let front_ptr = bufs.chunk().as_ptr(); let bytes = bufs.take_owned::(4); assert_eq!(bytes, vec![4u8; 4]); assert_eq!(bytes.as_ptr(), front_ptr); assert!(matches!(bufs.inner, IoBufsInner::Pair(_))); // A view keeps the chunk's native owner and converts in place once it is the last holder. drop(bufs); assert!(view.try_into_mut().is_ok()); // Spanning prefixes are copied once. An IoBuf lands in one native buffer (one data // allocation, plus the owner record that splitting the second chunk allocates), Bytes in // one shared buffer, and the remainder stays canonical. let mut bufs = IoBufs::from(chunk(1)); bufs.append(chunk(2)); let (copied, calls) = allocation_calls_during(|| bufs.take_owned::(12)); assert_eq!(calls, 2); assert_eq!(copied.as_ref(), [[1u8; 8].as_slice(), &[2u8; 4]].concat()); assert!(copied.try_into_mut().is_ok()); assert!(matches!(bufs.inner, IoBufsInner::Single(_))); assert_eq!(bufs.len(), 4); let mut bufs = IoBufs::from(chunk(1)); bufs.append(chunk(2)); let (copied, calls) = allocation_calls_during(|| bufs.take_owned::(12)); assert_eq!(calls, 1); assert_eq!(copied, [[1u8; 8].as_slice(), &[2u8; 4]].concat()); assert_eq!(bufs.len(), 4); // Zero-length takes are empty views that leave the container untouched. let empty = bufs.take_owned::(0); assert!(empty.is_empty()); assert_eq!(bufs.len(), 4); assert!(IoBufs::default().take_owned::(0).is_empty()); } /// An unshared native buffer allocates only its data until its first share, which allocates /// its owner record once. It lives here, next to the tracking allocator that counts the calls. #[test] fn test_iobuf_owner_record_allocations() { // Allocate, write, freeze, recover, write, freeze, and drop: one data allocation. let ((), calls) = allocation_calls_during(|| { let mut buf = IoBufMut::with_capacity(64); buf.put_slice(b"payload"); let mut buf = buf.freeze().try_into_mut().unwrap(); buf.put_slice(b"!"); drop(buf.freeze()); }); assert_eq!(calls, 1); // Moves, whole splits, and empty views allocate nothing. let mut source = IoBuf::copy_from_slice(&[7u8; 64]); let (mut buf, calls) = allocation_calls_during(|| { let mut moved = std::mem::take(&mut source); assert!(moved.slice(0..0).is_empty()); assert!(moved.split_to(0).is_empty()); moved.split_to(64) }); assert_eq!(calls, 0); // The first share allocates the record, and later shares reuse it. let (first, calls) = allocation_calls_during(|| buf.clone()); assert_eq!(calls, 1); let (later, calls) = allocation_calls_during(|| (buf.slice(1..3), buf.split_to(8), buf.clone())); assert_eq!(calls, 0); drop((first, later)); // Spare capacity too large for the handle allocates the record at freeze instead. let mut big = IoBufMut::with_capacity((1 << 17) + 1); big.put_u8(1); let (frozen, calls) = allocation_calls_during(|| big.freeze()); assert_eq!(calls, 1); let (_view, calls) = allocation_calls_during(|| frozen.clone()); assert_eq!(calls, 0); // Whole conversions and drains into Bytes move the descriptor into the Bytes box without // promoting it, so each allocates only that box. let mut advanced = IoBuf::copy_from_slice(&[7u8; 16]); advanced.advance(2); let (_bytes, calls) = allocation_calls_during(|| Bytes::from(advanced)); assert_eq!(calls, 1); let mut drained = IoBuf::copy_from_slice(&[7u8; 16]); let (_bytes, calls) = allocation_calls_during(|| drained.copy_to_bytes(16)); assert_eq!(calls, 1); let mut drained = IoBufMut::with_capacity(16); drained.put_slice(&[7u8; 8]); let (_bytes, calls) = allocation_calls_during(|| drained.copy_to_bytes(8)); assert_eq!(calls, 1); } fn assert_encode_iobufs_matches_encode(value: &T) { let mut encoded = value.encode_iobufs(); let baseline = value.encode(); let mut encoded_bytes = vec![0u8; encoded.remaining()]; encoded.copy_to_slice(&mut encoded_bytes); assert_eq!(encoded_bytes, baseline.as_ref()); } #[test] fn test_iobufs_from_iter_and_array_shapes() { fn check() { let data = b"abcde"; let chunks: [IoBuf; N] = std::array::from_fn(|i| IoBuf::from(&data[i..i + 1])); for bufs in [IoBufs::from(chunks.clone()), chunks.into_iter().collect()] { assert!(matches!( (&bufs.inner, N), (IoBufsInner::Single(_), 0 | 1) | (IoBufsInner::Pair(_), 2) | (IoBufsInner::Triple(_), 3) | (IoBufsInner::Chunked(_), 4..) )); assert_eq!(bufs.chunk_count(), N); for index in 0..N { assert_eq!( bufs.chunk_at(index).unwrap().as_ptr(), data[index..].as_ptr() ); } assert_eq!(bufs.coalesce().as_ref(), &data[..N]); } } check::<0>(); check::<1>(); check::<2>(); check::<3>(); check::<4>(); check::<5>(); } #[test] fn test_iobufs_from_iter_and_array_drop_empty_buffers() { fn check(chunks: [&'static [u8]; N], expected: &[u8]) { let chunks = chunks.map(IoBuf::from); for bufs in [IoBufs::from(chunks.clone()), chunks.into_iter().collect()] { // Each nonempty input contains one byte. assert_eq!(bufs.chunk_count(), expected.len()); assert!(matches!( (&bufs.inner, expected.len()), (IoBufsInner::Single(_), 0 | 1) | (IoBufsInner::Pair(_), 2) | (IoBufsInner::Triple(_), 3) | (IoBufsInner::Chunked(_), 4..) )); assert_eq!(bufs.coalesce().as_ref(), expected); } } check([b"", b"", b""], b""); check([b"", b"a", b""], b"a"); check([b"", b"a", b"b"], b"ab"); check([b"a", b"", b"b"], b"ab"); check([b"a", b"b", b""], b"ab"); check([b"a", b"", b"b", b"c"], b"abc"); check([b"a", b"b", b"", b"c", b"d"], b"abcd"); check([b"", b"a", b"", b"b", b"c", b"d", b"e", b""], b"abcde"); } #[test] fn test_iobufs_from_iter_error_releases_buffers() { const BUFFER_COUNT: usize = 6; // Cover failures before and after the transition to deque storage, // always leaving at least one unconsumed buffer in the iterator. for fail_at in 0..=4 { let owners: Vec> = (0..BUFFER_COUNT).map(|_| Arc::from(&b"x"[..])).collect(); let bufs: Vec<_> = owners .iter() .map(|owner| IoBuf::from(Bytes::from_owner(owner.clone()))) .collect(); // The test and the buffer each hold one reference to the owner. assert!(owners.iter().all(|owner| Arc::strong_count(owner) == 2)); let mut items_pulled = 0; let result = bufs .into_iter() .enumerate() .map(|(index, buf)| { items_pulled += 1; if index == fail_at { Err("injected") } else { Ok(buf) } }) .collect::>(); assert_eq!(result.unwrap_err(), "injected"); assert_eq!(items_pulled, fail_at + 1); // Only the test's reference remains, including for unconsumed buffers. for owner in owners { assert_eq!(Arc::strong_count(&owner), 1); } } } #[test] fn test_iobufsmut_from_iter_and_array_shapes() { fn check() { let chunks = || -> [IoBufMut; N] { std::array::from_fn(|_| IoBufMut::with_capacity(1)) }; for bufs in [IoBufsMut::from(chunks()), chunks().into_iter().collect()] { assert!(matches!( (&bufs.inner, N), (IoBufsMutInner::Single(_), 0 | 1) | (IoBufsMutInner::Pair(_), 2) | (IoBufsMutInner::Triple(_), 3) | (IoBufsMutInner::Chunked(_), 4..) )); assert!(bufs.is_empty()); assert_eq!(bufs.capacity(), N); assert_eq!(bufs.remaining_mut(), N); } } check::<0>(); check::<1>(); check::<2>(); check::<3>(); check::<4>(); check::<5>(); } #[test] fn test_iobufsmut_from_iter_and_array_preserve_capacity() { fn check(convert: impl FnOnce([IoBufMut; 4]) -> IoBufsMut) { let empty = IoBufMut::with_capacity(2); let mut partial = IoBufMut::with_capacity(2); partial.put_slice(b"a"); let full = IoBufMut::from(b"bc"); let pointers = [ empty.as_ref().as_ptr(), partial.as_ref().as_ptr(), full.as_ref().as_ptr(), ]; let mut bufs = convert([IoBufMut::default(), empty, partial, full]); assert!(matches!(bufs.inner, IoBufsMutInner::Triple(_))); assert_eq!(bufs.len(), 3); assert_eq!(bufs.capacity(), 6); assert_eq!(bufs.remaining_mut(), 3); // Fill the empty buffer with XY and the partial buffer's tail with Z. bufs.put_slice(b"XYZ"); assert_eq!(bufs.capacity(), 6); assert_eq!(bufs.remaining_mut(), 0); let frozen = bufs.freeze(); for (index, pointer) in pointers.into_iter().enumerate() { assert_eq!(frozen.chunk_at(index).unwrap().as_ptr(), pointer); } assert_eq!(frozen.coalesce(), b"XYaZbc"); } check(IoBufsMut::from); check(|chunks| chunks.into_iter().collect()); } #[test] fn test_iobufs_emptiness_after_consumption() { fn assert_state(bufs: &IoBufs, expected: &[u8]) { assert_eq!(bufs.remaining(), expected.len()); assert_eq!(bufs.is_empty(), expected.is_empty()); assert_eq!(bufs.has_remaining(), !expected.is_empty()); let mut actual = Vec::new(); bufs.for_each_chunk(|chunk| actual.extend_from_slice(chunk)); assert_eq!(actual, expected); } let data = b"abcdefgh"; let source = IoBufs::from(vec![ IoBuf::default(), IoBuf::from(b"ab"), IoBuf::default(), IoBuf::from(b"cd"), IoBuf::from(b"ef"), IoBuf::default(), IoBuf::from(b"gh"), IoBuf::default(), ]); // Prefixes exercise every representation, including an empty Single. for end in 0..=data.len() { let mut suffix = source.clone(); let prefix = suffix.split_to(end); assert_state(&prefix, &data[..end]); assert_state(&suffix, &data[end..]); for step in [1, 2, 3, data.len()] { let mut advanced = prefix.clone(); let mut copied = prefix.clone(); let mut offset = 0; while offset < end { let count = step.min(end - offset); advanced.advance(count); assert_eq!(copied.copy_to_bytes(count), &data[offset..offset + count]); offset += count; assert_state(&advanced, &data[offset..end]); assert_state(&copied, &data[offset..end]); } // Empty inputs preserve emptiness, and new bytes make a drained buffer readable. advanced.append(IoBuf::default()); advanced.prepend(IoBuf::default()); assert_state(&advanced, b""); advanced.append(IoBuf::from(b"z")); advanced.prepend(IoBuf::from(b"y")); assert_state(&advanced, b"yz"); } let mut input = prefix; let mut output = IoBufMut::with_capacity(data.len()); output.put(&mut input); assert_eq!(output.as_ref(), &data[..end]); assert_state(&input, b""); } } #[test] fn test_iobufs_for_each_chunk_single_and_empty() { // The Single arm is bypassed by coalesce (which returns the buffer // directly), so exercise the public method on both Single shapes. let mut seen = Vec::new(); IoBufs::from(b"hello").for_each_chunk(|chunk| seen.push(chunk.to_vec())); assert_eq!(seen, vec![b"hello".to_vec()]); let mut count = 0; IoBufs::default().for_each_chunk(|_| count += 1); assert_eq!(count, 0); } #[test] fn test_iobufs_shapes_and_read_paths() { // Empty construction normalizes to an empty single chunk. let empty = IoBufs::from(Vec::::new()); assert!(empty.is_empty()); assert!(empty.is_single()); assert!(empty.as_single().is_some()); // Single-buffer read path. let mut single = IoBufs::from(b"hello world"); assert!(single.is_single()); assert_eq!(single.chunk(), b"hello world"); single.advance(6); assert_eq!(single.chunk(), b"world"); assert_eq!(single.copy_to_bytes(5).as_ref(), b"world"); assert_eq!(single.remaining(), 0); // Fast-path shapes (Pair/Triple/Chunked). let mut pair = IoBufs::from(IoBuf::from(b"a")); pair.append(IoBuf::from(b"b")); assert!(matches!(pair.inner, IoBufsInner::Pair(_))); assert!(pair.as_single().is_none()); let mut triple = IoBufs::from(IoBuf::from(b"a")); triple.append(IoBuf::from(b"b")); triple.append(IoBuf::from(b"c")); assert!(matches!(triple.inner, IoBufsInner::Triple(_))); let mut chunked = IoBufs::from(IoBuf::from(b"a")); chunked.append(IoBuf::from(b"b")); chunked.append(IoBuf::from(b"c")); chunked.append(IoBuf::from(b"d")); assert!(matches!(chunked.inner, IoBufsInner::Chunked(_))); // prepend + append preserve ordering. let mut joined = IoBufs::from(b"middle"); joined.prepend(IoBuf::from(b"start ")); joined.append(IoBuf::from(b" end")); assert_eq!(joined.coalesce(), b"start middle end"); // prepending empty is a no-op, and prepending into pair upgrades to triple. let mut prepend_noop = IoBufs::from(b"x"); prepend_noop.prepend(IoBuf::default()); assert_eq!(prepend_noop.coalesce(), b"x"); // Prepending into an empty aggregate should stay on the single-buffer fast path. let mut prepend_into_empty = IoBufs::default(); prepend_into_empty.prepend(IoBuf::from(b"z")); assert!(prepend_into_empty.is_single()); assert_eq!(prepend_into_empty.coalesce(), b"z"); let mut prepend_pair = IoBufs::from(vec![IoBuf::from(b"b"), IoBuf::from(b"c")]); prepend_pair.prepend(IoBuf::from(b"a")); assert!(matches!(prepend_pair.inner, IoBufsInner::Triple(_))); assert_eq!(prepend_pair.coalesce(), b"abc"); // canonicalizing a non-empty single should keep the same representation. let mut canonical_single = IoBufs::from(b"q"); canonical_single.canonicalize(); assert!(canonical_single.is_single()); assert_eq!(canonical_single.coalesce(), b"q"); } #[test] fn test_iobufs_split_to_cases() { // Zero and full split on a single chunk. let mut bufs = IoBufs::from(b"hello"); let empty = bufs.split_to(0); assert!(empty.is_empty()); assert_eq!(bufs.coalesce(), b"hello"); let mut bufs = IoBufs::from(b"hello"); let all = bufs.split_to(5); assert_eq!(all.coalesce(), b"hello"); assert!(bufs.is_single()); assert!(bufs.is_empty()); // Single split in the middle. let mut single_mid = IoBufs::from(b"hello"); let single_prefix = single_mid.split_to(2); assert!(single_prefix.is_single()); assert_eq!(single_prefix.coalesce(), b"he"); assert_eq!(single_mid.coalesce(), b"llo"); // Pair split paths: in-first, boundary-after-first, crossing-into-second. let mut pair = IoBufs::from(vec![IoBuf::from(b"ab"), IoBuf::from(b"cd")]); let pair_prefix = pair.split_to(1); assert!(pair_prefix.is_single()); assert_eq!(pair_prefix.coalesce(), b"a"); assert!(matches!(pair.inner, IoBufsInner::Pair(_))); assert_eq!(pair.coalesce(), b"bcd"); let mut pair = IoBufs::from(vec![IoBuf::from(b"ab"), IoBuf::from(b"cd")]); let pair_prefix = pair.split_to(2); assert!(pair_prefix.is_single()); assert_eq!(pair_prefix.coalesce(), b"ab"); assert!(pair.is_single()); assert_eq!(pair.coalesce(), b"cd"); let mut pair = IoBufs::from(vec![IoBuf::from(b"ab"), IoBuf::from(b"cd")]); let pair_prefix = pair.split_to(3); assert!(matches!(pair_prefix.inner, IoBufsInner::Pair(_))); assert_eq!(pair_prefix.coalesce(), b"abc"); assert!(pair.is_single()); assert_eq!(pair.coalesce(), b"d"); // Triple split paths: in-first, boundary-after-first, in-second, boundary-after-second, // and reaching into third. let mut triple = IoBufs::from(vec![ IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), ]); let triple_prefix = triple.split_to(1); assert!(triple_prefix.is_single()); assert_eq!(triple_prefix.coalesce(), b"a"); assert!(matches!(triple.inner, IoBufsInner::Triple(_))); assert_eq!(triple.coalesce(), b"bcdef"); let mut triple = IoBufs::from(vec![ IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), ]); let triple_prefix = triple.split_to(2); assert!(triple_prefix.is_single()); assert_eq!(triple_prefix.coalesce(), b"ab"); assert!(matches!(triple.inner, IoBufsInner::Pair(_))); assert_eq!(triple.coalesce(), b"cdef"); let mut triple = IoBufs::from(vec![ IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), ]); let triple_prefix = triple.split_to(3); assert!(matches!(triple_prefix.inner, IoBufsInner::Pair(_))); assert_eq!(triple_prefix.coalesce(), b"abc"); assert!(matches!(triple.inner, IoBufsInner::Pair(_))); assert_eq!(triple.coalesce(), b"def"); let mut triple = IoBufs::from(vec![ IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), ]); let triple_prefix = triple.split_to(4); assert!(matches!(triple_prefix.inner, IoBufsInner::Pair(_))); assert_eq!(triple_prefix.coalesce(), b"abcd"); assert!(triple.is_single()); assert_eq!(triple.coalesce(), b"ef"); let mut triple = IoBufs::from(vec![ IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), ]); let triple_prefix = triple.split_to(5); assert!(matches!(triple_prefix.inner, IoBufsInner::Triple(_))); assert_eq!(triple_prefix.coalesce(), b"abcde"); assert!(triple.is_single()); assert_eq!(triple.coalesce(), b"f"); // Chunked split can canonicalize remainder/prefix shapes. let mut bufs = IoBufs::from(vec![ IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), IoBuf::from(b"gh"), ]); let prefix = bufs.split_to(4); assert!(matches!(prefix.inner, IoBufsInner::Pair(_))); assert_eq!(prefix.coalesce(), b"abcd"); assert!(matches!(bufs.inner, IoBufsInner::Pair(_))); assert_eq!(bufs.coalesce(), b"efgh"); // Chunked split inside a chunk. let mut bufs = IoBufs::from(vec![ IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), IoBuf::from(b"gh"), ]); let prefix = bufs.split_to(5); assert!(matches!(prefix.inner, IoBufsInner::Triple(_))); assert_eq!(prefix.coalesce(), b"abcde"); assert!(matches!(bufs.inner, IoBufsInner::Pair(_))); assert_eq!(bufs.coalesce(), b"fgh"); // Chunked split can remain chunked on both sides when both have >= 4 chunks. let mut bufs = IoBufs::from(vec![ IoBuf::from(b"a"), IoBuf::from(b"b"), IoBuf::from(b"c"), IoBuf::from(b"d"), IoBuf::from(b"e"), IoBuf::from(b"f"), IoBuf::from(b"g"), IoBuf::from(b"h"), ]); let prefix = bufs.split_to(4); assert!(matches!(prefix.inner, IoBufsInner::Chunked(_))); assert_eq!(prefix.coalesce(), b"abcd"); assert!(matches!(bufs.inner, IoBufsInner::Chunked(_))); assert_eq!(bufs.coalesce(), b"efgh"); // Chunked split inside a chunk that leaves both sides chunked. let mut bufs: IoBufs = (b'a'..b'i') .map(|byte| IoBuf::from(vec![byte; 2])) .collect(); let prefix = bufs.split_to(9); assert!(matches!(prefix.inner, IoBufsInner::Chunked(_))); assert_eq!(prefix.chunk_count(), 5); assert_eq!(prefix.coalesce(), b"aabbccdde"); assert!(matches!(bufs.inner, IoBufsInner::Chunked(_))); assert_eq!(bufs.chunk_count(), 4); assert_eq!(bufs.coalesce(), b"effgghh"); // Defensive path: tolerate accidental empty chunks in non-canonical chunked input. let mut bufs = IoBufs { inner: IoBufsInner::Chunked(VecDeque::from([ IoBuf::default(), IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), IoBuf::from(b"gh"), ])), }; let prefix = bufs.split_to(3); assert_eq!(prefix.coalesce(), b"abc"); assert_eq!(bufs.coalesce(), b"defgh"); } #[test] #[should_panic(expected = "split_to out of bounds")] fn test_iobufs_split_to_out_of_bounds() { let mut bufs = IoBufs::from(b"abc"); let _ = bufs.split_to(4); } /// Splits at every position of every representation keep the bytes in order and leave both /// sides canonical. The positions cover prefixes within the front chunk and across chunks, /// and chunked splits that move out either side. #[test] fn test_iobufs_split_positions() { for count in 1..10 { for rotation in 0..count { for at in 0..=2 * count { // Rotation wraps a chunked deque around the end of its buffer, so the drains // also run over a deque stored in two slices. Four or more chunks keep that // deque as their chunked representation. let mut chunks: VecDeque = (0..count as u8) .map(|byte| IoBuf::copy_from_slice(&[byte; 2])) .collect(); chunks.rotate_left(rotation); let expected: Vec = chunks .iter() .flat_map(|buf| buf.as_ref().to_vec()) .collect(); let mut bufs = if count >= 4 { IoBufs { inner: IoBufsInner::Chunked(chunks), } } else { IoBufs::from_chunks_iter(chunks) }; let prefix = bufs.split_to(at); // Each side stores only non-empty chunks, in the representation its chunk // count calls for, and the two sides hold the original bytes in order. for side in [&prefix, &bufs] { let chunks = side.chunk_count(); assert_eq!(side.iter().filter(|buf| !buf.is_empty()).count(), chunks); assert_eq!(side.is_single(), chunks <= 1); assert_eq!(matches!(side.inner, IoBufsInner::Chunked(_)), chunks >= 4); } assert_eq!(prefix.coalesce().as_ref(), &expected[..at]); assert_eq!(bufs.coalesce().as_ref(), &expected[at..]); } } } } /// A split of a chunked container moves out the side with fewer chunks, counting the piece /// toward the prefix, and keeps the deque for the other, so taking a small prefix or leaving a /// small suffix allocates nothing however many chunks remain. The chunks are static, so /// splitting one in place allocates nothing either. #[test] fn test_iobufs_chunked_split_allocations() { // Every split folds over the remaining chunks, so the count stays small enough for Miri // while keeping the container chunked throughout. const CHUNKS: usize = 256; // Splitting inside the first of many chunks. let mut bufs: IoBufs = (0..CHUNKS).map(|_| IoBuf::from(b"ab")).collect(); let (prefix, calls) = allocation_calls_during(|| bufs.split_to(1)); assert_eq!(calls, 0); assert_eq!(prefix.coalesce(), b"a"); assert_eq!(bufs.len(), 2 * CHUNKS - 1); // A split across the first two chunks moves out only those two. let (prefix, calls) = allocation_calls_during(|| bufs.split_to(2)); assert_eq!(calls, 0); assert_eq!(prefix.coalesce(), b"ba"); // A split that leaves only the last two chunks moves out only those two. let at = bufs.len() - 3; let (_, calls) = allocation_calls_during(|| bufs.split_to(at)); assert_eq!(calls, 0); assert_eq!(bufs.coalesce(), b"bab"); // A split inside the fourth of six chunks puts four chunks in the prefix and three in the // suffix, so the suffix moves out and the deque stays. let mut bufs: IoBufs = (0..6).map(|_| IoBuf::from(b"ab")).collect(); let (prefix, calls) = allocation_calls_during(|| bufs.split_to(7)); assert_eq!(calls, 0); assert_eq!(prefix.coalesce(), b"abababa"); assert_eq!(bufs.coalesce(), b"babab"); // Extracting the chunks one at a time. let mut bufs: IoBufs = (0..CHUNKS).map(|_| IoBuf::from(b"ab")).collect(); let ((), calls) = allocation_calls_during(|| { while bufs.chunk_count() >= 4 { assert_eq!(bufs.split_to(2).coalesce(), b"ab"); } }); assert_eq!(calls, 0); } #[test] fn test_iobufs_chunk_count() { assert_eq!(IoBufs::default().chunk_count(), 0); assert_eq!(IoBufs::from(IoBuf::from(b"a")).chunk_count(), 1); assert_eq!( IoBufs::from(vec![IoBuf::from(b"b"), IoBuf::from(b"c")]).chunk_count(), 2 ); assert_eq!( IoBufs::from(vec![ IoBuf::from(b"a"), IoBuf::from(b"b"), IoBuf::from(b"c") ]) .chunk_count(), 3 ); assert_eq!( IoBufs::from(vec![ IoBuf::from(b"a"), IoBuf::from(b"b"), IoBuf::from(b"c"), IoBuf::from(b"d") ]) .chunk_count(), 4 ); } #[test] fn test_iobufs_chunk_at() { let chunks = [b"ab".as_slice(), b"cd", b"ef", b"gh", b"ij"]; for count in 0..=chunks.len() { let mut bufs = IoBufs::from( chunks[..count] .iter() .map(|chunk| IoBuf::from(*chunk)) .collect::>(), ); for (index, chunk) in chunks[..count].iter().enumerate() { assert_eq!(bufs.chunk_at(index), Some(*chunk)); } assert_eq!(bufs.chunk_at(count), None); assert_eq!(bufs.chunk_at(usize::MAX), None); assert_eq!(bufs.len(), count * 2); // Indexing follows the readable chunks after partial and whole // chunks are consumed, including representation transitions. if count > 0 { bufs.advance(1); assert_eq!(bufs.chunk_at(0), Some(b"b".as_slice())); bufs.advance(1); assert_eq!(bufs.chunk_at(0), chunks[1..count].first().copied()); } } } #[test] fn test_iobufs_iter_borrows_every_owner() { let chunks = [b"ab".as_slice(), b"cd", b"ef", b"gh", b"ij"]; for count in 0..=chunks.len() { let bufs = IoBufs::from( chunks[..count] .iter() .map(|chunk| IoBuf::from(*chunk)) .collect::>(), ); let iter = bufs.iter(); assert_eq!(iter.len(), count); assert_eq!( iter.clone().map(AsRef::as_ref).collect::>(), chunks[..count] ); assert_eq!(iter.count(), count); } } #[test] fn test_iobufs_coalesce_after_advance() { let mut bufs = IoBufs::from(IoBuf::from(b"hello")); bufs.append(IoBuf::from(b" world")); assert_eq!(bufs.len(), 11); bufs.advance(3); assert_eq!(bufs.len(), 8); assert_eq!(bufs.coalesce(), b"lo world"); } #[test] fn test_iobufs_coalesce_shapes() { // Single buffer: zero-copy (same pointer) let buf = IoBuf::from(vec![1u8, 2, 3, 4, 5]); let original_ptr = buf.as_ptr(); let bufs = IoBufs::from(buf); let coalesced = bufs.coalesce(); assert_eq!(coalesced, [1, 2, 3, 4, 5]); assert_eq!(coalesced.as_ptr(), original_ptr); // Multiple buffers: merged into one allocation let mut bufs = IoBufs::from(IoBuf::from(b"hello")); bufs.append(IoBuf::from(b" world")); let coalesced = bufs.coalesce(); assert_eq!(coalesced, b"hello world"); // Multiple buffers after advance: only remaining data coalesced let mut bufs = IoBufs::from(IoBuf::from(b"hello")); bufs.append(IoBuf::from(b" world")); bufs.advance(3); let coalesced = bufs.coalesce(); assert_eq!(coalesced, b"lo world"); // Empty buffers in the middle let mut bufs = IoBufs::from(IoBuf::from(b"hello")); bufs.append(IoBuf::default()); bufs.append(IoBuf::from(b" world")); let coalesced = bufs.coalesce(); assert_eq!(coalesced, b"hello world"); // Empty IoBufs let bufs = IoBufs::default(); let coalesced = bufs.coalesce(); assert!(coalesced.is_empty()); // 4+ buffers: exercise the chunked coalesce path. let bufs = IoBufs::from(vec![ IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), IoBuf::from(b"gh"), ]); let coalesced = bufs.coalesce(); assert_eq!(coalesced, b"abcdefgh"); } #[test] fn test_iobufsmut_decode_preserves_byte_fields() { let first = IoBufMut::from(&b"\x02\x03abc"[..]); let second = IoBufMut::from(&b"\x03def"[..]); let ranges = [ first.as_ref().as_ptr_range(), second.as_ref().as_ptr_range(), ]; let source = IoBufsMut::from(vec![first, second]); let decoded = Vec::::decode_cfg(source, &((..).into(), (..).into())).unwrap(); assert_eq!( decoded, [Bytes::from_static(b"abc"), Bytes::from_static(b"def")] ); for (field, range) in decoded.iter().zip(ranges) { assert!(range.contains(&field.as_ptr())); } } #[test] fn test_iobufs_empty_chunks_and_copy_to_bytes_paths() { // Empty chunks are skipped while reading across multiple chunks. let mut bufs = IoBufs::default(); bufs.append(IoBuf::from(b"hello")); bufs.append(IoBuf::default()); bufs.append(IoBuf::from(b" ")); bufs.append(IoBuf::default()); bufs.append(IoBuf::from(b"world")); assert_eq!(bufs.len(), 11); assert_eq!(bufs.chunk(), b"hello"); bufs.advance(5); assert_eq!(bufs.chunk(), b" "); bufs.advance(1); assert_eq!(bufs.chunk(), b"world"); // Single-buffer copy_to_bytes path. let mut single = IoBufs::from(b"hello world"); assert_eq!(single.copy_to_bytes(5).as_ref(), b"hello"); assert_eq!(single.remaining(), 6); // Multi-buffer copy_to_bytes path across boundaries. let mut multi = IoBufs::from(b"hello"); multi.prepend(IoBuf::from(b"say ")); assert_eq!(multi.copy_to_bytes(7).as_ref(), b"say hel"); assert_eq!(multi.copy_to_bytes(2).as_ref(), b"lo"); } #[test] fn test_iobufs_copy_to_bytes_pair_and_triple() { // Pair: crossing one boundary should collapse to the trailing single chunk. let mut pair = IoBufs::from(IoBuf::from(b"ab")); pair.append(IoBuf::from(b"cd")); let first = pair.copy_to_bytes(3); assert_eq!(&first[..], b"abc"); assert!(pair.is_single()); assert_eq!(pair.chunk(), b"d"); // Triple: draining across two chunks leaves the final chunk readable. let mut triple = IoBufs::from(IoBuf::from(b"ab")); triple.append(IoBuf::from(b"cd")); triple.append(IoBuf::from(b"ef")); let first = triple.copy_to_bytes(5); assert_eq!(&first[..], b"abcde"); assert!(triple.is_single()); assert_eq!(triple.chunk(), b"f"); } #[test] fn test_iobufs_copy_to_bytes_chunked_four_plus() { let mut bufs = IoBufs::from(vec![ IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), IoBuf::from(b"gh"), ]); // Chunked fast-path: first chunk alone satisfies request. let first = bufs.copy_to_bytes(1); assert_eq!(&first[..], b"a"); // Chunked slow-path: request crosses chunk boundaries. let second = bufs.copy_to_bytes(4); assert_eq!(&second[..], b"bcde"); let rest = bufs.copy_to_bytes(3); assert_eq!(&rest[..], b"fgh"); assert_eq!(bufs.remaining(), 0); } #[test] fn test_iobufs_copy_to_bytes_edge_cases() { // A non-canonical leading empty chunk in the deque path is popped // without affecting the copied payload. Canonical construction never // stores empties, so the state is built directly. let mut iobufs = IoBufs { inner: IoBufsInner::Chunked(VecDeque::from([ IoBuf::default(), IoBuf::from(b"hel"), IoBuf::from(b"lo"), IoBuf::from(b" world"), ])), }; assert_eq!(iobufs.copy_to_bytes(5).as_ref(), b"hello"); assert_eq!(iobufs.remaining(), 6); // Boundary-aligned reads should return exact chunk payloads in-order. let mut boundary = IoBufs::from(IoBuf::from(b"hello")); boundary.append(IoBuf::from(b"world")); assert_eq!(boundary.copy_to_bytes(5).as_ref(), b"hello"); assert_eq!(boundary.copy_to_bytes(5).as_ref(), b"world"); assert_eq!(boundary.remaining(), 0); } #[test] #[should_panic(expected = "cannot advance past end of buffer")] fn test_iobufs_advance_past_end() { let mut bufs = IoBufs::from(b"hel"); bufs.append(IoBuf::from(b"lo")); bufs.advance(10); } #[test] #[should_panic(expected = "not enough data")] fn test_iobufs_copy_to_bytes_past_end() { let mut bufs = IoBufs::from(b"hel"); bufs.append(IoBuf::from(b"lo")); bufs.copy_to_bytes(10); } #[test] fn test_iobufs_matches_bytes_chain() { let b1 = Bytes::from_static(b"hello"); let b2 = Bytes::from_static(b" "); let b3 = Bytes::from_static(b"world"); // Buf parity for remaining/chunk/advance should match `Bytes::chain`. let mut chain = b1.clone().chain(b2.clone()).chain(b3.clone()); let mut iobufs = IoBufs::from(IoBuf::from(b1.clone())); iobufs.append(IoBuf::from(b2.clone())); iobufs.append(IoBuf::from(b3.clone())); assert_eq!(chain.remaining(), iobufs.remaining()); assert_eq!(chain.chunk(), iobufs.chunk()); chain.advance(3); iobufs.advance(3); assert_eq!(chain.remaining(), iobufs.remaining()); assert_eq!(chain.chunk(), iobufs.chunk()); chain.advance(3); iobufs.advance(3); assert_eq!(chain.remaining(), iobufs.remaining()); assert_eq!(chain.chunk(), iobufs.chunk()); // Test copy_to_bytes let mut chain = b1.clone().chain(b2.clone()).chain(b3.clone()); let mut iobufs = IoBufs::from(IoBuf::from(b1)); iobufs.append(IoBuf::from(b2)); iobufs.append(IoBuf::from(b3)); assert_eq!(chain.copy_to_bytes(3), iobufs.copy_to_bytes(3)); assert_eq!(chain.copy_to_bytes(4), iobufs.copy_to_bytes(4)); assert_eq!( chain.copy_to_bytes(chain.remaining()), iobufs.copy_to_bytes(iobufs.remaining()) ); assert_eq!(chain.remaining(), 0); assert_eq!(iobufs.remaining(), 0); } #[test] fn test_iobufs_try_into_single() { let single = IoBufs::from(IoBuf::from(b"hello")); let single = single.try_into_single().expect("single expected"); assert_eq!(single, b"hello"); let multi = IoBufs::from(vec![IoBuf::from(b"ab"), IoBuf::from(b"cd")]); let multi = multi.try_into_single().expect_err("multi expected"); assert_eq!(multi.coalesce(), b"abcd"); } #[test] fn test_iobufs_chunks_vectored_multiple_slices() { // Single non-empty buffers should export exactly one slice. let single = IoBufs::from(IoBuf::from(b"xy")); let mut single_dst = [IoSlice::new(&[]); 2]; let count = single.chunks_vectored(&mut single_dst); assert_eq!(count, 1); assert_eq!(&single_dst[0][..], b"xy"); // Single empty buffers should export no slices. let empty_single = IoBufs::default(); let mut empty_single_dst = [IoSlice::new(&[]); 1]; assert_eq!(empty_single.chunks_vectored(&mut empty_single_dst), 0); let bufs = IoBufs::from(vec![ IoBuf::from(b"ab"), IoBuf::from(b"cd"), IoBuf::from(b"ef"), IoBuf::from(b"gh"), ]); // Destination capacity should cap how many chunks we export. let mut small = [IoSlice::new(&[]); 2]; let count = bufs.chunks_vectored(&mut small); assert_eq!(count, 2); assert_eq!(&small[0][..], b"ab"); assert_eq!(&small[1][..], b"cd"); // Larger destination should include every readable chunk. let mut large = [IoSlice::new(&[]); 8]; let count = bufs.chunks_vectored(&mut large); assert_eq!(count, 4); assert_eq!(&large[0][..], b"ab"); assert_eq!(&large[1][..], b"cd"); assert_eq!(&large[2][..], b"ef"); assert_eq!(&large[3][..], b"gh"); // Empty destination cannot accept any slices. let mut empty_dst: [IoSlice<'_>; 0] = []; assert_eq!(bufs.chunks_vectored(&mut empty_dst), 0); // Non-canonical shapes should skip empty leading chunks. let sparse = IoBufs { inner: IoBufsInner::Pair([IoBuf::default(), IoBuf::from(b"x")]), }; let mut dst = [IoSlice::new(&[]); 2]; let count = sparse.chunks_vectored(&mut dst); assert_eq!(count, 1); assert_eq!(&dst[0][..], b"x"); // Triple should skip empty chunks and preserve readable order. let sparse_triple = IoBufs { inner: IoBufsInner::Triple([IoBuf::default(), IoBuf::from(b"y"), IoBuf::from(b"z")]), }; let mut dst = [IoSlice::new(&[]); 3]; let count = sparse_triple.chunks_vectored(&mut dst); assert_eq!(count, 2); assert_eq!(&dst[0][..], b"y"); assert_eq!(&dst[1][..], b"z"); // Chunked shapes with only empty buffers should export no slices. let empty_chunked = IoBufs { inner: IoBufsInner::Chunked(VecDeque::from([IoBuf::default(), IoBuf::default()])), }; let mut dst = [IoSlice::new(&[]); 2]; assert_eq!(empty_chunked.chunks_vectored(&mut dst), 0); } #[test] fn test_iobufsmut_freeze_chunked() { // Multiple non-empty buffers stay multi-chunk. let buf1 = IoBufMut::from(b"hello".as_ref()); let buf2 = IoBufMut::from(b" world".as_ref()); let bufs = IoBufsMut::from(vec![buf1, buf2]); let mut frozen = bufs.freeze(); assert!(!frozen.is_single()); assert_eq!(frozen.chunk(), b"hello"); frozen.advance(5); assert_eq!(frozen.chunk(), b" world"); frozen.advance(6); assert_eq!(frozen.remaining(), 0); // Empty buffers are filtered out. let buf1 = IoBufMut::from(b"hello".as_ref()); let empty = IoBufMut::default(); let buf2 = IoBufMut::from(b" world".as_ref()); let bufs = IoBufsMut::from(vec![buf1, empty, buf2]); let mut frozen = bufs.freeze(); assert!(!frozen.is_single()); assert_eq!(frozen.chunk(), b"hello"); frozen.advance(5); assert_eq!(frozen.chunk(), b" world"); frozen.advance(6); assert_eq!(frozen.remaining(), 0); // Collapses to Single when one non-empty buffer remains let empty1 = IoBufMut::default(); let buf = IoBufMut::from(b"only one".as_ref()); let empty2 = IoBufMut::default(); let bufs = IoBufsMut::from(vec![empty1, buf, empty2]); let frozen = bufs.freeze(); assert!(frozen.is_single()); assert_eq!(frozen.coalesce(), b"only one"); // All empty buffers -> Single with empty buffer let empty1 = IoBufMut::default(); let empty2 = IoBufMut::default(); let bufs = IoBufsMut::from(vec![empty1, empty2]); let frozen = bufs.freeze(); assert!(frozen.is_single()); assert!(frozen.is_empty()); } #[test] fn test_iobufsmut_from_vec() { // Empty Vec becomes Single with empty buffer let bufs = IoBufsMut::from(Vec::::new()); assert!(bufs.is_single()); assert!(bufs.is_empty()); // Vec with one element becomes Single let buf = IoBufMut::from(b"test"); let bufs = IoBufsMut::from(vec![buf]); assert!(bufs.is_single()); assert_eq!(bufs.chunk(), b"test"); // Vec with multiple elements becomes multi-chunk. let buf1 = IoBufMut::from(b"hello"); let buf2 = IoBufMut::from(b" world"); let bufs = IoBufsMut::from(vec![buf1, buf2]); assert!(!bufs.is_single()); } #[test] fn test_iobufsmut_fast_path_shapes() { let pair = IoBufsMut::from(vec![IoBufMut::from(b"a"), IoBufMut::from(b"b")]); assert!(matches!(pair.inner, IoBufsMutInner::Pair(_))); let triple = IoBufsMut::from(vec![ IoBufMut::from(b"a"), IoBufMut::from(b"b"), IoBufMut::from(b"c"), ]); assert!(matches!(triple.inner, IoBufsMutInner::Triple(_))); let chunked = IoBufsMut::from(vec![ IoBufMut::from(b"a"), IoBufMut::from(b"b"), IoBufMut::from(b"c"), IoBufMut::from(b"d"), ]); assert!(matches!(chunked.inner, IoBufsMutInner::Chunked(_))); } #[test] fn test_iobufsmut_default() { // Default IoBufsMut should be a single empty chunk. let bufs = IoBufsMut::default(); assert!(bufs.is_single()); assert!(bufs.is_empty()); assert_eq!(bufs.len(), 0); } #[test] fn test_iobufsmut_from_array() { // From<[u8; N]> should create a single-chunk container with the array data. let bufs = IoBufsMut::from([1u8, 2, 3, 4, 5]); assert!(bufs.is_single()); assert_eq!(bufs.len(), 5); assert_eq!(bufs.chunk(), &[1, 2, 3, 4, 5]); } #[test] fn test_iobufsmut_buf_trait_chunked() { let buf1 = IoBufMut::from(b"hello"); let buf2 = IoBufMut::from(b" "); let buf3 = IoBufMut::from(b"world"); let mut bufs = IoBufsMut::from(vec![buf1, buf2, buf3]); assert_eq!(bufs.remaining(), 11); assert_eq!(bufs.chunk(), b"hello"); // Advance within first buffer bufs.advance(3); assert_eq!(bufs.remaining(), 8); assert_eq!(bufs.chunk(), b"lo"); // Advance past first buffer (should pop_front) bufs.advance(2); assert_eq!(bufs.remaining(), 6); assert_eq!(bufs.chunk(), b" "); // Advance exactly one buffer bufs.advance(1); assert_eq!(bufs.remaining(), 5); assert_eq!(bufs.chunk(), b"world"); // Advance to end bufs.advance(5); assert_eq!(bufs.remaining(), 0); } #[test] #[should_panic(expected = "cannot advance past end of buffer")] fn test_iobufsmut_advance_past_end() { let buf1 = IoBufMut::from(b"hello"); let buf2 = IoBufMut::from(b" world"); let mut bufs = IoBufsMut::from(vec![buf1, buf2]); bufs.advance(20); } #[test] fn test_iobufsmut_bufmut_trait_single() { let mut bufs = IoBufsMut::from(IoBufMut::with_capacity(20)); assert_eq!(bufs.remaining_mut(), 20); bufs.put_slice(b"hello"); assert_eq!(bufs.chunk(), b"hello"); assert_eq!(bufs.len(), 5); assert_eq!(bufs.remaining_mut(), 15); bufs.put_slice(b" world"); assert_eq!(bufs.coalesce(), b"hello world"); } #[test] fn test_iobufsmut_zeroed_write() { // Use zeroed buffers which have a fixed length let bufs = IoBufsMut::from(IoBufMut::zeroed(20)); assert_eq!(bufs.len(), 20); // Can write using as_mut on coalesced buffer let mut coalesced = bufs.coalesce(); coalesced.as_mut()[..5].copy_from_slice(b"hello"); assert_eq!(&coalesced.as_ref()[..5], b"hello"); } #[test] fn test_iobufsmut_bufmut_put_slice() { // Test writing across multiple buffers let buf1 = IoBufMut::with_capacity(5); let buf2 = IoBufMut::with_capacity(6); let mut bufs = IoBufsMut::from(vec![buf1, buf2]); // Write data bufs.put_slice(b"hello"); bufs.put_slice(b" world"); assert_eq!(bufs.coalesce(), b"hello world"); } #[test] fn test_iobufs_advance_drains_buffers() { let mut bufs = IoBufs::from(IoBuf::from(b"hello")); bufs.append(IoBuf::from(b" ")); bufs.append(IoBuf::from(b"world")); // Advance exactly past first buffer bufs.advance(5); assert_eq!(bufs.remaining(), 6); assert_eq!(bufs.chunk(), b" "); // Advance across multiple buffers bufs.advance(4); assert_eq!(bufs.remaining(), 2); assert_eq!(bufs.chunk(), b"ld"); } #[test] fn test_iobufs_advance_exactly_to_boundary() { let mut bufs = IoBufs::from(IoBuf::from(b"abc")); bufs.append(IoBuf::from(b"def")); // Advance exactly to first buffer boundary bufs.advance(3); assert_eq!(bufs.remaining(), 3); assert_eq!(bufs.chunk(), b"def"); // Advance exactly to end bufs.advance(3); assert_eq!(bufs.remaining(), 0); } #[test] fn test_iobufs_advance_canonicalizes_pair_to_single() { let mut bufs = IoBufs::from(IoBuf::from(b"ab")); bufs.append(IoBuf::from(b"cd")); bufs.advance(2); assert!(bufs.is_single()); assert_eq!(bufs.chunk(), b"cd"); } #[test] fn test_iobufsmut_emptiness_with_reserved_capacity() { for count in 0..=8 { for readable in 0..=count { let mut chunks = Vec::new(); for index in 0..count { let mut chunk = IoBufMut::with_capacity(4); if index == readable { chunk.put_u8(7); } chunks.push(chunk); } let mut bufs = IoBufsMut::from(chunks); assert_eq!(bufs.is_empty(), readable == count); assert_eq!(bufs.has_remaining(), readable != count); if readable != count { bufs.advance(1); assert!(bufs.is_empty()); assert!(!bufs.has_remaining()); } if bufs.remaining_mut() > 0 { bufs.put_slice(&[9]); assert!(!bufs.is_empty()); assert!(bufs.has_remaining()); } } } } #[test] fn test_iobufsmut_with_empty_buffers() { let buf1 = IoBufMut::from(b"hello"); let buf2 = IoBufMut::default(); let buf3 = IoBufMut::from(b" world"); let mut bufs = IoBufsMut::from(vec![buf1, buf2, buf3]); assert_eq!(bufs.remaining(), 11); assert_eq!(bufs.chunk(), b"hello"); // Advance past first buffer bufs.advance(5); // Empty buffer should be skipped assert_eq!(bufs.chunk(), b" world"); assert_eq!(bufs.remaining(), 6); } #[test] fn test_iobufsmut_advance_skips_leading_writable_empty_chunk() { // A leading chunk with capacity but no readable bytes (len == 0) should // be skipped during advance, reaching the next readable chunk. let empty_writable = IoBufMut::with_capacity(4); let payload = IoBufMut::from(b"xy"); let mut bufs = IoBufsMut::from(vec![empty_writable, payload]); bufs.advance(1); assert_eq!(bufs.chunk(), b"y"); assert_eq!(bufs.remaining(), 1); } #[test] fn test_iobufsmut_read_ops_preserve_writable_capacity() { // Draining the filled first chunk must keep the never-filled second // chunk's reserved capacity, so reads leave remaining_mut unchanged. let mut a = IoBufMut::with_capacity(8); a.put_slice(&[1u8; 8]); let b = IoBufMut::with_capacity(8); let mut bufs = IoBufsMut::from(vec![a, b]); assert_eq!(bufs.remaining(), 8); assert_eq!(bufs.remaining_mut(), 8); bufs.advance(8); assert_eq!(bufs.remaining(), 0); assert_eq!(bufs.remaining_mut(), 8); bufs.put_slice(&[2u8; 8]); assert_eq!(bufs.copy_to_bytes(8).as_ref(), &[2u8; 8]); // copy_to_bytes drains must preserve capacity the same way. let mut a = IoBufMut::with_capacity(8); a.put_slice(&[3u8; 8]); let b = IoBufMut::with_capacity(8); let mut bufs = IoBufsMut::from(vec![a, b]); assert_eq!(bufs.copy_to_bytes(8).as_ref(), &[3u8; 8]); assert_eq!(bufs.remaining_mut(), 8); // Zero-length drains do not disturb chunk state at all. assert!(bufs.copy_to_bytes(0).is_empty()); assert_eq!(bufs.remaining_mut(), 8); } #[test] fn test_iobufsmut_from_vec_u8_preserves_capacity() { // The vec's allocation is adopted and its reserved capacity is // preserved exactly. let mut vec = Vec::with_capacity(128); vec.extend_from_slice(b"abc"); let cap = vec.capacity(); let mut bufs = IoBufsMut::from(vec); assert_eq!(bufs.remaining(), 3); assert_eq!(bufs.capacity(), cap); bufs.put_slice(b"d"); assert_eq!(bufs.copy_to_bytes(4).as_ref(), b"abcd"); // An empty vec keeps its full reservation, so `read_at_buf` can fill a // `Vec::with_capacity(len)` in place through `set_len(len)`. let vec = Vec::::with_capacity(64); let len = vec.capacity(); let mut bufs = IoBufsMut::from(vec); assert!(bufs.is_empty()); assert_eq!(bufs.capacity(), len); // SAFETY: all `len` bytes are initialized by copy_from_slice below. unsafe { bufs.set_len(len) }; bufs.copy_from_slice(&vec![9u8; len]); assert_eq!(bufs.copy_to_bytes(len).as_ref(), vec![9u8; len].as_slice()); // Exactly-sized vec: adopted with contents intact. let bufs = IoBufsMut::from(vec![1u8, 2, 3]); assert_eq!(bufs.remaining(), 3); assert_eq!(bufs.capacity(), 3); assert_eq!(bufs.chunk(), &[1, 2, 3]); } #[test] fn test_iobufsmut_coalesce_after_advance() { // Advance mid-chunk: advance 3 of 11 bytes let buf1 = IoBufMut::from(b"hello"); let buf2 = IoBufMut::from(b" world"); let mut bufs = IoBufsMut::from(vec![buf1, buf2]); bufs.advance(3); assert_eq!(bufs.coalesce(), b"lo world"); // Advance to exact chunk boundary: advance 5 of 11 bytes let buf1 = IoBufMut::from(b"hello"); let buf2 = IoBufMut::from(b" world"); let mut bufs = IoBufsMut::from(vec![buf1, buf2]); bufs.advance(5); assert_eq!(bufs.coalesce(), b" world"); } #[test] fn test_iobufsmut_copy_to_bytes() { let buf1 = IoBufMut::from(b"hello"); let buf2 = IoBufMut::from(b" world"); let mut bufs = IoBufsMut::from(vec![buf1, buf2]); // First read spans chunks and leaves unread suffix. let first = bufs.copy_to_bytes(7); assert_eq!(&first[..], b"hello w"); assert_eq!(bufs.remaining(), 4); // Second read drains the remainder. let rest = bufs.copy_to_bytes(4); assert_eq!(&rest[..], b"orld"); assert_eq!(bufs.remaining(), 0); } #[test] fn test_iobufsmut_copy_to_bytes_chunked_four_plus() { let mut bufs = IoBufsMut::from(vec![ IoBufMut::from(b"ab"), IoBufMut::from(b"cd"), IoBufMut::from(b"ef"), IoBufMut::from(b"gh"), ]); // Exercise chunked advance path before copy_to_bytes. bufs.advance(1); assert_eq!(bufs.chunk(), b"b"); bufs.advance(1); assert_eq!(bufs.chunk(), b"cd"); // Chunked fast-path: first chunk alone satisfies request. let first = bufs.copy_to_bytes(1); assert_eq!(&first[..], b"c"); // Chunked slow-path: request crosses chunk boundaries. let second = bufs.copy_to_bytes(4); assert_eq!(&second[..], b"defg"); let rest = bufs.copy_to_bytes(1); assert_eq!(&rest[..], b"h"); assert_eq!(bufs.remaining(), 0); // Enter copy_to_bytes while still in chunked representation. let mut bufs = IoBufsMut::from(vec![ IoBufMut::from(b"a"), IoBufMut::from(b"b"), IoBufMut::from(b"c"), IoBufMut::from(b"d"), IoBufMut::from(b"e"), ]); assert!(matches!(bufs.inner, IoBufsMutInner::Chunked(_))); let first = bufs.copy_to_bytes(1); assert_eq!(&first[..], b"a"); // Stay chunked while consuming across multiple tiny chunks. let next = bufs.copy_to_bytes(3); assert_eq!(&next[..], b"bcd"); assert_eq!(bufs.chunk(), b"e"); assert_eq!(bufs.remaining(), 1); } #[test] fn test_iobufsmut_copy_to_bytes_canonicalizes_pair() { let mut bufs = IoBufsMut::from(vec![IoBufMut::from(b"ab"), IoBufMut::from(b"cd")]); assert!(matches!(bufs.inner, IoBufsMutInner::Pair(_))); let first = bufs.copy_to_bytes(2); assert_eq!(&first[..], b"ab"); assert!(bufs.is_single()); assert_eq!(bufs.chunk(), b"cd"); assert_eq!(bufs.remaining(), 2); } #[test] fn test_iobufsmut_copy_from_slice_single() { let mut bufs = IoBufsMut::from(IoBufMut::zeroed(11)); bufs.copy_from_slice(b"hello world"); assert_eq!(bufs.coalesce(), b"hello world"); } #[test] fn test_iobufsmut_copy_from_slice_chunked() { let buf1 = IoBufMut::zeroed(5); let buf2 = IoBufMut::zeroed(6); let mut bufs = IoBufsMut::from(vec![buf1, buf2]); bufs.copy_from_slice(b"hello world"); // Verify each chunk was filled correctly. assert_eq!(bufs.chunk(), b"hello"); bufs.advance(5); assert_eq!(bufs.chunk(), b" world"); bufs.advance(6); assert_eq!(bufs.remaining(), 0); } #[test] #[should_panic(expected = "source slice length must match buffer length")] fn test_iobufsmut_copy_from_slice_wrong_length() { let mut bufs = IoBufsMut::from(IoBufMut::zeroed(5)); bufs.copy_from_slice(b"hello world"); // 11 bytes into 5-byte buffer } #[test] fn test_iobufsmut_matches_bytesmut_chain() { // Create three BytesMut with capacity for final content comparison. let mut bm1 = BytesMut::with_capacity(5); let mut bm2 = BytesMut::with_capacity(6); let mut bm3 = BytesMut::with_capacity(7); // Create matching IoBufsMut let mut iobufs = IoBufsMut::from(vec![ IoBufMut::with_capacity(5), IoBufMut::with_capacity(6), IoBufMut::with_capacity(7), ]); // Fixed-capacity IoBufsMut exposes the current writable chunk. assert_eq!(iobufs.chunk_mut().len(), 5); // Write some data (&mut bm1) .chain_mut(&mut bm2) .chain_mut(&mut bm3) .put_slice(b"hel"); iobufs.put_slice(b"hel"); assert_eq!(iobufs.chunk_mut().len(), 2); // Write more data (&mut bm1) .chain_mut(&mut bm2) .chain_mut(&mut bm3) .put_slice(b"lo world!"); iobufs.put_slice(b"lo world!"); assert_eq!(iobufs.chunk_mut().len(), 6); // Verify final content matches let frozen = iobufs.freeze().coalesce(); let mut chain_content = bm1.to_vec(); chain_content.extend_from_slice(&bm2); chain_content.extend_from_slice(&bm3); assert_eq!(frozen, chain_content.as_slice()); assert_eq!(frozen, b"hello world!"); } #[test] fn test_iobufsmut_buf_matches_bytes_chain() { // Create pre-filled Bytes buffers let mut b1 = Bytes::from_static(b"hello"); let mut b2 = Bytes::from_static(b" world"); let b3 = Bytes::from_static(b"!"); // Create matching IoBufsMut let mut iobufs = IoBufsMut::from(vec![ IoBufMut::from(b"hello"), IoBufMut::from(b" world"), IoBufMut::from(b"!"), ]); // Test Buf::remaining matches let chain_remaining = b1.clone().chain(b2.clone()).chain(b3.clone()).remaining(); assert_eq!(chain_remaining, iobufs.remaining()); // Test Buf::chunk matches let chain_chunk = b1 .clone() .chain(b2.clone()) .chain(b3.clone()) .chunk() .to_vec(); assert_eq!(chain_chunk, iobufs.chunk().to_vec()); // Advance and test again b1.advance(3); iobufs.advance(3); let chain_remaining = b1.clone().chain(b2.clone()).chain(b3.clone()).remaining(); assert_eq!(chain_remaining, iobufs.remaining()); let chain_chunk = b1 .clone() .chain(b2.clone()) .chain(b3.clone()) .chunk() .to_vec(); assert_eq!(chain_chunk, iobufs.chunk().to_vec()); // Advance past first buffer boundary into second b1.advance(2); iobufs.advance(2); let chain_remaining = b1.clone().chain(b2.clone()).chain(b3.clone()).remaining(); assert_eq!(chain_remaining, iobufs.remaining()); // Now we should be in the second buffer let chain_chunk = b1 .clone() .chain(b2.clone()) .chain(b3.clone()) .chunk() .to_vec(); assert_eq!(chain_chunk, iobufs.chunk().to_vec()); // Advance past second buffer boundary into third b2.advance(6); iobufs.advance(6); let chain_remaining = b1.clone().chain(b2.clone()).chain(b3.clone()).remaining(); assert_eq!(chain_remaining, iobufs.remaining()); // Now we should be in the third buffer let chain_chunk = b1.chain(b2).chain(b3).chunk().to_vec(); assert_eq!(chain_chunk, iobufs.chunk().to_vec()); // Test copy_to_bytes let b1 = Bytes::from_static(b"hello"); let b2 = Bytes::from_static(b" world"); let b3 = Bytes::from_static(b"!"); let mut iobufs = IoBufsMut::from(vec![ IoBufMut::from(b"hello"), IoBufMut::from(b" world"), IoBufMut::from(b"!"), ]); let chain_bytes = b1.chain(b2).chain(b3).copy_to_bytes(8); let iobufs_bytes = iobufs.copy_to_bytes(8); assert_eq!(chain_bytes, iobufs_bytes); assert_eq!(chain_bytes.as_ref(), b"hello wo"); } #[test] fn test_iobufsmut_chunks_vectored_multiple_slices() { // Single non-empty buffers should export exactly one slice. let single = IoBufsMut::from(IoBufMut::from(b"xy")); let mut single_dst = [IoSlice::new(&[]); 2]; let count = single.chunks_vectored(&mut single_dst); assert_eq!(count, 1); assert_eq!(&single_dst[0][..], b"xy"); // Single empty buffers should export no slices. let empty_single = IoBufsMut::default(); let mut empty_single_dst = [IoSlice::new(&[]); 1]; assert_eq!(empty_single.chunks_vectored(&mut empty_single_dst), 0); let bufs = IoBufsMut::from(vec![ IoBufMut::from(b"ab"), IoBufMut::from(b"cd"), IoBufMut::from(b"ef"), IoBufMut::from(b"gh"), ]); // Destination capacity should cap how many chunks we export. let mut small = [IoSlice::new(&[]); 2]; let count = bufs.chunks_vectored(&mut small); assert_eq!(count, 2); assert_eq!(&small[0][..], b"ab"); assert_eq!(&small[1][..], b"cd"); // Larger destination should include every readable chunk. let mut large = [IoSlice::new(&[]); 8]; let count = bufs.chunks_vectored(&mut large); assert_eq!(count, 4); assert_eq!(&large[0][..], b"ab"); assert_eq!(&large[1][..], b"cd"); assert_eq!(&large[2][..], b"ef"); assert_eq!(&large[3][..], b"gh"); // Empty destination cannot accept any slices. let mut empty_dst: [IoSlice<'_>; 0] = []; assert_eq!(bufs.chunks_vectored(&mut empty_dst), 0); // Non-canonical shapes should skip empty leading chunks. let sparse = IoBufsMut { inner: IoBufsMutInner::Pair([IoBufMut::default(), IoBufMut::from(b"y")]), }; let mut dst = [IoSlice::new(&[]); 2]; let count = sparse.chunks_vectored(&mut dst); assert_eq!(count, 1); assert_eq!(&dst[0][..], b"y"); // Triple should skip empty chunks and preserve readable order. let sparse_triple = IoBufsMut { inner: IoBufsMutInner::Triple([ IoBufMut::default(), IoBufMut::from(b"z"), IoBufMut::from(b"w"), ]), }; let mut dst = [IoSlice::new(&[]); 3]; let count = sparse_triple.chunks_vectored(&mut dst); assert_eq!(count, 2); assert_eq!(&dst[0][..], b"z"); assert_eq!(&dst[1][..], b"w"); // Chunked shapes with only empty buffers should export no slices. let empty_chunked = IoBufsMut { inner: IoBufsMutInner::Chunked(VecDeque::from([ IoBufMut::default(), IoBufMut::default(), ])), }; let mut dst = [IoSlice::new(&[]); 2]; assert_eq!(empty_chunked.chunks_vectored(&mut dst), 0); } #[test] fn test_iobufsmut_try_into_single() { let single = IoBufsMut::from(IoBufMut::from(b"hello")); let single = single.try_into_single().expect("single expected"); assert_eq!(single, b"hello"); let multi = IoBufsMut::from(vec![IoBufMut::from(b"ab"), IoBufMut::from(b"cd")]); let multi = multi.try_into_single().expect_err("multi expected"); assert_eq!(multi.coalesce(), b"abcd"); } #[test] fn test_iobufsmut_freeze_after_advance() { // Partial advance: advance 3 of 11 bytes let buf1 = IoBufMut::from(b"hello"); let buf2 = IoBufMut::from(b" world"); let mut bufs = IoBufsMut::from(vec![buf1, buf2]); bufs.advance(3); assert_eq!(bufs.len(), 8); let frozen = bufs.freeze(); assert_eq!(frozen.len(), 8); assert_eq!(frozen.coalesce(), b"lo world"); // Exact boundary advance: advance 5 of 11 bytes (first buf is 5 bytes) let buf1 = IoBufMut::from(b"hello"); let buf2 = IoBufMut::from(b" world"); let mut bufs = IoBufsMut::from(vec![buf1, buf2]); bufs.advance(5); assert_eq!(bufs.len(), 6); // First buffer should be fully consumed (empty after advance) // freeze() filters empty buffers, so result should be Single let frozen = bufs.freeze(); assert!(frozen.is_single()); assert_eq!(frozen.coalesce(), b" world"); } #[test] fn test_iobufsmut_coalesce_shapes() { // A single buffer is returned as is. let mut buf = IoBufMut::from(b"hello"); let original_ptr = buf.as_mut_ptr(); let bufs = IoBufsMut::from(buf); let coalesced = bufs.coalesce(); assert_eq!(coalesced, b"hello"); assert_eq!(coalesced.as_ref().as_ptr(), original_ptr); // Pair, triple, and chunked (4+) shapes copy into one allocation of // exactly the readable length. let pair = IoBufsMut::from(vec![IoBufMut::from(b"ab"), IoBufMut::from(b"cd")]); let pair = pair.coalesce(); assert_eq!(pair, b"abcd"); assert_eq!(pair.capacity(), 4); let triple = IoBufsMut::from(vec![ IoBufMut::from(b"a"), IoBufMut::from(b"b"), IoBufMut::from(b"c"), ]); let triple = triple.coalesce(); assert_eq!(triple, b"abc"); assert_eq!(triple.capacity(), 3); let chunked = IoBufsMut::from(vec![ IoBufMut::from(b"1"), IoBufMut::from(b"2"), IoBufMut::from(b"3"), IoBufMut::from(b"4"), ]); let chunked = chunked.coalesce(); assert_eq!(chunked, b"1234"); assert_eq!(chunked.capacity(), 4); } #[test] fn test_iobufs_additional_shape_and_conversion_paths() { // Constructor coverage for mutable/immutable/slice-backed inputs. let from_mut = IoBufs::from(IoBufMut::from(b"m")); assert_eq!(from_mut.chunk(), b"m"); let from_bytes = IoBufs::from(Bytes::from_static(b"b")); assert_eq!(from_bytes.chunk(), b"b"); let from_bytesmut = IoBufs::from(BytesMut::from(&b"bm"[..])); assert_eq!(from_bytesmut.chunk(), b"bm"); let from_vec = IoBufs::from(vec![1u8, 2u8]); assert_eq!(from_vec.chunk(), &[1u8, 2]); let static_slice: &'static [u8] = b"slice"; let from_static = IoBufs::from(static_slice); assert_eq!(from_static.chunk(), b"slice"); // Canonicalizing an already-empty buffer remains a single empty chunk. let mut single_empty = IoBufs::default(); single_empty.canonicalize(); assert!(single_empty.is_single()); // Triple path: prepend/append can promote into chunked while preserving order. let mut triple = IoBufs::from(vec![ IoBuf::from(b"a".to_vec()), IoBuf::from(b"b".to_vec()), IoBuf::from(b"c".to_vec()), ]); assert!(triple.as_single().is_none()); triple.prepend(IoBuf::from(vec![b'0'])); triple.prepend(IoBuf::from(vec![b'1'])); triple.append(IoBuf::from(vec![b'2'])); assert_eq!(triple.copy_to_bytes(triple.remaining()).as_ref(), b"10abc2"); // Appending to an existing triple keeps byte order stable. let mut triple_append = IoBufs::from(vec![ IoBuf::from(b"x".to_vec()), IoBuf::from(b"y".to_vec()), IoBuf::from(b"z".to_vec()), ]); triple_append.append(IoBuf::from(vec![b'w'])); assert_eq!(triple_append.coalesce(), b"xyzw"); // coalesce on a triple should preserve contents. let triple = IoBufs::from(vec![ IoBuf::from(b"a".to_vec()), IoBuf::from(b"b".to_vec()), IoBuf::from(b"c".to_vec()), ]); assert_eq!(triple.coalesce(), b"abc"); // Coalescing after advancing a four-chunk input preserves the remaining bytes. let mut chunked_advanced = IoBufs::from(vec![ IoBuf::from(b"a".to_vec()), IoBuf::from(b"b".to_vec()), IoBuf::from(b"c".to_vec()), IoBuf::from(b"d".to_vec()), ]); assert_eq!(chunked_advanced.remaining(), 4); chunked_advanced.advance(1); assert_eq!(chunked_advanced.coalesce(), b"bcd"); // Non-canonical Pair/Triple/Chunked shapes should still expose the first readable chunk. let pair_second = IoBufs { inner: IoBufsInner::Pair([IoBuf::default(), IoBuf::from(vec![1u8])]), }; assert_eq!(pair_second.chunk(), &[1u8]); let pair_empty = IoBufs { inner: IoBufsInner::Pair([IoBuf::default(), IoBuf::default()]), }; assert_eq!(pair_empty.chunk(), b""); let triple_third = IoBufs { inner: IoBufsInner::Triple([ IoBuf::default(), IoBuf::default(), IoBuf::from(vec![3u8]), ]), }; assert_eq!(triple_third.chunk(), &[3u8]); let triple_second = IoBufs { inner: IoBufsInner::Triple([ IoBuf::default(), IoBuf::from(vec![2u8]), IoBuf::default(), ]), }; assert_eq!(triple_second.chunk(), &[2u8]); let triple_empty = IoBufs { inner: IoBufsInner::Triple([IoBuf::default(), IoBuf::default(), IoBuf::default()]), }; assert_eq!(triple_empty.chunk(), b""); let chunked_second = IoBufs { inner: IoBufsInner::Chunked(VecDeque::from([IoBuf::default(), IoBuf::from(vec![9u8])])), }; assert_eq!(chunked_second.chunk(), &[9u8]); let chunked_empty = IoBufs { inner: IoBufsInner::Chunked(VecDeque::from([IoBuf::default()])), }; assert_eq!(chunked_empty.chunk(), b""); } #[test] fn test_iobufsmut_additional_shape_and_conversion_paths() { // `as_single` accessors should work only for single-shape containers. let mut single = IoBufsMut::from(IoBufMut::from(b"x")); assert!(single.as_single().is_some()); assert!(single.as_single_mut().is_some()); single.canonicalize(); assert!(single.is_single()); let mut pair = IoBufsMut::from(vec![IoBufMut::from(b"a"), IoBufMut::from(b"b")]); assert!(pair.as_single().is_none()); assert!(pair.as_single_mut().is_none()); // Constructor coverage for raw vec and BytesMut sources. let from_vec = IoBufsMut::from(vec![1u8, 2u8]); assert_eq!(from_vec.chunk(), &[1u8, 2]); let from_bytesmut = IoBufsMut::from(BytesMut::from(&b"cd"[..])); assert_eq!(from_bytesmut.chunk(), b"cd"); // Chunked write path: set_len + copy_from_slice + freeze round-trip. let mut chunked = IoBufsMut::from(vec![ IoBufMut::with_capacity(1), IoBufMut::with_capacity(1), IoBufMut::with_capacity(1), IoBufMut::with_capacity(1), ]); // SAFETY: We only write/read initialized bytes after `copy_from_slice`. unsafe { chunked.set_len(4) }; chunked.copy_from_slice(b"wxyz"); assert_eq!(chunked.capacity(), 4); assert_eq!(chunked.remaining(), 4); let frozen = chunked.freeze(); assert_eq!(frozen.coalesce(), b"wxyz"); } #[test] fn test_iobufsmut_noncanonical_chunk_and_chunk_mut_paths() { fn no_spare_capacity_buf() -> IoBufMut { let mut buf = IoBufMut::with_capacity(1); let cap = buf.capacity(); // SAFETY: We never read from this buffer in this helper. unsafe { buf.set_len(cap) }; buf } // `chunk()` should skip empty front buffers across all shapes. let pair_second = IoBufsMut { inner: IoBufsMutInner::Pair([IoBufMut::default(), IoBufMut::from(b"b")]), }; assert_eq!(pair_second.chunk(), b"b"); let pair_empty = IoBufsMut { inner: IoBufsMutInner::Pair([IoBufMut::default(), IoBufMut::default()]), }; assert_eq!(pair_empty.chunk(), b""); let triple_third = IoBufsMut { inner: IoBufsMutInner::Triple([ IoBufMut::default(), IoBufMut::default(), IoBufMut::from(b"c"), ]), }; assert_eq!(triple_third.chunk(), b"c"); let triple_second = IoBufsMut { inner: IoBufsMutInner::Triple([ IoBufMut::default(), IoBufMut::from(b"b"), IoBufMut::default(), ]), }; assert_eq!(triple_second.chunk(), b"b"); let triple_empty = IoBufsMut { inner: IoBufsMutInner::Triple([ IoBufMut::default(), IoBufMut::default(), IoBufMut::default(), ]), }; assert_eq!(triple_empty.chunk(), b""); let chunked_second = IoBufsMut { inner: IoBufsMutInner::Chunked(VecDeque::from([ IoBufMut::default(), IoBufMut::from(b"d"), ])), }; assert_eq!(chunked_second.chunk(), b"d"); let chunked_empty = IoBufsMut { inner: IoBufsMutInner::Chunked(VecDeque::from([IoBufMut::default()])), }; assert_eq!(chunked_empty.chunk(), b""); // `chunk_mut()` should skip non-writable fronts and return first writable chunk. let mut pair_chunk_mut = IoBufsMut { inner: IoBufsMutInner::Pair([no_spare_capacity_buf(), IoBufMut::with_capacity(2)]), }; assert!(pair_chunk_mut.chunk_mut().len() >= 2); let mut pair_chunk_mut_empty = IoBufsMut { inner: IoBufsMutInner::Pair([no_spare_capacity_buf(), no_spare_capacity_buf()]), }; assert_eq!(pair_chunk_mut_empty.chunk_mut().len(), 0); let mut triple_chunk_mut = IoBufsMut { inner: IoBufsMutInner::Triple([ no_spare_capacity_buf(), no_spare_capacity_buf(), IoBufMut::with_capacity(3), ]), }; assert!(triple_chunk_mut.chunk_mut().len() >= 3); let mut triple_chunk_mut_second = IoBufsMut { inner: IoBufsMutInner::Triple([ no_spare_capacity_buf(), IoBufMut::with_capacity(2), no_spare_capacity_buf(), ]), }; assert!(triple_chunk_mut_second.chunk_mut().len() >= 2); let mut triple_chunk_mut_empty = IoBufsMut { inner: IoBufsMutInner::Triple([ no_spare_capacity_buf(), no_spare_capacity_buf(), no_spare_capacity_buf(), ]), }; assert_eq!(triple_chunk_mut_empty.chunk_mut().len(), 0); let mut chunked_chunk_mut = IoBufsMut { inner: IoBufsMutInner::Chunked(VecDeque::from([ IoBufMut::default(), IoBufMut::with_capacity(4), ])), }; assert!(chunked_chunk_mut.chunk_mut().len() >= 4); let mut chunked_chunk_mut_empty = IoBufsMut { inner: IoBufsMutInner::Chunked(VecDeque::from([no_spare_capacity_buf()])), }; assert_eq!(chunked_chunk_mut_empty.chunk_mut().len(), 0); } #[test] fn test_iobuf_internal_chunk_helpers() { // `copy_to_bytes_chunked` drops leading empties on zero-length reads // and asks for canonicalization so the emptied deque collapses back // to the Single representation. let mut empty_with_leading = VecDeque::from([IoBuf::default()]); let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut empty_with_leading, 0, "x"); assert!(bytes.is_empty()); assert!(needs_canonicalize); assert!(empty_with_leading.is_empty()); // Fast path: front chunk can fully satisfy the request. let mut fast = VecDeque::from([ IoBuf::from(b"ab".to_vec()), IoBuf::from(b"cd".to_vec()), IoBuf::from(b"ef".to_vec()), IoBuf::from(b"gh".to_vec()), ]); let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut fast, 2, "x"); assert_eq!(bytes.as_ref(), b"ab"); assert!(needs_canonicalize); assert_eq!(fast.front().expect("front exists").as_ref(), b"cd"); // Slow path: request spans multiple chunks. let mut slow = VecDeque::from([ IoBuf::from(b"a".to_vec()), IoBuf::from(b"bc".to_vec()), IoBuf::from(b"d".to_vec()), IoBuf::from(b"e".to_vec()), ]); let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut slow, 3, "x"); assert_eq!(bytes.as_ref(), b"abc"); assert!(needs_canonicalize); let mut empty_with_leading_mut = VecDeque::from([IoBufMut::default()]); let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut empty_with_leading_mut, 0, "x"); assert!(bytes.is_empty()); assert!(needs_canonicalize); assert!(empty_with_leading_mut.is_empty()); // Mirror the fast/slow chunked helper paths for mutable chunks too. let mut fast_mut = VecDeque::from([ IoBufMut::from(b"ab"), IoBufMut::from(b"cd"), IoBufMut::from(b"ef"), IoBufMut::from(b"gh"), ]); let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut fast_mut, 2, "x"); assert_eq!(bytes.as_ref(), b"ab"); assert!(needs_canonicalize); assert_eq!(fast_mut.front().expect("front exists").as_ref(), b"cd"); let mut slow_mut = VecDeque::from([ IoBufMut::from(b"a"), IoBufMut::from(b"bc"), IoBufMut::from(b"de"), IoBufMut::from(b"f"), ]); let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut slow_mut, 4, "x"); assert_eq!(bytes.as_ref(), b"abcd"); assert!(needs_canonicalize); assert_eq!(slow_mut.front().expect("front exists").as_ref(), b"e"); // `advance_chunked_front` should skip empties and drain in linear order. let mut advance_chunked = VecDeque::from([ IoBuf::default(), IoBuf::from(b"abc".to_vec()), IoBuf::from(b"d".to_vec()), ]); advance_chunked_front(&mut advance_chunked, 2); assert_eq!( advance_chunked.front().expect("front exists").as_ref(), b"c" ); advance_chunked_front(&mut advance_chunked, 2); assert!(advance_chunked.is_empty()); // The front-advance helper also has a separate mutable monomorphization. let mut advance_chunked_mut = VecDeque::from([ IoBufMut::default(), IoBufMut::from(b"abc"), IoBufMut::from(b"d"), ]); advance_chunked_front(&mut advance_chunked_mut, 2); assert_eq!( advance_chunked_mut.front().expect("front exists").as_ref(), b"c" ); advance_chunked_front(&mut advance_chunked_mut, 2); assert!(advance_chunked_mut.is_empty()); // `advance_small_chunks` signals canonicalization when front chunks are exhausted. let mut small = [IoBuf::default(), IoBuf::from(b"abc".to_vec())]; let needs_canonicalize = advance_small_chunks(&mut small, 2); assert!(needs_canonicalize); assert_eq!(small[1].as_ref(), b"c"); let mut small_exact = [ IoBuf::from(b"a".to_vec()), IoBuf::from(b"b".to_vec()), IoBuf::from(b"c".to_vec()), ]; let needs_canonicalize = advance_small_chunks(&mut small_exact, 3); assert!(needs_canonicalize); assert_eq!(small_exact[0].remaining(), 0); assert_eq!(small_exact[1].remaining(), 0); assert_eq!(small_exact[2].remaining(), 0); // Small-chunk copy canonicalization is also instantiated for mutable chunks. let mut small_mut = [ IoBufMut::from(b"a"), IoBufMut::from(b"bc"), IoBufMut::from(b"d"), ]; let (bytes, needs_canonicalize) = copy_to_bytes_small_chunks(&mut small_mut, 3, "x"); assert_eq!(bytes.as_ref(), b"abc"); assert!(needs_canonicalize); assert_eq!(small_mut[2].as_ref(), b"d"); // `advance_mut_in_chunks` returns whether the request fully fit in writable chunks. let mut writable = [IoBufMut::with_capacity(2), IoBufMut::with_capacity(1)]; let mut remaining = 3usize; // SAFETY: We do not read from advanced bytes in this test. let all_advanced = unsafe { advance_mut_in_chunks(&mut writable, &mut remaining) }; assert!(all_advanced); assert_eq!(remaining, 0); // `advance_mut_in_chunks` should skip non-writable chunks. let mut full = IoBufMut::with_capacity(1); // SAFETY: We only mark initialized capacity and never read bytes. unsafe { full.set_len(full.capacity()) }; let mut writable_after_full = [full, IoBufMut::with_capacity(2)]; let mut remaining = 2usize; // SAFETY: We do not read from advanced bytes in this test. let all_advanced = unsafe { advance_mut_in_chunks(&mut writable_after_full, &mut remaining) }; assert!(all_advanced); assert_eq!(remaining, 0); let mut writable_short = [IoBufMut::with_capacity(1), IoBufMut::with_capacity(1)]; let mut remaining = 3usize; // SAFETY: We do not read from advanced bytes in this test. let all_advanced = unsafe { advance_mut_in_chunks(&mut writable_short, &mut remaining) }; assert!(!all_advanced); assert_eq!(remaining, 1); } #[test] fn test_iobufsmut_advance_mut_success_paths() { // Pair path. let mut pair = IoBufsMut { inner: IoBufsMutInner::Pair([IoBufMut::with_capacity(2), IoBufMut::with_capacity(2)]), }; // SAFETY: We only verify cursor movement (`remaining`) and do not read bytes. unsafe { pair.advance_mut(3) }; assert_eq!(pair.remaining(), 3); // Triple path. let mut triple = IoBufsMut { inner: IoBufsMutInner::Triple([ IoBufMut::with_capacity(1), IoBufMut::with_capacity(1), IoBufMut::with_capacity(1), ]), }; // SAFETY: We only verify cursor movement (`remaining`) and do not read bytes. unsafe { triple.advance_mut(2) }; assert_eq!(triple.remaining(), 2); // Chunked wrapped-VecDeque path. let mut wrapped = VecDeque::with_capacity(5); wrapped.push_back(IoBufMut::with_capacity(1)); wrapped.push_back(IoBufMut::with_capacity(1)); wrapped.push_back(IoBufMut::with_capacity(1)); wrapped.push_back(IoBufMut::with_capacity(1)); wrapped.push_back(IoBufMut::with_capacity(1)); let _ = wrapped.pop_front(); wrapped.push_back(IoBufMut::with_capacity(1)); let (first, second) = wrapped.as_slices(); assert!(!first.is_empty()); assert!(!second.is_empty()); // Force `advance_mut` to consume across the wrapped second slice as well. let to_advance = first.len() + 1; let mut chunked = IoBufsMut { inner: IoBufsMutInner::Chunked(wrapped), }; let before = chunked.remaining_mut(); // SAFETY: We only verify cursor movement (`remaining`) and do not read bytes. unsafe { chunked.advance_mut(to_advance) }; assert_eq!(chunked.remaining(), to_advance); assert_eq!(chunked.remaining_mut(), before - to_advance); } #[test] fn test_iobufsmut_advance_mut_zero_noop_when_full() { fn full_chunk() -> IoBufMut { // IoBufMut capacity is fixed, so setting len to capacity leaves // remaining_mut() at zero. let mut buf = IoBufMut::with_capacity(1); let cap = buf.capacity(); // SAFETY: We never read from this buffer in this test. unsafe { buf.set_len(cap) }; buf } // Pair path: fully-written chunks should allow advance_mut(0) as a no-op. let mut pair = IoBufsMut::from(vec![full_chunk(), full_chunk()]); assert!(matches!(pair.inner, IoBufsMutInner::Pair(_))); assert_eq!(pair.remaining_mut(), 0); let before = pair.remaining(); // SAFETY: Advancing by 0 does not expose uninitialized bytes. unsafe { pair.advance_mut(0) }; assert_eq!(pair.remaining(), before); // Triple path: same no-op behavior. let mut triple = IoBufsMut::from(vec![full_chunk(), full_chunk(), full_chunk()]); assert!(matches!(triple.inner, IoBufsMutInner::Triple(_))); assert_eq!(triple.remaining_mut(), 0); let before = triple.remaining(); // SAFETY: Advancing by 0 does not expose uninitialized bytes. unsafe { triple.advance_mut(0) }; assert_eq!(triple.remaining(), before); // Chunked path: 4+ fully-written chunks should also no-op. let mut chunked = IoBufsMut::from(vec![full_chunk(), full_chunk(), full_chunk(), full_chunk()]); assert!(matches!(chunked.inner, IoBufsMutInner::Chunked(_))); assert_eq!(chunked.remaining_mut(), 0); let before = chunked.remaining(); // SAFETY: Advancing by 0 does not expose uninitialized bytes. unsafe { chunked.advance_mut(0) }; assert_eq!(chunked.remaining(), before); } #[test] #[should_panic(expected = "cannot advance past end of buffer")] fn test_iobufsmut_advance_mut_past_end_pair() { let mut pair = IoBufsMut { inner: IoBufsMutInner::Pair([IoBufMut::with_capacity(1), IoBufMut::with_capacity(1)]), }; // SAFETY: Intentional panic path coverage. unsafe { pair.advance_mut(3) }; } #[test] #[should_panic(expected = "cannot advance past end of buffer")] fn test_iobufsmut_advance_mut_past_end_triple() { let mut triple = IoBufsMut { inner: IoBufsMutInner::Triple([ IoBufMut::with_capacity(1), IoBufMut::with_capacity(1), IoBufMut::with_capacity(1), ]), }; // SAFETY: Intentional panic path coverage. unsafe { triple.advance_mut(4) }; } #[test] #[should_panic(expected = "cannot advance past end of buffer")] fn test_iobufsmut_advance_mut_past_end_chunked() { let mut chunked = IoBufsMut { inner: IoBufsMutInner::Chunked(VecDeque::from([ IoBufMut::with_capacity(1), IoBufMut::with_capacity(1), IoBufMut::with_capacity(1), IoBufMut::with_capacity(1), ])), }; // SAFETY: Intentional panic path coverage. unsafe { chunked.advance_mut(5) }; } #[test] fn test_iobufsmut_set_len() { // SAFETY: we don't read the uninitialized bytes. unsafe { // Single buffer let mut bufs = IoBufsMut::from(IoBufMut::with_capacity(16)); bufs.set_len(10); assert_eq!(bufs.len(), 10); // Chunked: distributes across chunks [cap 5, cap 10], set 12 -> [5, 7] let mut bufs = IoBufsMut::from(vec![ IoBufMut::with_capacity(5), IoBufMut::with_capacity(10), ]); bufs.set_len(12); assert_eq!(bufs.len(), 12); assert_eq!(bufs.chunk().len(), 5); bufs.advance(5); assert_eq!(bufs.chunk().len(), 7); bufs.advance(7); assert_eq!(bufs.remaining(), 0); // Uneven capacities [3, 20, 2], set 18 -> [3, 15, 0]. let mut bufs = IoBufsMut::from(vec![ IoBufMut::with_capacity(3), IoBufMut::with_capacity(20), IoBufMut::with_capacity(2), ]); bufs.set_len(18); assert_eq!(bufs.chunk().len(), 3); bufs.advance(3); assert_eq!(bufs.chunk().len(), 15); bufs.advance(15); assert_eq!(bufs.remaining(), 0); // Exact total capacity [4, 4], set 8 -> [4, 4] let mut bufs = IoBufsMut::from(vec![IoBufMut::with_capacity(4), IoBufMut::with_capacity(4)]); bufs.set_len(8); assert_eq!(bufs.chunk().len(), 4); bufs.advance(4); assert_eq!(bufs.chunk().len(), 4); bufs.advance(4); assert_eq!(bufs.remaining(), 0); // Zero length preserves caller-provided layout. let mut bufs = IoBufsMut::from(vec![IoBufMut::with_capacity(4), IoBufMut::with_capacity(4)]); bufs.set_len(0); assert_eq!(bufs.len(), 0); assert_eq!(bufs.chunk(), b""); } } #[test] #[should_panic(expected = "set_len(9) exceeds capacity(8)")] fn test_iobufsmut_set_len_overflow() { let mut bufs = IoBufsMut::from(vec![IoBufMut::with_capacity(4), IoBufMut::with_capacity(4)]); // SAFETY: this will panic before any read. unsafe { bufs.set_len(9) }; } #[test] fn test_encode_iobufs_matches_encode() { let value = vec![1u8, 2, 3, 4, 5, 6]; assert_encode_iobufs_matches_encode(&value); } /// Reports a larger encoding size than `write` produces, so /// [`EncodeExt::encode_iobufs`] fails its size assert. struct UnderWriter; impl Write for UnderWriter { fn write(&self, buf: &mut impl BufMut) { buf.put_slice(b"ab"); } } impl EncodeSize for UnderWriter { fn encode_size(&self) -> usize { 4 } } #[test] #[should_panic(expected = "write_bufs() must write exactly encode_size() bytes")] fn test_encode_iobufs_rejects_short_write() { let _ = UnderWriter.encode_iobufs(); } #[test] fn test_iobuf_encode_iobufs_matches_encode() { let value = IoBuf::from(vec![0xAB; 512]); assert_encode_iobufs_matches_encode(&value); } /// A pushed native [`IoBuf`] becomes its own chunk as is: once the caller's /// handle is gone, the chunk recovers the original allocation as an [`IoBufMut`]. #[test] fn test_encode_iobufs_keeps_pushed_iobuf() { let value = IoBuf::from(vec![0xAB; 512]); let ptr = value.as_ptr(); let mut encoded = value.encode_iobufs(); drop(value); let prefix = 512usize.encode_size(); drop(encoded.split_to(prefix)); let chunk = encoded .try_into_single() .expect("the payload must be one chunk"); let recovered = chunk .try_into_mut() .expect("the chunk must be the pushed native buffer"); assert_eq!(recovered.as_ref().as_ptr(), ptr); } #[test] fn test_nested_container_encode_iobufs_matches_encode() { let value = ( Some(Bytes::from(vec![0xAA; 256])), vec![Bytes::from(vec![0xBB; 128]), Bytes::from(vec![0xCC; 64])], ); assert_encode_iobufs_matches_encode(&value); } #[test] fn test_map_encode_iobufs_matches_encode() { let mut btree = BTreeMap::new(); btree.insert(2u8, Bytes::from(vec![0xDD; 96])); btree.insert(1u8, Bytes::from(vec![0xEE; 48])); assert_encode_iobufs_matches_encode(&btree); let mut hash = HashMap::new(); hash.insert(2u8, Bytes::from(vec![0x11; 96])); hash.insert(1u8, Bytes::from(vec![0x22; 48])); assert_encode_iobufs_matches_encode(&hash); } #[test] fn test_lazy_encode_iobufs_matches_encode() { let value = Lazy::new(Bytes::from(vec![0x44; 200])); assert_encode_iobufs_matches_encode(&value); // A deferred value pushes its pending encoding as the only chunk and // writes no inline bytes. let mut pending = 42u64.encode(); let deferred = Lazy::::deferred(&mut pending, ()); assert_eq!(deferred.encode_inline_size(), 0); assert_encode_iobufs_matches_encode(&deferred); assert_eq!(deferred.encode_iobufs().chunk_count(), 1); } /// A value with neither inline bytes nor pushed chunks encodes without /// allocating. #[test] fn test_encode_iobufs_without_inline_bytes_does_not_allocate() { let (bufs, calls) = allocation_calls_during(|| [0u8; 0].encode_iobufs()); assert_eq!(calls, 0); assert_eq!(bufs.remaining(), 0); assert_eq!(bufs.chunk_count(), 0); } #[test] fn test_non_empty_range_encode_iobufs_matches_encode() { let range = NonEmptyRange::new(Bytes::from(vec![0x10; 32])..Bytes::from(vec![0x20; 48])).unwrap(); assert_encode_iobufs_matches_encode(&range); } mod builder_tests { use super::*; use commonware_codec::{BufsMut, Encode, Write}; // Only inline writes, no pushes. #[test] fn test_inline_only() { let mut b = Builder::new(64); b.put_u32(42); b.put_u8(7); let mut r = b.finish(); assert_eq!(r.remaining(), 5); assert_eq!(r.get_u32(), 42); assert_eq!(r.get_u8(), 7); } // Only zero-copy pushes, no inline writes. #[test] fn test_push_only() { let mut b = Builder::new(64); let data = Bytes::from(vec![0xAA; 1024]); b.push(data.clone()); let mut r = b.finish(); assert_eq!(r.remaining(), 1024); assert_eq!(r.copy_to_bytes(1024), data); } // Pushed Bytes appear in the output without a payload copy. #[test] fn test_push_is_zero_copy() { let mut b = Builder::new(64); b.put_u16(99); let payload = Bytes::from(vec![0xDD; 1024]); b.push(payload.clone()); let mut r = b.finish(); r.advance(2); assert_eq!(r.chunk().as_ptr(), payload.as_ptr()); } // Interleaved: inline header, zero-copy push, inline trailer. #[test] fn test_inline_push_inline() { let mut b = Builder::new(64); b.put_u16(99); let payload = Bytes::from(vec![0xBB; 512]); b.push(payload.clone()); b.put_u8(1); let mut r = b.finish(); assert_eq!(r.remaining(), 2 + 512 + 1); assert_eq!(r.get_u16(), 99); assert_eq!(r.copy_to_bytes(512), payload); assert_eq!(r.get_u8(), 1); } // Bytes::write_bufs produces identical wire format to Bytes::write. #[test] fn test_write_bufs_matches_write() { let data = Bytes::from(vec![0xCC; 256]); let mut b = Builder::new(64); data.write_bufs(&mut b); let mut bufs = b.finish(); let mut out = vec![0u8; bufs.remaining()]; bufs.copy_to_slice(&mut out); assert_eq!(out, data.encode().as_ref()); } // Finishing an unused builder produces empty IoBufs. #[test] fn test_empty() { let bufs = Builder::new(64).finish(); assert_eq!(bufs.remaining(), 0); } // Inline writes exceeding capacity panic. `Builder` does not override // the `BufMut` write methods, so the panic (and message) comes from // bytes' trait defaults checking `remaining_mut`. #[test] #[should_panic(expected = "advance out of bounds")] fn test_inline_overflow_panics() { let mut b = Builder::new(1); let cap = b.remaining_mut(); b.put_slice(&vec![0xFF; cap]); b.put_u8(1); // exceeds capacity } // Pushing empty Bytes is a no-op. #[test] fn test_empty_push_ignored() { let mut b = Builder::new(64); b.push(Bytes::new()); b.put_u8(1); let bufs = b.finish(); assert_eq!(bufs.remaining(), 1); } // Consecutive pushes without inline writes between them. #[test] fn test_multiple_pushes() { let mut b = Builder::new(64); let a = Bytes::from(vec![0xAA; 100]); let c = Bytes::from(vec![0xCC; 200]); b.push(a.clone()); b.push(c.clone()); let mut r = b.finish(); assert_eq!(r.remaining(), 300); assert_eq!(r.copy_to_bytes(100), a); assert_eq!(r.copy_to_bytes(200), c); } // A zero-capacity builder allocates no inline buffer and emits the // pushed chunks alone, with no empty inline slices between them. #[test] fn test_zero_capacity_push_only() { let (bufs, calls) = allocation_calls_during(|| Builder::new(0).finish()); assert_eq!(calls, 0); assert!(bufs.is_empty()); let mut b = Builder::new(0); assert_eq!(b.remaining_mut(), 0); let a = Bytes::from(vec![0xAA; 100]); let c = Bytes::from(vec![0xCC; 200]); b.push(a.clone()); b.push(Bytes::new()); b.push(c.clone()); let mut r = b.finish(); assert_eq!(r.chunk_count(), 2); assert_eq!(r.chunk().as_ptr(), a.as_ptr()); assert_eq!(r.copy_to_bytes(100), a); assert_eq!(r.chunk().as_ptr(), c.as_ptr()); assert_eq!(r.copy_to_bytes(200), c); } // The inline segments around a push share the inline buffer, so // finishing allocates only that buffer's owner record. #[test] fn test_inline_segments_share_one_record() { let mut b = Builder::new(64); b.put_slice(b"head"); let pushed = IoBuf::from(b"pushed"); b.push(pushed.clone()); b.put_slice(b"tail"); let (mut r, calls) = allocation_calls_during(|| b.finish()); assert_eq!(calls, 1); assert_eq!(r.chunk_count(), 3); let head = r.chunk().as_ptr(); assert_eq!(r.copy_to_bytes(4), &b"head"[..]); assert_eq!(r.chunk().as_ptr(), pushed.as_ptr()); assert_eq!(r.copy_to_bytes(6), &b"pushed"[..]); // The tail segment starts right after the head in the same buffer. assert_eq!(r.chunk().as_ptr(), head.wrapping_add(4)); assert_eq!(r.copy_to_bytes(4), &b"tail"[..]); } // put() exceeding capacity panics (in bytes' trait default). #[test] #[should_panic(expected = "advance out of bounds")] fn test_put_exceeding_capacity_panics() { let mut b = Builder::new(1); let cap = b.remaining_mut(); let src = Bytes::from(vec![0xAB; cap + 1]); b.put(src); } // put_slice() exceeding capacity panics (in bytes' trait default). #[test] #[should_panic(expected = "advance out of bounds")] fn test_put_slice_exceeding_capacity_panics() { let mut b = Builder::new(1); let cap = b.remaining_mut(); b.put_slice(&vec![0xFE; cap + 1]); } // Simulates a multi-field struct: [u16 | Bytes (via push) | u32]. // Verifies write_bufs produces identical wire format to write. #[test] fn test_multi_field_struct_equivalence() { let header: u16 = 0xCAFE; let payload = Bytes::from(vec![0xDD; 1024]); let trailer: u32 = 0xDEADBEEF; // Flat encoding via write. let size = header.encode_size() + payload.encode_size() + trailer.encode_size(); let mut flat = BytesMut::with_capacity(size); header.write(&mut flat); payload.write(&mut flat); trailer.write(&mut flat); // Multi-buffer encoding via write_bufs. let mut b = Builder::new(64); header.write(&mut b); payload.write_bufs(&mut b); trailer.write(&mut b); let mut bufs = b.finish(); let mut out = vec![0u8; bufs.remaining()]; bufs.copy_to_slice(&mut out); assert_eq!(out, flat.as_ref()); } // encode_iobufs (Builder path) matches encode (flat BytesMut path). #[test] fn test_encode_iobufs_matches_encode() { let data = Bytes::from(vec![0xEE; 500]); let mut encoded = data.encode_iobufs(); let baseline = data.encode(); let mut out = vec![0u8; encoded.remaining()]; encoded.copy_to_slice(&mut out); assert_eq!(out, baseline.as_ref()); } // Exercise remaining_mut, chunk_mut, and advance_mut directly. #[test] fn test_chunk_mut_and_advance_mut() { let mut b = Builder::new(64); let initial = b.remaining_mut(); assert!(initial >= 64); let chunk = b.chunk_mut(); chunk[0..1].copy_from_slice(&[0xAB]); // SAFETY: We just wrote 1 byte into chunk_mut above. unsafe { b.advance_mut(1) }; assert_eq!(b.remaining_mut(), initial - 1); let mut r = b.finish(); assert_eq!(r.remaining(), 1); assert_eq!(r.get_u8(), 0xAB); } // Writing past a full buffer panics (fixed capacity, with the panic coming // from bytes' trait default). #[test] #[should_panic(expected = "advance out of bounds")] fn test_write_past_full_panics() { let mut b = Builder::new(1); let cap = b.remaining_mut(); b.put_slice(&vec![0xFF; cap]); // fill the buffer completely assert_eq!(b.remaining_mut(), 0); b.put_u8(0x42); // panics } // Push at offset 0 with inline trailer exercises finish branch // where offset == pos (no inline prefix before push). #[test] fn test_push_at_start_with_trailer() { let mut b = Builder::new(64); let payload = Bytes::from(vec![0xCC; 32]); b.push(payload.clone()); b.put_u8(0x01); let mut r = b.finish(); assert_eq!(r.remaining(), 33); assert_eq!(r.copy_to_bytes(32), payload); assert_eq!(r.get_u8(), 0x01); } } }