bytes/buf/
take.rs

Help
1use crate::Buf;
2
3use core::cmp;
4
5#[cfg(feature = "std")]
6use std::io::IoSlice;
7
8/// A `Buf` adapter which limits the bytes read from an underlying buffer.
9///
10/// This struct is generally created by calling `take()` on `Buf`. See
11/// documentation of [`take()`](Buf::take) for more details.
12#[derive(Debug)]
13pub struct Take<T> {
14    inner: T,
15    limit: usize,
16}
17
18pub fn new<T>(inner: T, limit: usize) -> Take<T> {
19    Take { inner, limit }
20}
21
22impl<T> Take<T> {
23    /// Consumes this `Take`, returning the underlying value.
24    ///
25    /// # Examples
26    ///
27    /// ```rust
28    /// use bytes::{Buf, BufMut};
29    ///
30    /// let mut buf = b"hello world".take(2);
31    /// let mut dst = vec![];
32    ///
33    /// dst.put(&mut buf);
34    /// assert_eq!(*dst, b"he"[..]);
35    ///
36    /// let mut buf = buf.into_inner();
37    ///
38    /// dst.clear();
39    /// dst.put(&mut buf);
40    /// assert_eq!(*dst, b"llo world"[..]);
41    /// ```
42    pub fn into_inner(self) -> T {
43        self.inner
44    }
45
46    /// Gets a reference to the underlying `Buf`.
47    ///
48    /// It is inadvisable to directly read from the underlying `Buf`.
49    ///
50    /// # Examples
51    ///
52    /// ```rust
53    /// use bytes::Buf;
54    ///
55    /// let buf = b"hello world".take(2);
56    ///
57    /// assert_eq!(11, buf.get_ref().remaining());
58    /// ```
59    pub fn get_ref(&self) -> &T {
60        &self.inner
61    }
62
63    /// Gets a mutable reference to the underlying `Buf`.
64    ///
65    /// It is inadvisable to directly read from the underlying `Buf`.
66    ///
67    /// # Examples
68    ///
69    /// ```rust
70    /// use bytes::{Buf, BufMut};
71    ///
72    /// let mut buf = b"hello world".take(2);
73    /// let mut dst = vec![];
74    ///
75    /// buf.get_mut().advance(2);
76    ///
77    /// dst.put(&mut buf);
78    /// assert_eq!(*dst, b"ll"[..]);
79    /// ```
80    pub fn get_mut(&mut self) -> &mut T {
81        &mut self.inner
82    }
83
84    /// Returns the maximum number of bytes that can be read.
85    ///
86    /// # Note
87    ///
88    /// If the inner `Buf` has fewer bytes than indicated by this method then
89    /// that is the actual number of available bytes.
90    ///
91    /// # Examples
92    ///
93    /// ```rust
94    /// use bytes::Buf;
95    ///
96    /// let mut buf = b"hello world".take(2);
97    ///
98    /// assert_eq!(2, buf.limit());
99    /// assert_eq!(b'h', buf.get_u8());
100    /// assert_eq!(1, buf.limit());
101    /// ```
102    pub fn limit(&self) -> usize {
103        self.limit
104    }
105
106    /// Sets the maximum number of bytes that can be read.
107    ///
108    /// # Note
109    ///
110    /// If the inner `Buf` has fewer bytes than `lim` then that is the actual
111    /// number of available bytes.
112    ///
113    /// # Examples
114    ///
115    /// ```rust
116    /// use bytes::{Buf, BufMut};
117    ///
118    /// let mut buf = b"hello world".take(2);
119    /// let mut dst = vec![];
120    ///
121    /// dst.put(&mut buf);
122    /// assert_eq!(*dst, b"he"[..]);
123    ///
124    /// dst.clear();
125    ///
126    /// buf.set_limit(3);
127    /// dst.put(&mut buf);
128    /// assert_eq!(*dst, b"llo"[..]);
129    /// ```
130    pub fn set_limit(&mut self, lim: usize) {
131        self.limit = lim
132    }
133}
134
135impl<T: Buf> Buf for Take<T> {
136    fn remaining(&self) -> usize {
137        cmp::min(self.inner.remaining(), self.limit)
138    }
139
140    fn chunk(&self) -> &[u8] {
141        let bytes = self.inner.chunk();
142        &bytes[..cmp::min(bytes.len(), self.limit)]
143    }
144
145    fn advance(&mut self, cnt: usize) {
146        assert!(cnt <= self.limit);
147        self.inner.advance(cnt);
148        self.limit -= cnt;
149    }
150
151    fn copy_to_bytes(&mut self, len: usize) -> crate::Bytes {
152        assert!(len <= self.remaining(), "`len` greater than remaining");
153
154        let r = self.inner.copy_to_bytes(len);
155        self.limit -= len;
156        r
157    }
158
159    #[cfg(feature = "std")]
160    fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
161        if self.limit == 0 {
162            return 0;
163        }
164
165        const LEN: usize = 16;
166        let mut slices: [IoSlice<'a>; LEN] = [
167            IoSlice::new(&[]),
168            IoSlice::new(&[]),
169            IoSlice::new(&[]),
170            IoSlice::new(&[]),
171            IoSlice::new(&[]),
172            IoSlice::new(&[]),
173            IoSlice::new(&[]),
174            IoSlice::new(&[]),
175            IoSlice::new(&[]),
176            IoSlice::new(&[]),
177            IoSlice::new(&[]),
178            IoSlice::new(&[]),
179            IoSlice::new(&[]),
180            IoSlice::new(&[]),
181            IoSlice::new(&[]),
182            IoSlice::new(&[]),
183        ];
184
185        let cnt = self
186            .inner
187            .chunks_vectored(&mut slices[..dst.len().min(LEN)]);
188        let mut limit = self.limit;
189        for (i, (dst, slice)) in dst[..cnt].iter_mut().zip(slices.iter()).enumerate() {
190            if let Some(buf) = slice.get(..limit) {
191                // SAFETY: We could do this safely with `IoSlice::advance` if we had a larger MSRV.
192                let buf = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(buf) };
193                *dst = IoSlice::new(buf);
194                return i + 1;
195            } else {
196                // SAFETY: We could do this safely with `IoSlice::advance` if we had a larger MSRV.
197                let buf = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(slice) };
198                *dst = IoSlice::new(buf);
199                limit -= slice.len();
200            }
201        }
202        cnt
203    }
204}
bytes/buf/take.rs

bytes/buf/
take.rs
