1 use crate::{Buf, BufMut};
2 use crate::buf::IntoIter;
3 
4 use core::mem::MaybeUninit;
5 
6 #[cfg(feature = "std")]
7 use std::io::{IoSlice};
8 #[cfg(feature = "std")]
9 use crate::buf::IoSliceMut;
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 {
45             a,
46             b,
47         }
48     }
49 
50     /// Gets a reference to the first underlying `Buf`.
51     ///
52     /// # Examples
53     ///
54     /// ```
55     /// use bytes::buf::BufExt;
56     ///
57     /// let buf = (&b"hello"[..])
58     ///     .chain(&b"world"[..]);
59     ///
60     /// assert_eq!(buf.first_ref()[..], b"hello"[..]);
61     /// ```
first_ref(&self) -> &T62     pub fn first_ref(&self) -> &T {
63         &self.a
64     }
65 
66     /// Gets a mutable reference to the first underlying `Buf`.
67     ///
68     /// # Examples
69     ///
70     /// ```
71     /// use bytes::{Buf, buf::BufExt};
72     ///
73     /// let mut buf = (&b"hello"[..])
74     ///     .chain(&b"world"[..]);
75     ///
76     /// buf.first_mut().advance(1);
77     ///
78     /// let full = buf.to_bytes();
79     /// assert_eq!(full, b"elloworld"[..]);
80     /// ```
first_mut(&mut self) -> &mut T81     pub fn first_mut(&mut self) -> &mut T {
82         &mut self.a
83     }
84 
85     /// Gets a reference to the last underlying `Buf`.
86     ///
87     /// # Examples
88     ///
89     /// ```
90     /// use bytes::buf::BufExt;
91     ///
92     /// let buf = (&b"hello"[..])
93     ///     .chain(&b"world"[..]);
94     ///
95     /// assert_eq!(buf.last_ref()[..], b"world"[..]);
96     /// ```
last_ref(&self) -> &U97     pub fn last_ref(&self) -> &U {
98         &self.b
99     }
100 
101     /// Gets a mutable reference to the last underlying `Buf`.
102     ///
103     /// # Examples
104     ///
105     /// ```
106     /// use bytes::{Buf, buf::BufExt};
107     ///
108     /// let mut buf = (&b"hello "[..])
109     ///     .chain(&b"world"[..]);
110     ///
111     /// buf.last_mut().advance(1);
112     ///
113     /// let full = buf.to_bytes();
114     /// assert_eq!(full, b"hello orld"[..]);
115     /// ```
last_mut(&mut self) -> &mut U116     pub fn last_mut(&mut self) -> &mut U {
117         &mut self.b
118     }
119 
120     /// Consumes this `Chain`, returning the underlying values.
121     ///
122     /// # Examples
123     ///
124     /// ```
125     /// use bytes::buf::BufExt;
126     ///
127     /// let chain = (&b"hello"[..])
128     ///     .chain(&b"world"[..]);
129     ///
130     /// let (first, last) = chain.into_inner();
131     /// assert_eq!(first[..], b"hello"[..]);
132     /// assert_eq!(last[..], b"world"[..]);
133     /// ```
into_inner(self) -> (T, U)134     pub fn into_inner(self) -> (T, U) {
135         (self.a, self.b)
136     }
137 }
138 
139 impl<T, U> Buf for Chain<T, U>
140     where T: Buf,
141           U: Buf,
142 {
remaining(&self) -> usize143     fn remaining(&self) -> usize {
144         self.a.remaining() + self.b.remaining()
145     }
146 
bytes(&self) -> &[u8]147     fn bytes(&self) -> &[u8] {
148         if self.a.has_remaining() {
149             self.a.bytes()
150         } else {
151             self.b.bytes()
152         }
153     }
154 
advance(&mut self, mut cnt: usize)155     fn advance(&mut self, mut cnt: usize) {
156         let a_rem = self.a.remaining();
157 
158         if a_rem != 0 {
159             if a_rem >= cnt {
160                 self.a.advance(cnt);
161                 return;
162             }
163 
164             // Consume what is left of a
165             self.a.advance(a_rem);
166 
167             cnt -= a_rem;
168         }
169 
170         self.b.advance(cnt);
171     }
172 
173     #[cfg(feature = "std")]
bytes_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize174     fn bytes_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
175         let mut n = self.a.bytes_vectored(dst);
176         n += self.b.bytes_vectored(&mut dst[n..]);
177         n
178     }
179 }
180 
181 impl<T, U> BufMut for Chain<T, U>
182     where T: BufMut,
183           U: BufMut,
184 {
remaining_mut(&self) -> usize185     fn remaining_mut(&self) -> usize {
186         self.a.remaining_mut() + self.b.remaining_mut()
187     }
188 
bytes_mut(&mut self) -> &mut [MaybeUninit<u8>]189     fn bytes_mut(&mut self) -> &mut [MaybeUninit<u8>] {
190         if self.a.has_remaining_mut() {
191             self.a.bytes_mut()
192         } else {
193             self.b.bytes_mut()
194         }
195     }
196 
advance_mut(&mut self, mut cnt: usize)197     unsafe fn advance_mut(&mut self, mut cnt: usize) {
198         let a_rem = self.a.remaining_mut();
199 
200         if a_rem != 0 {
201             if a_rem >= cnt {
202                 self.a.advance_mut(cnt);
203                 return;
204             }
205 
206             // Consume what is left of a
207             self.a.advance_mut(a_rem);
208 
209             cnt -= a_rem;
210         }
211 
212         self.b.advance_mut(cnt);
213     }
214 
215     #[cfg(feature = "std")]
bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize216     fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize {
217         let mut n = self.a.bytes_vectored_mut(dst);
218         n += self.b.bytes_vectored_mut(&mut dst[n..]);
219         n
220     }
221 }
222 
223 impl<T, U> IntoIterator for Chain<T, U>
224 where
225     T: Buf,
226     U: Buf,
227 {
228     type Item = u8;
229     type IntoIter = IntoIter<Chain<T, U>>;
230 
into_iter(self) -> Self::IntoIter231     fn into_iter(self) -> Self::IntoIter {
232         IntoIter::new(self)
233     }
234 }
235