1 | use crate::buf::UninitSlice; |
2 | use crate::BufMut; |
3 | |
4 | use core::cmp; |
5 | |
6 | /// A `BufMut` adapter which limits the amount of bytes that can be written |
7 | /// to an underlying buffer. |
8 | #[derive(Debug)] |
9 | pub struct Limit<T> { |
10 | inner: T, |
11 | limit: usize, |
12 | } |
13 | |
14 | pub(super) fn new<T>(inner: T, limit: usize) -> Limit<T> { |
15 | Limit { inner, limit } |
16 | } |
17 | |
18 | impl<T> Limit<T> { |
19 | /// Consumes this `Limit`, returning the underlying value. |
20 | pub fn into_inner(self) -> T { |
21 | self.inner |
22 | } |
23 | |
24 | /// Gets a reference to the underlying `BufMut`. |
25 | /// |
26 | /// It is inadvisable to directly write to the underlying `BufMut`. |
27 | pub fn get_ref(&self) -> &T { |
28 | &self.inner |
29 | } |
30 | |
31 | /// Gets a mutable reference to the underlying `BufMut`. |
32 | /// |
33 | /// It is inadvisable to directly write to the underlying `BufMut`. |
34 | pub fn get_mut(&mut self) -> &mut T { |
35 | &mut self.inner |
36 | } |
37 | |
38 | /// Returns the maximum number of bytes that can be written |
39 | /// |
40 | /// # Note |
41 | /// |
42 | /// If the inner `BufMut` has fewer bytes than indicated by this method then |
43 | /// that is the actual number of available bytes. |
44 | pub fn limit(&self) -> usize { |
45 | self.limit |
46 | } |
47 | |
48 | /// Sets the maximum number of bytes that can be written. |
49 | /// |
50 | /// # Note |
51 | /// |
52 | /// If the inner `BufMut` has fewer bytes than `lim` then that is the actual |
53 | /// number of available bytes. |
54 | pub fn set_limit(&mut self, lim: usize) { |
55 | self.limit = lim |
56 | } |
57 | } |
58 | |
59 | unsafe impl<T: BufMut> BufMut for Limit<T> { |
60 | fn remaining_mut(&self) -> usize { |
61 | cmp::min(self.inner.remaining_mut(), self.limit) |
62 | } |
63 | |
64 | fn chunk_mut(&mut self) -> &mut UninitSlice { |
65 | let bytes = self.inner.chunk_mut(); |
66 | let end = cmp::min(bytes.len(), self.limit); |
67 | &mut bytes[..end] |
68 | } |
69 | |
70 | unsafe fn advance_mut(&mut self, cnt: usize) { |
71 | assert!(cnt <= self.limit); |
72 | self.inner.advance_mut(cnt); |
73 | self.limit -= cnt; |
74 | } |
75 | } |
76 | |