1 use crate::buf::{IntoIter, UninitSlice};
2 use crate::{Buf, BufMut, Bytes};
3 
4 #[cfg(feature = "std")]
5 use std::io::IoSlice;
6 
7 /// A `Chain` sequences two buffers.
8 ///
9 /// `Chain` is an adapter that links two underlying buffers and provides a
10 /// continuous view across both buffers. It is able to sequence either immutable
11 /// buffers ([`Buf`] values) or mutable buffers ([`BufMut`] values).
12 ///
13 /// This struct is generally created by calling [`Buf::chain`]. Please see that
14 /// function's documentation for more detail.
15 ///
16 /// # Examples
17 ///
18 /// ```
19 /// use bytes::{Bytes, Buf};
20 ///
21 /// let mut buf = (&b"hello "[..])
22 ///     .chain(&b"world"[..]);
23 ///
24 /// let full: Bytes = buf.copy_to_bytes(11);
25 /// assert_eq!(full[..], b"hello world"[..]);
26 /// ```
27 ///
28 /// [`Buf::chain`]: trait.Buf.html#method.chain
29 /// [`Buf`]: trait.Buf.html
30 /// [`BufMut`]: trait.BufMut.html
31 #[derive(Debug)]
32 pub struct Chain<T, U> {
33     a: T,
34     b: U,
35 }
36 
37 impl<T, U> Chain<T, U> {
38     /// Creates a new `Chain` sequencing the provided values.
new(a: T, b: U) -> Chain<T, U>39     pub(crate) fn new(a: T, b: U) -> Chain<T, U> {
40         Chain { a, b }
41     }
42 
43     /// Gets a reference to the first underlying `Buf`.
44     ///
45     /// # Examples
46     ///
47     /// ```
48     /// use bytes::Buf;
49     ///
50     /// let buf = (&b"hello"[..])
51     ///     .chain(&b"world"[..]);
52     ///
53     /// assert_eq!(buf.first_ref()[..], b"hello"[..]);
54     /// ```
first_ref(&self) -> &T55     pub fn first_ref(&self) -> &T {
56         &self.a
57     }
58 
59     /// Gets a mutable reference to the first underlying `Buf`.
60     ///
61     /// # Examples
62     ///
63     /// ```
64     /// use bytes::Buf;
65     ///
66     /// let mut buf = (&b"hello"[..])
67     ///     .chain(&b"world"[..]);
68     ///
69     /// buf.first_mut().advance(1);
70     ///
71     /// let full = buf.copy_to_bytes(9);
72     /// assert_eq!(full, b"elloworld"[..]);
73     /// ```
first_mut(&mut self) -> &mut T74     pub fn first_mut(&mut self) -> &mut T {
75         &mut self.a
76     }
77 
78     /// Gets a reference to the last underlying `Buf`.
79     ///
80     /// # Examples
81     ///
82     /// ```
83     /// use bytes::Buf;
84     ///
85     /// let buf = (&b"hello"[..])
86     ///     .chain(&b"world"[..]);
87     ///
88     /// assert_eq!(buf.last_ref()[..], b"world"[..]);
89     /// ```
last_ref(&self) -> &U90     pub fn last_ref(&self) -> &U {
91         &self.b
92     }
93 
94     /// Gets a mutable reference to the last underlying `Buf`.
95     ///
96     /// # Examples
97     ///
98     /// ```
99     /// use bytes::Buf;
100     ///
101     /// let mut buf = (&b"hello "[..])
102     ///     .chain(&b"world"[..]);
103     ///
104     /// buf.last_mut().advance(1);
105     ///
106     /// let full = buf.copy_to_bytes(10);
107     /// assert_eq!(full, b"hello orld"[..]);
108     /// ```
last_mut(&mut self) -> &mut U109     pub fn last_mut(&mut self) -> &mut U {
110         &mut self.b
111     }
112 
113     /// Consumes this `Chain`, returning the underlying values.
114     ///
115     /// # Examples
116     ///
117     /// ```
118     /// use bytes::Buf;
119     ///
120     /// let chain = (&b"hello"[..])
121     ///     .chain(&b"world"[..]);
122     ///
123     /// let (first, last) = chain.into_inner();
124     /// assert_eq!(first[..], b"hello"[..]);
125     /// assert_eq!(last[..], b"world"[..]);
126     /// ```
into_inner(self) -> (T, U)127     pub fn into_inner(self) -> (T, U) {
128         (self.a, self.b)
129     }
130 }
131 
132 impl<T, U> Buf for Chain<T, U>
133 where
134     T: Buf,
135     U: Buf,
136 {
remaining(&self) -> usize137     fn remaining(&self) -> usize {
138         self.a.remaining().checked_add(self.b.remaining()).unwrap()
139     }
140 
chunk(&self) -> &[u8]141     fn chunk(&self) -> &[u8] {
142         if self.a.has_remaining() {
143             self.a.chunk()
144         } else {
145             self.b.chunk()
146         }
147     }
148 
advance(&mut self, mut cnt: usize)149     fn advance(&mut self, mut cnt: usize) {
150         let a_rem = self.a.remaining();
151 
152         if a_rem != 0 {
153             if a_rem >= cnt {
154                 self.a.advance(cnt);
155                 return;
156             }
157 
158             // Consume what is left of a
159             self.a.advance(a_rem);
160 
161             cnt -= a_rem;
162         }
163 
164         self.b.advance(cnt);
165     }
166 
167     #[cfg(feature = "std")]
chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize168     fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
169         let mut n = self.a.chunks_vectored(dst);
170         n += self.b.chunks_vectored(&mut dst[n..]);
171         n
172     }
173 
copy_to_bytes(&mut self, len: usize) -> Bytes174     fn copy_to_bytes(&mut self, len: usize) -> Bytes {
175         let a_rem = self.a.remaining();
176         if a_rem >= len {
177             self.a.copy_to_bytes(len)
178         } else if a_rem == 0 {
179             self.b.copy_to_bytes(len)
180         } else {
181             assert!(
182                 len - a_rem <= self.b.remaining(),
183                 "`len` greater than remaining"
184             );
185             let mut ret = crate::BytesMut::with_capacity(len);
186             ret.put(&mut self.a);
187             ret.put((&mut self.b).take(len - a_rem));
188             ret.freeze()
189         }
190     }
191 }
192 
193 unsafe impl<T, U> BufMut for Chain<T, U>
194 where
195     T: BufMut,
196     U: BufMut,
197 {
remaining_mut(&self) -> usize198     fn remaining_mut(&self) -> usize {
199         self.a
200             .remaining_mut()
201             .checked_add(self.b.remaining_mut())
202             .unwrap()
203     }
204 
chunk_mut(&mut self) -> &mut UninitSlice205     fn chunk_mut(&mut self) -> &mut UninitSlice {
206         if self.a.has_remaining_mut() {
207             self.a.chunk_mut()
208         } else {
209             self.b.chunk_mut()
210         }
211     }
212 
advance_mut(&mut self, mut cnt: usize)213     unsafe fn advance_mut(&mut self, mut cnt: usize) {
214         let a_rem = self.a.remaining_mut();
215 
216         if a_rem != 0 {
217             if a_rem >= cnt {
218                 self.a.advance_mut(cnt);
219                 return;
220             }
221 
222             // Consume what is left of a
223             self.a.advance_mut(a_rem);
224 
225             cnt -= a_rem;
226         }
227 
228         self.b.advance_mut(cnt);
229     }
230 }
231 
232 impl<T, U> IntoIterator for Chain<T, U>
233 where
234     T: Buf,
235     U: Buf,
236 {
237     type Item = u8;
238     type IntoIter = IntoIter<Chain<T, U>>;
239 
into_iter(self) -> Self::IntoIter240     fn into_iter(self) -> Self::IntoIter {
241         IntoIter::new(self)
242     }
243 }
244