diff --git a/src/buf/ext/limit.rs b/src/buf/ext/limit.rs new file mode 100644 index 000000000..b16c049c7 --- /dev/null +++ b/src/buf/ext/limit.rs @@ -0,0 +1,36 @@ +use crate::BufMut; + +use core::{cmp, mem::MaybeUninit}; + +/// A `BufMut` adapter which limits the amount of bytes that can be written +/// to an underlying buffer. +#[derive(Debug)] +pub struct Limit { + inner: T, + limit: usize, +} + +pub(super) fn new(inner: T, limit: usize) -> Limit { + Limit { + inner, + limit, + } +} + +impl BufMut for Limit { + fn remaining_mut(&self) -> usize { + cmp::min(self.inner.remaining_mut(), self.limit) + } + + fn bytes_mut(&mut self) -> &mut [MaybeUninit] { + let bytes = self.inner.bytes_mut(); + let end = cmp::min(bytes.len(), self.limit); + &mut bytes[..end] + } + + unsafe fn advance_mut(&mut self, cnt: usize) { + assert!(cnt <= self.limit); + self.inner.advance_mut(cnt); + self.limit -= cnt; + } +} diff --git a/src/buf/ext/mod.rs b/src/buf/ext/mod.rs index 10b9a34c0..0b9fdcce9 100644 --- a/src/buf/ext/mod.rs +++ b/src/buf/ext/mod.rs @@ -3,12 +3,14 @@ use super::{Buf, BufMut}; mod chain; +mod limit; #[cfg(feature = "std")] mod reader; mod take; #[cfg(feature = "std")] mod writer; +use self::limit::Limit; use self::take::Take; use self::chain::Chain; @@ -98,6 +100,25 @@ impl BufExt for B {} /// Extra methods for implementations of `BufMut`. pub trait BufMutExt: BufMut { + /// Creates an adaptor which can write at most `limit` bytes to `self`. + /// + /// # Examples + /// + /// ``` + /// use bytes::{BufMut, buf::BufMutExt}; + /// + /// let arr = &mut [0u8; 128][..]; + /// assert_eq!(arr.remaining_mut(), 128); + /// + /// let dst = arr.limit(10); + /// assert_eq!(dst.remaining_mut(), 10); + /// ``` + fn limit(self, limit: usize) -> Limit + where Self: Sized + { + limit::new(self, limit) + } + /// Creates an adaptor which implements the `Write` trait for `self`. /// /// This function returns a new value which implements `Write` by adapting