Skip to main content

commonware_runtime/iobuf/
bufs.rs

1//! Immutable and mutable collections of I/O buffers.
2//!
3//! [`IoBufs`] and [`IoBufsMut`] present multiple buffers as one logical byte
4//! stream. [`Builder`] assembles an [`IoBufs`] from encoded bytes and existing
5//! byte segments, and [`EncodeExt`] encodes values into [`IoBufs`] through a
6//! [`Builder`].
7
8use super::{
9    buf::{IoBuf, IoBufMut},
10    panic_advance,
11};
12use bytes::{Buf as _, BufMut, Bytes, BytesMut};
13use commonware_codec::{Buf, BufsMut, EncodeSize, Input, OwnedBytes, Write};
14use std::{collections::VecDeque, io::IoSlice};
15
16/// Container for one or more immutable buffers.
17#[derive(Clone, Debug)]
18pub struct IoBufs {
19    inner: IoBufsInner,
20}
21
22/// Internal immutable representation.
23///
24/// - Representation is canonical and minimal for readable data:
25///   - `Single` is the only representation for empty data and one-chunk data.
26///   - `Chunked` is used only when four or more readable chunks remain.
27/// - `Pair`, `Triple`, and `Chunked` never store empty chunks.
28#[derive(Debug)]
29enum IoBufsInner {
30    /// Single buffer (fast path).
31    Single(IoBuf),
32    /// Two buffers (fast path).
33    Pair([IoBuf; 2]),
34    /// Three buffers (fast path).
35    Triple([IoBuf; 3]),
36    /// Four or more buffers.
37    Chunked(VecDeque<IoBuf>),
38}
39
40// Clones `Pair` and `Triple` chunks one at a time so each stays in registers
41// until it is stored into the result. The derived clone of a chunk array builds
42// a temporary array, to drop the chunks already cloned if a later clone
43// unwinds, and then copies it into place. That copy reloads the array with
44// loads wider than the stores that filled it and stalls store forwarding.
45impl Clone for IoBufsInner {
46    #[inline]
47    fn clone(&self) -> Self {
48        match self {
49            Self::Single(buf) => Self::Single(buf.clone()),
50            Self::Pair([a, b]) => Self::Pair([a.clone(), b.clone()]),
51            Self::Triple([a, b, c]) => Self::Triple([a.clone(), b.clone(), c.clone()]),
52            Self::Chunked(bufs) => Self::Chunked(bufs.clone()),
53        }
54    }
55}
56
57impl Default for IoBufs {
58    fn default() -> Self {
59        Self {
60            inner: IoBufsInner::Single(IoBuf::default()),
61        }
62    }
63}
64
65impl IoBufs {
66    /// Build canonical immutable chunk storage from readable chunks.
67    ///
68    /// Empty chunks are removed before representation selection.
69    fn from_chunks_iter(chunks: impl IntoIterator<Item = IoBuf>) -> Self {
70        let mut iter = chunks.into_iter().filter(|buf| !buf.is_empty());
71        let first = match iter.next() {
72            Some(first) => first,
73            None => return Self::default(),
74        };
75        let second = match iter.next() {
76            Some(second) => second,
77            None => {
78                return Self {
79                    inner: IoBufsInner::Single(first),
80                };
81            }
82        };
83        let third = match iter.next() {
84            Some(third) => third,
85            None => {
86                return Self {
87                    inner: IoBufsInner::Pair([first, second]),
88                };
89            }
90        };
91        let fourth = match iter.next() {
92            Some(fourth) => fourth,
93            None => {
94                return Self {
95                    inner: IoBufsInner::Triple([first, second, third]),
96                };
97            }
98        };
99
100        let mut bufs = VecDeque::with_capacity(4);
101        bufs.push_back(first);
102        bufs.push_back(second);
103        bufs.push_back(third);
104        bufs.push_back(fourth);
105        bufs.extend(iter);
106
107        Self {
108            inner: IoBufsInner::Chunked(bufs),
109        }
110    }
111
112    /// Re-establish canonical immutable representation invariants.
113    fn canonicalize(&mut self) {
114        let inner = std::mem::replace(&mut self.inner, IoBufsInner::Single(IoBuf::default()));
115        self.inner = match inner {
116            IoBufsInner::Single(buf) => {
117                if buf.is_empty() {
118                    IoBufsInner::Single(IoBuf::default())
119                } else {
120                    IoBufsInner::Single(buf)
121                }
122            }
123            IoBufsInner::Pair([a, b]) => Self::from_chunks_iter([a, b]).inner,
124            IoBufsInner::Triple([a, b, c]) => Self::from_chunks_iter([a, b, c]).inner,
125            IoBufsInner::Chunked(bufs) => Self::from_chunks_iter(bufs).inner,
126        };
127    }
128
129    /// Returns a reference to the single contiguous buffer, if present.
130    ///
131    /// Returns `Some` only when all remaining data is in one contiguous buffer.
132    pub const fn as_single(&self) -> Option<&IoBuf> {
133        match &self.inner {
134            IoBufsInner::Single(buf) => Some(buf),
135            _ => None,
136        }
137    }
138
139    /// Consume this container and return the single buffer if present.
140    ///
141    /// Returns `Ok(IoBuf)` only when all remaining data is already contained in
142    /// a single chunk. Returns `Err(Self)` with the original container
143    /// otherwise.
144    pub fn try_into_single(self) -> Result<IoBuf, Self> {
145        match self.inner {
146            IoBufsInner::Single(buf) => Ok(buf),
147            inner => Err(Self { inner }),
148        }
149    }
150
151    /// Number of bytes remaining across all buffers.
152    #[inline]
153    pub fn len(&self) -> usize {
154        self.remaining()
155    }
156
157    /// Number of non-empty readable chunks.
158    #[inline]
159    pub fn chunk_count(&self) -> usize {
160        // This assumes canonical form.
161        match &self.inner {
162            IoBufsInner::Single(buf) => {
163                if buf.is_empty() {
164                    0
165                } else {
166                    1
167                }
168            }
169            IoBufsInner::Pair(_) => 2,
170            IoBufsInner::Triple(_) => 3,
171            IoBufsInner::Chunked(bufs) => bufs.len(),
172        }
173    }
174
175    /// Returns the readable chunk at the given index without advancing the buffers.
176    ///
177    /// Returns `None` when `index` is greater than or equal to [`Self::chunk_count`].
178    pub fn chunk_at(&self, index: usize) -> Option<&[u8]> {
179        match &self.inner {
180            IoBufsInner::Single(buf) => (index == 0 && !buf.is_empty()).then(|| buf.as_ref()),
181            IoBufsInner::Pair(bufs) => bufs.get(index).map(AsRef::as_ref),
182            IoBufsInner::Triple(bufs) => bufs.get(index).map(AsRef::as_ref),
183            IoBufsInner::Chunked(bufs) => bufs.get(index).map(AsRef::as_ref),
184        }
185    }
186
187    /// Iterates over the readable buffer owners without coalescing them.
188    pub(crate) fn iter(&self) -> impl ExactSizeIterator<Item = &IoBuf> + Clone {
189        (0..self.chunk_count()).map(|index| match &self.inner {
190            IoBufsInner::Single(buf) => {
191                assert_eq!(index, 0);
192                buf
193            }
194            IoBufsInner::Pair(bufs) => &bufs[index],
195            IoBufsInner::Triple(bufs) => &bufs[index],
196            IoBufsInner::Chunked(bufs) => &bufs[index],
197        })
198    }
199
200    /// Whether all buffers are empty.
201    #[inline]
202    pub const fn is_empty(&self) -> bool {
203        // Only Single can be empty. All other variants contain readable bytes.
204        match &self.inner {
205            IoBufsInner::Single(buf) => buf.is_empty(),
206            IoBufsInner::Pair(_) | IoBufsInner::Triple(_) | IoBufsInner::Chunked(_) => false,
207        }
208    }
209
210    /// Whether this contains a single contiguous buffer.
211    ///
212    /// When true, `chunk()` returns all remaining bytes.
213    #[inline]
214    pub const fn is_single(&self) -> bool {
215        matches!(self.inner, IoBufsInner::Single(_))
216    }
217
218    /// Visit each readable chunk in order without coalescing.
219    #[inline]
220    pub fn for_each_chunk(&self, mut f: impl FnMut(&[u8])) {
221        match &self.inner {
222            IoBufsInner::Single(buf) => {
223                let chunk = buf.as_ref();
224                if !chunk.is_empty() {
225                    f(chunk);
226                }
227            }
228            IoBufsInner::Pair(pair) => {
229                for buf in pair {
230                    let chunk = buf.as_ref();
231                    if !chunk.is_empty() {
232                        f(chunk);
233                    }
234                }
235            }
236            IoBufsInner::Triple(triple) => {
237                for buf in triple {
238                    let chunk = buf.as_ref();
239                    if !chunk.is_empty() {
240                        f(chunk);
241                    }
242                }
243            }
244            IoBufsInner::Chunked(bufs) => {
245                for buf in bufs {
246                    let chunk = buf.as_ref();
247                    if !chunk.is_empty() {
248                        f(chunk);
249                    }
250                }
251            }
252        }
253    }
254
255    /// Prepend a buffer to the front.
256    ///
257    /// Empty input buffers are ignored.
258    pub fn prepend(&mut self, buf: IoBuf) {
259        if buf.is_empty() {
260            return;
261        }
262        let inner = std::mem::replace(&mut self.inner, IoBufsInner::Single(IoBuf::default()));
263        self.inner = match inner {
264            IoBufsInner::Single(existing) if existing.is_empty() => IoBufsInner::Single(buf),
265            IoBufsInner::Single(existing) => IoBufsInner::Pair([buf, existing]),
266            IoBufsInner::Pair([a, b]) => IoBufsInner::Triple([buf, a, b]),
267            IoBufsInner::Triple([a, b, c]) => {
268                let mut bufs = VecDeque::with_capacity(4);
269                bufs.push_back(buf);
270                bufs.push_back(a);
271                bufs.push_back(b);
272                bufs.push_back(c);
273                IoBufsInner::Chunked(bufs)
274            }
275            IoBufsInner::Chunked(mut bufs) => {
276                bufs.push_front(buf);
277                IoBufsInner::Chunked(bufs)
278            }
279        };
280    }
281
282    /// Append a buffer to the back.
283    ///
284    /// Empty input buffers are ignored.
285    pub fn append(&mut self, buf: IoBuf) {
286        if buf.is_empty() {
287            return;
288        }
289        let inner = std::mem::replace(&mut self.inner, IoBufsInner::Single(IoBuf::default()));
290        self.inner = match inner {
291            IoBufsInner::Single(existing) if existing.is_empty() => IoBufsInner::Single(buf),
292            IoBufsInner::Single(existing) => IoBufsInner::Pair([existing, buf]),
293            IoBufsInner::Pair([a, b]) => IoBufsInner::Triple([a, b, buf]),
294            IoBufsInner::Triple([a, b, c]) => {
295                let mut bufs = VecDeque::with_capacity(4);
296                bufs.push_back(a);
297                bufs.push_back(b);
298                bufs.push_back(c);
299                bufs.push_back(buf);
300                IoBufsInner::Chunked(bufs)
301            }
302            IoBufsInner::Chunked(mut bufs) => {
303                bufs.push_back(buf);
304                IoBufsInner::Chunked(bufs)
305            }
306        };
307    }
308
309    /// Splits the buffer(s) into two at the given index.
310    ///
311    /// Afterwards `self` contains bytes `[at, len)`, and the returned
312    /// [`IoBufs`] contains bytes `[0, at)`.
313    ///
314    /// Whole chunks are moved without copying. If the split point lands inside
315    /// a chunk, the chunk is split zero-copy via [`IoBuf::split_to`].
316    ///
317    /// # Panics
318    ///
319    /// Panics if `at > len`.
320    pub fn split_to(&mut self, at: usize) -> Self {
321        if at == 0 {
322            return Self::default();
323        }
324
325        let remaining = self.remaining();
326        assert!(
327            at <= remaining,
328            "split_to out of bounds: {:?} <= {:?}",
329            at,
330            remaining,
331        );
332
333        if at == remaining {
334            return std::mem::take(self);
335        }
336
337        // A prefix within the front chunk of a chunked representation splits
338        // that chunk in place, without moving the deque out and back. The
339        // smaller representations handle that case in their own arms below.
340        if matches!(self.inner, IoBufsInner::Chunked(_))
341            && let Some(prefix) = self.split_front(at)
342        {
343            return Self::from(prefix);
344        }
345
346        let inner = std::mem::replace(&mut self.inner, IoBufsInner::Single(IoBuf::default()));
347        match inner {
348            IoBufsInner::Single(mut buf) => {
349                // Delegate directly and keep remainder as single
350                let prefix = buf.split_to(at);
351                self.inner = IoBufsInner::Single(buf);
352                Self::from(prefix)
353            }
354            IoBufsInner::Pair([mut a, mut b]) => {
355                let a_len = a.remaining();
356                if at < a_len {
357                    // Split stays entirely in chunk `a`.
358                    let prefix = a.split_to(at);
359                    self.inner = IoBufsInner::Pair([a, b]);
360                    return Self::from(prefix);
361                }
362                if at == a_len {
363                    // Exact chunk boundary: move `a` out, keep `b`.
364                    self.inner = IoBufsInner::Single(b);
365                    return Self::from(a);
366                }
367
368                // Split crosses from `a` into `b`.
369                let b_prefix_len = at - a_len;
370                let b_prefix = b.split_to(b_prefix_len);
371                self.inner = IoBufsInner::Single(b);
372                Self {
373                    inner: IoBufsInner::Pair([a, b_prefix]),
374                }
375            }
376            IoBufsInner::Triple([mut a, mut b, mut c]) => {
377                let a_len = a.remaining();
378                if at < a_len {
379                    // Split stays entirely in chunk `a`.
380                    let prefix = a.split_to(at);
381                    self.inner = IoBufsInner::Triple([a, b, c]);
382                    return Self::from(prefix);
383                }
384                if at == a_len {
385                    // Exact boundary after `a`.
386                    self.inner = IoBufsInner::Pair([b, c]);
387                    return Self::from(a);
388                }
389
390                let mut remaining = at - a_len;
391                let b_len = b.remaining();
392                if remaining < b_len {
393                    // Split lands inside `b`.
394                    let b_prefix = b.split_to(remaining);
395                    self.inner = IoBufsInner::Pair([b, c]);
396                    return Self {
397                        inner: IoBufsInner::Pair([a, b_prefix]),
398                    };
399                }
400                if remaining == b_len {
401                    // Exact boundary after `b`.
402                    self.inner = IoBufsInner::Single(c);
403                    return Self {
404                        inner: IoBufsInner::Pair([a, b]),
405                    };
406                }
407
408                // Split reaches into `c`.
409                remaining -= b_len;
410                let c_prefix = c.split_to(remaining);
411                self.inner = IoBufsInner::Single(c);
412                Self {
413                    inner: IoBufsInner::Triple([a, b, c_prefix]),
414                }
415            }
416            IoBufsInner::Chunked(mut bufs) => {
417                // Count the chunks the prefix takes whole. The split lands
418                // inside the next one when bytes remain.
419                let mut remaining = at;
420                let mut whole = 0;
421                for buf in &bufs {
422                    let avail = buf.remaining();
423                    if remaining == 0 || remaining < avail {
424                        break;
425                    }
426                    remaining -= avail;
427                    whole += 1;
428                }
429
430                // Split in place the chunk that the prefix ends inside. The
431                // bounds check above guarantees that chunk exists.
432                let piece = (remaining > 0).then(|| bufs[whole].split_to(remaining));
433
434                // Move out the side with fewer chunks, counting the piece
435                // toward the prefix, and keep the deque for the other, so the
436                // copying and any new deque scale with the shorter side. A
437                // prefix of fewer than four chunks from a large container
438                // needs no new deque.
439                let prefix_chunks = whole + usize::from(piece.is_some());
440                let suffix_chunks = bufs.len() - whole;
441                if prefix_chunks <= suffix_chunks {
442                    let prefix =
443                        Self::from_exact_chunks(bufs.drain(..whole).chain(piece), prefix_chunks);
444                    self.inner = Self::from_chunk_deque(bufs).inner;
445                    prefix
446                } else {
447                    self.inner = Self::from_exact_chunks(bufs.drain(whole..), suffix_chunks).inner;
448                    bufs.extend(piece);
449                    Self::from_chunk_deque(bufs)
450                }
451            }
452        }
453    }
454
455    /// Builds canonical storage from a deque of non-empty chunks and keeps the
456    /// deque when it holds four or more.
457    fn from_chunk_deque(bufs: VecDeque<IoBuf>) -> Self {
458        if bufs.len() >= 4 {
459            Self {
460                inner: IoBufsInner::Chunked(bufs),
461            }
462        } else {
463            Self::from_chunks_iter(bufs)
464        }
465    }
466
467    /// Builds canonical storage from `len` non-empty chunks and allocates a
468    /// deque only when there are four or more.
469    fn from_exact_chunks(chunks: impl Iterator<Item = IoBuf>, len: usize) -> Self {
470        if len >= 4 {
471            Self {
472                inner: IoBufsInner::Chunked(chunks.collect()),
473            }
474        } else {
475            Self::from_chunks_iter(chunks)
476        }
477    }
478
479    /// Splits off the first `len` bytes as a view of the front chunk when that chunk holds them
480    /// all, and returns `None` when they span chunks.
481    ///
482    /// A prefix that leaves part of the front chunk returns right after its split. One that
483    /// consumes the whole front chunk of a multi-chunk representation restores the canonical form
484    /// through the outlined [`Self::remove_front`].
485    #[inline(always)]
486    fn split_front(&mut self, len: usize) -> Option<IoBuf> {
487        let front = match &mut self.inner {
488            IoBufsInner::Single(buf) => return (buf.len() >= len).then(|| buf.split_to(len)),
489            IoBufsInner::Pair([front, ..]) | IoBufsInner::Triple([front, ..]) => front,
490            IoBufsInner::Chunked(bufs) => bufs.front_mut()?,
491        };
492        if front.len() > len {
493            return Some(front.split_to(len));
494        }
495        if front.len() < len {
496            return None;
497        }
498        let prefix = std::mem::take(front);
499        self.remove_front();
500        Some(prefix)
501    }
502
503    /// Removes the empty front chunk of a multi-chunk representation and
504    /// restores the canonical form.
505    fn remove_front(&mut self) {
506        match &mut self.inner {
507            IoBufsInner::Single(_) => {}
508            IoBufsInner::Pair(_) | IoBufsInner::Triple(_) => self.canonicalize(),
509            IoBufsInner::Chunked(bufs) => {
510                bufs.pop_front();
511                if bufs.len() <= 3 {
512                    self.canonicalize();
513                }
514            }
515        }
516    }
517
518    /// Takes a prefix that spans chunks, copying it once.
519    // Kept out of line so the inlined front-chunk fast path stays small.
520    #[inline(never)]
521    fn take_spanning<T: OwnedBytes>(&mut self, len: usize) -> T {
522        // A prefix that spans chunks is copied once, into one `IoBuf` when `T` claims `IoBuf`
523        // views and into `Bytes` otherwise. Claiming an empty view tells which case applies.
524        if T::claim(IoBuf::default()).is_ok() {
525            T::claim(self.split_to(len).coalesce())
526                .unwrap_or_else(|view| T::from_bytes(view.into()))
527        } else {
528            T::from_bytes(self.copy_to_bytes(len))
529        }
530    }
531
532    /// Coalesce all remaining bytes into a single contiguous [`IoBuf`].
533    ///
534    /// Zero-copy if only one buffer. Copies into one native heap allocation
535    /// if multiple buffers, so the result supports zero-copy
536    /// [`IoBuf::try_into_mut`].
537    #[inline]
538    pub fn coalesce(self) -> IoBuf {
539        match self.inner {
540            IoBufsInner::Single(buf) => buf,
541            inner => {
542                let bufs = Self { inner };
543                let mut out = IoBufMut::with_capacity(bufs.remaining());
544                bufs.for_each_chunk(|chunk| out.put_slice(chunk));
545                out.freeze()
546            }
547        }
548    }
549}
550
551impl Buf for IoBufs {
552    /// Returns the prefix as a view of the front chunk when that chunk holds it all, and copies it
553    /// once when it spans chunks. A `T` that claims [`IoBuf`] views receives the [`IoBuf`], and
554    /// any other `T` is built from its bytes.
555    // Inlined with `split_front` so a front-chunk view is built in registers at the call site. A
556    // view returned through a stack slot can be reloaded with loads wider than the stores that
557    // filled it, and that stalls store forwarding.
558    #[inline(always)]
559    fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
560        if let Some(view) = self.split_front(len) {
561            return T::claim(view).unwrap_or_else(|view| T::from_bytes(view.into()));
562        }
563        self.take_spanning(len)
564    }
565}
566
567impl bytes::Buf for IoBufs {
568    #[inline]
569    fn has_remaining(&self) -> bool {
570        !self.is_empty()
571    }
572
573    #[inline]
574    fn remaining(&self) -> usize {
575        match &self.inner {
576            IoBufsInner::Single(buf) => buf.remaining(),
577            IoBufsInner::Pair([a, b]) => a.remaining().saturating_add(b.remaining()),
578            IoBufsInner::Triple([a, b, c]) => a
579                .remaining()
580                .saturating_add(b.remaining())
581                .saturating_add(c.remaining()),
582            IoBufsInner::Chunked(bufs) => bufs
583                .iter()
584                .map(|b| b.remaining())
585                .fold(0, usize::saturating_add),
586        }
587    }
588
589    #[inline]
590    fn chunk(&self) -> &[u8] {
591        match &self.inner {
592            IoBufsInner::Single(buf) => buf.chunk(),
593            IoBufsInner::Pair([a, b]) => {
594                if a.remaining() > 0 {
595                    a.chunk()
596                } else if b.remaining() > 0 {
597                    b.chunk()
598                } else {
599                    &[]
600                }
601            }
602            IoBufsInner::Triple([a, b, c]) => {
603                if a.remaining() > 0 {
604                    a.chunk()
605                } else if b.remaining() > 0 {
606                    b.chunk()
607                } else if c.remaining() > 0 {
608                    c.chunk()
609                } else {
610                    &[]
611                }
612            }
613            IoBufsInner::Chunked(bufs) => {
614                for buf in bufs.iter() {
615                    if buf.remaining() > 0 {
616                        return buf.chunk();
617                    }
618                }
619                &[]
620            }
621        }
622    }
623
624    #[inline]
625    fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
626        if dst.is_empty() {
627            return 0;
628        }
629
630        match &self.inner {
631            IoBufsInner::Single(buf) => {
632                let chunk = buf.chunk();
633                if !chunk.is_empty() {
634                    dst[0] = IoSlice::new(chunk);
635                    return 1;
636                }
637                0
638            }
639            IoBufsInner::Pair([a, b]) => fill_vectored_from_chunks(dst, [a.chunk(), b.chunk()]),
640            IoBufsInner::Triple([a, b, c]) => {
641                fill_vectored_from_chunks(dst, [a.chunk(), b.chunk(), c.chunk()])
642            }
643            IoBufsInner::Chunked(bufs) => {
644                fill_vectored_from_chunks(dst, bufs.iter().map(|buf| buf.chunk()))
645            }
646        }
647    }
648
649    #[inline]
650    fn advance(&mut self, cnt: usize) {
651        let should_canonicalize = match &mut self.inner {
652            IoBufsInner::Single(buf) => {
653                buf.advance(cnt);
654                false
655            }
656            IoBufsInner::Pair(pair) => advance_small_chunks(pair.as_mut_slice(), cnt),
657            IoBufsInner::Triple(triple) => advance_small_chunks(triple.as_mut_slice(), cnt),
658            IoBufsInner::Chunked(bufs) => {
659                advance_chunked_front(bufs, cnt);
660                bufs.len() <= 3
661            }
662        };
663
664        if should_canonicalize {
665            self.canonicalize();
666        }
667    }
668
669    #[inline]
670    fn copy_to_bytes(&mut self, len: usize) -> Bytes {
671        let (result, needs_canonicalize) = match &mut self.inner {
672            IoBufsInner::Single(buf) => return buf.copy_to_bytes(len),
673            IoBufsInner::Pair(pair) => {
674                copy_to_bytes_small_chunks(pair, len, "IoBufs::copy_to_bytes: not enough data")
675            }
676            IoBufsInner::Triple(triple) => {
677                copy_to_bytes_small_chunks(triple, len, "IoBufs::copy_to_bytes: not enough data")
678            }
679            IoBufsInner::Chunked(bufs) => {
680                copy_to_bytes_chunked(bufs, len, "IoBufs::copy_to_bytes: not enough data")
681            }
682        };
683
684        if needs_canonicalize {
685            self.canonicalize();
686        }
687
688        result
689    }
690}
691
692/// Zero-copy: wraps the buffer as the single chunk.
693impl From<IoBuf> for IoBufs {
694    fn from(buf: IoBuf) -> Self {
695        Self {
696            inner: IoBufsInner::Single(buf),
697        }
698    }
699}
700
701/// Zero-copy: freezes the buffer and wraps it as the single chunk.
702impl From<IoBufMut> for IoBufs {
703    fn from(buf: IoBufMut) -> Self {
704        Self {
705            inner: IoBufsInner::Single(buf.freeze()),
706        }
707    }
708}
709
710/// Zero-copy via `From<Bytes> for IoBuf`.
711impl From<Bytes> for IoBufs {
712    fn from(bytes: Bytes) -> Self {
713        Self::from(IoBuf::from(bytes))
714    }
715}
716
717/// Zero-copy via `From<BytesMut> for IoBuf`.
718impl From<BytesMut> for IoBufs {
719    fn from(bytes: BytesMut) -> Self {
720        Self::from(IoBuf::from(bytes))
721    }
722}
723
724/// Zero-copy via `From<Vec<u8>> for IoBuf`.
725impl From<Vec<u8>> for IoBufs {
726    fn from(vec: Vec<u8>) -> Self {
727        Self::from(IoBuf::from(vec))
728    }
729}
730
731/// Zero-copy: collects the chunks, dropping empty ones.
732impl From<Vec<IoBuf>> for IoBufs {
733    fn from(bufs: Vec<IoBuf>) -> Self {
734        Self::from_iter(bufs)
735    }
736}
737
738/// Zero-copy: collects the chunks, dropping empty ones.
739impl<const N: usize> From<[IoBuf; N]> for IoBufs {
740    fn from(bufs: [IoBuf; N]) -> Self {
741        Self::from_iter(bufs)
742    }
743}
744
745/// Zero-copy: collects the chunks, dropping empty ones.
746impl FromIterator<IoBuf> for IoBufs {
747    fn from_iter<T: IntoIterator<Item = IoBuf>>(bufs: T) -> Self {
748        Self::from_chunks_iter(bufs)
749    }
750}
751
752/// Zero-copy: creates a static view with no owner.
753impl<const N: usize> From<&'static [u8; N]> for IoBufs {
754    fn from(array: &'static [u8; N]) -> Self {
755        Self::from(IoBuf::from(array))
756    }
757}
758
759/// Zero-copy: creates a static view with no owner.
760impl From<&'static [u8]> for IoBufs {
761    fn from(slice: &'static [u8]) -> Self {
762        Self::from(IoBuf::from(slice))
763    }
764}
765
766/// Container for one or more mutable buffers.
767///
768/// # Capacity retention
769///
770/// The intended usage is fill-then-read: write into the container (through
771/// [`BufMut`], [`Self::copy_from_slice`], or
772/// [`Blob::read_at_buf`](crate::Blob::read_at_buf)), then consume it through
773/// [`bytes::Buf`]. Caller-reserved write capacity generally survives read
774/// operations, with three exceptions:
775/// - The deque-backed read paths (four or more chunks) skip past a chunk
776///   with no readable bytes by popping it, so a never-filled chunk ordered
777///   before readable data loses its capacity when a read crosses it.
778/// - The same paths pop any chunk whose readable bytes a read fully
779///   consumes, releasing its writable tail. The two and three-chunk shapes
780///   advance in place and retain such tails, so `remaining_mut()` after an
781///   identical read can differ by shape.
782/// - A nonzero `copy_to_bytes` that exactly drains the front chunk's
783///   readable bytes consumes that chunk's whole handle on every shape (see
784///   [`IoBufMut`]'s `copy_to_bytes` doc), so spare capacity behind an
785///   exactly-drained chunk is released rather than retained.
786///
787/// Interleaving writes after reads is outside this contract. A read can
788/// leave a drained chunk's surviving tail ordered before later readable
789/// bytes, and a subsequent [`BufMut`] write fills that tail first.
790#[derive(Debug)]
791pub struct IoBufsMut {
792    inner: IoBufsMutInner,
793}
794
795/// Internal mutable representation.
796///
797/// Construction and canonicalization keep every chunk that still owns
798/// storage (`capacity() > 0`), readable or not, so caller-reserved write
799/// capacity generally survives read operations. The reader-facing rules and
800/// the accepted exceptions are documented on [`IoBufsMut`] under "Capacity
801/// retention". Only fully-drained chunks (capacity consumed by `advance`)
802/// and empty defaults are removed as the shape collapses.
803#[derive(Debug)]
804enum IoBufsMutInner {
805    /// Single buffer (common case, no allocation).
806    Single(IoBufMut),
807    /// Two buffers (fast path, no VecDeque allocation).
808    Pair([IoBufMut; 2]),
809    /// Three buffers (fast path, no VecDeque allocation).
810    Triple([IoBufMut; 3]),
811    /// Four or more buffers.
812    Chunked(VecDeque<IoBufMut>),
813}
814
815impl Default for IoBufsMut {
816    fn default() -> Self {
817        Self {
818            inner: IoBufsMutInner::Single(IoBufMut::default()),
819        }
820    }
821}
822
823impl IoBufsMut {
824    /// Build mutable chunk storage from already-filtered chunks.
825    ///
826    /// This helper intentionally does not filter. Callers route through
827    /// [`Self::from_writable_chunks_iter`] so storage-owning chunks are kept.
828    fn from_chunks_iter(chunks: impl IntoIterator<Item = IoBufMut>) -> Self {
829        let mut iter = chunks.into_iter();
830        let first = match iter.next() {
831            Some(first) => first,
832            None => return Self::default(),
833        };
834        let second = match iter.next() {
835            Some(second) => second,
836            None => {
837                return Self {
838                    inner: IoBufsMutInner::Single(first),
839                };
840            }
841        };
842        let third = match iter.next() {
843            Some(third) => third,
844            None => {
845                return Self {
846                    inner: IoBufsMutInner::Pair([first, second]),
847                };
848            }
849        };
850        let fourth = match iter.next() {
851            Some(fourth) => fourth,
852            None => {
853                return Self {
854                    inner: IoBufsMutInner::Triple([first, second, third]),
855                };
856            }
857        };
858
859        let mut bufs = VecDeque::with_capacity(4);
860        bufs.push_back(first);
861        bufs.push_back(second);
862        bufs.push_back(third);
863        bufs.push_back(fourth);
864        bufs.extend(iter);
865        Self {
866            inner: IoBufsMutInner::Chunked(bufs),
867        }
868    }
869
870    /// Build canonical mutable chunk storage from writable chunks.
871    ///
872    /// Keeps chunks that still own storage, readable or not (`capacity()`
873    /// covers both readable bytes and the writable tail), so a never-filled
874    /// chunk's reserved capacity is not discarded. Fully-drained chunks
875    /// (capacity consumed by `advance`) and empty defaults are removed.
876    fn from_writable_chunks_iter(chunks: impl IntoIterator<Item = IoBufMut>) -> Self {
877        Self::from_chunks_iter(chunks.into_iter().filter(|buf| buf.capacity() > 0))
878    }
879
880    /// Re-establish canonical mutable representation invariants.
881    ///
882    /// Uses the same storage-keeping filter as construction, so chunks that
883    /// were drained of readable bytes but still own writable capacity survive.
884    /// Reads then leave `remaining_mut()` unchanged, apart from the exceptions
885    /// listed under "Capacity retention" on [`IoBufsMut`].
886    fn canonicalize(&mut self) {
887        let inner = std::mem::replace(&mut self.inner, IoBufsMutInner::Single(IoBufMut::default()));
888        self.inner = match inner {
889            IoBufsMutInner::Single(buf) => IoBufsMutInner::Single(buf),
890            IoBufsMutInner::Pair([a, b]) => Self::from_writable_chunks_iter([a, b]).inner,
891            IoBufsMutInner::Triple([a, b, c]) => Self::from_writable_chunks_iter([a, b, c]).inner,
892            IoBufsMutInner::Chunked(bufs) => Self::from_writable_chunks_iter(bufs).inner,
893        };
894    }
895
896    #[inline]
897    fn for_each_chunk_mut(&mut self, mut f: impl FnMut(&mut IoBufMut)) {
898        match &mut self.inner {
899            IoBufsMutInner::Single(buf) => f(buf),
900            IoBufsMutInner::Pair(pair) => {
901                for buf in pair.iter_mut() {
902                    f(buf);
903                }
904            }
905            IoBufsMutInner::Triple(triple) => {
906                for buf in triple.iter_mut() {
907                    f(buf);
908                }
909            }
910            IoBufsMutInner::Chunked(bufs) => {
911                for buf in bufs.iter_mut() {
912                    f(buf);
913                }
914            }
915        }
916    }
917
918    /// Returns a reference to the single contiguous buffer, if present.
919    ///
920    /// Returns `Some` only when this is currently represented as one chunk.
921    pub const fn as_single(&self) -> Option<&IoBufMut> {
922        match &self.inner {
923            IoBufsMutInner::Single(buf) => Some(buf),
924            _ => None,
925        }
926    }
927
928    /// Returns a mutable reference to the single contiguous buffer, if present.
929    ///
930    /// Returns `Some` only when this is currently represented as one chunk.
931    pub const fn as_single_mut(&mut self) -> Option<&mut IoBufMut> {
932        match &mut self.inner {
933            IoBufsMutInner::Single(buf) => Some(buf),
934            _ => None,
935        }
936    }
937
938    /// Consume this container and return the single buffer if present.
939    ///
940    /// Returns `Ok(IoBufMut)` only when readable data is represented as one
941    /// chunk. Returns `Err(Self)` with the original container otherwise.
942    pub fn try_into_single(self) -> Result<IoBufMut, Self> {
943        match self.inner {
944            IoBufsMutInner::Single(buf) => Ok(buf),
945            inner => Err(Self { inner }),
946        }
947    }
948
949    /// Number of bytes remaining across all buffers.
950    #[inline]
951    pub fn len(&self) -> usize {
952        self.remaining()
953    }
954
955    /// Whether all buffers are empty.
956    #[inline]
957    pub fn is_empty(&self) -> bool {
958        match &self.inner {
959            IoBufsMutInner::Single(buf) => buf.is_empty(),
960            IoBufsMutInner::Pair([a, b]) => a.is_empty() && b.is_empty(),
961            IoBufsMutInner::Triple([a, b, c]) => a.is_empty() && b.is_empty() && c.is_empty(),
962            IoBufsMutInner::Chunked(bufs) => bufs.iter().all(IoBufMut::is_empty),
963        }
964    }
965
966    /// Whether this contains a single contiguous buffer.
967    ///
968    /// When true, `chunk()` returns all remaining bytes.
969    #[inline]
970    pub const fn is_single(&self) -> bool {
971        matches!(self.inner, IoBufsMutInner::Single(_))
972    }
973
974    /// Freeze into immutable [`IoBufs`].
975    ///
976    /// Each chunk freezes through [`IoBufMut::freeze`], so a chunk whose spare
977    /// capacity is too large to record in its handle allocates its owner record.
978    pub fn freeze(self) -> IoBufs {
979        match self.inner {
980            IoBufsMutInner::Single(buf) => IoBufs::from(buf.freeze()),
981            IoBufsMutInner::Pair([a, b]) => IoBufs::from_chunks_iter([a.freeze(), b.freeze()]),
982            IoBufsMutInner::Triple([a, b, c]) => {
983                IoBufs::from_chunks_iter([a.freeze(), b.freeze(), c.freeze()])
984            }
985            IoBufsMutInner::Chunked(bufs) => {
986                IoBufs::from_chunks_iter(bufs.into_iter().map(IoBufMut::freeze))
987            }
988        }
989    }
990
991    /// Coalesce all buffers into a single contiguous [`IoBufMut`].
992    ///
993    /// Zero-copy if only one buffer. Otherwise copies into one allocation of
994    /// exactly the readable length.
995    pub fn coalesce(self) -> IoBufMut {
996        match self.inner {
997            IoBufsMutInner::Single(buf) => buf,
998            inner => {
999                let mut bufs = Self { inner };
1000                let mut out = IoBufMut::with_capacity(bufs.len());
1001                bufs.for_each_chunk_mut(|buf| out.put_slice(buf.as_ref()));
1002                out
1003            }
1004        }
1005    }
1006
1007    /// Returns the total capacity across all buffers.
1008    pub fn capacity(&self) -> usize {
1009        match &self.inner {
1010            IoBufsMutInner::Single(buf) => buf.capacity(),
1011            IoBufsMutInner::Pair([a, b]) => a.capacity().saturating_add(b.capacity()),
1012            IoBufsMutInner::Triple([a, b, c]) => a
1013                .capacity()
1014                .saturating_add(b.capacity())
1015                .saturating_add(c.capacity()),
1016            IoBufsMutInner::Chunked(bufs) => bufs
1017                .iter()
1018                .map(|b| b.capacity())
1019                .fold(0, usize::saturating_add),
1020        }
1021    }
1022
1023    /// Sets the length of the buffer(s) to `len`, distributing across chunks
1024    /// while preserving the current chunk layout.
1025    ///
1026    /// This is useful for APIs that must fill caller-provided buffer structure
1027    /// in place (for example [`Blob::read_at_buf`](crate::Blob::read_at_buf)).
1028    ///
1029    /// # Safety
1030    ///
1031    /// Caller must initialize all `len` bytes before the buffer is read.
1032    ///
1033    /// # Panics
1034    ///
1035    /// Panics if `len` exceeds total capacity.
1036    pub(crate) unsafe fn set_len(&mut self, len: usize) {
1037        let capacity = self.capacity();
1038        assert!(
1039            len <= capacity,
1040            "set_len({len}) exceeds capacity({capacity})"
1041        );
1042        let mut remaining = len;
1043        self.for_each_chunk_mut(|buf| {
1044            let cap = buf.capacity();
1045            let to_set = remaining.min(cap);
1046            // SAFETY: forwarded from this method's contract. The caller
1047            // initializes all `len` bytes before any read.
1048            unsafe { buf.set_len(to_set) };
1049            remaining -= to_set;
1050        });
1051    }
1052
1053    /// Copy data from a slice into the buffers.
1054    ///
1055    /// # Panics
1056    ///
1057    /// Panics if the slice length doesn't match the total buffer length.
1058    pub fn copy_from_slice(&mut self, src: &[u8]) {
1059        assert_eq!(
1060            src.len(),
1061            self.len(),
1062            "source slice length must match buffer length"
1063        );
1064        let mut offset = 0;
1065        self.for_each_chunk_mut(|buf| {
1066            let len = buf.len();
1067            buf.as_mut().copy_from_slice(&src[offset..offset + len]);
1068            offset += len;
1069        });
1070    }
1071}
1072
1073impl Input for IoBufsMut {
1074    type Buf = IoBufs;
1075}
1076
1077/// Freezes every chunk (see [`IoBufsMut::freeze`]).
1078impl From<IoBufsMut> for IoBufs {
1079    fn from(bufs: IoBufsMut) -> Self {
1080        bufs.freeze()
1081    }
1082}
1083
1084impl bytes::Buf for IoBufsMut {
1085    #[inline]
1086    fn has_remaining(&self) -> bool {
1087        !self.is_empty()
1088    }
1089
1090    #[inline]
1091    fn remaining(&self) -> usize {
1092        match &self.inner {
1093            IoBufsMutInner::Single(buf) => buf.remaining(),
1094            IoBufsMutInner::Pair([a, b]) => a.remaining().saturating_add(b.remaining()),
1095            IoBufsMutInner::Triple([a, b, c]) => a
1096                .remaining()
1097                .saturating_add(b.remaining())
1098                .saturating_add(c.remaining()),
1099            IoBufsMutInner::Chunked(bufs) => bufs
1100                .iter()
1101                .map(|b| b.remaining())
1102                .fold(0, usize::saturating_add),
1103        }
1104    }
1105
1106    #[inline]
1107    fn chunk(&self) -> &[u8] {
1108        match &self.inner {
1109            IoBufsMutInner::Single(buf) => buf.chunk(),
1110            IoBufsMutInner::Pair([a, b]) => {
1111                if a.remaining() > 0 {
1112                    a.chunk()
1113                } else if b.remaining() > 0 {
1114                    b.chunk()
1115                } else {
1116                    &[]
1117                }
1118            }
1119            IoBufsMutInner::Triple([a, b, c]) => {
1120                if a.remaining() > 0 {
1121                    a.chunk()
1122                } else if b.remaining() > 0 {
1123                    b.chunk()
1124                } else if c.remaining() > 0 {
1125                    c.chunk()
1126                } else {
1127                    &[]
1128                }
1129            }
1130            IoBufsMutInner::Chunked(bufs) => {
1131                for buf in bufs.iter() {
1132                    if buf.remaining() > 0 {
1133                        return buf.chunk();
1134                    }
1135                }
1136                &[]
1137            }
1138        }
1139    }
1140
1141    #[inline]
1142    fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
1143        if dst.is_empty() {
1144            return 0;
1145        }
1146
1147        match &self.inner {
1148            IoBufsMutInner::Single(buf) => {
1149                let chunk = buf.chunk();
1150                if !chunk.is_empty() {
1151                    dst[0] = IoSlice::new(chunk);
1152                    return 1;
1153                }
1154                0
1155            }
1156            IoBufsMutInner::Pair([a, b]) => fill_vectored_from_chunks(dst, [a.chunk(), b.chunk()]),
1157            IoBufsMutInner::Triple([a, b, c]) => {
1158                fill_vectored_from_chunks(dst, [a.chunk(), b.chunk(), c.chunk()])
1159            }
1160            IoBufsMutInner::Chunked(bufs) => {
1161                fill_vectored_from_chunks(dst, bufs.iter().map(|buf| buf.chunk()))
1162            }
1163        }
1164    }
1165
1166    #[inline]
1167    fn advance(&mut self, cnt: usize) {
1168        let should_canonicalize = match &mut self.inner {
1169            IoBufsMutInner::Single(buf) => {
1170                buf.advance(cnt);
1171                false
1172            }
1173            IoBufsMutInner::Pair(pair) => advance_small_chunks(pair.as_mut_slice(), cnt),
1174            IoBufsMutInner::Triple(triple) => advance_small_chunks(triple.as_mut_slice(), cnt),
1175            IoBufsMutInner::Chunked(bufs) => {
1176                advance_chunked_front(bufs, cnt);
1177                bufs.len() <= 3
1178            }
1179        };
1180
1181        if should_canonicalize {
1182            self.canonicalize();
1183        }
1184    }
1185
1186    #[inline]
1187    fn copy_to_bytes(&mut self, len: usize) -> Bytes {
1188        // Zero-length drains must not disturb chunk state. The deque-backed
1189        // path pops chunks with no readable bytes, and popping a never-filled
1190        // chunk would discard its reserved capacity.
1191        if len == 0 {
1192            return Bytes::new();
1193        }
1194        let (result, needs_canonicalize) = match &mut self.inner {
1195            IoBufsMutInner::Single(buf) => return buf.copy_to_bytes(len),
1196            IoBufsMutInner::Pair(pair) => {
1197                copy_to_bytes_small_chunks(pair, len, "IoBufsMut::copy_to_bytes: not enough data")
1198            }
1199            IoBufsMutInner::Triple(triple) => {
1200                copy_to_bytes_small_chunks(triple, len, "IoBufsMut::copy_to_bytes: not enough data")
1201            }
1202            IoBufsMutInner::Chunked(bufs) => {
1203                copy_to_bytes_chunked(bufs, len, "IoBufsMut::copy_to_bytes: not enough data")
1204            }
1205        };
1206
1207        if needs_canonicalize {
1208            self.canonicalize();
1209        }
1210
1211        result
1212    }
1213}
1214
1215// SAFETY: Delegates to IoBufMut which implements BufMut safely.
1216unsafe impl BufMut for IoBufsMut {
1217    #[inline]
1218    fn remaining_mut(&self) -> usize {
1219        match &self.inner {
1220            IoBufsMutInner::Single(buf) => buf.remaining_mut(),
1221            IoBufsMutInner::Pair([a, b]) => a.remaining_mut().saturating_add(b.remaining_mut()),
1222            IoBufsMutInner::Triple([a, b, c]) => a
1223                .remaining_mut()
1224                .saturating_add(b.remaining_mut())
1225                .saturating_add(c.remaining_mut()),
1226            IoBufsMutInner::Chunked(bufs) => bufs
1227                .iter()
1228                .map(|b| b.remaining_mut())
1229                .fold(0, usize::saturating_add),
1230        }
1231    }
1232
1233    #[inline]
1234    unsafe fn advance_mut(&mut self, cnt: usize) {
1235        // SAFETY: The caller guarantees `cnt <= self.remaining_mut()`, so each
1236        // delegated advance remains within the writable chunks.
1237        unsafe {
1238            // On failure, every writable byte was consumed before the chunks ran
1239            // out, so the advanced amount (`cnt - remaining`) is exactly what was
1240            // available.
1241            let mut remaining = cnt;
1242            let advanced = match &mut self.inner {
1243                IoBufsMutInner::Single(buf) => {
1244                    buf.advance_mut(cnt);
1245                    return;
1246                }
1247                IoBufsMutInner::Pair(pair) => advance_mut_in_chunks(pair, &mut remaining),
1248                IoBufsMutInner::Triple(triple) => advance_mut_in_chunks(triple, &mut remaining),
1249                IoBufsMutInner::Chunked(bufs) => {
1250                    let (first, second) = bufs.as_mut_slices();
1251                    advance_mut_in_chunks(first, &mut remaining)
1252                        || advance_mut_in_chunks(second, &mut remaining)
1253                }
1254            };
1255            if !advanced {
1256                panic_advance(cnt, cnt - remaining);
1257            }
1258        }
1259    }
1260
1261    #[inline]
1262    fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
1263        match &mut self.inner {
1264            IoBufsMutInner::Single(buf) => buf.chunk_mut(),
1265            IoBufsMutInner::Pair(pair) => {
1266                if pair[0].remaining_mut() > 0 {
1267                    pair[0].chunk_mut()
1268                } else if pair[1].remaining_mut() > 0 {
1269                    pair[1].chunk_mut()
1270                } else {
1271                    bytes::buf::UninitSlice::new(&mut [])
1272                }
1273            }
1274            IoBufsMutInner::Triple(triple) => {
1275                if triple[0].remaining_mut() > 0 {
1276                    triple[0].chunk_mut()
1277                } else if triple[1].remaining_mut() > 0 {
1278                    triple[1].chunk_mut()
1279                } else if triple[2].remaining_mut() > 0 {
1280                    triple[2].chunk_mut()
1281                } else {
1282                    bytes::buf::UninitSlice::new(&mut [])
1283                }
1284            }
1285            IoBufsMutInner::Chunked(bufs) => {
1286                for buf in bufs.iter_mut() {
1287                    if buf.remaining_mut() > 0 {
1288                        return buf.chunk_mut();
1289                    }
1290                }
1291                bytes::buf::UninitSlice::new(&mut [])
1292            }
1293        }
1294    }
1295}
1296
1297/// Zero-copy: wraps the buffer as the single chunk.
1298impl From<IoBufMut> for IoBufsMut {
1299    fn from(buf: IoBufMut) -> Self {
1300        Self {
1301            inner: IoBufsMutInner::Single(buf),
1302        }
1303    }
1304}
1305
1306/// Convert a [`Vec<u8>`] into a single-buffer [`IoBufsMut`].
1307///
1308/// Adopts the vec's allocation and capacity via `From<Vec<u8>> for IoBufMut`.
1309impl From<Vec<u8>> for IoBufsMut {
1310    fn from(vec: Vec<u8>) -> Self {
1311        Self {
1312            inner: IoBufsMutInner::Single(IoBufMut::from(vec)),
1313        }
1314    }
1315}
1316
1317/// Reuses the allocation via `From<BytesMut> for IoBufMut` when no other
1318/// handle shares it and copies otherwise. Either way at least the reserved
1319/// capacity is kept.
1320impl From<BytesMut> for IoBufsMut {
1321    fn from(bytes: BytesMut) -> Self {
1322        Self {
1323            inner: IoBufsMutInner::Single(IoBufMut::from(bytes)),
1324        }
1325    }
1326}
1327
1328/// Zero-copy: collects the chunks, dropping zero-capacity ones.
1329impl From<Vec<IoBufMut>> for IoBufsMut {
1330    fn from(bufs: Vec<IoBufMut>) -> Self {
1331        Self::from_iter(bufs)
1332    }
1333}
1334
1335/// Zero-copy: collects the chunks, dropping zero-capacity ones.
1336impl<const N: usize> From<[IoBufMut; N]> for IoBufsMut {
1337    fn from(bufs: [IoBufMut; N]) -> Self {
1338        Self::from_iter(bufs)
1339    }
1340}
1341
1342/// Zero-copy: collects the chunks, dropping zero-capacity ones.
1343impl FromIterator<IoBufMut> for IoBufsMut {
1344    fn from_iter<T: IntoIterator<Item = IoBufMut>>(bufs: T) -> Self {
1345        Self::from_writable_chunks_iter(bufs)
1346    }
1347}
1348
1349/// Copies via `From<[u8; N]> for IoBufMut`.
1350impl<const N: usize> From<[u8; N]> for IoBufsMut {
1351    fn from(array: [u8; N]) -> Self {
1352        Self {
1353            inner: IoBufsMutInner::Single(IoBufMut::from(array)),
1354        }
1355    }
1356}
1357
1358/// Drain `len` readable bytes from a small fixed chunk array (`Pair`/`Triple`).
1359///
1360/// Returns drained bytes plus whether the caller should canonicalize afterward.
1361#[inline]
1362fn copy_to_bytes_small_chunks<B: bytes::Buf, const N: usize>(
1363    chunks: &mut [B; N],
1364    len: usize,
1365    not_enough_data_msg: &str,
1366) -> (Bytes, bool) {
1367    let total = chunks
1368        .iter()
1369        .map(|buf| buf.remaining())
1370        .fold(0, usize::saturating_add);
1371    assert!(total >= len, "{not_enough_data_msg}");
1372
1373    if chunks[0].remaining() >= len {
1374        let bytes = chunks[0].copy_to_bytes(len);
1375        return (bytes, chunks[0].remaining() == 0);
1376    }
1377
1378    let mut out = BytesMut::with_capacity(len);
1379    let mut remaining = len;
1380    for buf in chunks.iter_mut() {
1381        if remaining == 0 {
1382            break;
1383        }
1384        let to_copy = remaining.min(buf.remaining());
1385        out.extend_from_slice(&buf.chunk()[..to_copy]);
1386        buf.advance(to_copy);
1387        remaining -= to_copy;
1388    }
1389
1390    // Slow path always consumes past chunk 0, so canonicalization is required.
1391    (out.freeze(), true)
1392}
1393
1394/// Drain `len` readable bytes from a deque-backed chunk representation.
1395///
1396/// Returns drained bytes plus whether the caller should canonicalize afterward.
1397#[inline]
1398fn copy_to_bytes_chunked<B: bytes::Buf>(
1399    bufs: &mut VecDeque<B>,
1400    len: usize,
1401    not_enough_data_msg: &str,
1402) -> (Bytes, bool) {
1403    while bufs.front().is_some_and(|buf| buf.remaining() == 0) {
1404        bufs.pop_front();
1405    }
1406
1407    if bufs.front().is_none() {
1408        assert_eq!(len, 0, "{not_enough_data_msg}");
1409        // The deque is empty now, so the container must collapse back to the
1410        // canonical Single representation.
1411        return (Bytes::new(), true);
1412    }
1413
1414    if bufs.front().is_some_and(|front| front.remaining() >= len) {
1415        let front = bufs.front_mut().expect("front checked above");
1416        let bytes = front.copy_to_bytes(len);
1417        if front.remaining() == 0 {
1418            bufs.pop_front();
1419        }
1420        return (bytes, bufs.len() <= 3);
1421    }
1422
1423    let total = bufs
1424        .iter()
1425        .map(|buf| buf.remaining())
1426        .fold(0, usize::saturating_add);
1427    assert!(total >= len, "{not_enough_data_msg}");
1428
1429    let mut out = BytesMut::with_capacity(len);
1430    let mut remaining = len;
1431    while remaining > 0 {
1432        let front = bufs
1433            .front_mut()
1434            .expect("remaining > 0 implies non-empty bufs");
1435        let to_copy = remaining.min(front.remaining());
1436        out.extend_from_slice(&front.chunk()[..to_copy]);
1437        front.advance(to_copy);
1438        if front.remaining() == 0 {
1439            bufs.pop_front();
1440        }
1441        remaining -= to_copy;
1442    }
1443
1444    (out.freeze(), bufs.len() <= 3)
1445}
1446
1447/// Advance across a [`VecDeque`] of chunks by consuming from the front.
1448#[inline]
1449fn advance_chunked_front<B: bytes::Buf>(bufs: &mut VecDeque<B>, mut cnt: usize) {
1450    while cnt > 0 {
1451        let front = bufs.front_mut().expect("cannot advance past end of buffer");
1452        let avail = front.remaining();
1453        if avail == 0 {
1454            bufs.pop_front();
1455            continue;
1456        }
1457        if cnt < avail {
1458            front.advance(cnt);
1459            break;
1460        }
1461        front.advance(avail);
1462        bufs.pop_front();
1463        cnt -= avail;
1464    }
1465}
1466
1467/// Advance across a small fixed set of chunks (`Pair`/`Triple`).
1468///
1469/// Returns `true` when one or more chunks became (or were) empty, so callers
1470/// can canonicalize once after the operation.
1471#[inline]
1472fn advance_small_chunks<B: bytes::Buf>(chunks: &mut [B], mut cnt: usize) -> bool {
1473    let mut idx = 0;
1474    let mut needs_canonicalize = false;
1475
1476    while cnt > 0 {
1477        let chunk = chunks
1478            .get_mut(idx)
1479            .expect("cannot advance past end of buffer");
1480        let avail = chunk.remaining();
1481        if avail == 0 {
1482            idx += 1;
1483            needs_canonicalize = true;
1484            continue;
1485        }
1486        if cnt < avail {
1487            chunk.advance(cnt);
1488            return needs_canonicalize;
1489        }
1490        chunk.advance(avail);
1491        cnt -= avail;
1492        idx += 1;
1493        needs_canonicalize = true;
1494    }
1495
1496    needs_canonicalize
1497}
1498
1499/// Advance writable cursors across `chunks` by up to `*remaining` bytes.
1500///
1501/// Returns `true` when the full request has been satisfied.
1502///
1503/// # Safety
1504///
1505/// Forwards to [`BufMut::advance_mut`], so callers must ensure the advanced
1506/// region has been initialized according to [`BufMut`]'s contract.
1507#[inline]
1508unsafe fn advance_mut_in_chunks<B: BufMut>(chunks: &mut [B], remaining: &mut usize) -> bool {
1509    if *remaining == 0 {
1510        return true;
1511    }
1512
1513    for buf in chunks.iter_mut() {
1514        let avail = buf.chunk_mut().len();
1515        if avail == 0 {
1516            continue;
1517        }
1518        if *remaining <= avail {
1519            // SAFETY: Upheld by this function's safety contract.
1520            unsafe { buf.advance_mut(*remaining) };
1521            *remaining = 0;
1522            return true;
1523        }
1524        // SAFETY: Upheld by this function's safety contract.
1525        unsafe { buf.advance_mut(avail) };
1526        *remaining -= avail;
1527    }
1528    false
1529}
1530
1531/// Fill `dst` with `IoSlice`s built from `chunks`.
1532///
1533/// Empty chunks are skipped. At most `dst.len()` slices are written.
1534/// Returns the number of slices written.
1535#[inline]
1536fn fill_vectored_from_chunks<'a, I>(dst: &mut [IoSlice<'a>], chunks: I) -> usize
1537where
1538    I: IntoIterator<Item = &'a [u8]>,
1539{
1540    let mut written = 0;
1541    for chunk in chunks
1542        .into_iter()
1543        .filter(|chunk| !chunk.is_empty())
1544        .take(dst.len())
1545    {
1546        dst[written] = IoSlice::new(chunk);
1547        written += 1;
1548    }
1549    written
1550}
1551
1552/// Assembles [`IoBufs`] from a mix of inline writes and zero-copy pieces.
1553///
1554/// All inline writes go into a single buffer, and [`BufsMut::push`] records
1555/// the inline offset at which each pushed piece belongs. [`Builder::finish`]
1556/// freezes the buffer once, splits it at those offsets with
1557/// [`IoBuf::split_to`], and interleaves the pushed chunks with the inline
1558/// segments.
1559///
1560/// A pushed [`IoBuf`] becomes a chunk as is. Any other pushed value converts
1561/// into [`Bytes`], and the chunk owns that `Bytes` without copying.
1562///
1563/// The inline buffer has a fixed capacity set at construction and will not
1564/// grow. Callers must ensure the capacity accounts for all inline
1565/// (non-pushed) bytes that will be written. Exceeding it will panic.
1566///
1567/// ```text
1568/// builder.put_u16(99);                   // inline
1569/// builder.push(shard_payload.clone());   // zero-copy: shares the Bytes
1570/// builder.put_u32(checksum);             // inline
1571/// let output = builder.finish();
1572///
1573/// // inline buffer  [ 99 | checksum ]           one allocation, frozen once
1574/// // pushed Bytes   [ ----- 1 MB shard ----- ]  the caller's, never copied
1575/// //
1576/// // output         [ 99 ][ ----- 1 MB shard ----- ][ checksum ]
1577/// //                 inline      pushed Bytes           inline
1578/// ```
1579pub struct Builder {
1580    // Single working buffer for all inline writes.
1581    buf: IoBufMut,
1582    // Each entry is (offset_in_buf, pushed_chunk) recording where a push
1583    // interrupts the inline byte stream.
1584    pushes: Vec<(usize, IoBuf)>,
1585}
1586
1587impl Builder {
1588    /// Creates a builder whose inline buffer holds exactly `capacity` bytes.
1589    ///
1590    /// A zero capacity allocates nothing, so the builder accepts only pushed
1591    /// chunks.
1592    ///
1593    /// # Panics
1594    ///
1595    /// Panics if `capacity` exceeds `isize::MAX`. Writing more than `capacity`
1596    /// inline bytes panics at the write.
1597    pub fn new(capacity: usize) -> Self {
1598        Self {
1599            buf: IoBufMut::with_capacity(capacity),
1600            pushes: Vec::new(),
1601        }
1602    }
1603
1604    /// Freezes the inline buffer and assembles [`IoBufs`] by splitting it at
1605    /// the recorded push boundaries.
1606    pub fn finish(self) -> IoBufs {
1607        if self.pushes.is_empty() {
1608            return IoBufs::from(self.buf.freeze());
1609        }
1610
1611        // Segments split off the owned frozen buffer. A first split of a
1612        // descriptor stores the shared record with a plain write, and the
1613        // remainder moves into the result without taking another reference.
1614        let mut frozen = self.buf.freeze();
1615        let mut result = IoBufs::default();
1616        let mut pos = 0;
1617
1618        for (offset, pushed) in self.pushes {
1619            if offset > pos {
1620                result.append(frozen.split_to(offset - pos));
1621            }
1622            result.append(pushed);
1623            pos = offset;
1624        }
1625
1626        result.append(frozen);
1627        result
1628    }
1629}
1630
1631// SAFETY: All methods delegate directly to `self.buf`, an `IoBufMut` with a
1632// sound `BufMut` implementation. The inline buffer has fixed capacity. Writes
1633// that exceed it panic in bytes' `BufMut` trait defaults (which check
1634// `remaining_mut`) or, for a direct `advance_mut`, in `IoBufMut::advance_mut`.
1635unsafe impl BufMut for Builder {
1636    #[inline]
1637    fn remaining_mut(&self) -> usize {
1638        self.buf.remaining_mut()
1639    }
1640
1641    #[inline]
1642    unsafe fn advance_mut(&mut self, cnt: usize) {
1643        // SAFETY: The caller guarantees that `cnt` is within the writable chunk.
1644        unsafe {
1645            self.buf.advance_mut(cnt);
1646        }
1647    }
1648
1649    #[inline]
1650    fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
1651        self.buf.chunk_mut()
1652    }
1653}
1654
1655impl BufsMut for Builder {
1656    fn push(&mut self, chunk: impl Into<Bytes> + 'static) {
1657        let chunk = IoBuf::claim(chunk).unwrap_or_else(|chunk| IoBuf::from(chunk.into()));
1658        if !chunk.is_empty() {
1659            self.pushes.push((self.buf.len(), chunk));
1660        }
1661    }
1662}
1663
1664/// Extension trait for encoding values that implement [`Write`] and
1665/// [`EncodeSize`] into [`IoBufs`].
1666pub trait EncodeExt: EncodeSize + Write {
1667    /// Encode this value into [`IoBufs`] through a [`Builder`].
1668    ///
1669    /// Override [`Write::write_bufs`] to avoid copying large byte fields.
1670    ///
1671    /// # Panics
1672    ///
1673    /// Panics if [`Write::write_bufs`] writes more inline bytes than
1674    /// [`EncodeSize::encode_inline_size`], or if [`EncodeSize::encode_size`]
1675    /// does not match the total bytes written.
1676    fn encode_iobufs(&self) -> IoBufs {
1677        let len = self.encode_size();
1678        let mut builder = Builder::new(self.encode_inline_size());
1679        self.write_bufs(&mut builder);
1680        let bufs = builder.finish();
1681        assert_eq!(
1682            bufs.remaining(),
1683            len,
1684            "write_bufs() must write exactly encode_size() bytes"
1685        );
1686        bufs
1687    }
1688}
1689
1690impl<T: EncodeSize + Write> EncodeExt for T {}
1691
1692#[cfg(all(test, not(feature = "loom")))]
1693#[commonware_macros::test_group("miri")]
1694mod tests {
1695    use super::*;
1696    use bytes::{Bytes, BytesMut};
1697    use commonware_codec::{Decode, Encode, types::lazy::Lazy};
1698    use commonware_utils::range::NonEmptyRange;
1699    use std::{
1700        alloc::{GlobalAlloc, Layout, System},
1701        cell::Cell,
1702        collections::{BTreeMap, HashMap},
1703        sync::Arc,
1704    };
1705
1706    /// Counts this thread's allocations during [`allocation_calls_during`] and delegates every
1707    /// call to the system allocator.
1708    struct TrackingAllocator;
1709
1710    std::thread_local! {
1711        /// This thread's allocation count, or `None` while tracking is off.
1712        static ALLOCATION_CALLS: Cell<Option<usize>> = const { Cell::new(None) };
1713    }
1714
1715    #[global_allocator]
1716    static ALLOCATOR: TrackingAllocator = TrackingAllocator;
1717
1718    /// Counts an allocation when tracking is enabled on this thread.
1719    fn record_allocation() {
1720        ALLOCATION_CALLS.with(|calls| {
1721            if let Some(count) = calls.get() {
1722                calls.set(Some(count + 1));
1723            }
1724        });
1725    }
1726
1727    // SAFETY: Every operation delegates to the system allocator with unchanged
1728    // arguments. Recording an allocation only updates a thread-local cell.
1729    unsafe impl GlobalAlloc for TrackingAllocator {
1730        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
1731            record_allocation();
1732            // SAFETY: The caller supplies the layout required by GlobalAlloc.
1733            unsafe { System.alloc(layout) }
1734        }
1735
1736        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
1737            record_allocation();
1738            // SAFETY: The caller supplies the layout required by GlobalAlloc.
1739            unsafe { System.alloc_zeroed(layout) }
1740        }
1741
1742        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
1743            // SAFETY: This wrapper delegates every allocation to System, so the
1744            // pointer came from System with this layout.
1745            unsafe { System.dealloc(ptr, layout) }
1746        }
1747
1748        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
1749            record_allocation();
1750            // SAFETY: The pointer and layout came from System through this wrapper,
1751            // and the caller supplies the new size required by GlobalAlloc.
1752            unsafe { System.realloc(ptr, layout, new_size) }
1753        }
1754    }
1755
1756    /// Runs `f` and returns its result with the number of allocations it made on this thread.
1757    fn allocation_calls_during<T>(f: impl FnOnce() -> T) -> (T, usize) {
1758        ALLOCATION_CALLS.with(|calls| assert!(calls.replace(Some(0)).is_none()));
1759        let result = f();
1760        let calls = ALLOCATION_CALLS.with(|calls| {
1761            calls
1762                .replace(None)
1763                .expect("allocation tracking must stay enabled")
1764        });
1765        (result, calls)
1766    }
1767
1768    /// A prefix held by the front chunk is handed out as a view of that chunk, in every
1769    /// representation, and consuming the chunk keeps the container canonical. Only the first
1770    /// partial view of an unshared chunk allocates, for the chunk's owner record. A prefix that
1771    /// spans chunks is copied once into the requested type.
1772    #[test]
1773    fn test_take_owned_front_chunk() {
1774        let chunk = |byte: u8| IoBuf::copy_from_slice(&[byte; 8]);
1775        let mut bufs = IoBufs::default();
1776        for byte in 1..=5 {
1777            bufs.append(chunk(byte));
1778        }
1779        assert!(matches!(bufs.inner, IoBufsInner::Chunked(_)));
1780
1781        // Chunked: a partial take leaves the chunk in place, a take of the remainder removes it.
1782        // The partial take shares the unshared chunk, which allocates its owner record.
1783        let front_ptr = bufs.chunk().as_ptr();
1784        let (view, calls) = allocation_calls_during(|| bufs.take_owned::<IoBuf>(3));
1785        assert_eq!(calls, 1);
1786        assert_eq!(view, [1u8; 3]);
1787        assert_eq!(view.as_ptr(), front_ptr);
1788        let (view, calls) = allocation_calls_during(|| bufs.take_owned::<IoBuf>(5));
1789        assert_eq!(calls, 0);
1790        assert_eq!(view, [1u8; 5]);
1791        assert!(matches!(bufs.inner, IoBufsInner::Chunked(ref chunks) if chunks.len() == 4));
1792
1793        // Consuming the next chunk collapses to the small representation.
1794        let (view, calls) = allocation_calls_during(|| bufs.take_owned::<IoBuf>(8));
1795        assert_eq!(calls, 0);
1796        assert_eq!(view, [2u8; 8]);
1797        assert!(matches!(bufs.inner, IoBufsInner::Triple(_)));
1798
1799        // Triple and Pair front takes are views too, and a `Bytes` field gets the view's bytes.
1800        let (view, calls) = allocation_calls_during(|| bufs.take_owned::<IoBuf>(8));
1801        assert_eq!(calls, 0);
1802        assert_eq!(view, [3u8; 8]);
1803        assert!(matches!(bufs.inner, IoBufsInner::Pair(_)));
1804        let front_ptr = bufs.chunk().as_ptr();
1805        let bytes = bufs.take_owned::<Bytes>(4);
1806        assert_eq!(bytes, vec![4u8; 4]);
1807        assert_eq!(bytes.as_ptr(), front_ptr);
1808        assert!(matches!(bufs.inner, IoBufsInner::Pair(_)));
1809
1810        // A view keeps the chunk's native owner and converts in place once it is the last holder.
1811        drop(bufs);
1812        assert!(view.try_into_mut().is_ok());
1813
1814        // Spanning prefixes are copied once. An IoBuf lands in one native buffer (one data
1815        // allocation, plus the owner record that splitting the second chunk allocates), Bytes in
1816        // one shared buffer, and the remainder stays canonical.
1817        let mut bufs = IoBufs::from(chunk(1));
1818        bufs.append(chunk(2));
1819        let (copied, calls) = allocation_calls_during(|| bufs.take_owned::<IoBuf>(12));
1820        assert_eq!(calls, 2);
1821        assert_eq!(copied.as_ref(), [[1u8; 8].as_slice(), &[2u8; 4]].concat());
1822        assert!(copied.try_into_mut().is_ok());
1823        assert!(matches!(bufs.inner, IoBufsInner::Single(_)));
1824        assert_eq!(bufs.len(), 4);
1825
1826        let mut bufs = IoBufs::from(chunk(1));
1827        bufs.append(chunk(2));
1828        let (copied, calls) = allocation_calls_during(|| bufs.take_owned::<Bytes>(12));
1829        assert_eq!(calls, 1);
1830        assert_eq!(copied, [[1u8; 8].as_slice(), &[2u8; 4]].concat());
1831        assert_eq!(bufs.len(), 4);
1832
1833        // Zero-length takes are empty views that leave the container untouched.
1834        let empty = bufs.take_owned::<IoBuf>(0);
1835        assert!(empty.is_empty());
1836        assert_eq!(bufs.len(), 4);
1837        assert!(IoBufs::default().take_owned::<IoBuf>(0).is_empty());
1838    }
1839
1840    /// An unshared native buffer allocates only its data until its first share, which allocates
1841    /// its owner record once. It lives here, next to the tracking allocator that counts the calls.
1842    #[test]
1843    fn test_iobuf_owner_record_allocations() {
1844        // Allocate, write, freeze, recover, write, freeze, and drop: one data allocation.
1845        let ((), calls) = allocation_calls_during(|| {
1846            let mut buf = IoBufMut::with_capacity(64);
1847            buf.put_slice(b"payload");
1848            let mut buf = buf.freeze().try_into_mut().unwrap();
1849            buf.put_slice(b"!");
1850            drop(buf.freeze());
1851        });
1852        assert_eq!(calls, 1);
1853
1854        // Moves, whole splits, and empty views allocate nothing.
1855        let mut source = IoBuf::copy_from_slice(&[7u8; 64]);
1856        let (mut buf, calls) = allocation_calls_during(|| {
1857            let mut moved = std::mem::take(&mut source);
1858            assert!(moved.slice(0..0).is_empty());
1859            assert!(moved.split_to(0).is_empty());
1860            moved.split_to(64)
1861        });
1862        assert_eq!(calls, 0);
1863
1864        // The first share allocates the record, and later shares reuse it.
1865        let (first, calls) = allocation_calls_during(|| buf.clone());
1866        assert_eq!(calls, 1);
1867        let (later, calls) =
1868            allocation_calls_during(|| (buf.slice(1..3), buf.split_to(8), buf.clone()));
1869        assert_eq!(calls, 0);
1870        drop((first, later));
1871
1872        // Spare capacity too large for the handle allocates the record at freeze instead.
1873        let mut big = IoBufMut::with_capacity((1 << 17) + 1);
1874        big.put_u8(1);
1875        let (frozen, calls) = allocation_calls_during(|| big.freeze());
1876        assert_eq!(calls, 1);
1877        let (_view, calls) = allocation_calls_during(|| frozen.clone());
1878        assert_eq!(calls, 0);
1879
1880        // Whole conversions and drains into Bytes move the descriptor into the Bytes box without
1881        // promoting it, so each allocates only that box.
1882        let mut advanced = IoBuf::copy_from_slice(&[7u8; 16]);
1883        advanced.advance(2);
1884        let (_bytes, calls) = allocation_calls_during(|| Bytes::from(advanced));
1885        assert_eq!(calls, 1);
1886        let mut drained = IoBuf::copy_from_slice(&[7u8; 16]);
1887        let (_bytes, calls) = allocation_calls_during(|| drained.copy_to_bytes(16));
1888        assert_eq!(calls, 1);
1889        let mut drained = IoBufMut::with_capacity(16);
1890        drained.put_slice(&[7u8; 8]);
1891        let (_bytes, calls) = allocation_calls_during(|| drained.copy_to_bytes(8));
1892        assert_eq!(calls, 1);
1893    }
1894
1895    fn assert_encode_iobufs_matches_encode<T: Encode + EncodeExt>(value: &T) {
1896        let mut encoded = value.encode_iobufs();
1897        let baseline = value.encode();
1898        let mut encoded_bytes = vec![0u8; encoded.remaining()];
1899        encoded.copy_to_slice(&mut encoded_bytes);
1900        assert_eq!(encoded_bytes, baseline.as_ref());
1901    }
1902
1903    #[test]
1904    fn test_iobufs_from_iter_and_array_shapes() {
1905        fn check<const N: usize>() {
1906            let data = b"abcde";
1907            let chunks: [IoBuf; N] = std::array::from_fn(|i| IoBuf::from(&data[i..i + 1]));
1908            for bufs in [IoBufs::from(chunks.clone()), chunks.into_iter().collect()] {
1909                assert!(matches!(
1910                    (&bufs.inner, N),
1911                    (IoBufsInner::Single(_), 0 | 1)
1912                        | (IoBufsInner::Pair(_), 2)
1913                        | (IoBufsInner::Triple(_), 3)
1914                        | (IoBufsInner::Chunked(_), 4..)
1915                ));
1916                assert_eq!(bufs.chunk_count(), N);
1917                for index in 0..N {
1918                    assert_eq!(
1919                        bufs.chunk_at(index).unwrap().as_ptr(),
1920                        data[index..].as_ptr()
1921                    );
1922                }
1923                assert_eq!(bufs.coalesce().as_ref(), &data[..N]);
1924            }
1925        }
1926
1927        check::<0>();
1928        check::<1>();
1929        check::<2>();
1930        check::<3>();
1931        check::<4>();
1932        check::<5>();
1933    }
1934
1935    #[test]
1936    fn test_iobufs_from_iter_and_array_drop_empty_buffers() {
1937        fn check<const N: usize>(chunks: [&'static [u8]; N], expected: &[u8]) {
1938            let chunks = chunks.map(IoBuf::from);
1939            for bufs in [IoBufs::from(chunks.clone()), chunks.into_iter().collect()] {
1940                // Each nonempty input contains one byte.
1941                assert_eq!(bufs.chunk_count(), expected.len());
1942                assert!(matches!(
1943                    (&bufs.inner, expected.len()),
1944                    (IoBufsInner::Single(_), 0 | 1)
1945                        | (IoBufsInner::Pair(_), 2)
1946                        | (IoBufsInner::Triple(_), 3)
1947                        | (IoBufsInner::Chunked(_), 4..)
1948                ));
1949                assert_eq!(bufs.coalesce().as_ref(), expected);
1950            }
1951        }
1952
1953        check([b"", b"", b""], b"");
1954        check([b"", b"a", b""], b"a");
1955        check([b"", b"a", b"b"], b"ab");
1956        check([b"a", b"", b"b"], b"ab");
1957        check([b"a", b"b", b""], b"ab");
1958        check([b"a", b"", b"b", b"c"], b"abc");
1959        check([b"a", b"b", b"", b"c", b"d"], b"abcd");
1960        check([b"", b"a", b"", b"b", b"c", b"d", b"e", b""], b"abcde");
1961    }
1962
1963    #[test]
1964    fn test_iobufs_from_iter_error_releases_buffers() {
1965        const BUFFER_COUNT: usize = 6;
1966
1967        // Cover failures before and after the transition to deque storage,
1968        // always leaving at least one unconsumed buffer in the iterator.
1969        for fail_at in 0..=4 {
1970            let owners: Vec<Arc<[u8]>> = (0..BUFFER_COUNT).map(|_| Arc::from(&b"x"[..])).collect();
1971            let bufs: Vec<_> = owners
1972                .iter()
1973                .map(|owner| IoBuf::from(Bytes::from_owner(owner.clone())))
1974                .collect();
1975            // The test and the buffer each hold one reference to the owner.
1976            assert!(owners.iter().all(|owner| Arc::strong_count(owner) == 2));
1977
1978            let mut items_pulled = 0;
1979            let result = bufs
1980                .into_iter()
1981                .enumerate()
1982                .map(|(index, buf)| {
1983                    items_pulled += 1;
1984                    if index == fail_at {
1985                        Err("injected")
1986                    } else {
1987                        Ok(buf)
1988                    }
1989                })
1990                .collect::<Result<IoBufs, _>>();
1991
1992            assert_eq!(result.unwrap_err(), "injected");
1993            assert_eq!(items_pulled, fail_at + 1);
1994            // Only the test's reference remains, including for unconsumed buffers.
1995            for owner in owners {
1996                assert_eq!(Arc::strong_count(&owner), 1);
1997            }
1998        }
1999    }
2000
2001    #[test]
2002    fn test_iobufsmut_from_iter_and_array_shapes() {
2003        fn check<const N: usize>() {
2004            let chunks =
2005                || -> [IoBufMut; N] { std::array::from_fn(|_| IoBufMut::with_capacity(1)) };
2006            for bufs in [IoBufsMut::from(chunks()), chunks().into_iter().collect()] {
2007                assert!(matches!(
2008                    (&bufs.inner, N),
2009                    (IoBufsMutInner::Single(_), 0 | 1)
2010                        | (IoBufsMutInner::Pair(_), 2)
2011                        | (IoBufsMutInner::Triple(_), 3)
2012                        | (IoBufsMutInner::Chunked(_), 4..)
2013                ));
2014                assert!(bufs.is_empty());
2015                assert_eq!(bufs.capacity(), N);
2016                assert_eq!(bufs.remaining_mut(), N);
2017            }
2018        }
2019
2020        check::<0>();
2021        check::<1>();
2022        check::<2>();
2023        check::<3>();
2024        check::<4>();
2025        check::<5>();
2026    }
2027
2028    #[test]
2029    fn test_iobufsmut_from_iter_and_array_preserve_capacity() {
2030        fn check(convert: impl FnOnce([IoBufMut; 4]) -> IoBufsMut) {
2031            let empty = IoBufMut::with_capacity(2);
2032            let mut partial = IoBufMut::with_capacity(2);
2033            partial.put_slice(b"a");
2034            let full = IoBufMut::from(b"bc");
2035            let pointers = [
2036                empty.as_ref().as_ptr(),
2037                partial.as_ref().as_ptr(),
2038                full.as_ref().as_ptr(),
2039            ];
2040
2041            let mut bufs = convert([IoBufMut::default(), empty, partial, full]);
2042            assert!(matches!(bufs.inner, IoBufsMutInner::Triple(_)));
2043            assert_eq!(bufs.len(), 3);
2044            assert_eq!(bufs.capacity(), 6);
2045            assert_eq!(bufs.remaining_mut(), 3);
2046
2047            // Fill the empty buffer with XY and the partial buffer's tail with Z.
2048            bufs.put_slice(b"XYZ");
2049            assert_eq!(bufs.capacity(), 6);
2050            assert_eq!(bufs.remaining_mut(), 0);
2051            let frozen = bufs.freeze();
2052            for (index, pointer) in pointers.into_iter().enumerate() {
2053                assert_eq!(frozen.chunk_at(index).unwrap().as_ptr(), pointer);
2054            }
2055            assert_eq!(frozen.coalesce(), b"XYaZbc");
2056        }
2057
2058        check(IoBufsMut::from);
2059        check(|chunks| chunks.into_iter().collect());
2060    }
2061
2062    #[test]
2063    fn test_iobufs_emptiness_after_consumption() {
2064        fn assert_state(bufs: &IoBufs, expected: &[u8]) {
2065            assert_eq!(bufs.remaining(), expected.len());
2066            assert_eq!(bufs.is_empty(), expected.is_empty());
2067            assert_eq!(bufs.has_remaining(), !expected.is_empty());
2068            let mut actual = Vec::new();
2069            bufs.for_each_chunk(|chunk| actual.extend_from_slice(chunk));
2070            assert_eq!(actual, expected);
2071        }
2072
2073        let data = b"abcdefgh";
2074        let source = IoBufs::from(vec![
2075            IoBuf::default(),
2076            IoBuf::from(b"ab"),
2077            IoBuf::default(),
2078            IoBuf::from(b"cd"),
2079            IoBuf::from(b"ef"),
2080            IoBuf::default(),
2081            IoBuf::from(b"gh"),
2082            IoBuf::default(),
2083        ]);
2084
2085        // Prefixes exercise every representation, including an empty Single.
2086        for end in 0..=data.len() {
2087            let mut suffix = source.clone();
2088            let prefix = suffix.split_to(end);
2089            assert_state(&prefix, &data[..end]);
2090            assert_state(&suffix, &data[end..]);
2091
2092            for step in [1, 2, 3, data.len()] {
2093                let mut advanced = prefix.clone();
2094                let mut copied = prefix.clone();
2095                let mut offset = 0;
2096                while offset < end {
2097                    let count = step.min(end - offset);
2098                    advanced.advance(count);
2099                    assert_eq!(copied.copy_to_bytes(count), &data[offset..offset + count]);
2100                    offset += count;
2101                    assert_state(&advanced, &data[offset..end]);
2102                    assert_state(&copied, &data[offset..end]);
2103                }
2104
2105                // Empty inputs preserve emptiness, and new bytes make a drained buffer readable.
2106                advanced.append(IoBuf::default());
2107                advanced.prepend(IoBuf::default());
2108                assert_state(&advanced, b"");
2109                advanced.append(IoBuf::from(b"z"));
2110                advanced.prepend(IoBuf::from(b"y"));
2111                assert_state(&advanced, b"yz");
2112            }
2113
2114            let mut input = prefix;
2115            let mut output = IoBufMut::with_capacity(data.len());
2116            output.put(&mut input);
2117            assert_eq!(output.as_ref(), &data[..end]);
2118            assert_state(&input, b"");
2119        }
2120    }
2121
2122    #[test]
2123    fn test_iobufs_for_each_chunk_single_and_empty() {
2124        // The Single arm is bypassed by coalesce (which returns the buffer
2125        // directly), so exercise the public method on both Single shapes.
2126        let mut seen = Vec::new();
2127        IoBufs::from(b"hello").for_each_chunk(|chunk| seen.push(chunk.to_vec()));
2128        assert_eq!(seen, vec![b"hello".to_vec()]);
2129
2130        let mut count = 0;
2131        IoBufs::default().for_each_chunk(|_| count += 1);
2132        assert_eq!(count, 0);
2133    }
2134
2135    #[test]
2136    fn test_iobufs_shapes_and_read_paths() {
2137        // Empty construction normalizes to an empty single chunk.
2138        let empty = IoBufs::from(Vec::<u8>::new());
2139        assert!(empty.is_empty());
2140        assert!(empty.is_single());
2141        assert!(empty.as_single().is_some());
2142
2143        // Single-buffer read path.
2144        let mut single = IoBufs::from(b"hello world");
2145        assert!(single.is_single());
2146        assert_eq!(single.chunk(), b"hello world");
2147        single.advance(6);
2148        assert_eq!(single.chunk(), b"world");
2149        assert_eq!(single.copy_to_bytes(5).as_ref(), b"world");
2150        assert_eq!(single.remaining(), 0);
2151
2152        // Fast-path shapes (Pair/Triple/Chunked).
2153        let mut pair = IoBufs::from(IoBuf::from(b"a"));
2154        pair.append(IoBuf::from(b"b"));
2155        assert!(matches!(pair.inner, IoBufsInner::Pair(_)));
2156        assert!(pair.as_single().is_none());
2157
2158        let mut triple = IoBufs::from(IoBuf::from(b"a"));
2159        triple.append(IoBuf::from(b"b"));
2160        triple.append(IoBuf::from(b"c"));
2161        assert!(matches!(triple.inner, IoBufsInner::Triple(_)));
2162
2163        let mut chunked = IoBufs::from(IoBuf::from(b"a"));
2164        chunked.append(IoBuf::from(b"b"));
2165        chunked.append(IoBuf::from(b"c"));
2166        chunked.append(IoBuf::from(b"d"));
2167        assert!(matches!(chunked.inner, IoBufsInner::Chunked(_)));
2168
2169        // prepend + append preserve ordering.
2170        let mut joined = IoBufs::from(b"middle");
2171        joined.prepend(IoBuf::from(b"start "));
2172        joined.append(IoBuf::from(b" end"));
2173        assert_eq!(joined.coalesce(), b"start middle end");
2174
2175        // prepending empty is a no-op, and prepending into pair upgrades to triple.
2176        let mut prepend_noop = IoBufs::from(b"x");
2177        prepend_noop.prepend(IoBuf::default());
2178        assert_eq!(prepend_noop.coalesce(), b"x");
2179
2180        // Prepending into an empty aggregate should stay on the single-buffer fast path.
2181        let mut prepend_into_empty = IoBufs::default();
2182        prepend_into_empty.prepend(IoBuf::from(b"z"));
2183        assert!(prepend_into_empty.is_single());
2184        assert_eq!(prepend_into_empty.coalesce(), b"z");
2185
2186        let mut prepend_pair = IoBufs::from(vec![IoBuf::from(b"b"), IoBuf::from(b"c")]);
2187        prepend_pair.prepend(IoBuf::from(b"a"));
2188        assert!(matches!(prepend_pair.inner, IoBufsInner::Triple(_)));
2189        assert_eq!(prepend_pair.coalesce(), b"abc");
2190
2191        // canonicalizing a non-empty single should keep the same representation.
2192        let mut canonical_single = IoBufs::from(b"q");
2193        canonical_single.canonicalize();
2194        assert!(canonical_single.is_single());
2195        assert_eq!(canonical_single.coalesce(), b"q");
2196    }
2197
2198    #[test]
2199    fn test_iobufs_split_to_cases() {
2200        // Zero and full split on a single chunk.
2201        let mut bufs = IoBufs::from(b"hello");
2202
2203        let empty = bufs.split_to(0);
2204        assert!(empty.is_empty());
2205        assert_eq!(bufs.coalesce(), b"hello");
2206
2207        let mut bufs = IoBufs::from(b"hello");
2208        let all = bufs.split_to(5);
2209        assert_eq!(all.coalesce(), b"hello");
2210        assert!(bufs.is_single());
2211        assert!(bufs.is_empty());
2212
2213        // Single split in the middle.
2214        let mut single_mid = IoBufs::from(b"hello");
2215        let single_prefix = single_mid.split_to(2);
2216        assert!(single_prefix.is_single());
2217        assert_eq!(single_prefix.coalesce(), b"he");
2218        assert_eq!(single_mid.coalesce(), b"llo");
2219
2220        // Pair split paths: in-first, boundary-after-first, crossing-into-second.
2221        let mut pair = IoBufs::from(vec![IoBuf::from(b"ab"), IoBuf::from(b"cd")]);
2222        let pair_prefix = pair.split_to(1);
2223        assert!(pair_prefix.is_single());
2224        assert_eq!(pair_prefix.coalesce(), b"a");
2225        assert!(matches!(pair.inner, IoBufsInner::Pair(_)));
2226        assert_eq!(pair.coalesce(), b"bcd");
2227
2228        let mut pair = IoBufs::from(vec![IoBuf::from(b"ab"), IoBuf::from(b"cd")]);
2229        let pair_prefix = pair.split_to(2);
2230        assert!(pair_prefix.is_single());
2231        assert_eq!(pair_prefix.coalesce(), b"ab");
2232        assert!(pair.is_single());
2233        assert_eq!(pair.coalesce(), b"cd");
2234
2235        let mut pair = IoBufs::from(vec![IoBuf::from(b"ab"), IoBuf::from(b"cd")]);
2236        let pair_prefix = pair.split_to(3);
2237        assert!(matches!(pair_prefix.inner, IoBufsInner::Pair(_)));
2238        assert_eq!(pair_prefix.coalesce(), b"abc");
2239        assert!(pair.is_single());
2240        assert_eq!(pair.coalesce(), b"d");
2241
2242        // Triple split paths: in-first, boundary-after-first, in-second, boundary-after-second,
2243        // and reaching into third.
2244        let mut triple = IoBufs::from(vec![
2245            IoBuf::from(b"ab"),
2246            IoBuf::from(b"cd"),
2247            IoBuf::from(b"ef"),
2248        ]);
2249        let triple_prefix = triple.split_to(1);
2250        assert!(triple_prefix.is_single());
2251        assert_eq!(triple_prefix.coalesce(), b"a");
2252        assert!(matches!(triple.inner, IoBufsInner::Triple(_)));
2253        assert_eq!(triple.coalesce(), b"bcdef");
2254
2255        let mut triple = IoBufs::from(vec![
2256            IoBuf::from(b"ab"),
2257            IoBuf::from(b"cd"),
2258            IoBuf::from(b"ef"),
2259        ]);
2260        let triple_prefix = triple.split_to(2);
2261        assert!(triple_prefix.is_single());
2262        assert_eq!(triple_prefix.coalesce(), b"ab");
2263        assert!(matches!(triple.inner, IoBufsInner::Pair(_)));
2264        assert_eq!(triple.coalesce(), b"cdef");
2265
2266        let mut triple = IoBufs::from(vec![
2267            IoBuf::from(b"ab"),
2268            IoBuf::from(b"cd"),
2269            IoBuf::from(b"ef"),
2270        ]);
2271        let triple_prefix = triple.split_to(3);
2272        assert!(matches!(triple_prefix.inner, IoBufsInner::Pair(_)));
2273        assert_eq!(triple_prefix.coalesce(), b"abc");
2274        assert!(matches!(triple.inner, IoBufsInner::Pair(_)));
2275        assert_eq!(triple.coalesce(), b"def");
2276
2277        let mut triple = IoBufs::from(vec![
2278            IoBuf::from(b"ab"),
2279            IoBuf::from(b"cd"),
2280            IoBuf::from(b"ef"),
2281        ]);
2282        let triple_prefix = triple.split_to(4);
2283        assert!(matches!(triple_prefix.inner, IoBufsInner::Pair(_)));
2284        assert_eq!(triple_prefix.coalesce(), b"abcd");
2285        assert!(triple.is_single());
2286        assert_eq!(triple.coalesce(), b"ef");
2287
2288        let mut triple = IoBufs::from(vec![
2289            IoBuf::from(b"ab"),
2290            IoBuf::from(b"cd"),
2291            IoBuf::from(b"ef"),
2292        ]);
2293        let triple_prefix = triple.split_to(5);
2294        assert!(matches!(triple_prefix.inner, IoBufsInner::Triple(_)));
2295        assert_eq!(triple_prefix.coalesce(), b"abcde");
2296        assert!(triple.is_single());
2297        assert_eq!(triple.coalesce(), b"f");
2298
2299        // Chunked split can canonicalize remainder/prefix shapes.
2300        let mut bufs = IoBufs::from(vec![
2301            IoBuf::from(b"ab"),
2302            IoBuf::from(b"cd"),
2303            IoBuf::from(b"ef"),
2304            IoBuf::from(b"gh"),
2305        ]);
2306        let prefix = bufs.split_to(4);
2307        assert!(matches!(prefix.inner, IoBufsInner::Pair(_)));
2308        assert_eq!(prefix.coalesce(), b"abcd");
2309        assert!(matches!(bufs.inner, IoBufsInner::Pair(_)));
2310        assert_eq!(bufs.coalesce(), b"efgh");
2311
2312        // Chunked split inside a chunk.
2313        let mut bufs = IoBufs::from(vec![
2314            IoBuf::from(b"ab"),
2315            IoBuf::from(b"cd"),
2316            IoBuf::from(b"ef"),
2317            IoBuf::from(b"gh"),
2318        ]);
2319        let prefix = bufs.split_to(5);
2320        assert!(matches!(prefix.inner, IoBufsInner::Triple(_)));
2321        assert_eq!(prefix.coalesce(), b"abcde");
2322        assert!(matches!(bufs.inner, IoBufsInner::Pair(_)));
2323        assert_eq!(bufs.coalesce(), b"fgh");
2324
2325        // Chunked split can remain chunked on both sides when both have >= 4 chunks.
2326        let mut bufs = IoBufs::from(vec![
2327            IoBuf::from(b"a"),
2328            IoBuf::from(b"b"),
2329            IoBuf::from(b"c"),
2330            IoBuf::from(b"d"),
2331            IoBuf::from(b"e"),
2332            IoBuf::from(b"f"),
2333            IoBuf::from(b"g"),
2334            IoBuf::from(b"h"),
2335        ]);
2336        let prefix = bufs.split_to(4);
2337        assert!(matches!(prefix.inner, IoBufsInner::Chunked(_)));
2338        assert_eq!(prefix.coalesce(), b"abcd");
2339        assert!(matches!(bufs.inner, IoBufsInner::Chunked(_)));
2340        assert_eq!(bufs.coalesce(), b"efgh");
2341
2342        // Chunked split inside a chunk that leaves both sides chunked.
2343        let mut bufs: IoBufs = (b'a'..b'i')
2344            .map(|byte| IoBuf::from(vec![byte; 2]))
2345            .collect();
2346        let prefix = bufs.split_to(9);
2347        assert!(matches!(prefix.inner, IoBufsInner::Chunked(_)));
2348        assert_eq!(prefix.chunk_count(), 5);
2349        assert_eq!(prefix.coalesce(), b"aabbccdde");
2350        assert!(matches!(bufs.inner, IoBufsInner::Chunked(_)));
2351        assert_eq!(bufs.chunk_count(), 4);
2352        assert_eq!(bufs.coalesce(), b"effgghh");
2353
2354        // Defensive path: tolerate accidental empty chunks in non-canonical chunked input.
2355        let mut bufs = IoBufs {
2356            inner: IoBufsInner::Chunked(VecDeque::from([
2357                IoBuf::default(),
2358                IoBuf::from(b"ab"),
2359                IoBuf::from(b"cd"),
2360                IoBuf::from(b"ef"),
2361                IoBuf::from(b"gh"),
2362            ])),
2363        };
2364        let prefix = bufs.split_to(3);
2365        assert_eq!(prefix.coalesce(), b"abc");
2366        assert_eq!(bufs.coalesce(), b"defgh");
2367    }
2368
2369    #[test]
2370    #[should_panic(expected = "split_to out of bounds")]
2371    fn test_iobufs_split_to_out_of_bounds() {
2372        let mut bufs = IoBufs::from(b"abc");
2373        let _ = bufs.split_to(4);
2374    }
2375
2376    /// Splits at every position of every representation keep the bytes in order and leave both
2377    /// sides canonical. The positions cover prefixes within the front chunk and across chunks,
2378    /// and chunked splits that move out either side.
2379    #[test]
2380    fn test_iobufs_split_positions() {
2381        for count in 1..10 {
2382            for rotation in 0..count {
2383                for at in 0..=2 * count {
2384                    // Rotation wraps a chunked deque around the end of its buffer, so the drains
2385                    // also run over a deque stored in two slices. Four or more chunks keep that
2386                    // deque as their chunked representation.
2387                    let mut chunks: VecDeque<IoBuf> = (0..count as u8)
2388                        .map(|byte| IoBuf::copy_from_slice(&[byte; 2]))
2389                        .collect();
2390                    chunks.rotate_left(rotation);
2391                    let expected: Vec<u8> = chunks
2392                        .iter()
2393                        .flat_map(|buf| buf.as_ref().to_vec())
2394                        .collect();
2395                    let mut bufs = if count >= 4 {
2396                        IoBufs {
2397                            inner: IoBufsInner::Chunked(chunks),
2398                        }
2399                    } else {
2400                        IoBufs::from_chunks_iter(chunks)
2401                    };
2402
2403                    let prefix = bufs.split_to(at);
2404
2405                    // Each side stores only non-empty chunks, in the representation its chunk
2406                    // count calls for, and the two sides hold the original bytes in order.
2407                    for side in [&prefix, &bufs] {
2408                        let chunks = side.chunk_count();
2409                        assert_eq!(side.iter().filter(|buf| !buf.is_empty()).count(), chunks);
2410                        assert_eq!(side.is_single(), chunks <= 1);
2411                        assert_eq!(matches!(side.inner, IoBufsInner::Chunked(_)), chunks >= 4);
2412                    }
2413                    assert_eq!(prefix.coalesce().as_ref(), &expected[..at]);
2414                    assert_eq!(bufs.coalesce().as_ref(), &expected[at..]);
2415                }
2416            }
2417        }
2418    }
2419
2420    /// A split of a chunked container moves out the side with fewer chunks, counting the piece
2421    /// toward the prefix, and keeps the deque for the other, so taking a small prefix or leaving a
2422    /// small suffix allocates nothing however many chunks remain. The chunks are static, so
2423    /// splitting one in place allocates nothing either.
2424    #[test]
2425    fn test_iobufs_chunked_split_allocations() {
2426        // Every split folds over the remaining chunks, so the count stays small enough for Miri
2427        // while keeping the container chunked throughout.
2428        const CHUNKS: usize = 256;
2429
2430        // Splitting inside the first of many chunks.
2431        let mut bufs: IoBufs = (0..CHUNKS).map(|_| IoBuf::from(b"ab")).collect();
2432        let (prefix, calls) = allocation_calls_during(|| bufs.split_to(1));
2433        assert_eq!(calls, 0);
2434        assert_eq!(prefix.coalesce(), b"a");
2435        assert_eq!(bufs.len(), 2 * CHUNKS - 1);
2436
2437        // A split across the first two chunks moves out only those two.
2438        let (prefix, calls) = allocation_calls_during(|| bufs.split_to(2));
2439        assert_eq!(calls, 0);
2440        assert_eq!(prefix.coalesce(), b"ba");
2441
2442        // A split that leaves only the last two chunks moves out only those two.
2443        let at = bufs.len() - 3;
2444        let (_, calls) = allocation_calls_during(|| bufs.split_to(at));
2445        assert_eq!(calls, 0);
2446        assert_eq!(bufs.coalesce(), b"bab");
2447
2448        // A split inside the fourth of six chunks puts four chunks in the prefix and three in the
2449        // suffix, so the suffix moves out and the deque stays.
2450        let mut bufs: IoBufs = (0..6).map(|_| IoBuf::from(b"ab")).collect();
2451        let (prefix, calls) = allocation_calls_during(|| bufs.split_to(7));
2452        assert_eq!(calls, 0);
2453        assert_eq!(prefix.coalesce(), b"abababa");
2454        assert_eq!(bufs.coalesce(), b"babab");
2455
2456        // Extracting the chunks one at a time.
2457        let mut bufs: IoBufs = (0..CHUNKS).map(|_| IoBuf::from(b"ab")).collect();
2458        let ((), calls) = allocation_calls_during(|| {
2459            while bufs.chunk_count() >= 4 {
2460                assert_eq!(bufs.split_to(2).coalesce(), b"ab");
2461            }
2462        });
2463        assert_eq!(calls, 0);
2464    }
2465
2466    #[test]
2467    fn test_iobufs_chunk_count() {
2468        assert_eq!(IoBufs::default().chunk_count(), 0);
2469        assert_eq!(IoBufs::from(IoBuf::from(b"a")).chunk_count(), 1);
2470        assert_eq!(
2471            IoBufs::from(vec![IoBuf::from(b"b"), IoBuf::from(b"c")]).chunk_count(),
2472            2
2473        );
2474        assert_eq!(
2475            IoBufs::from(vec![
2476                IoBuf::from(b"a"),
2477                IoBuf::from(b"b"),
2478                IoBuf::from(b"c")
2479            ])
2480            .chunk_count(),
2481            3
2482        );
2483        assert_eq!(
2484            IoBufs::from(vec![
2485                IoBuf::from(b"a"),
2486                IoBuf::from(b"b"),
2487                IoBuf::from(b"c"),
2488                IoBuf::from(b"d")
2489            ])
2490            .chunk_count(),
2491            4
2492        );
2493    }
2494
2495    #[test]
2496    fn test_iobufs_chunk_at() {
2497        let chunks = [b"ab".as_slice(), b"cd", b"ef", b"gh", b"ij"];
2498        for count in 0..=chunks.len() {
2499            let mut bufs = IoBufs::from(
2500                chunks[..count]
2501                    .iter()
2502                    .map(|chunk| IoBuf::from(*chunk))
2503                    .collect::<Vec<_>>(),
2504            );
2505            for (index, chunk) in chunks[..count].iter().enumerate() {
2506                assert_eq!(bufs.chunk_at(index), Some(*chunk));
2507            }
2508            assert_eq!(bufs.chunk_at(count), None);
2509            assert_eq!(bufs.chunk_at(usize::MAX), None);
2510            assert_eq!(bufs.len(), count * 2);
2511
2512            // Indexing follows the readable chunks after partial and whole
2513            // chunks are consumed, including representation transitions.
2514            if count > 0 {
2515                bufs.advance(1);
2516                assert_eq!(bufs.chunk_at(0), Some(b"b".as_slice()));
2517                bufs.advance(1);
2518                assert_eq!(bufs.chunk_at(0), chunks[1..count].first().copied());
2519            }
2520        }
2521    }
2522
2523    #[test]
2524    fn test_iobufs_iter_borrows_every_owner() {
2525        let chunks = [b"ab".as_slice(), b"cd", b"ef", b"gh", b"ij"];
2526        for count in 0..=chunks.len() {
2527            let bufs = IoBufs::from(
2528                chunks[..count]
2529                    .iter()
2530                    .map(|chunk| IoBuf::from(*chunk))
2531                    .collect::<Vec<_>>(),
2532            );
2533            let iter = bufs.iter();
2534            assert_eq!(iter.len(), count);
2535            assert_eq!(
2536                iter.clone().map(AsRef::as_ref).collect::<Vec<_>>(),
2537                chunks[..count]
2538            );
2539            assert_eq!(iter.count(), count);
2540        }
2541    }
2542
2543    #[test]
2544    fn test_iobufs_coalesce_after_advance() {
2545        let mut bufs = IoBufs::from(IoBuf::from(b"hello"));
2546        bufs.append(IoBuf::from(b" world"));
2547
2548        assert_eq!(bufs.len(), 11);
2549
2550        bufs.advance(3);
2551        assert_eq!(bufs.len(), 8);
2552
2553        assert_eq!(bufs.coalesce(), b"lo world");
2554    }
2555
2556    #[test]
2557    fn test_iobufs_coalesce_shapes() {
2558        // Single buffer: zero-copy (same pointer)
2559        let buf = IoBuf::from(vec![1u8, 2, 3, 4, 5]);
2560        let original_ptr = buf.as_ptr();
2561        let bufs = IoBufs::from(buf);
2562        let coalesced = bufs.coalesce();
2563        assert_eq!(coalesced, [1, 2, 3, 4, 5]);
2564        assert_eq!(coalesced.as_ptr(), original_ptr);
2565
2566        // Multiple buffers: merged into one allocation
2567        let mut bufs = IoBufs::from(IoBuf::from(b"hello"));
2568        bufs.append(IoBuf::from(b" world"));
2569        let coalesced = bufs.coalesce();
2570        assert_eq!(coalesced, b"hello world");
2571
2572        // Multiple buffers after advance: only remaining data coalesced
2573        let mut bufs = IoBufs::from(IoBuf::from(b"hello"));
2574        bufs.append(IoBuf::from(b" world"));
2575        bufs.advance(3);
2576        let coalesced = bufs.coalesce();
2577        assert_eq!(coalesced, b"lo world");
2578
2579        // Empty buffers in the middle
2580        let mut bufs = IoBufs::from(IoBuf::from(b"hello"));
2581        bufs.append(IoBuf::default());
2582        bufs.append(IoBuf::from(b" world"));
2583        let coalesced = bufs.coalesce();
2584        assert_eq!(coalesced, b"hello world");
2585
2586        // Empty IoBufs
2587        let bufs = IoBufs::default();
2588        let coalesced = bufs.coalesce();
2589        assert!(coalesced.is_empty());
2590
2591        // 4+ buffers: exercise the chunked coalesce path.
2592        let bufs = IoBufs::from(vec![
2593            IoBuf::from(b"ab"),
2594            IoBuf::from(b"cd"),
2595            IoBuf::from(b"ef"),
2596            IoBuf::from(b"gh"),
2597        ]);
2598        let coalesced = bufs.coalesce();
2599        assert_eq!(coalesced, b"abcdefgh");
2600    }
2601
2602    #[test]
2603    fn test_iobufsmut_decode_preserves_byte_fields() {
2604        let first = IoBufMut::from(&b"\x02\x03abc"[..]);
2605        let second = IoBufMut::from(&b"\x03def"[..]);
2606        let ranges = [
2607            first.as_ref().as_ptr_range(),
2608            second.as_ref().as_ptr_range(),
2609        ];
2610        let source = IoBufsMut::from(vec![first, second]);
2611        let decoded = Vec::<Bytes>::decode_cfg(source, &((..).into(), (..).into())).unwrap();
2612        assert_eq!(
2613            decoded,
2614            [Bytes::from_static(b"abc"), Bytes::from_static(b"def")]
2615        );
2616        for (field, range) in decoded.iter().zip(ranges) {
2617            assert!(range.contains(&field.as_ptr()));
2618        }
2619    }
2620
2621    #[test]
2622    fn test_iobufs_empty_chunks_and_copy_to_bytes_paths() {
2623        // Empty chunks are skipped while reading across multiple chunks.
2624        let mut bufs = IoBufs::default();
2625        bufs.append(IoBuf::from(b"hello"));
2626        bufs.append(IoBuf::default());
2627        bufs.append(IoBuf::from(b" "));
2628        bufs.append(IoBuf::default());
2629        bufs.append(IoBuf::from(b"world"));
2630        assert_eq!(bufs.len(), 11);
2631        assert_eq!(bufs.chunk(), b"hello");
2632        bufs.advance(5);
2633        assert_eq!(bufs.chunk(), b" ");
2634        bufs.advance(1);
2635        assert_eq!(bufs.chunk(), b"world");
2636
2637        // Single-buffer copy_to_bytes path.
2638        let mut single = IoBufs::from(b"hello world");
2639        assert_eq!(single.copy_to_bytes(5).as_ref(), b"hello");
2640        assert_eq!(single.remaining(), 6);
2641
2642        // Multi-buffer copy_to_bytes path across boundaries.
2643        let mut multi = IoBufs::from(b"hello");
2644        multi.prepend(IoBuf::from(b"say "));
2645        assert_eq!(multi.copy_to_bytes(7).as_ref(), b"say hel");
2646        assert_eq!(multi.copy_to_bytes(2).as_ref(), b"lo");
2647    }
2648
2649    #[test]
2650    fn test_iobufs_copy_to_bytes_pair_and_triple() {
2651        // Pair: crossing one boundary should collapse to the trailing single chunk.
2652        let mut pair = IoBufs::from(IoBuf::from(b"ab"));
2653        pair.append(IoBuf::from(b"cd"));
2654        let first = pair.copy_to_bytes(3);
2655        assert_eq!(&first[..], b"abc");
2656        assert!(pair.is_single());
2657        assert_eq!(pair.chunk(), b"d");
2658
2659        // Triple: draining across two chunks leaves the final chunk readable.
2660        let mut triple = IoBufs::from(IoBuf::from(b"ab"));
2661        triple.append(IoBuf::from(b"cd"));
2662        triple.append(IoBuf::from(b"ef"));
2663        let first = triple.copy_to_bytes(5);
2664        assert_eq!(&first[..], b"abcde");
2665        assert!(triple.is_single());
2666        assert_eq!(triple.chunk(), b"f");
2667    }
2668
2669    #[test]
2670    fn test_iobufs_copy_to_bytes_chunked_four_plus() {
2671        let mut bufs = IoBufs::from(vec![
2672            IoBuf::from(b"ab"),
2673            IoBuf::from(b"cd"),
2674            IoBuf::from(b"ef"),
2675            IoBuf::from(b"gh"),
2676        ]);
2677
2678        // Chunked fast-path: first chunk alone satisfies request.
2679        let first = bufs.copy_to_bytes(1);
2680        assert_eq!(&first[..], b"a");
2681
2682        // Chunked slow-path: request crosses chunk boundaries.
2683        let second = bufs.copy_to_bytes(4);
2684        assert_eq!(&second[..], b"bcde");
2685
2686        let rest = bufs.copy_to_bytes(3);
2687        assert_eq!(&rest[..], b"fgh");
2688        assert_eq!(bufs.remaining(), 0);
2689    }
2690
2691    #[test]
2692    fn test_iobufs_copy_to_bytes_edge_cases() {
2693        // A non-canonical leading empty chunk in the deque path is popped
2694        // without affecting the copied payload. Canonical construction never
2695        // stores empties, so the state is built directly.
2696        let mut iobufs = IoBufs {
2697            inner: IoBufsInner::Chunked(VecDeque::from([
2698                IoBuf::default(),
2699                IoBuf::from(b"hel"),
2700                IoBuf::from(b"lo"),
2701                IoBuf::from(b" world"),
2702            ])),
2703        };
2704        assert_eq!(iobufs.copy_to_bytes(5).as_ref(), b"hello");
2705        assert_eq!(iobufs.remaining(), 6);
2706
2707        // Boundary-aligned reads should return exact chunk payloads in-order.
2708        let mut boundary = IoBufs::from(IoBuf::from(b"hello"));
2709        boundary.append(IoBuf::from(b"world"));
2710        assert_eq!(boundary.copy_to_bytes(5).as_ref(), b"hello");
2711        assert_eq!(boundary.copy_to_bytes(5).as_ref(), b"world");
2712        assert_eq!(boundary.remaining(), 0);
2713    }
2714
2715    #[test]
2716    #[should_panic(expected = "cannot advance past end of buffer")]
2717    fn test_iobufs_advance_past_end() {
2718        let mut bufs = IoBufs::from(b"hel");
2719        bufs.append(IoBuf::from(b"lo"));
2720        bufs.advance(10);
2721    }
2722
2723    #[test]
2724    #[should_panic(expected = "not enough data")]
2725    fn test_iobufs_copy_to_bytes_past_end() {
2726        let mut bufs = IoBufs::from(b"hel");
2727        bufs.append(IoBuf::from(b"lo"));
2728        bufs.copy_to_bytes(10);
2729    }
2730
2731    #[test]
2732    fn test_iobufs_matches_bytes_chain() {
2733        let b1 = Bytes::from_static(b"hello");
2734        let b2 = Bytes::from_static(b" ");
2735        let b3 = Bytes::from_static(b"world");
2736
2737        // Buf parity for remaining/chunk/advance should match `Bytes::chain`.
2738        let mut chain = b1.clone().chain(b2.clone()).chain(b3.clone());
2739        let mut iobufs = IoBufs::from(IoBuf::from(b1.clone()));
2740        iobufs.append(IoBuf::from(b2.clone()));
2741        iobufs.append(IoBuf::from(b3.clone()));
2742
2743        assert_eq!(chain.remaining(), iobufs.remaining());
2744        assert_eq!(chain.chunk(), iobufs.chunk());
2745
2746        chain.advance(3);
2747        iobufs.advance(3);
2748        assert_eq!(chain.remaining(), iobufs.remaining());
2749        assert_eq!(chain.chunk(), iobufs.chunk());
2750
2751        chain.advance(3);
2752        iobufs.advance(3);
2753        assert_eq!(chain.remaining(), iobufs.remaining());
2754        assert_eq!(chain.chunk(), iobufs.chunk());
2755
2756        // Test copy_to_bytes
2757        let mut chain = b1.clone().chain(b2.clone()).chain(b3.clone());
2758        let mut iobufs = IoBufs::from(IoBuf::from(b1));
2759        iobufs.append(IoBuf::from(b2));
2760        iobufs.append(IoBuf::from(b3));
2761
2762        assert_eq!(chain.copy_to_bytes(3), iobufs.copy_to_bytes(3));
2763        assert_eq!(chain.copy_to_bytes(4), iobufs.copy_to_bytes(4));
2764        assert_eq!(
2765            chain.copy_to_bytes(chain.remaining()),
2766            iobufs.copy_to_bytes(iobufs.remaining())
2767        );
2768        assert_eq!(chain.remaining(), 0);
2769        assert_eq!(iobufs.remaining(), 0);
2770    }
2771
2772    #[test]
2773    fn test_iobufs_try_into_single() {
2774        let single = IoBufs::from(IoBuf::from(b"hello"));
2775        let single = single.try_into_single().expect("single expected");
2776        assert_eq!(single, b"hello");
2777
2778        let multi = IoBufs::from(vec![IoBuf::from(b"ab"), IoBuf::from(b"cd")]);
2779        let multi = multi.try_into_single().expect_err("multi expected");
2780        assert_eq!(multi.coalesce(), b"abcd");
2781    }
2782
2783    #[test]
2784    fn test_iobufs_chunks_vectored_multiple_slices() {
2785        // Single non-empty buffers should export exactly one slice.
2786        let single = IoBufs::from(IoBuf::from(b"xy"));
2787        let mut single_dst = [IoSlice::new(&[]); 2];
2788        let count = single.chunks_vectored(&mut single_dst);
2789        assert_eq!(count, 1);
2790        assert_eq!(&single_dst[0][..], b"xy");
2791
2792        // Single empty buffers should export no slices.
2793        let empty_single = IoBufs::default();
2794        let mut empty_single_dst = [IoSlice::new(&[]); 1];
2795        assert_eq!(empty_single.chunks_vectored(&mut empty_single_dst), 0);
2796
2797        let bufs = IoBufs::from(vec![
2798            IoBuf::from(b"ab"),
2799            IoBuf::from(b"cd"),
2800            IoBuf::from(b"ef"),
2801            IoBuf::from(b"gh"),
2802        ]);
2803
2804        // Destination capacity should cap how many chunks we export.
2805        let mut small = [IoSlice::new(&[]); 2];
2806        let count = bufs.chunks_vectored(&mut small);
2807        assert_eq!(count, 2);
2808        assert_eq!(&small[0][..], b"ab");
2809        assert_eq!(&small[1][..], b"cd");
2810
2811        // Larger destination should include every readable chunk.
2812        let mut large = [IoSlice::new(&[]); 8];
2813        let count = bufs.chunks_vectored(&mut large);
2814        assert_eq!(count, 4);
2815        assert_eq!(&large[0][..], b"ab");
2816        assert_eq!(&large[1][..], b"cd");
2817        assert_eq!(&large[2][..], b"ef");
2818        assert_eq!(&large[3][..], b"gh");
2819
2820        // Empty destination cannot accept any slices.
2821        let mut empty_dst: [IoSlice<'_>; 0] = [];
2822        assert_eq!(bufs.chunks_vectored(&mut empty_dst), 0);
2823
2824        // Non-canonical shapes should skip empty leading chunks.
2825        let sparse = IoBufs {
2826            inner: IoBufsInner::Pair([IoBuf::default(), IoBuf::from(b"x")]),
2827        };
2828        let mut dst = [IoSlice::new(&[]); 2];
2829        let count = sparse.chunks_vectored(&mut dst);
2830        assert_eq!(count, 1);
2831        assert_eq!(&dst[0][..], b"x");
2832
2833        // Triple should skip empty chunks and preserve readable order.
2834        let sparse_triple = IoBufs {
2835            inner: IoBufsInner::Triple([IoBuf::default(), IoBuf::from(b"y"), IoBuf::from(b"z")]),
2836        };
2837        let mut dst = [IoSlice::new(&[]); 3];
2838        let count = sparse_triple.chunks_vectored(&mut dst);
2839        assert_eq!(count, 2);
2840        assert_eq!(&dst[0][..], b"y");
2841        assert_eq!(&dst[1][..], b"z");
2842
2843        // Chunked shapes with only empty buffers should export no slices.
2844        let empty_chunked = IoBufs {
2845            inner: IoBufsInner::Chunked(VecDeque::from([IoBuf::default(), IoBuf::default()])),
2846        };
2847        let mut dst = [IoSlice::new(&[]); 2];
2848        assert_eq!(empty_chunked.chunks_vectored(&mut dst), 0);
2849    }
2850
2851    #[test]
2852    fn test_iobufsmut_freeze_chunked() {
2853        // Multiple non-empty buffers stay multi-chunk.
2854        let buf1 = IoBufMut::from(b"hello".as_ref());
2855        let buf2 = IoBufMut::from(b" world".as_ref());
2856        let bufs = IoBufsMut::from(vec![buf1, buf2]);
2857        let mut frozen = bufs.freeze();
2858        assert!(!frozen.is_single());
2859        assert_eq!(frozen.chunk(), b"hello");
2860        frozen.advance(5);
2861        assert_eq!(frozen.chunk(), b" world");
2862        frozen.advance(6);
2863        assert_eq!(frozen.remaining(), 0);
2864
2865        // Empty buffers are filtered out.
2866        let buf1 = IoBufMut::from(b"hello".as_ref());
2867        let empty = IoBufMut::default();
2868        let buf2 = IoBufMut::from(b" world".as_ref());
2869        let bufs = IoBufsMut::from(vec![buf1, empty, buf2]);
2870        let mut frozen = bufs.freeze();
2871        assert!(!frozen.is_single());
2872        assert_eq!(frozen.chunk(), b"hello");
2873        frozen.advance(5);
2874        assert_eq!(frozen.chunk(), b" world");
2875        frozen.advance(6);
2876        assert_eq!(frozen.remaining(), 0);
2877
2878        // Collapses to Single when one non-empty buffer remains
2879        let empty1 = IoBufMut::default();
2880        let buf = IoBufMut::from(b"only one".as_ref());
2881        let empty2 = IoBufMut::default();
2882        let bufs = IoBufsMut::from(vec![empty1, buf, empty2]);
2883        let frozen = bufs.freeze();
2884        assert!(frozen.is_single());
2885        assert_eq!(frozen.coalesce(), b"only one");
2886
2887        // All empty buffers -> Single with empty buffer
2888        let empty1 = IoBufMut::default();
2889        let empty2 = IoBufMut::default();
2890        let bufs = IoBufsMut::from(vec![empty1, empty2]);
2891        let frozen = bufs.freeze();
2892        assert!(frozen.is_single());
2893        assert!(frozen.is_empty());
2894    }
2895
2896    #[test]
2897    fn test_iobufsmut_from_vec() {
2898        // Empty Vec becomes Single with empty buffer
2899        let bufs = IoBufsMut::from(Vec::<IoBufMut>::new());
2900        assert!(bufs.is_single());
2901        assert!(bufs.is_empty());
2902
2903        // Vec with one element becomes Single
2904        let buf = IoBufMut::from(b"test");
2905        let bufs = IoBufsMut::from(vec![buf]);
2906        assert!(bufs.is_single());
2907        assert_eq!(bufs.chunk(), b"test");
2908
2909        // Vec with multiple elements becomes multi-chunk.
2910        let buf1 = IoBufMut::from(b"hello");
2911        let buf2 = IoBufMut::from(b" world");
2912        let bufs = IoBufsMut::from(vec![buf1, buf2]);
2913        assert!(!bufs.is_single());
2914    }
2915
2916    #[test]
2917    fn test_iobufsmut_fast_path_shapes() {
2918        let pair = IoBufsMut::from(vec![IoBufMut::from(b"a"), IoBufMut::from(b"b")]);
2919        assert!(matches!(pair.inner, IoBufsMutInner::Pair(_)));
2920
2921        let triple = IoBufsMut::from(vec![
2922            IoBufMut::from(b"a"),
2923            IoBufMut::from(b"b"),
2924            IoBufMut::from(b"c"),
2925        ]);
2926        assert!(matches!(triple.inner, IoBufsMutInner::Triple(_)));
2927
2928        let chunked = IoBufsMut::from(vec![
2929            IoBufMut::from(b"a"),
2930            IoBufMut::from(b"b"),
2931            IoBufMut::from(b"c"),
2932            IoBufMut::from(b"d"),
2933        ]);
2934        assert!(matches!(chunked.inner, IoBufsMutInner::Chunked(_)));
2935    }
2936
2937    #[test]
2938    fn test_iobufsmut_default() {
2939        // Default IoBufsMut should be a single empty chunk.
2940        let bufs = IoBufsMut::default();
2941        assert!(bufs.is_single());
2942        assert!(bufs.is_empty());
2943        assert_eq!(bufs.len(), 0);
2944    }
2945
2946    #[test]
2947    fn test_iobufsmut_from_array() {
2948        // From<[u8; N]> should create a single-chunk container with the array data.
2949        let bufs = IoBufsMut::from([1u8, 2, 3, 4, 5]);
2950        assert!(bufs.is_single());
2951        assert_eq!(bufs.len(), 5);
2952        assert_eq!(bufs.chunk(), &[1, 2, 3, 4, 5]);
2953    }
2954
2955    #[test]
2956    fn test_iobufsmut_buf_trait_chunked() {
2957        let buf1 = IoBufMut::from(b"hello");
2958        let buf2 = IoBufMut::from(b" ");
2959        let buf3 = IoBufMut::from(b"world");
2960        let mut bufs = IoBufsMut::from(vec![buf1, buf2, buf3]);
2961
2962        assert_eq!(bufs.remaining(), 11);
2963        assert_eq!(bufs.chunk(), b"hello");
2964
2965        // Advance within first buffer
2966        bufs.advance(3);
2967        assert_eq!(bufs.remaining(), 8);
2968        assert_eq!(bufs.chunk(), b"lo");
2969
2970        // Advance past first buffer (should pop_front)
2971        bufs.advance(2);
2972        assert_eq!(bufs.remaining(), 6);
2973        assert_eq!(bufs.chunk(), b" ");
2974
2975        // Advance exactly one buffer
2976        bufs.advance(1);
2977        assert_eq!(bufs.remaining(), 5);
2978        assert_eq!(bufs.chunk(), b"world");
2979
2980        // Advance to end
2981        bufs.advance(5);
2982        assert_eq!(bufs.remaining(), 0);
2983    }
2984
2985    #[test]
2986    #[should_panic(expected = "cannot advance past end of buffer")]
2987    fn test_iobufsmut_advance_past_end() {
2988        let buf1 = IoBufMut::from(b"hello");
2989        let buf2 = IoBufMut::from(b" world");
2990        let mut bufs = IoBufsMut::from(vec![buf1, buf2]);
2991        bufs.advance(20);
2992    }
2993
2994    #[test]
2995    fn test_iobufsmut_bufmut_trait_single() {
2996        let mut bufs = IoBufsMut::from(IoBufMut::with_capacity(20));
2997        assert_eq!(bufs.remaining_mut(), 20);
2998
2999        bufs.put_slice(b"hello");
3000        assert_eq!(bufs.chunk(), b"hello");
3001        assert_eq!(bufs.len(), 5);
3002        assert_eq!(bufs.remaining_mut(), 15);
3003
3004        bufs.put_slice(b" world");
3005        assert_eq!(bufs.coalesce(), b"hello world");
3006    }
3007
3008    #[test]
3009    fn test_iobufsmut_zeroed_write() {
3010        // Use zeroed buffers which have a fixed length
3011        let bufs = IoBufsMut::from(IoBufMut::zeroed(20));
3012        assert_eq!(bufs.len(), 20);
3013
3014        // Can write using as_mut on coalesced buffer
3015        let mut coalesced = bufs.coalesce();
3016        coalesced.as_mut()[..5].copy_from_slice(b"hello");
3017        assert_eq!(&coalesced.as_ref()[..5], b"hello");
3018    }
3019
3020    #[test]
3021    fn test_iobufsmut_bufmut_put_slice() {
3022        // Test writing across multiple buffers
3023        let buf1 = IoBufMut::with_capacity(5);
3024        let buf2 = IoBufMut::with_capacity(6);
3025        let mut bufs = IoBufsMut::from(vec![buf1, buf2]);
3026
3027        // Write data
3028        bufs.put_slice(b"hello");
3029        bufs.put_slice(b" world");
3030        assert_eq!(bufs.coalesce(), b"hello world");
3031    }
3032
3033    #[test]
3034    fn test_iobufs_advance_drains_buffers() {
3035        let mut bufs = IoBufs::from(IoBuf::from(b"hello"));
3036        bufs.append(IoBuf::from(b" "));
3037        bufs.append(IoBuf::from(b"world"));
3038
3039        // Advance exactly past first buffer
3040        bufs.advance(5);
3041        assert_eq!(bufs.remaining(), 6);
3042        assert_eq!(bufs.chunk(), b" ");
3043
3044        // Advance across multiple buffers
3045        bufs.advance(4);
3046        assert_eq!(bufs.remaining(), 2);
3047        assert_eq!(bufs.chunk(), b"ld");
3048    }
3049
3050    #[test]
3051    fn test_iobufs_advance_exactly_to_boundary() {
3052        let mut bufs = IoBufs::from(IoBuf::from(b"abc"));
3053        bufs.append(IoBuf::from(b"def"));
3054
3055        // Advance exactly to first buffer boundary
3056        bufs.advance(3);
3057        assert_eq!(bufs.remaining(), 3);
3058        assert_eq!(bufs.chunk(), b"def");
3059
3060        // Advance exactly to end
3061        bufs.advance(3);
3062        assert_eq!(bufs.remaining(), 0);
3063    }
3064
3065    #[test]
3066    fn test_iobufs_advance_canonicalizes_pair_to_single() {
3067        let mut bufs = IoBufs::from(IoBuf::from(b"ab"));
3068        bufs.append(IoBuf::from(b"cd"));
3069        bufs.advance(2);
3070        assert!(bufs.is_single());
3071        assert_eq!(bufs.chunk(), b"cd");
3072    }
3073
3074    #[test]
3075    fn test_iobufsmut_emptiness_with_reserved_capacity() {
3076        for count in 0..=8 {
3077            for readable in 0..=count {
3078                let mut chunks = Vec::new();
3079                for index in 0..count {
3080                    let mut chunk = IoBufMut::with_capacity(4);
3081                    if index == readable {
3082                        chunk.put_u8(7);
3083                    }
3084                    chunks.push(chunk);
3085                }
3086                let mut bufs = IoBufsMut::from(chunks);
3087                assert_eq!(bufs.is_empty(), readable == count);
3088                assert_eq!(bufs.has_remaining(), readable != count);
3089                if readable != count {
3090                    bufs.advance(1);
3091                    assert!(bufs.is_empty());
3092                    assert!(!bufs.has_remaining());
3093                }
3094                if bufs.remaining_mut() > 0 {
3095                    bufs.put_slice(&[9]);
3096                    assert!(!bufs.is_empty());
3097                    assert!(bufs.has_remaining());
3098                }
3099            }
3100        }
3101    }
3102
3103    #[test]
3104    fn test_iobufsmut_with_empty_buffers() {
3105        let buf1 = IoBufMut::from(b"hello");
3106        let buf2 = IoBufMut::default();
3107        let buf3 = IoBufMut::from(b" world");
3108        let mut bufs = IoBufsMut::from(vec![buf1, buf2, buf3]);
3109
3110        assert_eq!(bufs.remaining(), 11);
3111        assert_eq!(bufs.chunk(), b"hello");
3112
3113        // Advance past first buffer
3114        bufs.advance(5);
3115        // Empty buffer should be skipped
3116        assert_eq!(bufs.chunk(), b" world");
3117        assert_eq!(bufs.remaining(), 6);
3118    }
3119
3120    #[test]
3121    fn test_iobufsmut_advance_skips_leading_writable_empty_chunk() {
3122        // A leading chunk with capacity but no readable bytes (len == 0) should
3123        // be skipped during advance, reaching the next readable chunk.
3124        let empty_writable = IoBufMut::with_capacity(4);
3125        let payload = IoBufMut::from(b"xy");
3126        let mut bufs = IoBufsMut::from(vec![empty_writable, payload]);
3127
3128        bufs.advance(1);
3129        assert_eq!(bufs.chunk(), b"y");
3130        assert_eq!(bufs.remaining(), 1);
3131    }
3132
3133    #[test]
3134    fn test_iobufsmut_read_ops_preserve_writable_capacity() {
3135        // Draining the filled first chunk must keep the never-filled second
3136        // chunk's reserved capacity, so reads leave remaining_mut unchanged.
3137        let mut a = IoBufMut::with_capacity(8);
3138        a.put_slice(&[1u8; 8]);
3139        let b = IoBufMut::with_capacity(8);
3140        let mut bufs = IoBufsMut::from(vec![a, b]);
3141        assert_eq!(bufs.remaining(), 8);
3142        assert_eq!(bufs.remaining_mut(), 8);
3143        bufs.advance(8);
3144        assert_eq!(bufs.remaining(), 0);
3145        assert_eq!(bufs.remaining_mut(), 8);
3146        bufs.put_slice(&[2u8; 8]);
3147        assert_eq!(bufs.copy_to_bytes(8).as_ref(), &[2u8; 8]);
3148
3149        // copy_to_bytes drains must preserve capacity the same way.
3150        let mut a = IoBufMut::with_capacity(8);
3151        a.put_slice(&[3u8; 8]);
3152        let b = IoBufMut::with_capacity(8);
3153        let mut bufs = IoBufsMut::from(vec![a, b]);
3154        assert_eq!(bufs.copy_to_bytes(8).as_ref(), &[3u8; 8]);
3155        assert_eq!(bufs.remaining_mut(), 8);
3156
3157        // Zero-length drains do not disturb chunk state at all.
3158        assert!(bufs.copy_to_bytes(0).is_empty());
3159        assert_eq!(bufs.remaining_mut(), 8);
3160    }
3161
3162    #[test]
3163    fn test_iobufsmut_from_vec_u8_preserves_capacity() {
3164        // The vec's allocation is adopted and its reserved capacity is
3165        // preserved exactly.
3166        let mut vec = Vec::with_capacity(128);
3167        vec.extend_from_slice(b"abc");
3168        let cap = vec.capacity();
3169        let mut bufs = IoBufsMut::from(vec);
3170        assert_eq!(bufs.remaining(), 3);
3171        assert_eq!(bufs.capacity(), cap);
3172        bufs.put_slice(b"d");
3173        assert_eq!(bufs.copy_to_bytes(4).as_ref(), b"abcd");
3174
3175        // An empty vec keeps its full reservation, so `read_at_buf` can fill a
3176        // `Vec::with_capacity(len)` in place through `set_len(len)`.
3177        let vec = Vec::<u8>::with_capacity(64);
3178        let len = vec.capacity();
3179        let mut bufs = IoBufsMut::from(vec);
3180        assert!(bufs.is_empty());
3181        assert_eq!(bufs.capacity(), len);
3182        // SAFETY: all `len` bytes are initialized by copy_from_slice below.
3183        unsafe { bufs.set_len(len) };
3184        bufs.copy_from_slice(&vec![9u8; len]);
3185        assert_eq!(bufs.copy_to_bytes(len).as_ref(), vec![9u8; len].as_slice());
3186
3187        // Exactly-sized vec: adopted with contents intact.
3188        let bufs = IoBufsMut::from(vec![1u8, 2, 3]);
3189        assert_eq!(bufs.remaining(), 3);
3190        assert_eq!(bufs.capacity(), 3);
3191        assert_eq!(bufs.chunk(), &[1, 2, 3]);
3192    }
3193
3194    #[test]
3195    fn test_iobufsmut_coalesce_after_advance() {
3196        // Advance mid-chunk: advance 3 of 11 bytes
3197        let buf1 = IoBufMut::from(b"hello");
3198        let buf2 = IoBufMut::from(b" world");
3199        let mut bufs = IoBufsMut::from(vec![buf1, buf2]);
3200
3201        bufs.advance(3);
3202        assert_eq!(bufs.coalesce(), b"lo world");
3203
3204        // Advance to exact chunk boundary: advance 5 of 11 bytes
3205        let buf1 = IoBufMut::from(b"hello");
3206        let buf2 = IoBufMut::from(b" world");
3207        let mut bufs = IoBufsMut::from(vec![buf1, buf2]);
3208
3209        bufs.advance(5);
3210        assert_eq!(bufs.coalesce(), b" world");
3211    }
3212
3213    #[test]
3214    fn test_iobufsmut_copy_to_bytes() {
3215        let buf1 = IoBufMut::from(b"hello");
3216        let buf2 = IoBufMut::from(b" world");
3217        let mut bufs = IoBufsMut::from(vec![buf1, buf2]);
3218
3219        // First read spans chunks and leaves unread suffix.
3220        let first = bufs.copy_to_bytes(7);
3221        assert_eq!(&first[..], b"hello w");
3222        assert_eq!(bufs.remaining(), 4);
3223
3224        // Second read drains the remainder.
3225        let rest = bufs.copy_to_bytes(4);
3226        assert_eq!(&rest[..], b"orld");
3227        assert_eq!(bufs.remaining(), 0);
3228    }
3229
3230    #[test]
3231    fn test_iobufsmut_copy_to_bytes_chunked_four_plus() {
3232        let mut bufs = IoBufsMut::from(vec![
3233            IoBufMut::from(b"ab"),
3234            IoBufMut::from(b"cd"),
3235            IoBufMut::from(b"ef"),
3236            IoBufMut::from(b"gh"),
3237        ]);
3238
3239        // Exercise chunked advance path before copy_to_bytes.
3240        bufs.advance(1);
3241        assert_eq!(bufs.chunk(), b"b");
3242        bufs.advance(1);
3243        assert_eq!(bufs.chunk(), b"cd");
3244
3245        // Chunked fast-path: first chunk alone satisfies request.
3246        let first = bufs.copy_to_bytes(1);
3247        assert_eq!(&first[..], b"c");
3248
3249        // Chunked slow-path: request crosses chunk boundaries.
3250        let second = bufs.copy_to_bytes(4);
3251        assert_eq!(&second[..], b"defg");
3252
3253        let rest = bufs.copy_to_bytes(1);
3254        assert_eq!(&rest[..], b"h");
3255        assert_eq!(bufs.remaining(), 0);
3256
3257        // Enter copy_to_bytes while still in chunked representation.
3258        let mut bufs = IoBufsMut::from(vec![
3259            IoBufMut::from(b"a"),
3260            IoBufMut::from(b"b"),
3261            IoBufMut::from(b"c"),
3262            IoBufMut::from(b"d"),
3263            IoBufMut::from(b"e"),
3264        ]);
3265        assert!(matches!(bufs.inner, IoBufsMutInner::Chunked(_)));
3266        let first = bufs.copy_to_bytes(1);
3267        assert_eq!(&first[..], b"a");
3268        // Stay chunked while consuming across multiple tiny chunks.
3269        let next = bufs.copy_to_bytes(3);
3270        assert_eq!(&next[..], b"bcd");
3271        assert_eq!(bufs.chunk(), b"e");
3272        assert_eq!(bufs.remaining(), 1);
3273    }
3274
3275    #[test]
3276    fn test_iobufsmut_copy_to_bytes_canonicalizes_pair() {
3277        let mut bufs = IoBufsMut::from(vec![IoBufMut::from(b"ab"), IoBufMut::from(b"cd")]);
3278        assert!(matches!(bufs.inner, IoBufsMutInner::Pair(_)));
3279
3280        let first = bufs.copy_to_bytes(2);
3281        assert_eq!(&first[..], b"ab");
3282
3283        assert!(bufs.is_single());
3284        assert_eq!(bufs.chunk(), b"cd");
3285        assert_eq!(bufs.remaining(), 2);
3286    }
3287
3288    #[test]
3289    fn test_iobufsmut_copy_from_slice_single() {
3290        let mut bufs = IoBufsMut::from(IoBufMut::zeroed(11));
3291        bufs.copy_from_slice(b"hello world");
3292        assert_eq!(bufs.coalesce(), b"hello world");
3293    }
3294
3295    #[test]
3296    fn test_iobufsmut_copy_from_slice_chunked() {
3297        let buf1 = IoBufMut::zeroed(5);
3298        let buf2 = IoBufMut::zeroed(6);
3299        let mut bufs = IoBufsMut::from(vec![buf1, buf2]);
3300
3301        bufs.copy_from_slice(b"hello world");
3302
3303        // Verify each chunk was filled correctly.
3304        assert_eq!(bufs.chunk(), b"hello");
3305        bufs.advance(5);
3306        assert_eq!(bufs.chunk(), b" world");
3307        bufs.advance(6);
3308        assert_eq!(bufs.remaining(), 0);
3309    }
3310
3311    #[test]
3312    #[should_panic(expected = "source slice length must match buffer length")]
3313    fn test_iobufsmut_copy_from_slice_wrong_length() {
3314        let mut bufs = IoBufsMut::from(IoBufMut::zeroed(5));
3315        bufs.copy_from_slice(b"hello world"); // 11 bytes into 5-byte buffer
3316    }
3317
3318    #[test]
3319    fn test_iobufsmut_matches_bytesmut_chain() {
3320        // Create three BytesMut with capacity for final content comparison.
3321        let mut bm1 = BytesMut::with_capacity(5);
3322        let mut bm2 = BytesMut::with_capacity(6);
3323        let mut bm3 = BytesMut::with_capacity(7);
3324
3325        // Create matching IoBufsMut
3326        let mut iobufs = IoBufsMut::from(vec![
3327            IoBufMut::with_capacity(5),
3328            IoBufMut::with_capacity(6),
3329            IoBufMut::with_capacity(7),
3330        ]);
3331
3332        // Fixed-capacity IoBufsMut exposes the current writable chunk.
3333        assert_eq!(iobufs.chunk_mut().len(), 5);
3334
3335        // Write some data
3336        (&mut bm1)
3337            .chain_mut(&mut bm2)
3338            .chain_mut(&mut bm3)
3339            .put_slice(b"hel");
3340        iobufs.put_slice(b"hel");
3341
3342        assert_eq!(iobufs.chunk_mut().len(), 2);
3343
3344        // Write more data
3345        (&mut bm1)
3346            .chain_mut(&mut bm2)
3347            .chain_mut(&mut bm3)
3348            .put_slice(b"lo world!");
3349        iobufs.put_slice(b"lo world!");
3350
3351        assert_eq!(iobufs.chunk_mut().len(), 6);
3352
3353        // Verify final content matches
3354        let frozen = iobufs.freeze().coalesce();
3355        let mut chain_content = bm1.to_vec();
3356        chain_content.extend_from_slice(&bm2);
3357        chain_content.extend_from_slice(&bm3);
3358        assert_eq!(frozen, chain_content.as_slice());
3359        assert_eq!(frozen, b"hello world!");
3360    }
3361
3362    #[test]
3363    fn test_iobufsmut_buf_matches_bytes_chain() {
3364        // Create pre-filled Bytes buffers
3365        let mut b1 = Bytes::from_static(b"hello");
3366        let mut b2 = Bytes::from_static(b" world");
3367        let b3 = Bytes::from_static(b"!");
3368
3369        // Create matching IoBufsMut
3370        let mut iobufs = IoBufsMut::from(vec![
3371            IoBufMut::from(b"hello"),
3372            IoBufMut::from(b" world"),
3373            IoBufMut::from(b"!"),
3374        ]);
3375
3376        // Test Buf::remaining matches
3377        let chain_remaining = b1.clone().chain(b2.clone()).chain(b3.clone()).remaining();
3378        assert_eq!(chain_remaining, iobufs.remaining());
3379
3380        // Test Buf::chunk matches
3381        let chain_chunk = b1
3382            .clone()
3383            .chain(b2.clone())
3384            .chain(b3.clone())
3385            .chunk()
3386            .to_vec();
3387        assert_eq!(chain_chunk, iobufs.chunk().to_vec());
3388
3389        // Advance and test again
3390        b1.advance(3);
3391        iobufs.advance(3);
3392
3393        let chain_remaining = b1.clone().chain(b2.clone()).chain(b3.clone()).remaining();
3394        assert_eq!(chain_remaining, iobufs.remaining());
3395
3396        let chain_chunk = b1
3397            .clone()
3398            .chain(b2.clone())
3399            .chain(b3.clone())
3400            .chunk()
3401            .to_vec();
3402        assert_eq!(chain_chunk, iobufs.chunk().to_vec());
3403
3404        // Advance past first buffer boundary into second
3405        b1.advance(2);
3406        iobufs.advance(2);
3407
3408        let chain_remaining = b1.clone().chain(b2.clone()).chain(b3.clone()).remaining();
3409        assert_eq!(chain_remaining, iobufs.remaining());
3410
3411        // Now we should be in the second buffer
3412        let chain_chunk = b1
3413            .clone()
3414            .chain(b2.clone())
3415            .chain(b3.clone())
3416            .chunk()
3417            .to_vec();
3418        assert_eq!(chain_chunk, iobufs.chunk().to_vec());
3419
3420        // Advance past second buffer boundary into third
3421        b2.advance(6);
3422        iobufs.advance(6);
3423
3424        let chain_remaining = b1.clone().chain(b2.clone()).chain(b3.clone()).remaining();
3425        assert_eq!(chain_remaining, iobufs.remaining());
3426
3427        // Now we should be in the third buffer
3428        let chain_chunk = b1.chain(b2).chain(b3).chunk().to_vec();
3429        assert_eq!(chain_chunk, iobufs.chunk().to_vec());
3430
3431        // Test copy_to_bytes
3432        let b1 = Bytes::from_static(b"hello");
3433        let b2 = Bytes::from_static(b" world");
3434        let b3 = Bytes::from_static(b"!");
3435        let mut iobufs = IoBufsMut::from(vec![
3436            IoBufMut::from(b"hello"),
3437            IoBufMut::from(b" world"),
3438            IoBufMut::from(b"!"),
3439        ]);
3440
3441        let chain_bytes = b1.chain(b2).chain(b3).copy_to_bytes(8);
3442        let iobufs_bytes = iobufs.copy_to_bytes(8);
3443        assert_eq!(chain_bytes, iobufs_bytes);
3444        assert_eq!(chain_bytes.as_ref(), b"hello wo");
3445    }
3446
3447    #[test]
3448    fn test_iobufsmut_chunks_vectored_multiple_slices() {
3449        // Single non-empty buffers should export exactly one slice.
3450        let single = IoBufsMut::from(IoBufMut::from(b"xy"));
3451        let mut single_dst = [IoSlice::new(&[]); 2];
3452        let count = single.chunks_vectored(&mut single_dst);
3453        assert_eq!(count, 1);
3454        assert_eq!(&single_dst[0][..], b"xy");
3455
3456        // Single empty buffers should export no slices.
3457        let empty_single = IoBufsMut::default();
3458        let mut empty_single_dst = [IoSlice::new(&[]); 1];
3459        assert_eq!(empty_single.chunks_vectored(&mut empty_single_dst), 0);
3460
3461        let bufs = IoBufsMut::from(vec![
3462            IoBufMut::from(b"ab"),
3463            IoBufMut::from(b"cd"),
3464            IoBufMut::from(b"ef"),
3465            IoBufMut::from(b"gh"),
3466        ]);
3467
3468        // Destination capacity should cap how many chunks we export.
3469        let mut small = [IoSlice::new(&[]); 2];
3470        let count = bufs.chunks_vectored(&mut small);
3471        assert_eq!(count, 2);
3472        assert_eq!(&small[0][..], b"ab");
3473        assert_eq!(&small[1][..], b"cd");
3474
3475        // Larger destination should include every readable chunk.
3476        let mut large = [IoSlice::new(&[]); 8];
3477        let count = bufs.chunks_vectored(&mut large);
3478        assert_eq!(count, 4);
3479        assert_eq!(&large[0][..], b"ab");
3480        assert_eq!(&large[1][..], b"cd");
3481        assert_eq!(&large[2][..], b"ef");
3482        assert_eq!(&large[3][..], b"gh");
3483
3484        // Empty destination cannot accept any slices.
3485        let mut empty_dst: [IoSlice<'_>; 0] = [];
3486        assert_eq!(bufs.chunks_vectored(&mut empty_dst), 0);
3487
3488        // Non-canonical shapes should skip empty leading chunks.
3489        let sparse = IoBufsMut {
3490            inner: IoBufsMutInner::Pair([IoBufMut::default(), IoBufMut::from(b"y")]),
3491        };
3492        let mut dst = [IoSlice::new(&[]); 2];
3493        let count = sparse.chunks_vectored(&mut dst);
3494        assert_eq!(count, 1);
3495        assert_eq!(&dst[0][..], b"y");
3496
3497        // Triple should skip empty chunks and preserve readable order.
3498        let sparse_triple = IoBufsMut {
3499            inner: IoBufsMutInner::Triple([
3500                IoBufMut::default(),
3501                IoBufMut::from(b"z"),
3502                IoBufMut::from(b"w"),
3503            ]),
3504        };
3505        let mut dst = [IoSlice::new(&[]); 3];
3506        let count = sparse_triple.chunks_vectored(&mut dst);
3507        assert_eq!(count, 2);
3508        assert_eq!(&dst[0][..], b"z");
3509        assert_eq!(&dst[1][..], b"w");
3510
3511        // Chunked shapes with only empty buffers should export no slices.
3512        let empty_chunked = IoBufsMut {
3513            inner: IoBufsMutInner::Chunked(VecDeque::from([
3514                IoBufMut::default(),
3515                IoBufMut::default(),
3516            ])),
3517        };
3518        let mut dst = [IoSlice::new(&[]); 2];
3519        assert_eq!(empty_chunked.chunks_vectored(&mut dst), 0);
3520    }
3521
3522    #[test]
3523    fn test_iobufsmut_try_into_single() {
3524        let single = IoBufsMut::from(IoBufMut::from(b"hello"));
3525        let single = single.try_into_single().expect("single expected");
3526        assert_eq!(single, b"hello");
3527
3528        let multi = IoBufsMut::from(vec![IoBufMut::from(b"ab"), IoBufMut::from(b"cd")]);
3529        let multi = multi.try_into_single().expect_err("multi expected");
3530        assert_eq!(multi.coalesce(), b"abcd");
3531    }
3532
3533    #[test]
3534    fn test_iobufsmut_freeze_after_advance() {
3535        // Partial advance: advance 3 of 11 bytes
3536        let buf1 = IoBufMut::from(b"hello");
3537        let buf2 = IoBufMut::from(b" world");
3538        let mut bufs = IoBufsMut::from(vec![buf1, buf2]);
3539
3540        bufs.advance(3);
3541        assert_eq!(bufs.len(), 8);
3542
3543        let frozen = bufs.freeze();
3544        assert_eq!(frozen.len(), 8);
3545        assert_eq!(frozen.coalesce(), b"lo world");
3546
3547        // Exact boundary advance: advance 5 of 11 bytes (first buf is 5 bytes)
3548        let buf1 = IoBufMut::from(b"hello");
3549        let buf2 = IoBufMut::from(b" world");
3550        let mut bufs = IoBufsMut::from(vec![buf1, buf2]);
3551
3552        bufs.advance(5);
3553        assert_eq!(bufs.len(), 6);
3554
3555        // First buffer should be fully consumed (empty after advance)
3556        // freeze() filters empty buffers, so result should be Single
3557        let frozen = bufs.freeze();
3558        assert!(frozen.is_single());
3559        assert_eq!(frozen.coalesce(), b" world");
3560    }
3561
3562    #[test]
3563    fn test_iobufsmut_coalesce_shapes() {
3564        // A single buffer is returned as is.
3565        let mut buf = IoBufMut::from(b"hello");
3566        let original_ptr = buf.as_mut_ptr();
3567        let bufs = IoBufsMut::from(buf);
3568        let coalesced = bufs.coalesce();
3569        assert_eq!(coalesced, b"hello");
3570        assert_eq!(coalesced.as_ref().as_ptr(), original_ptr);
3571
3572        // Pair, triple, and chunked (4+) shapes copy into one allocation of
3573        // exactly the readable length.
3574        let pair = IoBufsMut::from(vec![IoBufMut::from(b"ab"), IoBufMut::from(b"cd")]);
3575        let pair = pair.coalesce();
3576        assert_eq!(pair, b"abcd");
3577        assert_eq!(pair.capacity(), 4);
3578
3579        let triple = IoBufsMut::from(vec![
3580            IoBufMut::from(b"a"),
3581            IoBufMut::from(b"b"),
3582            IoBufMut::from(b"c"),
3583        ]);
3584        let triple = triple.coalesce();
3585        assert_eq!(triple, b"abc");
3586        assert_eq!(triple.capacity(), 3);
3587
3588        let chunked = IoBufsMut::from(vec![
3589            IoBufMut::from(b"1"),
3590            IoBufMut::from(b"2"),
3591            IoBufMut::from(b"3"),
3592            IoBufMut::from(b"4"),
3593        ]);
3594        let chunked = chunked.coalesce();
3595        assert_eq!(chunked, b"1234");
3596        assert_eq!(chunked.capacity(), 4);
3597    }
3598
3599    #[test]
3600    fn test_iobufs_additional_shape_and_conversion_paths() {
3601        // Constructor coverage for mutable/immutable/slice-backed inputs.
3602        let from_mut = IoBufs::from(IoBufMut::from(b"m"));
3603        assert_eq!(from_mut.chunk(), b"m");
3604        let from_bytes = IoBufs::from(Bytes::from_static(b"b"));
3605        assert_eq!(from_bytes.chunk(), b"b");
3606        let from_bytesmut = IoBufs::from(BytesMut::from(&b"bm"[..]));
3607        assert_eq!(from_bytesmut.chunk(), b"bm");
3608        let from_vec = IoBufs::from(vec![1u8, 2u8]);
3609        assert_eq!(from_vec.chunk(), &[1u8, 2]);
3610        let static_slice: &'static [u8] = b"slice";
3611        let from_static = IoBufs::from(static_slice);
3612        assert_eq!(from_static.chunk(), b"slice");
3613
3614        // Canonicalizing an already-empty buffer remains a single empty chunk.
3615        let mut single_empty = IoBufs::default();
3616        single_empty.canonicalize();
3617        assert!(single_empty.is_single());
3618
3619        // Triple path: prepend/append can promote into chunked while preserving order.
3620        let mut triple = IoBufs::from(vec![
3621            IoBuf::from(b"a".to_vec()),
3622            IoBuf::from(b"b".to_vec()),
3623            IoBuf::from(b"c".to_vec()),
3624        ]);
3625        assert!(triple.as_single().is_none());
3626        triple.prepend(IoBuf::from(vec![b'0']));
3627        triple.prepend(IoBuf::from(vec![b'1']));
3628        triple.append(IoBuf::from(vec![b'2']));
3629        assert_eq!(triple.copy_to_bytes(triple.remaining()).as_ref(), b"10abc2");
3630
3631        // Appending to an existing triple keeps byte order stable.
3632        let mut triple_append = IoBufs::from(vec![
3633            IoBuf::from(b"x".to_vec()),
3634            IoBuf::from(b"y".to_vec()),
3635            IoBuf::from(b"z".to_vec()),
3636        ]);
3637        triple_append.append(IoBuf::from(vec![b'w']));
3638        assert_eq!(triple_append.coalesce(), b"xyzw");
3639
3640        // coalesce on a triple should preserve contents.
3641        let triple = IoBufs::from(vec![
3642            IoBuf::from(b"a".to_vec()),
3643            IoBuf::from(b"b".to_vec()),
3644            IoBuf::from(b"c".to_vec()),
3645        ]);
3646        assert_eq!(triple.coalesce(), b"abc");
3647
3648        // Coalescing after advancing a four-chunk input preserves the remaining bytes.
3649        let mut chunked_advanced = IoBufs::from(vec![
3650            IoBuf::from(b"a".to_vec()),
3651            IoBuf::from(b"b".to_vec()),
3652            IoBuf::from(b"c".to_vec()),
3653            IoBuf::from(b"d".to_vec()),
3654        ]);
3655        assert_eq!(chunked_advanced.remaining(), 4);
3656        chunked_advanced.advance(1);
3657        assert_eq!(chunked_advanced.coalesce(), b"bcd");
3658
3659        // Non-canonical Pair/Triple/Chunked shapes should still expose the first readable chunk.
3660        let pair_second = IoBufs {
3661            inner: IoBufsInner::Pair([IoBuf::default(), IoBuf::from(vec![1u8])]),
3662        };
3663        assert_eq!(pair_second.chunk(), &[1u8]);
3664        let pair_empty = IoBufs {
3665            inner: IoBufsInner::Pair([IoBuf::default(), IoBuf::default()]),
3666        };
3667        assert_eq!(pair_empty.chunk(), b"");
3668
3669        let triple_third = IoBufs {
3670            inner: IoBufsInner::Triple([
3671                IoBuf::default(),
3672                IoBuf::default(),
3673                IoBuf::from(vec![3u8]),
3674            ]),
3675        };
3676        assert_eq!(triple_third.chunk(), &[3u8]);
3677        let triple_second = IoBufs {
3678            inner: IoBufsInner::Triple([
3679                IoBuf::default(),
3680                IoBuf::from(vec![2u8]),
3681                IoBuf::default(),
3682            ]),
3683        };
3684        assert_eq!(triple_second.chunk(), &[2u8]);
3685        let triple_empty = IoBufs {
3686            inner: IoBufsInner::Triple([IoBuf::default(), IoBuf::default(), IoBuf::default()]),
3687        };
3688        assert_eq!(triple_empty.chunk(), b"");
3689
3690        let chunked_second = IoBufs {
3691            inner: IoBufsInner::Chunked(VecDeque::from([IoBuf::default(), IoBuf::from(vec![9u8])])),
3692        };
3693        assert_eq!(chunked_second.chunk(), &[9u8]);
3694        let chunked_empty = IoBufs {
3695            inner: IoBufsInner::Chunked(VecDeque::from([IoBuf::default()])),
3696        };
3697        assert_eq!(chunked_empty.chunk(), b"");
3698    }
3699
3700    #[test]
3701    fn test_iobufsmut_additional_shape_and_conversion_paths() {
3702        // `as_single` accessors should work only for single-shape containers.
3703        let mut single = IoBufsMut::from(IoBufMut::from(b"x"));
3704        assert!(single.as_single().is_some());
3705        assert!(single.as_single_mut().is_some());
3706        single.canonicalize();
3707        assert!(single.is_single());
3708
3709        let mut pair = IoBufsMut::from(vec![IoBufMut::from(b"a"), IoBufMut::from(b"b")]);
3710        assert!(pair.as_single().is_none());
3711        assert!(pair.as_single_mut().is_none());
3712
3713        // Constructor coverage for raw vec and BytesMut sources.
3714        let from_vec = IoBufsMut::from(vec![1u8, 2u8]);
3715        assert_eq!(from_vec.chunk(), &[1u8, 2]);
3716        let from_bytesmut = IoBufsMut::from(BytesMut::from(&b"cd"[..]));
3717        assert_eq!(from_bytesmut.chunk(), b"cd");
3718
3719        // Chunked write path: set_len + copy_from_slice + freeze round-trip.
3720        let mut chunked = IoBufsMut::from(vec![
3721            IoBufMut::with_capacity(1),
3722            IoBufMut::with_capacity(1),
3723            IoBufMut::with_capacity(1),
3724            IoBufMut::with_capacity(1),
3725        ]);
3726        // SAFETY: We only write/read initialized bytes after `copy_from_slice`.
3727        unsafe { chunked.set_len(4) };
3728        chunked.copy_from_slice(b"wxyz");
3729        assert_eq!(chunked.capacity(), 4);
3730        assert_eq!(chunked.remaining(), 4);
3731        let frozen = chunked.freeze();
3732        assert_eq!(frozen.coalesce(), b"wxyz");
3733    }
3734
3735    #[test]
3736    fn test_iobufsmut_noncanonical_chunk_and_chunk_mut_paths() {
3737        fn no_spare_capacity_buf() -> IoBufMut {
3738            let mut buf = IoBufMut::with_capacity(1);
3739            let cap = buf.capacity();
3740            // SAFETY: We never read from this buffer in this helper.
3741            unsafe { buf.set_len(cap) };
3742            buf
3743        }
3744        // `chunk()` should skip empty front buffers across all shapes.
3745        let pair_second = IoBufsMut {
3746            inner: IoBufsMutInner::Pair([IoBufMut::default(), IoBufMut::from(b"b")]),
3747        };
3748        assert_eq!(pair_second.chunk(), b"b");
3749        let pair_empty = IoBufsMut {
3750            inner: IoBufsMutInner::Pair([IoBufMut::default(), IoBufMut::default()]),
3751        };
3752        assert_eq!(pair_empty.chunk(), b"");
3753
3754        let triple_third = IoBufsMut {
3755            inner: IoBufsMutInner::Triple([
3756                IoBufMut::default(),
3757                IoBufMut::default(),
3758                IoBufMut::from(b"c"),
3759            ]),
3760        };
3761        assert_eq!(triple_third.chunk(), b"c");
3762        let triple_second = IoBufsMut {
3763            inner: IoBufsMutInner::Triple([
3764                IoBufMut::default(),
3765                IoBufMut::from(b"b"),
3766                IoBufMut::default(),
3767            ]),
3768        };
3769        assert_eq!(triple_second.chunk(), b"b");
3770        let triple_empty = IoBufsMut {
3771            inner: IoBufsMutInner::Triple([
3772                IoBufMut::default(),
3773                IoBufMut::default(),
3774                IoBufMut::default(),
3775            ]),
3776        };
3777        assert_eq!(triple_empty.chunk(), b"");
3778
3779        let chunked_second = IoBufsMut {
3780            inner: IoBufsMutInner::Chunked(VecDeque::from([
3781                IoBufMut::default(),
3782                IoBufMut::from(b"d"),
3783            ])),
3784        };
3785        assert_eq!(chunked_second.chunk(), b"d");
3786        let chunked_empty = IoBufsMut {
3787            inner: IoBufsMutInner::Chunked(VecDeque::from([IoBufMut::default()])),
3788        };
3789        assert_eq!(chunked_empty.chunk(), b"");
3790
3791        // `chunk_mut()` should skip non-writable fronts and return first writable chunk.
3792        let mut pair_chunk_mut = IoBufsMut {
3793            inner: IoBufsMutInner::Pair([no_spare_capacity_buf(), IoBufMut::with_capacity(2)]),
3794        };
3795        assert!(pair_chunk_mut.chunk_mut().len() >= 2);
3796
3797        let mut pair_chunk_mut_empty = IoBufsMut {
3798            inner: IoBufsMutInner::Pair([no_spare_capacity_buf(), no_spare_capacity_buf()]),
3799        };
3800        assert_eq!(pair_chunk_mut_empty.chunk_mut().len(), 0);
3801
3802        let mut triple_chunk_mut = IoBufsMut {
3803            inner: IoBufsMutInner::Triple([
3804                no_spare_capacity_buf(),
3805                no_spare_capacity_buf(),
3806                IoBufMut::with_capacity(3),
3807            ]),
3808        };
3809        assert!(triple_chunk_mut.chunk_mut().len() >= 3);
3810        let mut triple_chunk_mut_second = IoBufsMut {
3811            inner: IoBufsMutInner::Triple([
3812                no_spare_capacity_buf(),
3813                IoBufMut::with_capacity(2),
3814                no_spare_capacity_buf(),
3815            ]),
3816        };
3817        assert!(triple_chunk_mut_second.chunk_mut().len() >= 2);
3818
3819        let mut triple_chunk_mut_empty = IoBufsMut {
3820            inner: IoBufsMutInner::Triple([
3821                no_spare_capacity_buf(),
3822                no_spare_capacity_buf(),
3823                no_spare_capacity_buf(),
3824            ]),
3825        };
3826        assert_eq!(triple_chunk_mut_empty.chunk_mut().len(), 0);
3827
3828        let mut chunked_chunk_mut = IoBufsMut {
3829            inner: IoBufsMutInner::Chunked(VecDeque::from([
3830                IoBufMut::default(),
3831                IoBufMut::with_capacity(4),
3832            ])),
3833        };
3834        assert!(chunked_chunk_mut.chunk_mut().len() >= 4);
3835
3836        let mut chunked_chunk_mut_empty = IoBufsMut {
3837            inner: IoBufsMutInner::Chunked(VecDeque::from([no_spare_capacity_buf()])),
3838        };
3839        assert_eq!(chunked_chunk_mut_empty.chunk_mut().len(), 0);
3840    }
3841
3842    #[test]
3843    fn test_iobuf_internal_chunk_helpers() {
3844        // `copy_to_bytes_chunked` drops leading empties on zero-length reads
3845        // and asks for canonicalization so the emptied deque collapses back
3846        // to the Single representation.
3847        let mut empty_with_leading = VecDeque::from([IoBuf::default()]);
3848        let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut empty_with_leading, 0, "x");
3849        assert!(bytes.is_empty());
3850        assert!(needs_canonicalize);
3851        assert!(empty_with_leading.is_empty());
3852
3853        // Fast path: front chunk can fully satisfy the request.
3854        let mut fast = VecDeque::from([
3855            IoBuf::from(b"ab".to_vec()),
3856            IoBuf::from(b"cd".to_vec()),
3857            IoBuf::from(b"ef".to_vec()),
3858            IoBuf::from(b"gh".to_vec()),
3859        ]);
3860        let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut fast, 2, "x");
3861        assert_eq!(bytes.as_ref(), b"ab");
3862        assert!(needs_canonicalize);
3863        assert_eq!(fast.front().expect("front exists").as_ref(), b"cd");
3864
3865        // Slow path: request spans multiple chunks.
3866        let mut slow = VecDeque::from([
3867            IoBuf::from(b"a".to_vec()),
3868            IoBuf::from(b"bc".to_vec()),
3869            IoBuf::from(b"d".to_vec()),
3870            IoBuf::from(b"e".to_vec()),
3871        ]);
3872        let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut slow, 3, "x");
3873        assert_eq!(bytes.as_ref(), b"abc");
3874        assert!(needs_canonicalize);
3875
3876        let mut empty_with_leading_mut = VecDeque::from([IoBufMut::default()]);
3877        let (bytes, needs_canonicalize) =
3878            copy_to_bytes_chunked(&mut empty_with_leading_mut, 0, "x");
3879        assert!(bytes.is_empty());
3880        assert!(needs_canonicalize);
3881        assert!(empty_with_leading_mut.is_empty());
3882
3883        // Mirror the fast/slow chunked helper paths for mutable chunks too.
3884        let mut fast_mut = VecDeque::from([
3885            IoBufMut::from(b"ab"),
3886            IoBufMut::from(b"cd"),
3887            IoBufMut::from(b"ef"),
3888            IoBufMut::from(b"gh"),
3889        ]);
3890        let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut fast_mut, 2, "x");
3891        assert_eq!(bytes.as_ref(), b"ab");
3892        assert!(needs_canonicalize);
3893        assert_eq!(fast_mut.front().expect("front exists").as_ref(), b"cd");
3894
3895        let mut slow_mut = VecDeque::from([
3896            IoBufMut::from(b"a"),
3897            IoBufMut::from(b"bc"),
3898            IoBufMut::from(b"de"),
3899            IoBufMut::from(b"f"),
3900        ]);
3901        let (bytes, needs_canonicalize) = copy_to_bytes_chunked(&mut slow_mut, 4, "x");
3902        assert_eq!(bytes.as_ref(), b"abcd");
3903        assert!(needs_canonicalize);
3904        assert_eq!(slow_mut.front().expect("front exists").as_ref(), b"e");
3905
3906        // `advance_chunked_front` should skip empties and drain in linear order.
3907        let mut advance_chunked = VecDeque::from([
3908            IoBuf::default(),
3909            IoBuf::from(b"abc".to_vec()),
3910            IoBuf::from(b"d".to_vec()),
3911        ]);
3912        advance_chunked_front(&mut advance_chunked, 2);
3913        assert_eq!(
3914            advance_chunked.front().expect("front exists").as_ref(),
3915            b"c"
3916        );
3917        advance_chunked_front(&mut advance_chunked, 2);
3918        assert!(advance_chunked.is_empty());
3919
3920        // The front-advance helper also has a separate mutable monomorphization.
3921        let mut advance_chunked_mut = VecDeque::from([
3922            IoBufMut::default(),
3923            IoBufMut::from(b"abc"),
3924            IoBufMut::from(b"d"),
3925        ]);
3926        advance_chunked_front(&mut advance_chunked_mut, 2);
3927        assert_eq!(
3928            advance_chunked_mut.front().expect("front exists").as_ref(),
3929            b"c"
3930        );
3931        advance_chunked_front(&mut advance_chunked_mut, 2);
3932        assert!(advance_chunked_mut.is_empty());
3933
3934        // `advance_small_chunks` signals canonicalization when front chunks are exhausted.
3935        let mut small = [IoBuf::default(), IoBuf::from(b"abc".to_vec())];
3936        let needs_canonicalize = advance_small_chunks(&mut small, 2);
3937        assert!(needs_canonicalize);
3938        assert_eq!(small[1].as_ref(), b"c");
3939
3940        let mut small_exact = [
3941            IoBuf::from(b"a".to_vec()),
3942            IoBuf::from(b"b".to_vec()),
3943            IoBuf::from(b"c".to_vec()),
3944        ];
3945        let needs_canonicalize = advance_small_chunks(&mut small_exact, 3);
3946        assert!(needs_canonicalize);
3947        assert_eq!(small_exact[0].remaining(), 0);
3948        assert_eq!(small_exact[1].remaining(), 0);
3949        assert_eq!(small_exact[2].remaining(), 0);
3950
3951        // Small-chunk copy canonicalization is also instantiated for mutable chunks.
3952        let mut small_mut = [
3953            IoBufMut::from(b"a"),
3954            IoBufMut::from(b"bc"),
3955            IoBufMut::from(b"d"),
3956        ];
3957        let (bytes, needs_canonicalize) = copy_to_bytes_small_chunks(&mut small_mut, 3, "x");
3958        assert_eq!(bytes.as_ref(), b"abc");
3959        assert!(needs_canonicalize);
3960        assert_eq!(small_mut[2].as_ref(), b"d");
3961
3962        // `advance_mut_in_chunks` returns whether the request fully fit in writable chunks.
3963        let mut writable = [IoBufMut::with_capacity(2), IoBufMut::with_capacity(1)];
3964        let mut remaining = 3usize;
3965        // SAFETY: We do not read from advanced bytes in this test.
3966        let all_advanced = unsafe { advance_mut_in_chunks(&mut writable, &mut remaining) };
3967        assert!(all_advanced);
3968        assert_eq!(remaining, 0);
3969
3970        // `advance_mut_in_chunks` should skip non-writable chunks.
3971        let mut full = IoBufMut::with_capacity(1);
3972        // SAFETY: We only mark initialized capacity and never read bytes.
3973        unsafe { full.set_len(full.capacity()) };
3974        let mut writable_after_full = [full, IoBufMut::with_capacity(2)];
3975        let mut remaining = 2usize;
3976        // SAFETY: We do not read from advanced bytes in this test.
3977        let all_advanced =
3978            unsafe { advance_mut_in_chunks(&mut writable_after_full, &mut remaining) };
3979        assert!(all_advanced);
3980        assert_eq!(remaining, 0);
3981
3982        let mut writable_short = [IoBufMut::with_capacity(1), IoBufMut::with_capacity(1)];
3983        let mut remaining = 3usize;
3984        // SAFETY: We do not read from advanced bytes in this test.
3985        let all_advanced = unsafe { advance_mut_in_chunks(&mut writable_short, &mut remaining) };
3986        assert!(!all_advanced);
3987        assert_eq!(remaining, 1);
3988    }
3989
3990    #[test]
3991    fn test_iobufsmut_advance_mut_success_paths() {
3992        // Pair path.
3993        let mut pair = IoBufsMut {
3994            inner: IoBufsMutInner::Pair([IoBufMut::with_capacity(2), IoBufMut::with_capacity(2)]),
3995        };
3996        // SAFETY: We only verify cursor movement (`remaining`) and do not read bytes.
3997        unsafe { pair.advance_mut(3) };
3998        assert_eq!(pair.remaining(), 3);
3999
4000        // Triple path.
4001        let mut triple = IoBufsMut {
4002            inner: IoBufsMutInner::Triple([
4003                IoBufMut::with_capacity(1),
4004                IoBufMut::with_capacity(1),
4005                IoBufMut::with_capacity(1),
4006            ]),
4007        };
4008        // SAFETY: We only verify cursor movement (`remaining`) and do not read bytes.
4009        unsafe { triple.advance_mut(2) };
4010        assert_eq!(triple.remaining(), 2);
4011
4012        // Chunked wrapped-VecDeque path.
4013        let mut wrapped = VecDeque::with_capacity(5);
4014        wrapped.push_back(IoBufMut::with_capacity(1));
4015        wrapped.push_back(IoBufMut::with_capacity(1));
4016        wrapped.push_back(IoBufMut::with_capacity(1));
4017        wrapped.push_back(IoBufMut::with_capacity(1));
4018        wrapped.push_back(IoBufMut::with_capacity(1));
4019        let _ = wrapped.pop_front();
4020        wrapped.push_back(IoBufMut::with_capacity(1));
4021        let (first, second) = wrapped.as_slices();
4022        assert!(!first.is_empty());
4023        assert!(!second.is_empty());
4024
4025        // Force `advance_mut` to consume across the wrapped second slice as well.
4026        let to_advance = first.len() + 1;
4027        let mut chunked = IoBufsMut {
4028            inner: IoBufsMutInner::Chunked(wrapped),
4029        };
4030        let before = chunked.remaining_mut();
4031        // SAFETY: We only verify cursor movement (`remaining`) and do not read bytes.
4032        unsafe { chunked.advance_mut(to_advance) };
4033        assert_eq!(chunked.remaining(), to_advance);
4034        assert_eq!(chunked.remaining_mut(), before - to_advance);
4035    }
4036
4037    #[test]
4038    fn test_iobufsmut_advance_mut_zero_noop_when_full() {
4039        fn full_chunk() -> IoBufMut {
4040            // IoBufMut capacity is fixed, so setting len to capacity leaves
4041            // remaining_mut() at zero.
4042            let mut buf = IoBufMut::with_capacity(1);
4043            let cap = buf.capacity();
4044            // SAFETY: We never read from this buffer in this test.
4045            unsafe { buf.set_len(cap) };
4046            buf
4047        }
4048
4049        // Pair path: fully-written chunks should allow advance_mut(0) as a no-op.
4050        let mut pair = IoBufsMut::from(vec![full_chunk(), full_chunk()]);
4051        assert!(matches!(pair.inner, IoBufsMutInner::Pair(_)));
4052        assert_eq!(pair.remaining_mut(), 0);
4053        let before = pair.remaining();
4054        // SAFETY: Advancing by 0 does not expose uninitialized bytes.
4055        unsafe { pair.advance_mut(0) };
4056        assert_eq!(pair.remaining(), before);
4057
4058        // Triple path: same no-op behavior.
4059        let mut triple = IoBufsMut::from(vec![full_chunk(), full_chunk(), full_chunk()]);
4060        assert!(matches!(triple.inner, IoBufsMutInner::Triple(_)));
4061        assert_eq!(triple.remaining_mut(), 0);
4062        let before = triple.remaining();
4063        // SAFETY: Advancing by 0 does not expose uninitialized bytes.
4064        unsafe { triple.advance_mut(0) };
4065        assert_eq!(triple.remaining(), before);
4066
4067        // Chunked path: 4+ fully-written chunks should also no-op.
4068        let mut chunked =
4069            IoBufsMut::from(vec![full_chunk(), full_chunk(), full_chunk(), full_chunk()]);
4070        assert!(matches!(chunked.inner, IoBufsMutInner::Chunked(_)));
4071        assert_eq!(chunked.remaining_mut(), 0);
4072        let before = chunked.remaining();
4073        // SAFETY: Advancing by 0 does not expose uninitialized bytes.
4074        unsafe { chunked.advance_mut(0) };
4075        assert_eq!(chunked.remaining(), before);
4076    }
4077
4078    #[test]
4079    #[should_panic(expected = "cannot advance past end of buffer")]
4080    fn test_iobufsmut_advance_mut_past_end_pair() {
4081        let mut pair = IoBufsMut {
4082            inner: IoBufsMutInner::Pair([IoBufMut::with_capacity(1), IoBufMut::with_capacity(1)]),
4083        };
4084        // SAFETY: Intentional panic path coverage.
4085        unsafe { pair.advance_mut(3) };
4086    }
4087
4088    #[test]
4089    #[should_panic(expected = "cannot advance past end of buffer")]
4090    fn test_iobufsmut_advance_mut_past_end_triple() {
4091        let mut triple = IoBufsMut {
4092            inner: IoBufsMutInner::Triple([
4093                IoBufMut::with_capacity(1),
4094                IoBufMut::with_capacity(1),
4095                IoBufMut::with_capacity(1),
4096            ]),
4097        };
4098        // SAFETY: Intentional panic path coverage.
4099        unsafe { triple.advance_mut(4) };
4100    }
4101
4102    #[test]
4103    #[should_panic(expected = "cannot advance past end of buffer")]
4104    fn test_iobufsmut_advance_mut_past_end_chunked() {
4105        let mut chunked = IoBufsMut {
4106            inner: IoBufsMutInner::Chunked(VecDeque::from([
4107                IoBufMut::with_capacity(1),
4108                IoBufMut::with_capacity(1),
4109                IoBufMut::with_capacity(1),
4110                IoBufMut::with_capacity(1),
4111            ])),
4112        };
4113        // SAFETY: Intentional panic path coverage.
4114        unsafe { chunked.advance_mut(5) };
4115    }
4116
4117    #[test]
4118    fn test_iobufsmut_set_len() {
4119        // SAFETY: we don't read the uninitialized bytes.
4120        unsafe {
4121            // Single buffer
4122            let mut bufs = IoBufsMut::from(IoBufMut::with_capacity(16));
4123            bufs.set_len(10);
4124            assert_eq!(bufs.len(), 10);
4125
4126            // Chunked: distributes across chunks [cap 5, cap 10], set 12 -> [5, 7]
4127            let mut bufs = IoBufsMut::from(vec![
4128                IoBufMut::with_capacity(5),
4129                IoBufMut::with_capacity(10),
4130            ]);
4131            bufs.set_len(12);
4132            assert_eq!(bufs.len(), 12);
4133            assert_eq!(bufs.chunk().len(), 5);
4134            bufs.advance(5);
4135            assert_eq!(bufs.chunk().len(), 7);
4136            bufs.advance(7);
4137            assert_eq!(bufs.remaining(), 0);
4138
4139            // Uneven capacities [3, 20, 2], set 18 -> [3, 15, 0].
4140            let mut bufs = IoBufsMut::from(vec![
4141                IoBufMut::with_capacity(3),
4142                IoBufMut::with_capacity(20),
4143                IoBufMut::with_capacity(2),
4144            ]);
4145            bufs.set_len(18);
4146            assert_eq!(bufs.chunk().len(), 3);
4147            bufs.advance(3);
4148            assert_eq!(bufs.chunk().len(), 15);
4149            bufs.advance(15);
4150            assert_eq!(bufs.remaining(), 0);
4151
4152            // Exact total capacity [4, 4], set 8 -> [4, 4]
4153            let mut bufs =
4154                IoBufsMut::from(vec![IoBufMut::with_capacity(4), IoBufMut::with_capacity(4)]);
4155            bufs.set_len(8);
4156            assert_eq!(bufs.chunk().len(), 4);
4157            bufs.advance(4);
4158            assert_eq!(bufs.chunk().len(), 4);
4159            bufs.advance(4);
4160            assert_eq!(bufs.remaining(), 0);
4161
4162            // Zero length preserves caller-provided layout.
4163            let mut bufs =
4164                IoBufsMut::from(vec![IoBufMut::with_capacity(4), IoBufMut::with_capacity(4)]);
4165            bufs.set_len(0);
4166            assert_eq!(bufs.len(), 0);
4167            assert_eq!(bufs.chunk(), b"");
4168        }
4169    }
4170
4171    #[test]
4172    #[should_panic(expected = "set_len(9) exceeds capacity(8)")]
4173    fn test_iobufsmut_set_len_overflow() {
4174        let mut bufs =
4175            IoBufsMut::from(vec![IoBufMut::with_capacity(4), IoBufMut::with_capacity(4)]);
4176        // SAFETY: this will panic before any read.
4177        unsafe { bufs.set_len(9) };
4178    }
4179
4180    #[test]
4181    fn test_encode_iobufs_matches_encode() {
4182        let value = vec![1u8, 2, 3, 4, 5, 6];
4183        assert_encode_iobufs_matches_encode(&value);
4184    }
4185
4186    /// Reports a larger encoding size than `write` produces, so
4187    /// [`EncodeExt::encode_iobufs`] fails its size assert.
4188    struct UnderWriter;
4189
4190    impl Write for UnderWriter {
4191        fn write(&self, buf: &mut impl BufMut) {
4192            buf.put_slice(b"ab");
4193        }
4194    }
4195
4196    impl EncodeSize for UnderWriter {
4197        fn encode_size(&self) -> usize {
4198            4
4199        }
4200    }
4201
4202    #[test]
4203    #[should_panic(expected = "write_bufs() must write exactly encode_size() bytes")]
4204    fn test_encode_iobufs_rejects_short_write() {
4205        let _ = UnderWriter.encode_iobufs();
4206    }
4207
4208    #[test]
4209    fn test_iobuf_encode_iobufs_matches_encode() {
4210        let value = IoBuf::from(vec![0xAB; 512]);
4211        assert_encode_iobufs_matches_encode(&value);
4212    }
4213
4214    /// A pushed native [`IoBuf`] becomes its own chunk as is: once the caller's
4215    /// handle is gone, the chunk recovers the original allocation as an [`IoBufMut`].
4216    #[test]
4217    fn test_encode_iobufs_keeps_pushed_iobuf() {
4218        let value = IoBuf::from(vec![0xAB; 512]);
4219        let ptr = value.as_ptr();
4220        let mut encoded = value.encode_iobufs();
4221        drop(value);
4222
4223        let prefix = 512usize.encode_size();
4224        drop(encoded.split_to(prefix));
4225        let chunk = encoded
4226            .try_into_single()
4227            .expect("the payload must be one chunk");
4228        let recovered = chunk
4229            .try_into_mut()
4230            .expect("the chunk must be the pushed native buffer");
4231        assert_eq!(recovered.as_ref().as_ptr(), ptr);
4232    }
4233
4234    #[test]
4235    fn test_nested_container_encode_iobufs_matches_encode() {
4236        let value = (
4237            Some(Bytes::from(vec![0xAA; 256])),
4238            vec![Bytes::from(vec![0xBB; 128]), Bytes::from(vec![0xCC; 64])],
4239        );
4240        assert_encode_iobufs_matches_encode(&value);
4241    }
4242
4243    #[test]
4244    fn test_map_encode_iobufs_matches_encode() {
4245        let mut btree = BTreeMap::new();
4246        btree.insert(2u8, Bytes::from(vec![0xDD; 96]));
4247        btree.insert(1u8, Bytes::from(vec![0xEE; 48]));
4248        assert_encode_iobufs_matches_encode(&btree);
4249
4250        let mut hash = HashMap::new();
4251        hash.insert(2u8, Bytes::from(vec![0x11; 96]));
4252        hash.insert(1u8, Bytes::from(vec![0x22; 48]));
4253        assert_encode_iobufs_matches_encode(&hash);
4254    }
4255
4256    #[test]
4257    fn test_lazy_encode_iobufs_matches_encode() {
4258        let value = Lazy::new(Bytes::from(vec![0x44; 200]));
4259        assert_encode_iobufs_matches_encode(&value);
4260
4261        // A deferred value pushes its pending encoding as the only chunk and
4262        // writes no inline bytes.
4263        let mut pending = 42u64.encode();
4264        let deferred = Lazy::<u64>::deferred(&mut pending, ());
4265        assert_eq!(deferred.encode_inline_size(), 0);
4266        assert_encode_iobufs_matches_encode(&deferred);
4267        assert_eq!(deferred.encode_iobufs().chunk_count(), 1);
4268    }
4269
4270    /// A value with neither inline bytes nor pushed chunks encodes without
4271    /// allocating.
4272    #[test]
4273    fn test_encode_iobufs_without_inline_bytes_does_not_allocate() {
4274        let (bufs, calls) = allocation_calls_during(|| [0u8; 0].encode_iobufs());
4275        assert_eq!(calls, 0);
4276        assert_eq!(bufs.remaining(), 0);
4277        assert_eq!(bufs.chunk_count(), 0);
4278    }
4279
4280    #[test]
4281    fn test_non_empty_range_encode_iobufs_matches_encode() {
4282        let range =
4283            NonEmptyRange::new(Bytes::from(vec![0x10; 32])..Bytes::from(vec![0x20; 48])).unwrap();
4284        assert_encode_iobufs_matches_encode(&range);
4285    }
4286
4287    mod builder_tests {
4288        use super::*;
4289        use commonware_codec::{BufsMut, Encode, Write};
4290
4291        // Only inline writes, no pushes.
4292        #[test]
4293        fn test_inline_only() {
4294            let mut b = Builder::new(64);
4295            b.put_u32(42);
4296            b.put_u8(7);
4297            let mut r = b.finish();
4298            assert_eq!(r.remaining(), 5);
4299            assert_eq!(r.get_u32(), 42);
4300            assert_eq!(r.get_u8(), 7);
4301        }
4302
4303        // Only zero-copy pushes, no inline writes.
4304        #[test]
4305        fn test_push_only() {
4306            let mut b = Builder::new(64);
4307            let data = Bytes::from(vec![0xAA; 1024]);
4308            b.push(data.clone());
4309            let mut r = b.finish();
4310            assert_eq!(r.remaining(), 1024);
4311            assert_eq!(r.copy_to_bytes(1024), data);
4312        }
4313
4314        // Pushed Bytes appear in the output without a payload copy.
4315        #[test]
4316        fn test_push_is_zero_copy() {
4317            let mut b = Builder::new(64);
4318            b.put_u16(99);
4319            let payload = Bytes::from(vec![0xDD; 1024]);
4320            b.push(payload.clone());
4321            let mut r = b.finish();
4322            r.advance(2);
4323            assert_eq!(r.chunk().as_ptr(), payload.as_ptr());
4324        }
4325
4326        // Interleaved: inline header, zero-copy push, inline trailer.
4327        #[test]
4328        fn test_inline_push_inline() {
4329            let mut b = Builder::new(64);
4330            b.put_u16(99);
4331            let payload = Bytes::from(vec![0xBB; 512]);
4332            b.push(payload.clone());
4333            b.put_u8(1);
4334            let mut r = b.finish();
4335            assert_eq!(r.remaining(), 2 + 512 + 1);
4336            assert_eq!(r.get_u16(), 99);
4337            assert_eq!(r.copy_to_bytes(512), payload);
4338            assert_eq!(r.get_u8(), 1);
4339        }
4340
4341        // Bytes::write_bufs produces identical wire format to Bytes::write.
4342        #[test]
4343        fn test_write_bufs_matches_write() {
4344            let data = Bytes::from(vec![0xCC; 256]);
4345            let mut b = Builder::new(64);
4346            data.write_bufs(&mut b);
4347            let mut bufs = b.finish();
4348
4349            let mut out = vec![0u8; bufs.remaining()];
4350            bufs.copy_to_slice(&mut out);
4351            assert_eq!(out, data.encode().as_ref());
4352        }
4353
4354        // Finishing an unused builder produces empty IoBufs.
4355        #[test]
4356        fn test_empty() {
4357            let bufs = Builder::new(64).finish();
4358            assert_eq!(bufs.remaining(), 0);
4359        }
4360
4361        // Inline writes exceeding capacity panic. `Builder` does not override
4362        // the `BufMut` write methods, so the panic (and message) comes from
4363        // bytes' trait defaults checking `remaining_mut`.
4364        #[test]
4365        #[should_panic(expected = "advance out of bounds")]
4366        fn test_inline_overflow_panics() {
4367            let mut b = Builder::new(1);
4368            let cap = b.remaining_mut();
4369            b.put_slice(&vec![0xFF; cap]);
4370            b.put_u8(1); // exceeds capacity
4371        }
4372
4373        // Pushing empty Bytes is a no-op.
4374        #[test]
4375        fn test_empty_push_ignored() {
4376            let mut b = Builder::new(64);
4377            b.push(Bytes::new());
4378            b.put_u8(1);
4379            let bufs = b.finish();
4380            assert_eq!(bufs.remaining(), 1);
4381        }
4382
4383        // Consecutive pushes without inline writes between them.
4384        #[test]
4385        fn test_multiple_pushes() {
4386            let mut b = Builder::new(64);
4387            let a = Bytes::from(vec![0xAA; 100]);
4388            let c = Bytes::from(vec![0xCC; 200]);
4389            b.push(a.clone());
4390            b.push(c.clone());
4391            let mut r = b.finish();
4392            assert_eq!(r.remaining(), 300);
4393            assert_eq!(r.copy_to_bytes(100), a);
4394            assert_eq!(r.copy_to_bytes(200), c);
4395        }
4396
4397        // A zero-capacity builder allocates no inline buffer and emits the
4398        // pushed chunks alone, with no empty inline slices between them.
4399        #[test]
4400        fn test_zero_capacity_push_only() {
4401            let (bufs, calls) = allocation_calls_during(|| Builder::new(0).finish());
4402            assert_eq!(calls, 0);
4403            assert!(bufs.is_empty());
4404
4405            let mut b = Builder::new(0);
4406            assert_eq!(b.remaining_mut(), 0);
4407            let a = Bytes::from(vec![0xAA; 100]);
4408            let c = Bytes::from(vec![0xCC; 200]);
4409            b.push(a.clone());
4410            b.push(Bytes::new());
4411            b.push(c.clone());
4412            let mut r = b.finish();
4413            assert_eq!(r.chunk_count(), 2);
4414            assert_eq!(r.chunk().as_ptr(), a.as_ptr());
4415            assert_eq!(r.copy_to_bytes(100), a);
4416            assert_eq!(r.chunk().as_ptr(), c.as_ptr());
4417            assert_eq!(r.copy_to_bytes(200), c);
4418        }
4419
4420        // The inline segments around a push share the inline buffer, so
4421        // finishing allocates only that buffer's owner record.
4422        #[test]
4423        fn test_inline_segments_share_one_record() {
4424            let mut b = Builder::new(64);
4425            b.put_slice(b"head");
4426            let pushed = IoBuf::from(b"pushed");
4427            b.push(pushed.clone());
4428            b.put_slice(b"tail");
4429            let (mut r, calls) = allocation_calls_during(|| b.finish());
4430            assert_eq!(calls, 1);
4431            assert_eq!(r.chunk_count(), 3);
4432            let head = r.chunk().as_ptr();
4433            assert_eq!(r.copy_to_bytes(4), &b"head"[..]);
4434            assert_eq!(r.chunk().as_ptr(), pushed.as_ptr());
4435            assert_eq!(r.copy_to_bytes(6), &b"pushed"[..]);
4436
4437            // The tail segment starts right after the head in the same buffer.
4438            assert_eq!(r.chunk().as_ptr(), head.wrapping_add(4));
4439            assert_eq!(r.copy_to_bytes(4), &b"tail"[..]);
4440        }
4441
4442        // put() exceeding capacity panics (in bytes' trait default).
4443        #[test]
4444        #[should_panic(expected = "advance out of bounds")]
4445        fn test_put_exceeding_capacity_panics() {
4446            let mut b = Builder::new(1);
4447            let cap = b.remaining_mut();
4448            let src = Bytes::from(vec![0xAB; cap + 1]);
4449            b.put(src);
4450        }
4451
4452        // put_slice() exceeding capacity panics (in bytes' trait default).
4453        #[test]
4454        #[should_panic(expected = "advance out of bounds")]
4455        fn test_put_slice_exceeding_capacity_panics() {
4456            let mut b = Builder::new(1);
4457            let cap = b.remaining_mut();
4458            b.put_slice(&vec![0xFE; cap + 1]);
4459        }
4460
4461        // Simulates a multi-field struct: [u16 | Bytes (via push) | u32].
4462        // Verifies write_bufs produces identical wire format to write.
4463        #[test]
4464        fn test_multi_field_struct_equivalence() {
4465            let header: u16 = 0xCAFE;
4466            let payload = Bytes::from(vec![0xDD; 1024]);
4467            let trailer: u32 = 0xDEADBEEF;
4468
4469            // Flat encoding via write.
4470            let size = header.encode_size() + payload.encode_size() + trailer.encode_size();
4471            let mut flat = BytesMut::with_capacity(size);
4472            header.write(&mut flat);
4473            payload.write(&mut flat);
4474            trailer.write(&mut flat);
4475
4476            // Multi-buffer encoding via write_bufs.
4477            let mut b = Builder::new(64);
4478            header.write(&mut b);
4479            payload.write_bufs(&mut b);
4480            trailer.write(&mut b);
4481            let mut bufs = b.finish();
4482
4483            let mut out = vec![0u8; bufs.remaining()];
4484            bufs.copy_to_slice(&mut out);
4485            assert_eq!(out, flat.as_ref());
4486        }
4487
4488        // encode_iobufs (Builder path) matches encode (flat BytesMut path).
4489        #[test]
4490        fn test_encode_iobufs_matches_encode() {
4491            let data = Bytes::from(vec![0xEE; 500]);
4492            let mut encoded = data.encode_iobufs();
4493            let baseline = data.encode();
4494            let mut out = vec![0u8; encoded.remaining()];
4495            encoded.copy_to_slice(&mut out);
4496            assert_eq!(out, baseline.as_ref());
4497        }
4498
4499        // Exercise remaining_mut, chunk_mut, and advance_mut directly.
4500        #[test]
4501        fn test_chunk_mut_and_advance_mut() {
4502            let mut b = Builder::new(64);
4503            let initial = b.remaining_mut();
4504            assert!(initial >= 64);
4505            let chunk = b.chunk_mut();
4506            chunk[0..1].copy_from_slice(&[0xAB]);
4507            // SAFETY: We just wrote 1 byte into chunk_mut above.
4508            unsafe { b.advance_mut(1) };
4509            assert_eq!(b.remaining_mut(), initial - 1);
4510            let mut r = b.finish();
4511            assert_eq!(r.remaining(), 1);
4512            assert_eq!(r.get_u8(), 0xAB);
4513        }
4514
4515        // Writing past a full buffer panics (fixed capacity, with the panic coming
4516        // from bytes' trait default).
4517        #[test]
4518        #[should_panic(expected = "advance out of bounds")]
4519        fn test_write_past_full_panics() {
4520            let mut b = Builder::new(1);
4521            let cap = b.remaining_mut();
4522            b.put_slice(&vec![0xFF; cap]); // fill the buffer completely
4523            assert_eq!(b.remaining_mut(), 0);
4524            b.put_u8(0x42); // panics
4525        }
4526
4527        // Push at offset 0 with inline trailer exercises finish branch
4528        // where offset == pos (no inline prefix before push).
4529        #[test]
4530        fn test_push_at_start_with_trailer() {
4531            let mut b = Builder::new(64);
4532            let payload = Bytes::from(vec![0xCC; 32]);
4533            b.push(payload.clone());
4534            b.put_u8(0x01);
4535            let mut r = b.finish();
4536            assert_eq!(r.remaining(), 33);
4537            assert_eq!(r.copy_to_bytes(32), payload);
4538            assert_eq!(r.get_u8(), 0x01);
4539        }
4540    }
4541}