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