rcbytes/buf/
take.rs

1// The code in this file is heavily based on [Carl Lerche's LRU implementation](https://github.com/tokio-rs/bytes).
2//
3// MIT License
4//
5// Copyright (c) 2022 Al Liu (https://github.com/al8n/rcbytes)
6//
7// Copyright (c) 2018 Carl Lerche (https://github.com/tokio-rs/bytes)
8//
9// Permission is hereby granted, free of charge, to any person obtaining a copy
10// of this software and associated documentation files (the "Software"), to deal
11// in the Software without restriction, including without limitation the rights
12// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
13// copies of the Software, and to permit persons to whom the Software is
14// furnished to do so, subject to the following conditions:
15
16// The above copyright notice and this permission notice shall be included in all
17// copies or substantial portions of the Software.
18
19// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
20// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
21// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
22// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
23// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
24// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
25// SOFTWARE.
26use crate::{Buf, Bytes};
27
28use core::cmp;
29
30/// A `Buf` adapter which limits the bytes read from an underlying buffer.
31///
32/// This struct is generally created by calling `take()` on `Buf`. See
33/// documentation of [`take()`](trait.Buf.html#method.take) for more details.
34#[derive(Debug)]
35pub struct Take<T> {
36    inner: T,
37    limit: usize,
38}
39
40pub fn new<T>(inner: T, limit: usize) -> Take<T> {
41    Take { inner, limit }
42}
43
44impl<T> Take<T> {
45    /// Consumes this `Take`, returning the underlying value.
46    ///
47    /// # Examples
48    ///
49    /// ```rust
50    /// use rcbytes::{Buf, BufMut};
51    ///
52    /// let mut buf = b"hello world".take(2);
53    /// let mut dst = vec![];
54    ///
55    /// dst.put(&mut buf);
56    /// assert_eq!(*dst, b"he"[..]);
57    ///
58    /// let mut buf = buf.into_inner();
59    ///
60    /// dst.clear();
61    /// dst.put(&mut buf);
62    /// assert_eq!(*dst, b"llo world"[..]);
63    /// ```
64    pub fn into_inner(self) -> T {
65        self.inner
66    }
67
68    /// Gets a reference to the underlying `Buf`.
69    ///
70    /// It is inadvisable to directly read from the underlying `Buf`.
71    ///
72    /// # Examples
73    ///
74    /// ```rust
75    /// use rcbytes::Buf;
76    ///
77    /// let buf = b"hello world".take(2);
78    ///
79    /// assert_eq!(11, buf.get_ref().remaining());
80    /// ```
81    pub fn get_ref(&self) -> &T {
82        &self.inner
83    }
84
85    /// Gets a mutable reference to the underlying `Buf`.
86    ///
87    /// It is inadvisable to directly read from the underlying `Buf`.
88    ///
89    /// # Examples
90    ///
91    /// ```rust
92    /// use rcbytes::{Buf, BufMut};
93    ///
94    /// let mut buf = b"hello world".take(2);
95    /// let mut dst = vec![];
96    ///
97    /// buf.get_mut().advance(2);
98    ///
99    /// dst.put(&mut buf);
100    /// assert_eq!(*dst, b"ll"[..]);
101    /// ```
102    pub fn get_mut(&mut self) -> &mut T {
103        &mut self.inner
104    }
105
106    /// Returns the maximum number of bytes that can be read.
107    ///
108    /// # Note
109    ///
110    /// If the inner `Buf` has fewer bytes than indicated by this method then
111    /// that is the actual number of available bytes.
112    ///
113    /// # Examples
114    ///
115    /// ```rust
116    /// use rcbytes::Buf;
117    ///
118    /// let mut buf = b"hello world".take(2);
119    ///
120    /// assert_eq!(2, buf.limit());
121    /// assert_eq!(b'h', buf.get_u8());
122    /// assert_eq!(1, buf.limit());
123    /// ```
124    pub fn limit(&self) -> usize {
125        self.limit
126    }
127
128    /// Sets the maximum number of bytes that can be read.
129    ///
130    /// # Note
131    ///
132    /// If the inner `Buf` has fewer bytes than `lim` then that is the actual
133    /// number of available bytes.
134    ///
135    /// # Examples
136    ///
137    /// ```rust
138    /// use rcbytes::{Buf, BufMut};
139    ///
140    /// let mut buf = b"hello world".take(2);
141    /// let mut dst = vec![];
142    ///
143    /// dst.put(&mut buf);
144    /// assert_eq!(*dst, b"he"[..]);
145    ///
146    /// dst.clear();
147    ///
148    /// buf.set_limit(3);
149    /// dst.put(&mut buf);
150    /// assert_eq!(*dst, b"llo"[..]);
151    /// ```
152    pub fn set_limit(&mut self, lim: usize) {
153        self.limit = lim
154    }
155}
156
157impl<T: Buf> Buf for Take<T> {
158    fn remaining(&self) -> usize {
159        cmp::min(self.inner.remaining(), self.limit)
160    }
161
162    fn chunk(&self) -> &[u8] {
163        let bytes = self.inner.chunk();
164        &bytes[..cmp::min(bytes.len(), self.limit)]
165    }
166
167    fn advance(&mut self, cnt: usize) {
168        assert!(cnt <= self.limit);
169        self.inner.advance(cnt);
170        self.limit -= cnt;
171    }
172
173    fn copy_to_bytes(&mut self, len: usize) -> Bytes {
174        assert!(len <= self.remaining(), "`len` greater than remaining");
175
176        let r = self.inner.copy_to_bytes(len);
177        self.limit -= len;
178        r
179    }
180}