add BufMutExt::limit

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<T> {
+    inner: T,
+    limit: usize,
+}
+
+pub(super) fn new<T>(inner: T, limit: usize) -> Limit<T> {
+    Limit {
+        inner,
+        limit,
+    }
+}
+
+impl<T: BufMut> BufMut for Limit<T> {
+    fn remaining_mut(&self) -> usize {
+        cmp::min(self.inner.remaining_mut(), self.limit)
+    }
+
+    fn bytes_mut(&mut self) -> &mut [MaybeUninit<u8>] {
+        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;
+    }
+}

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<B: Buf + ?Sized> 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<Self>
+        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