bytes/buf/
take.rs

1use crate::{Buf, Bytes};
2
3use core::cmp;
4
5/// A `Buf` adapter which limits the bytes read from an underlying buffer.
6///
7/// This struct is generally created by calling `take()` on `Buf`. See
8/// documentation of [`take()`](Buf::take) for more details.
9#[derive(Debug)]
10pub struct Take<T> {
11    inner: T,
12    limit: usize,
13}
14
15pub fn new<T>(inner: T, limit: usize) -> Take<T> {
16    Take { inner, limit }
17}
18
19impl<T> Take<T> {
20    /// Consumes this `Take`, returning the underlying value.
21    ///
22    /// # Examples
23    ///
24    /// ```rust
25    /// use bytes::{Buf, BufMut};
26    ///
27    /// let mut buf = b"hello world".take(2);
28    /// let mut dst = vec![];
29    ///
30    /// dst.put(&mut buf);
31    /// assert_eq!(*dst, b"he"[..]);
32    ///
33    /// let mut buf = buf.into_inner();
34    ///
35    /// dst.clear();
36    /// dst.put(&mut buf);
37    /// assert_eq!(*dst, b"llo world"[..]);
38    /// ```
39    pub fn into_inner(self) -> T {
40        self.inner
41    }
42
43    /// Gets a reference to the underlying `Buf`.
44    ///
45    /// It is inadvisable to directly read from the underlying `Buf`.
46    ///
47    /// # Examples
48    ///
49    /// ```rust
50    /// use bytes::Buf;
51    ///
52    /// let buf = b"hello world".take(2);
53    ///
54    /// assert_eq!(11, buf.get_ref().remaining());
55    /// ```
56    pub fn get_ref(&self) -> &T {
57        &self.inner
58    }
59
60    /// Gets a mutable reference to the underlying `Buf`.
61    ///
62    /// It is inadvisable to directly read from the underlying `Buf`.
63    ///
64    /// # Examples
65    ///
66    /// ```rust
67    /// use bytes::{Buf, BufMut};
68    ///
69    /// let mut buf = b"hello world".take(2);
70    /// let mut dst = vec![];
71    ///
72    /// buf.get_mut().advance(2);
73    ///
74    /// dst.put(&mut buf);
75    /// assert_eq!(*dst, b"ll"[..]);
76    /// ```
77    pub fn get_mut(&mut self) -> &mut T {
78        &mut self.inner
79    }
80
81    /// Returns the maximum number of bytes that can be read.
82    ///
83    /// # Note
84    ///
85    /// If the inner `Buf` has fewer bytes than indicated by this method then
86    /// that is the actual number of available bytes.
87    ///
88    /// # Examples
89    ///
90    /// ```rust
91    /// use bytes::Buf;
92    ///
93    /// let mut buf = b"hello world".take(2);
94    ///
95    /// assert_eq!(2, buf.limit());
96    /// assert_eq!(b'h', buf.get_u8());
97    /// assert_eq!(1, buf.limit());
98    /// ```
99    pub fn limit(&self) -> usize {
100        self.limit
101    }
102
103    /// Sets the maximum number of bytes that can be read.
104    ///
105    /// # Note
106    ///
107    /// If the inner `Buf` has fewer bytes than `lim` then that is the actual
108    /// number of available bytes.
109    ///
110    /// # Examples
111    ///
112    /// ```rust
113    /// use bytes::{Buf, BufMut};
114    ///
115    /// let mut buf = b"hello world".take(2);
116    /// let mut dst = vec![];
117    ///
118    /// dst.put(&mut buf);
119    /// assert_eq!(*dst, b"he"[..]);
120    ///
121    /// dst.clear();
122    ///
123    /// buf.set_limit(3);
124    /// dst.put(&mut buf);
125    /// assert_eq!(*dst, b"llo"[..]);
126    /// ```
127    pub fn set_limit(&mut self, lim: usize) {
128        self.limit = lim
129    }
130}
131
132impl<T: Buf> Buf for Take<T> {
133    fn remaining(&self) -> usize {
134        cmp::min(self.inner.remaining(), self.limit)
135    }
136
137    fn chunk(&self) -> &[u8] {
138        let bytes = self.inner.chunk();
139        &bytes[..cmp::min(bytes.len(), self.limit)]
140    }
141
142    fn advance(&mut self, cnt: usize) {
143        assert!(cnt <= self.limit);
144        self.inner.advance(cnt);
145        self.limit -= cnt;
146    }
147
148    fn copy_to_bytes(&mut self, len: usize) -> Bytes {
149        assert!(len <= self.remaining(), "`len` greater than remaining");
150
151        let r = self.inner.copy_to_bytes(len);
152        self.limit -= len;
153        r
154    }
155}