http_body_util/
util.rs

1use std::collections::VecDeque;
2use std::io::IoSlice;
3
4use bytes::{Buf, BufMut, Bytes, BytesMut};
5
6#[derive(Debug)]
7pub(crate) struct BufList<T> {
8    bufs: VecDeque<T>,
9}
10
11impl<T: Buf> BufList<T> {
12    #[inline]
13    pub(crate) fn push(&mut self, buf: T) {
14        debug_assert!(buf.has_remaining());
15        self.bufs.push_back(buf);
16    }
17
18    #[inline]
19    pub(crate) fn pop(&mut self) -> Option<T> {
20        self.bufs.pop_front()
21    }
22}
23
24impl<T: Buf> Buf for BufList<T> {
25    #[inline]
26    fn remaining(&self) -> usize {
27        self.bufs.iter().map(|buf| buf.remaining()).sum()
28    }
29
30    #[inline]
31    fn has_remaining(&self) -> bool {
32        self.bufs.iter().any(|buf| buf.has_remaining())
33    }
34
35    #[inline]
36    fn chunk(&self) -> &[u8] {
37        self.bufs.front().map(Buf::chunk).unwrap_or_default()
38    }
39
40    #[inline]
41    fn advance(&mut self, mut cnt: usize) {
42        while cnt > 0 {
43            {
44                let front = &mut self.bufs[0];
45                let rem = front.remaining();
46                if rem > cnt {
47                    front.advance(cnt);
48                    return;
49                } else {
50                    front.advance(rem);
51                    cnt -= rem;
52                }
53            }
54            self.bufs.pop_front();
55        }
56    }
57
58    #[inline]
59    fn chunks_vectored<'t>(&'t self, dst: &mut [IoSlice<'t>]) -> usize {
60        if dst.is_empty() {
61            return 0;
62        }
63        let mut vecs = 0;
64        for buf in &self.bufs {
65            vecs += buf.chunks_vectored(&mut dst[vecs..]);
66            if vecs == dst.len() {
67                break;
68            }
69        }
70        vecs
71    }
72
73    #[inline]
74    fn copy_to_bytes(&mut self, len: usize) -> Bytes {
75        // Our inner buffer may have an optimized version of copy_to_bytes, and if the whole
76        // request can be fulfilled by the front buffer, we can take advantage.
77        match self.bufs.front_mut() {
78            Some(front) if front.remaining() == len => {
79                let b = front.copy_to_bytes(len);
80                self.bufs.pop_front();
81                b
82            }
83            Some(front) if front.remaining() > len => front.copy_to_bytes(len),
84            _ => {
85                let rem = self.remaining();
86                assert!(len <= rem, "`len` greater than remaining");
87                let mut bm = BytesMut::with_capacity(len);
88                if rem == len {
89                    // .take() costs a lot more, so skip it if we don't need it
90                    bm.put(self);
91                } else {
92                    bm.put(self.take(len));
93                }
94                bm.freeze()
95            }
96        }
97    }
98}
99
100impl<T> Default for BufList<T> {
101    fn default() -> Self {
102        BufList {
103            bufs: VecDeque::new(),
104        }
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use std::ptr;
111
112    use super::*;
113
114    fn hello_world_buf() -> BufList<Bytes> {
115        BufList {
116            bufs: vec![Bytes::from("Hello"), Bytes::from(" "), Bytes::from("World")].into(),
117        }
118    }
119
120    #[test]
121    fn to_bytes_shorter() {
122        let mut bufs = hello_world_buf();
123        let old_ptr = bufs.chunk().as_ptr();
124        let start = bufs.copy_to_bytes(4);
125        assert_eq!(start, "Hell");
126        assert!(ptr::eq(old_ptr, start.as_ptr()));
127        assert_eq!(bufs.chunk(), b"o");
128        assert!(ptr::eq(old_ptr.wrapping_add(4), bufs.chunk().as_ptr()));
129        assert_eq!(bufs.remaining(), 7);
130    }
131
132    #[test]
133    fn to_bytes_eq() {
134        let mut bufs = hello_world_buf();
135        let old_ptr = bufs.chunk().as_ptr();
136        let start = bufs.copy_to_bytes(5);
137        assert_eq!(start, "Hello");
138        assert!(ptr::eq(old_ptr, start.as_ptr()));
139        assert_eq!(bufs.chunk(), b" ");
140        assert_eq!(bufs.remaining(), 6);
141    }
142
143    #[test]
144    fn to_bytes_longer() {
145        let mut bufs = hello_world_buf();
146        let start = bufs.copy_to_bytes(7);
147        assert_eq!(start, "Hello W");
148        assert_eq!(bufs.remaining(), 4);
149    }
150
151    #[test]
152    fn one_long_buf_to_bytes() {
153        let mut buf = BufList::default();
154        buf.push(b"Hello World" as &[_]);
155        assert_eq!(buf.copy_to_bytes(5), "Hello");
156        assert_eq!(buf.chunk(), b" World");
157    }
158
159    #[test]
160    #[should_panic(expected = "`len` greater than remaining")]
161    fn buf_to_bytes_too_many() {
162        hello_world_buf().copy_to_bytes(42);
163    }
164}