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`]: Buf::chain
29 #[derive(Debug)]
30 pub struct Chain<T, U> {
31     a: T,
32     b: U,
33 }
34 
35 impl<T, U> Chain<T, U> {
36     /// Creates a new `Chain` sequencing the provided values.
new(a: T, b: U) -> Chain<T, U>37     pub(crate) fn new(a: T, b: U) -> Chain<T, U> {
38         Chain { a, b }
39     }
40 
41     /// Gets a reference to the first underlying `Buf`.
42     ///
43     /// # Examples
44     ///
45     /// ```
46     /// use bytes::Buf;
47     ///
48     /// let buf = (&b"hello"[..])
49     ///     .chain(&b"world"[..]);
50     ///
51     /// assert_eq!(buf.first_ref()[..], b"hello"[..]);
52     /// ```
first_ref(&self) -> &T53     pub fn first_ref(&self) -> &T {
54         &self.a
55     }
56 
57     /// Gets a mutable reference to the first underlying `Buf`.
58     ///
59     /// # Examples
60     ///
61     /// ```
62     /// use bytes::Buf;
63     ///
64     /// let mut buf = (&b"hello"[..])
65     ///     .chain(&b"world"[..]);
66     ///
67     /// buf.first_mut().advance(1);
68     ///
69     /// let full = buf.copy_to_bytes(9);
70     /// assert_eq!(full, b"elloworld"[..]);
71     /// ```
first_mut(&mut self) -> &mut T72     pub fn first_mut(&mut self) -> &mut T {
73         &mut self.a
74     }
75 
76     /// Gets a reference to the last underlying `Buf`.
77     ///
78     /// # Examples
79     ///
80     /// ```
81     /// use bytes::Buf;
82     ///
83     /// let buf = (&b"hello"[..])
84     ///     .chain(&b"world"[..]);
85     ///
86     /// assert_eq!(buf.last_ref()[..], b"world"[..]);
87     /// ```
last_ref(&self) -> &U88     pub fn last_ref(&self) -> &U {
89         &self.b
90     }
91 
92     /// Gets a mutable reference to the last underlying `Buf`.
93     ///
94     /// # Examples
95     ///
96     /// ```
97     /// use bytes::Buf;
98     ///
99     /// let mut buf = (&b"hello "[..])
100     ///     .chain(&b"world"[..]);
101     ///
102     /// buf.last_mut().advance(1);
103     ///
104     /// let full = buf.copy_to_bytes(10);
105     /// assert_eq!(full, b"hello orld"[..]);
106     /// ```
last_mut(&mut self) -> &mut U107     pub fn last_mut(&mut self) -> &mut U {
108         &mut self.b
109     }
110 
111     /// Consumes this `Chain`, returning the underlying values.
112     ///
113     /// # Examples
114     ///
115     /// ```
116     /// use bytes::Buf;
117     ///
118     /// let chain = (&b"hello"[..])
119     ///     .chain(&b"world"[..]);
120     ///
121     /// let (first, last) = chain.into_inner();
122     /// assert_eq!(first[..], b"hello"[..]);
123     /// assert_eq!(last[..], b"world"[..]);
124     /// ```
into_inner(self) -> (T, U)125     pub fn into_inner(self) -> (T, U) {
126         (self.a, self.b)
127     }
128 }
129 
130 impl<T, U> Buf for Chain<T, U>
131 where
132     T: Buf,
133     U: Buf,
134 {
remaining(&self) -> usize135     fn remaining(&self) -> usize {
136         self.a.remaining().saturating_add(self.b.remaining())
137     }
138 
chunk(&self) -> &[u8]139     fn chunk(&self) -> &[u8] {
140         if self.a.has_remaining() {
141             self.a.chunk()
142         } else {
143             self.b.chunk()
144         }
145     }
146 
advance(&mut self, mut cnt: usize)147     fn advance(&mut self, mut cnt: usize) {
148         let a_rem = self.a.remaining();
149 
150         if a_rem != 0 {
151             if a_rem >= cnt {
152                 self.a.advance(cnt);
153                 return;
154             }
155 
156             // Consume what is left of a
157             self.a.advance(a_rem);
158 
159             cnt -= a_rem;
160         }
161 
162         self.b.advance(cnt);
163     }
164 
165     #[cfg(feature = "std")]
chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize166     fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
167         let mut n = self.a.chunks_vectored(dst);
168         n += self.b.chunks_vectored(&mut dst[n..]);
169         n
170     }
171 
copy_to_bytes(&mut self, len: usize) -> Bytes172     fn copy_to_bytes(&mut self, len: usize) -> Bytes {
173         let a_rem = self.a.remaining();
174         if a_rem >= len {
175             self.a.copy_to_bytes(len)
176         } else if a_rem == 0 {
177             self.b.copy_to_bytes(len)
178         } else {
179             assert!(
180                 len - a_rem <= self.b.remaining(),
181                 "`len` greater than remaining"
182             );
183             let mut ret = crate::BytesMut::with_capacity(len);
184             ret.put(&mut self.a);
185             ret.put((&mut self.b).take(len - a_rem));
186             ret.freeze()
187         }
188     }
189 }
190 
191 unsafe impl<T, U> BufMut for Chain<T, U>
192 where
193     T: BufMut,
194     U: BufMut,
195 {
remaining_mut(&self) -> usize196     fn remaining_mut(&self) -> usize {
197         self.a
198             .remaining_mut()
199             .saturating_add(self.b.remaining_mut())
200     }
201 
chunk_mut(&mut self) -> &mut UninitSlice202     fn chunk_mut(&mut self) -> &mut UninitSlice {
203         if self.a.has_remaining_mut() {
204             self.a.chunk_mut()
205         } else {
206             self.b.chunk_mut()
207         }
208     }
209 
advance_mut(&mut self, mut cnt: usize)210     unsafe fn advance_mut(&mut self, mut cnt: usize) {
211         let a_rem = self.a.remaining_mut();
212 
213         if a_rem != 0 {
214             if a_rem >= cnt {
215                 self.a.advance_mut(cnt);
216                 return;
217             }
218 
219             // Consume what is left of a
220             self.a.advance_mut(a_rem);
221 
222             cnt -= a_rem;
223         }
224 
225         self.b.advance_mut(cnt);
226     }
227 }
228 
229 impl<T, U> IntoIterator for Chain<T, U>
230 where
231     T: Buf,
232     U: Buf,
233 {
234     type Item = u8;
235     type IntoIter = IntoIter<Chain<T, U>>;
236 
into_iter(self) -> Self::IntoIter237     fn into_iter(self) -> Self::IntoIter {
238         IntoIter::new(self)
239     }
240 }
241