1use alloc::collections::VecDeque;
2use alloc::vec::Vec;
3use core::cmp;
4use std::io;
5use std::io::Read;
6
7/// This is a byte buffer that is built from a vector
8/// of byte vectors. This avoids extra copies when
9/// appending a new byte vector, at the expense of
10/// more complexity when reading out.
11pub(crate) struct ChunkVecBuffer {
12 chunks: VecDeque<Vec<u8>>,
13 limit: Option<usize>,
14}
15
16impl ChunkVecBuffer {
17 pub(crate) fn new(limit: Option<usize>) -> Self {
18 Self {
19 chunks: VecDeque::new(),
20 limit,
21 }
22 }
23
24 /// Sets the upper limit on how many bytes this
25 /// object can store.
26 ///
27 /// Setting a lower limit than the currently stored
28 /// data is not an error.
29 ///
30 /// A [`None`] limit is interpreted as no limit.
31 pub(crate) fn set_limit(&mut self, new_limit: Option<usize>) {
32 self.limit = new_limit;
33 }
34
35 /// If we're empty
36 pub(crate) fn is_empty(&self) -> bool {
37 self.chunks.is_empty()
38 }
39
40 pub(crate) fn is_full(&self) -> bool {
41 self.limit
42 .map(|limit| self.len() > limit)
43 .unwrap_or_default()
44 }
45
46 /// How many bytes we're storing
47 pub(crate) fn len(&self) -> usize {
48 let mut len = 0;
49 for ch in &self.chunks {
50 len += ch.len();
51 }
52 len
53 }
54
55 /// For a proposed append of `len` bytes, how many
56 /// bytes should we actually append to adhere to the
57 /// currently set `limit`?
58 pub(crate) fn apply_limit(&self, len: usize) -> usize {
59 if let Some(limit) = self.limit {
60 let space = limit.saturating_sub(self.len());
61 cmp::min(len, space)
62 } else {
63 len
64 }
65 }
66
67 /// Append a copy of `bytes`, perhaps a prefix if
68 /// we're near the limit.
69 pub(crate) fn append_limited_copy(&mut self, bytes: &[u8]) -> usize {
70 let take = self.apply_limit(bytes.len());
71 self.append(bytes[..take].to_vec());
72 take
73 }
74
75 /// Take and append the given `bytes`.
76 pub(crate) fn append(&mut self, bytes: Vec<u8>) -> usize {
77 let len = bytes.len();
78
79 if !bytes.is_empty() {
80 self.chunks.push_back(bytes);
81 }
82
83 len
84 }
85
86 /// Take one of the chunks from this object. This
87 /// function panics if the object `is_empty`.
88 pub(crate) fn pop(&mut self) -> Option<Vec<u8>> {
89 self.chunks.pop_front()
90 }
91
92 /// Read data out of this object, writing it into `buf`
93 /// and returning how many bytes were written there.
94 pub(crate) fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
95 let mut offs = 0;
96
97 while offs < buf.len() && !self.is_empty() {
98 let used = self.chunks[0]
99 .as_slice()
100 .read(&mut buf[offs..])?;
101
102 self.consume(used);
103 offs += used;
104 }
105
106 Ok(offs)
107 }
108
109 #[cfg(read_buf)]
110 /// Read data out of this object, writing it into `cursor`.
111 pub(crate) fn read_buf(&mut self, mut cursor: core::io::BorrowedCursor<'_>) -> io::Result<()> {
112 while !self.is_empty() && cursor.capacity() > 0 {
113 let chunk = self.chunks[0].as_slice();
114 let used = core::cmp::min(chunk.len(), cursor.capacity());
115 cursor.append(&chunk[..used]);
116 self.consume(used);
117 }
118
119 Ok(())
120 }
121
122 fn consume(&mut self, mut used: usize) {
123 while let Some(mut buf) = self.chunks.pop_front() {
124 if used < buf.len() {
125 buf.drain(..used);
126 self.chunks.push_front(buf);
127 break;
128 } else {
129 used -= buf.len();
130 }
131 }
132 }
133
134 /// Read data out of this object, passing it `wr`
135 pub(crate) fn write_to(&mut self, wr: &mut dyn io::Write) -> io::Result<usize> {
136 if self.is_empty() {
137 return Ok(0);
138 }
139
140 let mut bufs = [io::IoSlice::new(&[]); 64];
141 for (iov, chunk) in bufs.iter_mut().zip(self.chunks.iter()) {
142 *iov = io::IoSlice::new(chunk);
143 }
144 let len = cmp::min(bufs.len(), self.chunks.len());
145 let used = wr.write_vectored(&bufs[..len])?;
146 self.consume(used);
147 Ok(used)
148 }
149}
150
151#[cfg(test)]
152mod tests {
153 use super::ChunkVecBuffer;
154
155 #[test]
156 fn short_append_copy_with_limit() {
157 let mut cvb = ChunkVecBuffer::new(Some(12));
158 assert_eq!(cvb.append_limited_copy(b"hello"), 5);
159 assert_eq!(cvb.append_limited_copy(b"world"), 5);
160 assert_eq!(cvb.append_limited_copy(b"hello"), 2);
161 assert_eq!(cvb.append_limited_copy(b"world"), 0);
162
163 let mut buf = [0u8; 12];
164 assert_eq!(cvb.read(&mut buf).unwrap(), 12);
165 assert_eq!(buf.to_vec(), b"helloworldhe".to_vec());
166 }
167
168 #[cfg(read_buf)]
169 #[test]
170 fn read_buf() {
171 use core::io::BorrowedBuf;
172 use core::mem::MaybeUninit;
173
174 {
175 let mut cvb = ChunkVecBuffer::new(None);
176 cvb.append(b"test ".to_vec());
177 cvb.append(b"fixture ".to_vec());
178 cvb.append(b"data".to_vec());
179
180 let mut buf = [MaybeUninit::<u8>::uninit(); 8];
181 let mut buf: BorrowedBuf<'_> = buf.as_mut_slice().into();
182 cvb.read_buf(buf.unfilled()).unwrap();
183 assert_eq!(buf.filled(), b"test fix");
184 buf.clear();
185 cvb.read_buf(buf.unfilled()).unwrap();
186 assert_eq!(buf.filled(), b"ture dat");
187 buf.clear();
188 cvb.read_buf(buf.unfilled()).unwrap();
189 assert_eq!(buf.filled(), b"a");
190 }
191
192 {
193 let mut cvb = ChunkVecBuffer::new(None);
194 cvb.append(b"short message".to_vec());
195
196 let mut buf = [MaybeUninit::<u8>::uninit(); 1024];
197 let mut buf: BorrowedBuf<'_> = buf.as_mut_slice().into();
198 cvb.read_buf(buf.unfilled()).unwrap();
199 assert_eq!(buf.filled(), b"short message");
200 }
201 }
202}
203