1use crate::buf::{IntoIter, UninitSlice};
2use crate::{Buf, BufMut, Bytes};
3
4#[cfg(feature = "std")]
5use 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)]
30pub struct Chain<T, U> {
31 a: T,
32 b: U,
33}
34
35impl<T, U> Chain<T, U> {
36 /// Creates a new `Chain` sequencing the provided values.
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 /// ```
53 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 /// ```
72 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 /// ```
88 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 /// ```
107 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 /// ```
125 pub fn into_inner(self) -> (T, U) {
126 (self.a, self.b)
127 }
128}
129
130impl<T, U> Buf for Chain<T, U>
131where
132 T: Buf,
133 U: Buf,
134{
135 fn remaining(&self) -> usize {
136 self.a.remaining().saturating_add(self.b.remaining())
137 }
138
139 fn chunk(&self) -> &[u8] {
140 if self.a.has_remaining() {
141 self.a.chunk()
142 } else {
143 self.b.chunk()
144 }
145 }
146
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")]
166 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
172 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
191unsafe impl<T, U> BufMut for Chain<T, U>
192where
193 T: BufMut,
194 U: BufMut,
195{
196 fn remaining_mut(&self) -> usize {
197 self.a
198 .remaining_mut()
199 .saturating_add(self.b.remaining_mut())
200 }
201
202 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
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
229impl<T, U> IntoIterator for Chain<T, U>
230where
231 T: Buf,
232 U: Buf,
233{
234 type Item = u8;
235 type IntoIter = IntoIter<Chain<T, U>>;
236
237 fn into_iter(self) -> Self::IntoIter {
238 IntoIter::new(self)
239 }
240}
241