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`]: trait.Buf.html#method.chain
29/// [`Buf`]: trait.Buf.html
30/// [`BufMut`]: trait.BufMut.html
31#[derive(Debug)]
32pub struct Chain<T, U> {
33 a: T,
34 b: U,
35}
36
37impl<T, U> Chain<T, U> {
38 /// Creates a new `Chain` sequencing the provided values.
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 /// ```
55 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 /// ```
74 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 /// ```
90 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 /// ```
109 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 /// ```
127 pub fn into_inner(self) -> (T, U) {
128 (self.a, self.b)
129 }
130}
131
132impl<T, U> Buf for Chain<T, U>
133where
134 T: Buf,
135 U: Buf,
136{
137 fn remaining(&self) -> usize {
138 self.a.remaining().checked_add(self.b.remaining()).unwrap()
139 }
140
141 fn chunk(&self) -> &[u8] {
142 if self.a.has_remaining() {
143 self.a.chunk()
144 } else {
145 self.b.chunk()
146 }
147 }
148
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")]
168 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
174 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
193unsafe impl<T, U> BufMut for Chain<T, U>
194where
195 T: BufMut,
196 U: BufMut,
197{
198 fn remaining_mut(&self) -> usize {
199 self.a
200 .remaining_mut()
201 .saturating_add(self.b.remaining_mut())
202 }
203
204 fn chunk_mut(&mut self) -> &mut UninitSlice {
205 if self.a.has_remaining_mut() {
206 self.a.chunk_mut()
207 } else {
208 self.b.chunk_mut()
209 }
210 }
211
212 unsafe fn advance_mut(&mut self, mut cnt: usize) {
213 let a_rem = self.a.remaining_mut();
214
215 if a_rem != 0 {
216 if a_rem >= cnt {
217 self.a.advance_mut(cnt);
218 return;
219 }
220
221 // Consume what is left of a
222 self.a.advance_mut(a_rem);
223
224 cnt -= a_rem;
225 }
226
227 self.b.advance_mut(cnt);
228 }
229}
230
231impl<T, U> IntoIterator for Chain<T, U>
232where
233 T: Buf,
234 U: Buf,
235{
236 type Item = u8;
237 type IntoIter = IntoIter<Chain<T, U>>;
238
239 fn into_iter(self) -> Self::IntoIter {
240 IntoIter::new(self)
241 }
242}
243