diff --git a/src/buf/chain.rs b/src/buf/chain.rs index 97ac2eca5..b4227b0f3 100644 --- a/src/buf/chain.rs +++ b/src/buf/chain.rs @@ -1,5 +1,5 @@ use crate::buf::{IntoIter, UninitSlice}; -use crate::{Buf, BufMut, Bytes}; +use crate::{Buf, BufMut, Bytes, SeekBuf}; #[cfg(feature = "std")] use std::io::IoSlice; @@ -188,6 +188,28 @@ where } } +impl SeekBuf for Chain +where + T: SeekBuf, + U: SeekBuf, +{ + fn chunk_from(&self, start: usize) -> Option<&[u8]> { + if start < self.a.remaining() { + self.a.chunk_from(start) + } else { + self.b.chunk_from(start - self.a.remaining()) + } + } + + fn chunk_to(&self, end: usize) -> Option<&[u8]> { + if end <= self.a.remaining() { + self.a.chunk_to(end) + } else { + self.b.chunk_to(end - self.a.remaining()) + } + } +} + unsafe impl BufMut for Chain where T: BufMut, diff --git a/src/buf/cursor.rs b/src/buf/cursor.rs new file mode 100644 index 000000000..a554abb6d --- /dev/null +++ b/src/buf/cursor.rs @@ -0,0 +1,548 @@ +use crate::buf::SeekBuf; +use crate::Buf; + +use core::cell::Cell; +use core::iter::FusedIterator; +use core::num::NonZeroUsize; +use core::ops::{Bound, Range, RangeBounds}; + +/// Provides a non-mutating view of the bytes in a buffer. +/// +/// [BufCursor] implements [Iterator] for `&u8` items, and provides +/// compatibility for the full suite of compatible iterator helpers for any +/// buffers that support arbitrary position reads via [SeekBuf]. +/// +/// A cursor closely resembles a buffer itself; in fact, it also implements +/// [Buf] and [SeekBuf] like its parent. This also makes recursive cursor +/// instantiation possible, so sub-cursors can be consumed without mutating +/// the original cursor or original buffer. +/// +/// [BufCursor] aims for optimal performance even for buffers which may +/// introduce latency for retrieving chunks (such as `dyn SeekBuf` buffers). It +/// stores the most recently retrieved chunk for subsequent access up to the +/// length of that chunk, amortizing the cost of any buffer retrieval calls +/// during iteration across the size of returned chunks. +#[derive(Debug)] +pub struct BufCursor<'b, B: SeekBuf + ?Sized> { + buf: &'b B, + + // Offset from the start of the buffer to the end of the front chunk. + // May overlap with the back chunk. + front_chunk_offset: Cell, + + // Offset from the start of the buffer to the start of the back chunk. + // May overlap with the front chunk. + back_chunk_offset: Cell, + + front_chunk: Cell>, + back_chunk: Cell>, +} + +impl<'b, B: SeekBuf + ?Sized> BufCursor<'b, B> { + /// Creates a new [BufCursor] starting at the beginning of the provided + /// buffer and ending at the end of the buffer, as determined by the length + /// provided by [Buf::remaining]. + /// + /// # Examples + /// + /// ``` + /// use bytes::buf::BufCursor; + /// + /// let buf = b"Hello World!".as_slice(); + /// + /// let mut cursor = BufCursor::new(&buf); + /// + /// assert_eq!(cursor.next(), Some(&b'H')); + /// assert_eq!(cursor.next_back(), Some(&b'!')); + /// ``` + pub fn new(buf: &'b B) -> Self { + Self { + buf, + front_chunk_offset: Cell::new(0), + back_chunk_offset: Cell::new(buf.remaining()), + front_chunk: Cell::new(None), + back_chunk: Cell::new(None), + } + } + + /// Returns the absolute cursor position within the original buffer as a + /// range of bytes. + /// + /// # Notes + /// + /// This method may allow callers to understand the layout of the + /// underlying buffer while only provided with a sub-cursor. + /// + /// # Examples + /// + /// ``` + /// use bytes::SeekBufExt; + /// + /// let buf = b"Hello World!".as_slice(); + /// + /// let mut cursor = buf.cursor().seek(4..8).unwrap(); + /// + /// assert_eq!(cursor.next(), Some(&b'o')); + /// assert_eq!(cursor.next_back(), Some(&b'o')); + /// + /// assert_eq!(cursor.cursor_position(), 5..7); + /// ``` + #[inline] + pub fn cursor_position(&self) -> Range { + self.front_offset()..self.back_offset() + } + + /// Moves the cursor to the range specified and returns a new cursor at the + /// respective front and back offsets, consuming itself in the process. + /// + /// If the range provided moves the cursor out of its supported range, + /// then [None] is returned and the cursor (`self`) is destroyed. + /// + /// Should the current cursor need to remain valid after the call to + /// [Self::seek], call [Self::cursor] first to create a new sub-cursor at + /// the current cursor position, then call [Self::seek] on the new + /// sub-cursor. + /// + /// # Examples + /// + /// ``` + /// use bytes::SeekBufExt; + /// + /// let buf = b"<<< TEXT >>>".as_slice(); + /// + /// let cursor = buf.cursor().seek(4..8).unwrap(); + /// + /// let bytes = cursor.copied().collect::>(); + /// + /// assert_eq!(bytes.as_slice(), b"TEXT".as_slice()) + /// ``` + pub fn seek>(self, range: R) -> Option { + let remaining = self.remaining(); + + let relative_front_offset = match range.start_bound() { + Bound::Included(start_inclusive) => *start_inclusive, + // Exclusive range start bounds are unimplemented in Rust 2023, + // but my be implemented in the future. This line may be uncovered. + Bound::Excluded(start_exclusive) => *start_exclusive + 1, + Bound::Unbounded => 0, + }; + + let relative_back_offset = match range.end_bound() { + Bound::Included(end_inclusive) => *end_inclusive + 1, + Bound::Excluded(end_exclusive) => *end_exclusive, + Bound::Unbounded => remaining, + }; + + if relative_front_offset > relative_back_offset || relative_back_offset > remaining { + return None; + } + + let absolute_base_offset = self.front_offset(); + + let mut front_chunk_offset = absolute_base_offset + relative_front_offset; + let mut back_chunk_offset = absolute_base_offset + relative_back_offset; + + // Check if the existing front chunk is valid after the seek operation. + // That is, if the new `front_chunk_offset` lies between the start of + // the buffer and the existing `self.front_chunk_offset` bound. + let front_chunk = match self.front_chunk.get() { + Some(front_chunk) if front_chunk_offset < self.front_chunk_offset.get() => { + // Reuse existing front chunk after seek operation. + front_chunk_offset = self.front_chunk_offset.get(); + Some(&front_chunk[relative_front_offset..]) + } + _ => None, // New front chunk buf lookup required. + }; + + // Check if the existing back chunk is valid after the seek operation. + // That is, if the new `back_chunk_offset` lies between the end of the + // buffer and the existing `self.back_chunk_offset` bound. + let back_chunk = match self.back_chunk.get() { + Some(back_chunk) if back_chunk_offset > self.back_chunk_offset.get() => { + // Reuse existing back chunk after seek operation. + let back_chunk_remaining = back_chunk_offset - self.back_chunk_offset.get(); + back_chunk_offset = self.back_chunk_offset.get(); + Some(&back_chunk[..back_chunk_remaining]) + } + _ => None, // New back chunk buf lookup required. + }; + + Some(Self { + buf: self.buf, + front_chunk_offset: Cell::new(front_chunk_offset), + back_chunk_offset: Cell::new(back_chunk_offset), + front_chunk: Cell::new(front_chunk), + back_chunk: Cell::new(back_chunk), + }) + } + + /// Advances the cursor forward in the buffer by a set number of bytes. + /// + /// `advance_by(n)` will always return `Ok(())` if the cursor successfully + /// moves forward by `n` bytes. If the cursor front would exceed the cursor + /// back (or the buffer end), then `Err(NonZeroUsize)` with value `k` is + /// returned, where `k` is the number of bytes that were in excess of the + /// limit and were not advanced. + /// + /// Calling `advance_by(0)` can be used to opportunistically preload the + /// next front chunk from the buffer (if not already loaded) without moving + /// the cursor forward. + /// + /// # Examples + /// + /// ``` + /// use bytes::SeekBufExt; + /// + /// let buf = b"123456789".as_slice(); + /// + /// let mut cursor = buf.cursor(); + /// + /// cursor.advance_by(4).unwrap(); + /// + /// assert_eq!( + /// cursor.copied().collect::>().as_slice(), + /// b"56789".as_slice(), + /// ); + /// ``` + pub fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> { + let chunk_len = if let Some(chunk) = self.front_chunk.get() { + let chunk_len = chunk.len(); + + if n < chunk_len { + self.front_chunk.set(Some(&chunk[n..])); + return Ok(()); + } + + chunk_len + } else { + 0 + }; + + let remaining = self.remaining(); + + // Clear the current front chunk since the advance surpassed it. + self.front_chunk.set(None); + + // Invariant: Front chunk offset is always greater than or equal to the + // current chunk's length, since it is a sum of the lengths of all + // preceding chunks (including the current chunk). + debug_assert!( + self.front_chunk_offset.get() >= chunk_len, + "internal exception" + ); + + if n < remaining { + self.front_chunk_offset + .set(self.front_chunk_offset.get() + n - chunk_len); + + // Load next front chunk from buffer. + let _ = self.next_front_chunk(); + + Ok(()) + } else { + self.front_chunk_offset + .set(self.front_chunk_offset.get() + remaining - chunk_len); + + debug_assert!(self.remaining() == 0); + + match NonZeroUsize::new(n - remaining) { + None => Ok(()), + Some(n_remaining) => Err(n_remaining), + } + } + } + + /// Advances the cursor backwards in the buffer by a set number of bytes. + /// + /// `advance_back_by(n)` will always return `Ok(())` if the cursor + /// successfully moves backward by `n` bytes. If the cursor back would + /// exceed the cursor front (or the buffer start), then `Err(NonZeroUsize)` + /// with value `k` is returned, where `k` is the number of bytes that were + /// in excess of the limit and were not advanced. + /// + /// Calling `advance_back_by(0)` can be used to opportunistically preload + /// the next back chunk from the buffer (if not already loaded) without + /// moving the cursor backward. + /// + /// # Examples + /// + /// ``` + /// use bytes::SeekBufExt; + /// + /// let buf = b"123456789".as_slice(); + /// + /// let mut cursor = buf.cursor(); + /// + /// cursor.advance_back_by(4).unwrap(); + /// + /// assert_eq!( + /// cursor.copied().collect::>().as_slice(), + /// b"12345".as_slice(), + /// ); + /// ``` + pub fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> { + let chunk_len = if let Some(chunk) = self.back_chunk.get() { + let chunk_len = chunk.len(); + + if n < chunk_len { + self.back_chunk.set(Some(&chunk[..chunk_len - n])); + return Ok(()); + } + + chunk_len + } else { + 0 + }; + + let remaining = self.remaining(); + + // Clear the current back chunk since the advance surpassed it. + self.back_chunk.set(None); + + // Invariant: The back chunk offset + chunk_len is always more than the + // difference between the back offset and the front offset (remaining). + debug_assert!( + self.back_chunk_offset.get() + chunk_len >= remaining, + "internal exception" + ); + + if n < remaining { + self.back_chunk_offset + .set(self.back_chunk_offset.get() + chunk_len - n); + + // Load next back chunk from buffer. + let _ = self.next_back_chunk(); + + Ok(()) + } else { + self.back_chunk_offset + .set(self.back_chunk_offset.get() + chunk_len - remaining); + + debug_assert!(self.remaining() == 0); + + match NonZeroUsize::new(n - remaining) { + None => Ok(()), + Some(n_remaining) => Err(n_remaining), + } + } + } +} + +impl<'b, B: SeekBuf + ?Sized> BufCursor<'b, B> { + /// Returns the absolute offset from the beginning of the buffer + /// to the cursor's current front position (inclusive). + /// It can be thought of as `buf[front_offset..]`. + #[inline] + fn front_offset(&self) -> usize { + // Invariant: The front chunk size is always less than or equal to the + // front chunk offset, since the chunk offset is the total size of all + // encountered front chunks. + debug_assert!(self.front_chunk_len() <= self.front_chunk_offset.get()); + self.front_chunk_offset.get() - self.front_chunk_len() + } + + /// Returns the absolute offset from the beginning of the buffer + /// to the cursor's current back position (exclusive). + /// It can be thought of as `buf[..back_offset]`. + #[inline] + fn back_offset(&self) -> usize { + self.back_chunk_offset.get() + self.back_chunk_len() + } + + /// Current length of the active front chunk. It is not guaranteed that the + /// entire length of the chunk is valid for the current cursor, and must be + /// compared to the current back offset as well. + #[inline] + fn front_chunk_len(&self) -> usize { + self.front_chunk.get().map_or(0, |chunk| chunk.len()) + } + + /// Current length of the active back chunk. It is not guaranteed that the + /// entire length of the chunk is valid for the current cursor, and must be + /// compared to the current front offset as well. + #[inline] + fn back_chunk_len(&self) -> usize { + self.back_chunk.get().map_or(0, |chunk| chunk.len()) + } + + /// Returns the currently active front chunk, or retrieves a new front + /// chunk from the buffer if one is not active. + #[inline] + fn next_front_chunk(&self) -> Option<&'b [u8]> { + match self.front_chunk.get() { + // likely branch. + Some(chunk) if !chunk.is_empty() => Some(chunk), + _ => { + // unlikely branch. + self.load_next_front_chunk(); + self.front_chunk.get() + } + } + } + + /// Returns the currently active back chunk, or retrieves a new back + /// chunk from the buffer if one is not active. + #[inline] + fn next_back_chunk(&self) -> Option<&'b [u8]> { + match self.back_chunk.get() { + // likely branch. + Some(chunk) if !chunk.is_empty() => Some(chunk), + _ => { + // unlikely branch. + self.load_next_back_chunk(); + self.back_chunk.get() + } + } + } + + fn load_next_front_chunk(&self) { + let chunk = self.buf.chunk_from(self.front_chunk_offset.get()); + let chunk_len = chunk.map_or(0, |chunk| chunk.len()); + + self.front_chunk_offset + .set(self.front_chunk_offset.get() + chunk_len); + + self.front_chunk.set(chunk); + } + + fn load_next_back_chunk(&self) { + let chunk = self.buf.chunk_to(self.back_chunk_offset.get()); + let chunk_len = chunk.map_or(0, |chunk| chunk.len()); + + // This assertion checks that the buf is implemented correctly. + // A chunk should never exceed the back chunk offset, unless + // the buf's `remaining` method mismatched the total number of + // bytes available in the buffer when returned by `chunk` calls. + assert!( + self.back_chunk_offset.get() >= chunk_len, + "chunk length overflow" + ); + + self.back_chunk_offset + .set(self.back_chunk_offset.get() - chunk_len); + + self.back_chunk.set(chunk); + } +} + +impl<'b, B: SeekBuf + ?Sized> Buf for BufCursor<'b, B> { + #[inline] + fn remaining(&self) -> usize { + self.back_offset() - self.front_offset() + } + + #[inline] + fn chunk(&self) -> &[u8] { + self.next_front_chunk().unwrap_or(&[]) + } + + #[inline] + fn advance(&mut self, cnt: usize) { + self.advance_by(cnt).unwrap() + } +} + +impl<'b, B: SeekBuf + ?Sized> SeekBuf for BufCursor<'b, B> { + #[inline] + fn chunk_from(&self, start: usize) -> Option<&[u8]> { + let start_offset = self.front_offset() + start; + + if start_offset >= self.back_offset() { + return None; + } + + let chunk = self.buf.chunk_from(start_offset)?; + + let included_len = chunk.len().min(self.back_offset() - start_offset); + + Some(&chunk[..included_len]) + } + + #[inline] + fn chunk_to(&self, end: usize) -> Option<&[u8]> { + let end_offset = self.front_offset() + end; + + if end_offset > self.back_offset() { + return None; + } + + let chunk = self.buf.chunk_to(end_offset)?; + + let excluded_len = chunk.len().saturating_sub(end); + + Some(&chunk[excluded_len..]) + } +} + +impl<'b, B: SeekBuf + ?Sized> Iterator for BufCursor<'b, B> { + type Item = &'b u8; + + #[inline] + fn next(&mut self) -> Option { + let remaining = self.remaining(); + + if remaining == 0 { + return None; + } + + let chunk = self.next_front_chunk()?; + + let (next, front_chunk) = match chunk.len() { + // Most SeekBuf implementations will not need this line, but should + // an implementation return an empty slice chunk instead of None, + // this match case (chunk length of zero) may be hit. + 0 => (None, None), + 1 => (Some(&chunk[0]), None), + _ => (Some(&chunk[0]), Some(&chunk[1..])), + }; + + self.front_chunk.set(front_chunk); + + next + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.remaining(), Some(self.remaining())) + } +} + +impl<'b, B: SeekBuf + ?Sized> DoubleEndedIterator for BufCursor<'b, B> { + #[inline] + fn next_back(&mut self) -> Option { + if self.remaining() == 0 { + return None; + } + + let chunk = self.next_back_chunk()?; + + let (next, back_chunk) = match chunk.len() { + // Most SeekBuf implementations will not need this line, but should + // an implementation return an empty slice chunk instead of None, + // this match case (chunk length of zero) may be hit. + 0 => (None, None), + 1 => (Some(&chunk[0]), None), + len => (Some(&chunk[len - 1]), Some(&chunk[..len - 1])), + }; + + self.back_chunk.set(back_chunk); + + next + } +} + +impl<'b, B: SeekBuf + ?Sized> FusedIterator for BufCursor<'b, B> {} + +impl<'b, B: SeekBuf + ?Sized> ExactSizeIterator for BufCursor<'b, B> {} + +impl<'b, B: SeekBuf + ?Sized> Clone for BufCursor<'b, B> { + fn clone(&self) -> Self { + Self { + buf: self.buf, + front_chunk_offset: Cell::new(self.front_chunk_offset.get()), + back_chunk_offset: Cell::new(self.back_chunk_offset.get()), + front_chunk: Cell::new(self.front_chunk.get()), + back_chunk: Cell::new(self.back_chunk.get()), + } + } +} diff --git a/src/buf/mod.rs b/src/buf/mod.rs index 1bf0a47e8..eeb8ff5d4 100644 --- a/src/buf/mod.rs +++ b/src/buf/mod.rs @@ -17,10 +17,12 @@ mod buf_impl; mod buf_mut; mod chain; +mod cursor; mod iter; mod limit; #[cfg(feature = "std")] mod reader; +mod seek_buf; mod take; mod uninit_slice; mod vec_deque; @@ -30,8 +32,10 @@ mod writer; pub use self::buf_impl::Buf; pub use self::buf_mut::BufMut; pub use self::chain::Chain; +pub use self::cursor::BufCursor; pub use self::iter::IntoIter; pub use self::limit::Limit; +pub use self::seek_buf::{SeekBuf, SeekBufExt}; pub use self::take::Take; pub use self::uninit_slice::UninitSlice; diff --git a/src/buf/seek_buf.rs b/src/buf/seek_buf.rs new file mode 100644 index 000000000..8ab0f79e0 --- /dev/null +++ b/src/buf/seek_buf.rs @@ -0,0 +1,170 @@ +use crate::Buf; + +use crate::buf::cursor::BufCursor; +use alloc::boxed::Box; + +/// Read bytes from an arbitrary position within a buffer. +/// +/// This is an extension over the standard [Buf] type that allows seeking to +/// arbitrary positions within the buffer for reads. The buffer may return +/// chunks of bytes at each position of undetermined length, to allow for +/// buffer implementations that are formed using non-contiguous memory. +/// +/// It is always true that a [SeekBuf] is also a [Buf], but not always the +/// inverse. +/// +/// Many types that implement [Buf] may also implement [SeekBuf], including: +/// - `&[u8]` +/// - [alloc::collections::VecDeque] +/// +/// # Examples +/// +/// ``` +/// use bytes::{Buf, SeekBufExt}; +/// +/// let buf = b"try to find the T in the haystack".as_slice(); +/// +/// let remaining = buf.remaining(); +/// +/// assert!(buf.cursor().find(|&&b| b == b'Q').is_none()); +/// assert!(buf.cursor().find(|&&b| b == b'T').is_some()); +/// +/// // No bytes in the buffer were consumed while using the cursor. +/// assert_eq!(remaining, buf.remaining()); +/// ``` +pub trait SeekBuf: Buf { + /// Returns a chunk of unspecified length (but not exceeding + /// [Self::remaining]) starting at the specified inclusive index position. + /// + /// This method can be alternately thought of as equivalent to the + /// `[start..]` range indexing operation. + /// + /// # Implementer notes + /// + /// Implementations of [Self::chunk_from] should return [None] if the + /// `start` index is out of range for optimal performance. Implementations + /// that return values of empty slices are functionally equivalent, but may + /// hit slower code paths during use. + /// + /// Note that the `end` argument is an **inclusive** bound. + /// + /// # Examples + /// + /// ``` + /// use bytes::SeekBuf; + /// + /// let buf = b"hello world".as_slice(); + /// + /// assert_eq!(buf.chunk_from(6), Some(b"world".as_slice())); + /// assert_eq!(buf.chunk_from(100), None); + /// ``` + fn chunk_from(&self, start: usize) -> Option<&[u8]>; + + /// Returns a chunk of unspecified length (but not exceeding + /// [Self::remaining]) ending at the specified exclusive index position. + /// + /// This method can be alternately thought of as equivalent to the + /// `[..end]` range indexing operation. + /// + /// # Implementer notes + /// + /// Implementations of [Self::chunk_to] should return [None] if the `end` + /// index is out of range for optimal performance. Implementations that + /// return values of empty slices are functionally equivalent, but may hit + /// slower code paths during use. + /// + /// An identity of this function is that any call with an `end` argument of + /// zero will unconditionally return `Some(&[])`, rather than `None`, since + /// a sub-slice of zero length is a valid chunk of any buffer of any length. + /// + /// Note that the `end` argument is an **exclusive** bound. + /// + /// # Examples + /// + /// ``` + /// use bytes::SeekBuf; + /// + /// let buf = b"hello world".as_slice(); + /// + /// assert_eq!(buf.chunk_to(5), Some(b"hello".as_slice())); + /// assert_eq!(buf.chunk_to(100), None); + /// + /// // It may not be intuitive, but an identity of `chunk_to` is that when + /// // passed an `end` of zero, it will always return an empty slice instead + /// // of `None`. + /// assert_eq!([].as_slice().chunk_to(0), Some([].as_slice())); + /// ``` + fn chunk_to(&self, end: usize) -> Option<&[u8]>; +} + +/// SeekBufExt provides additional functionality for any type that implements +/// the [SeekBuf] trait. +/// +/// Methods within this trait are not implemented directly on the [SeekBuf] +/// trait in order to ensure that [SeekBuf] remains object-safe. +pub trait SeekBufExt: SeekBuf { + /// Returns a new [BufCursor] that can iterate over the current buffer. + /// [Self] is borrowed immutably while the cursor is active. + /// + /// ``` + /// use bytes::SeekBufExt; + /// + /// let buf = b"hello world".as_slice(); + /// + /// let mut cursor = buf.cursor(); + /// + /// assert_eq!(cursor.next(), Some(&b'h')); + /// assert_eq!(cursor.next(), Some(&b'e')); + /// assert_eq!(cursor.next(), Some(&b'l')); + /// assert_eq!(cursor.next(), Some(&b'l')); + /// assert_eq!(cursor.next(), Some(&b'o')); + /// + /// let mut sub_cursor = cursor.cursor(); + /// + /// assert_eq!(sub_cursor.next(), Some(&b' ')); + /// assert_eq!(cursor.next(), Some(&b' ')); + /// ``` + fn cursor(&self) -> BufCursor<'_, Self> { + BufCursor::new(self) + } +} + +impl SeekBufExt for T {} + +macro_rules! deref_forward_seek_buf { + () => { + #[inline] + fn chunk_from(&self, start: usize) -> Option<&[u8]> { + (**self).chunk_from(start) + } + + #[inline] + fn chunk_to(&self, end: usize) -> Option<&[u8]> { + (**self).chunk_to(end) + } + }; +} + +impl SeekBuf for &mut T { + deref_forward_seek_buf!(); +} + +impl SeekBuf for Box { + deref_forward_seek_buf!(); +} + +impl SeekBuf for &[u8] { + #[inline] + fn chunk_from(&self, start: usize) -> Option<&[u8]> { + self.get(start..) + } + + #[inline] + fn chunk_to(&self, end: usize) -> Option<&[u8]> { + self.get(..end) + } +} + +// The existence of this function makes the compiler catch if the SeekBuf +// trait is "object-safe" or not. +fn _assert_trait_object(_b: &dyn SeekBuf) {} diff --git a/src/buf/take.rs b/src/buf/take.rs index a16a434ee..c2cfa9bc3 100644 --- a/src/buf/take.rs +++ b/src/buf/take.rs @@ -1,4 +1,4 @@ -use crate::{Buf, Bytes}; +use crate::{Buf, Bytes, SeekBuf}; use core::cmp; @@ -153,3 +153,23 @@ impl Buf for Take { r } } + +impl SeekBuf for Take { + fn chunk_from(&self, start: usize) -> Option<&[u8]> { + if start < self.limit { + let remaining = self.limit - start; + let chunk = self.inner.chunk_from(start)?; + Some(&chunk[..chunk.len().min(remaining)]) + } else { + None + } + } + + fn chunk_to(&self, end: usize) -> Option<&[u8]> { + if end <= self.limit { + self.inner.chunk_to(end) + } else { + None + } + } +} diff --git a/src/buf/vec_deque.rs b/src/buf/vec_deque.rs index 263167e83..a32e5382a 100644 --- a/src/buf/vec_deque.rs +++ b/src/buf/vec_deque.rs @@ -1,6 +1,6 @@ use alloc::collections::VecDeque; -use super::Buf; +use super::{Buf, SeekBuf}; impl Buf for VecDeque { fn remaining(&self) -> usize { @@ -20,3 +20,29 @@ impl Buf for VecDeque { self.drain(..cnt); } } + +impl SeekBuf for VecDeque { + fn chunk_from(&self, start: usize) -> Option<&[u8]> { + let slices = self.as_slices(); + + if start < slices.0.len() { + Some(&slices.0[start..]) + } else if start - slices.0.len() < slices.1.len() { + Some(&slices.1[start - slices.0.len()..]) + } else { + None + } + } + + fn chunk_to(&self, end: usize) -> Option<&[u8]> { + let slices = self.as_slices(); + + if end <= slices.0.len() { + Some(&slices.0[..end]) + } else if end - slices.0.len() <= slices.1.len() { + Some(&slices.1[..end - slices.0.len()]) + } else { + None + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 1b3e6fc40..76639e5d3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -77,7 +77,7 @@ extern crate alloc; extern crate std; pub mod buf; -pub use crate::buf::{Buf, BufMut}; +pub use crate::buf::{Buf, BufMut, SeekBuf, SeekBufExt}; mod bytes; mod bytes_mut; diff --git a/tests/test_chain.rs b/tests/test_chain.rs index cfda6b8dc..254352345 100644 --- a/tests/test_chain.rs +++ b/tests/test_chain.rs @@ -1,6 +1,6 @@ #![warn(rust_2018_idioms)] -use bytes::{Buf, BufMut, Bytes}; +use bytes::{Buf, BufMut, Bytes, SeekBufExt}; #[cfg(feature = "std")] use std::io::IoSlice; @@ -175,3 +175,27 @@ fn chain_get_bytes() { // assert `get_bytes` did not allocate assert_eq!(cd_ptr.wrapping_offset(1), d.as_ptr()); } + +#[test] +fn chain_seek_buf_iterator() { + let buf1 = b"01234".as_slice(); + let buf2 = b"56789".as_slice(); + let buf3 = b"abcde".as_slice(); + + let chain = buf1.chain(buf2).chain(buf3); + + assert_eq!( + chain.cursor().copied().collect::>().as_slice(), + b"0123456789abcde".as_slice() + ); + + assert_eq!( + chain + .cursor() + .rev() + .copied() + .collect::>() + .as_slice(), + b"edcba9876543210".as_slice() + ); +} diff --git a/tests/test_cursor.rs b/tests/test_cursor.rs new file mode 100644 index 000000000..1a2c7d704 --- /dev/null +++ b/tests/test_cursor.rs @@ -0,0 +1,267 @@ +use bytes::{Buf, SeekBufExt}; +use std::collections::VecDeque; + +#[test] +fn test_iterator() { + let buf = b"Hello World!".as_slice(); + + let mut cursor = buf.cursor(); + + assert_eq!(cursor.next(), Some(&b'H')); + assert_eq!(cursor.next(), Some(&b'e')); + assert_eq!(cursor.next(), Some(&b'l')); + assert_eq!(cursor.next(), Some(&b'l')); + assert_eq!(cursor.next(), Some(&b'o')); + + assert_eq!(cursor.next_back(), Some(&b'!')); + assert_eq!(cursor.next_back(), Some(&b'd')); + assert_eq!(cursor.next_back(), Some(&b'l')); + assert_eq!(cursor.next_back(), Some(&b'r')); + assert_eq!(cursor.next_back(), Some(&b'o')); + assert_eq!(cursor.next_back(), Some(&b'W')); + + assert_eq!(cursor.next(), Some(&b' ')); + + assert_eq!(cursor.next(), None); + assert_eq!(cursor.next_back(), None); +} + +#[test] +fn test_clone() { + let buf = b"Hello World!".as_slice(); + + let mut cursor = buf.cursor(); + + assert_eq!(cursor.next(), Some(&b'H')); + assert_eq!(cursor.next_back(), Some(&b'!')); + + let mut cursor_clone = cursor.clone(); + + assert_eq!(cursor.next(), Some(&b'e')); + assert_eq!(cursor.next_back(), Some(&b'd')); + + assert_eq!(cursor_clone.next(), Some(&b'e')); + assert_eq!(cursor_clone.next_back(), Some(&b'd')); +} + +#[test] +fn test_seek() { + let buf = b"<<< TEXT >>>".as_slice(); + + let cursor = buf.cursor().seek(..).unwrap(); + + assert_eq!( + cursor.cursor().copied().collect::>().as_slice(), + b"<<< TEXT >>>".as_slice() + ); + + let cursor = buf.cursor().seek(4..8).unwrap(); + + assert_eq!( + cursor.cursor().copied().collect::>().as_slice(), + b"TEXT".as_slice() + ); + + let cursor = cursor.seek(0..=1).unwrap(); + + assert_eq!( + cursor.cursor().copied().collect::>().as_slice(), + b"TE".as_slice() + ); +} + +#[test] +fn test_invalid_seek() { + let buf = b"123".as_slice(); + + assert!(buf.cursor().seek(4..).is_none()); +} + +#[test] +fn test_size() { + let buf = b"123456789".as_slice(); + + let mut cursor = buf.cursor(); + + assert_eq!(cursor.size_hint(), (9, Some(9))); + + cursor.next(); + + assert_eq!(cursor.size_hint(), (8, Some(8))); +} + +#[test] +fn test_advance_by() { + let buf = b"123456789".as_slice(); + + let mut cursor = buf.cursor(); + + cursor.advance_by(4).unwrap(); + + assert_eq!( + cursor.cursor().copied().collect::>().as_slice(), + b"56789".as_slice() + ); + + cursor.advance_back_by(4).unwrap(); + + assert_eq!( + cursor.cursor().copied().collect::>().as_slice(), + b"5".as_slice() + ); +} + +#[test] +fn test_vec_deque_cursor() { + let mut buf = VecDeque::with_capacity(8); + + while buf.len() < buf.capacity() { + buf.push_back(b'0') + } + + for _ in 0..4 { + buf.pop_front(); + } + + for _ in 0..4 { + buf.push_back(b'1') + } + + let mut cursor = buf.cursor(); + + cursor.advance_by(1).unwrap(); + + assert_eq!( + cursor + .cursor() + .seek(..3) + .unwrap() + .copied() + .collect::>() + .as_slice(), + &[b'0', b'0', b'0'], + ); + + cursor.advance_back_by(1).unwrap(); + + assert_eq!( + cursor + .cursor() + .seek(buf.len() - 8..) + .unwrap() + .copied() + .collect::>() + .as_slice(), + &[b'0', b'0', b'0', b'1', b'1', b'1'], + ); + + { + // Advance forward through all remaining bytes in the cursor. + let mut cursor = cursor.cursor(); + cursor.advance_by(cursor.remaining()).unwrap(); + assert_eq!(cursor.remaining(), 0); + } + + { + // Advance backward through all remaining bytes in the cursor. + let mut cursor = cursor.cursor(); + cursor.advance_back_by(cursor.remaining()).unwrap(); + assert_eq!(cursor.remaining(), 0); + } +} + +/// PRNG implements a basic LCG random number generator with sane defaults. +struct PRNG { + state: u32, + multiplier: u32, + increment: u32, +} + +impl PRNG { + pub fn new(seed: u32) -> Self { + Self { + state: seed, + multiplier: 32310901, + increment: 12345, + } + } + + pub fn random(&mut self) -> u32 { + self.state = self + .multiplier + .overflowing_mul(self.state) + .0 + .overflowing_add(self.increment) + .0; + self.state + } +} + +#[test] +fn test_vec_deque_cursor_random_insert_and_drain() { + let mut prng = PRNG::new(7877 /* not important */); + + let mut buf = VecDeque::with_capacity(4096); + + for _ in 0..1000 { + let remaining_mut = buf.capacity() - buf.len(); + + buf.resize(buf.len() + (prng.random() as usize % remaining_mut), 0); + let _ = buf.drain(..prng.random() as usize & buf.len()); + + let cursor = buf.cursor(); + + assert_eq!( + cursor + .cursor() + .seek(..4.min(buf.len())) + .unwrap() + .copied() + .collect::>() + .as_slice(), + &[0, 0, 0, 0][..4.min(buf.len())], + ); + + if buf.len() > 16 { + let mut cursor = cursor.cursor().seek(4..12).unwrap(); + + cursor.advance_by(0).unwrap(); + cursor.advance_back_by(0).unwrap(); + + assert_eq!( + cursor + .seek(2..6) + .unwrap() + .copied() + .collect::>() + .as_slice(), + &[0, 0, 0, 0], + ); + } + + assert_eq!( + cursor + .cursor() + .rev() + .take(4) + .copied() + .collect::>() + .as_slice(), + &[0, 0, 0, 0][..4.min(buf.len())], + ); + + { + // Advance forward through all remaining bytes in the cursor. + let mut cursor = cursor.cursor(); + cursor.advance_by(cursor.remaining()).unwrap(); + assert_eq!(cursor.remaining(), 0); + } + + { + // Advance backward through all remaining bytes in the cursor. + let mut cursor = cursor.cursor(); + cursor.advance_back_by(cursor.remaining()).unwrap(); + assert_eq!(cursor.remaining(), 0); + } + } +} diff --git a/tests/test_seek_buf.rs b/tests/test_seek_buf.rs new file mode 100644 index 000000000..48b6b147d --- /dev/null +++ b/tests/test_seek_buf.rs @@ -0,0 +1,69 @@ +use bytes::{Buf, SeekBuf, SeekBufExt}; +use std::collections::VecDeque; + +#[test] +fn test_seek_buf_cursor() { + let buf = b"try to find the T in the haystack".as_slice(); + + let remaining = buf.remaining(); + + assert!(buf.cursor().find(|&&b| b == b'Q').is_none()); + assert!(buf.cursor().find(|&&b| b == b'T').is_some()); + + // No bytes in the buffer were consumed while using the cursor. + assert_eq!(remaining, buf.remaining()); +} + +#[test] +fn test_chunk_from() { + let buf = b"hello world".as_slice(); + + assert_eq!(buf.chunk_from(6), Some(b"world".as_slice())); + assert_eq!(buf.chunk_from(100), None); +} + +#[test] +fn test_chunk_to() { + let buf = b"hello world".as_slice(); + + assert_eq!(buf.chunk_to(5), Some(b"hello".as_slice())); + assert_eq!(buf.chunk_to(100), None); + + // It may not be intuitive, but an identity of `chunk_to` is that when + // passed an `end` of zero, it will always return an empty slice instead + // of `None`. + assert_eq!([].as_slice().chunk_to(0), Some([].as_slice())); +} + +#[test] +fn test_vec_deque() { + let mut buf = VecDeque::with_capacity(32); + + for i in 0..buf.capacity() { + buf.push_back((i % 256) as u8); + } + + assert_eq!(&buf.chunk()[..4], [0, 1, 2, 3].as_slice()); + assert_eq!(buf.chunk_to(2), Some([0, 1].as_slice())); + assert_eq!( + &buf.chunk_from(28).unwrap()[..4], + [28, 29, 30, 31].as_slice() + ); + + for _ in 0..16 { + buf.pop_front(); + } + + for _ in 0..15 { + buf.push_back(0); + } + + assert_eq!(&buf.chunk_from(0).unwrap()[..2], [16, 17].as_slice()); + + buf.push_back(255); + + { + let last_chunk = buf.chunk_to(buf.remaining()).unwrap(); + assert_eq!(&last_chunk[last_chunk.len() - 1..], [255].as_slice()); + } +} diff --git a/tests/test_take.rs b/tests/test_take.rs index 51df91d14..6768f4301 100644 --- a/tests/test_take.rs +++ b/tests/test_take.rs @@ -1,7 +1,7 @@ #![warn(rust_2018_idioms)] use bytes::buf::Buf; -use bytes::Bytes; +use bytes::{Bytes, SeekBuf}; #[test] fn long_take() { @@ -12,6 +12,19 @@ fn long_take() { assert_eq!(b"hello world", buf.chunk()); } +#[test] +fn take_from_seek_buf() { + let buf = b"hello world".take(5); + + assert_eq!(buf.chunk_from(0), Some(b"hello".as_slice())); + assert_eq!(buf.chunk_from(3), Some(b"lo".as_slice())); + assert_eq!(buf.chunk_to(5), Some(b"hello".as_slice())); + assert_eq!(buf.chunk_to(2), Some(b"he".as_slice())); + + assert_eq!(buf.chunk_from(6), None); + assert_eq!(buf.chunk_to(6), None); +} + #[test] fn take_copy_to_bytes() { let mut abcd = Bytes::copy_from_slice(b"abcd");