1 use crate::iter::{DoubleEndedIterator, FusedIterator, Iterator, TrustedLen}; 2 use crate::ops::Try; 3 4 /// An iterator that links two iterators together, in a chain. 5 /// 6 /// This `struct` is created by [`Iterator::chain`]. See its documentation 7 /// for more. 8 /// 9 /// # Examples 10 /// 11 /// ``` 12 /// use std::iter::Chain; 13 /// use std::slice::Iter; 14 /// 15 /// let a1 = [1, 2, 3]; 16 /// let a2 = [4, 5, 6]; 17 /// let iter: Chain<Iter<_>, Iter<_>> = a1.iter().chain(a2.iter()); 18 /// ``` 19 #[derive(Clone, Debug)] 20 #[must_use = "iterators are lazy and do nothing unless consumed"] 21 #[stable(feature = "rust1", since = "1.0.0")] 22 pub struct Chain<A, B> { 23 // These are "fused" with `Option` so we don't need separate state to track which part is 24 // already exhausted, and we may also get niche layout for `None`. We don't use the real `Fuse` 25 // adapter because its specialization for `FusedIterator` unconditionally descends into the 26 // iterator, and that could be expensive to keep revisiting stuff like nested chains. It also 27 // hurts compiler performance to add more iterator layers to `Chain`. 28 // 29 // Only the "first" iterator is actually set `None` when exhausted, depending on whether you 30 // iterate forward or backward. If you mix directions, then both sides may be `None`. 31 a: Option<A>, 32 b: Option<B>, 33 } 34 impl<A, B> Chain<A, B> { new(a: A, b: B) -> Chain<A, B>35 pub(in super::super) fn new(a: A, b: B) -> Chain<A, B> { 36 Chain { a: Some(a), b: Some(b) } 37 } 38 } 39 40 /// Fuse the iterator if the expression is `None`. 41 macro_rules! fuse { 42 ($self:ident . $iter:ident . $($call:tt)+) => { 43 match $self.$iter { 44 Some(ref mut iter) => match iter.$($call)+ { 45 None => { 46 $self.$iter = None; 47 None 48 } 49 item => item, 50 }, 51 None => None, 52 } 53 }; 54 } 55 56 /// Try an iterator method without fusing, 57 /// like an inline `.as_mut().and_then(...)` 58 macro_rules! maybe { 59 ($self:ident . $iter:ident . $($call:tt)+) => { 60 match $self.$iter { 61 Some(ref mut iter) => iter.$($call)+, 62 None => None, 63 } 64 }; 65 } 66 67 #[stable(feature = "rust1", since = "1.0.0")] 68 impl<A, B> Iterator for Chain<A, B> 69 where 70 A: Iterator, 71 B: Iterator<Item = A::Item>, 72 { 73 type Item = A::Item; 74 75 #[inline] next(&mut self) -> Option<A::Item>76 fn next(&mut self) -> Option<A::Item> { 77 match fuse!(self.a.next()) { 78 None => maybe!(self.b.next()), 79 item => item, 80 } 81 } 82 83 #[inline] 84 #[rustc_inherit_overflow_checks] count(self) -> usize85 fn count(self) -> usize { 86 let a_count = match self.a { 87 Some(a) => a.count(), 88 None => 0, 89 }; 90 let b_count = match self.b { 91 Some(b) => b.count(), 92 None => 0, 93 }; 94 a_count + b_count 95 } 96 try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R where Self: Sized, F: FnMut(Acc, Self::Item) -> R, R: Try<Output = Acc>,97 fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R 98 where 99 Self: Sized, 100 F: FnMut(Acc, Self::Item) -> R, 101 R: Try<Output = Acc>, 102 { 103 if let Some(ref mut a) = self.a { 104 acc = a.try_fold(acc, &mut f)?; 105 self.a = None; 106 } 107 if let Some(ref mut b) = self.b { 108 acc = b.try_fold(acc, f)?; 109 // we don't fuse the second iterator 110 } 111 try { acc } 112 } 113 fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc where F: FnMut(Acc, Self::Item) -> Acc,114 fn fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc 115 where 116 F: FnMut(Acc, Self::Item) -> Acc, 117 { 118 if let Some(a) = self.a { 119 acc = a.fold(acc, &mut f); 120 } 121 if let Some(b) = self.b { 122 acc = b.fold(acc, f); 123 } 124 acc 125 } 126 127 #[inline] advance_by(&mut self, n: usize) -> Result<(), usize>128 fn advance_by(&mut self, n: usize) -> Result<(), usize> { 129 let mut rem = n; 130 131 if let Some(ref mut a) = self.a { 132 match a.advance_by(rem) { 133 Ok(()) => return Ok(()), 134 Err(k) => rem -= k, 135 } 136 self.a = None; 137 } 138 139 if let Some(ref mut b) = self.b { 140 match b.advance_by(rem) { 141 Ok(()) => return Ok(()), 142 Err(k) => rem -= k, 143 } 144 // we don't fuse the second iterator 145 } 146 147 if rem == 0 { Ok(()) } else { Err(n - rem) } 148 } 149 150 #[inline] nth(&mut self, mut n: usize) -> Option<Self::Item>151 fn nth(&mut self, mut n: usize) -> Option<Self::Item> { 152 if let Some(ref mut a) = self.a { 153 match a.advance_by(n) { 154 Ok(()) => match a.next() { 155 None => n = 0, 156 x => return x, 157 }, 158 Err(k) => n -= k, 159 } 160 161 self.a = None; 162 } 163 164 maybe!(self.b.nth(n)) 165 } 166 167 #[inline] find<P>(&mut self, mut predicate: P) -> Option<Self::Item> where P: FnMut(&Self::Item) -> bool,168 fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item> 169 where 170 P: FnMut(&Self::Item) -> bool, 171 { 172 match fuse!(self.a.find(&mut predicate)) { 173 None => maybe!(self.b.find(predicate)), 174 item => item, 175 } 176 } 177 178 #[inline] last(self) -> Option<A::Item>179 fn last(self) -> Option<A::Item> { 180 // Must exhaust a before b. 181 let a_last = match self.a { 182 Some(a) => a.last(), 183 None => None, 184 }; 185 let b_last = match self.b { 186 Some(b) => b.last(), 187 None => None, 188 }; 189 b_last.or(a_last) 190 } 191 192 #[inline] size_hint(&self) -> (usize, Option<usize>)193 fn size_hint(&self) -> (usize, Option<usize>) { 194 match self { 195 Chain { a: Some(a), b: Some(b) } => { 196 let (a_lower, a_upper) = a.size_hint(); 197 let (b_lower, b_upper) = b.size_hint(); 198 199 let lower = a_lower.saturating_add(b_lower); 200 201 let upper = match (a_upper, b_upper) { 202 (Some(x), Some(y)) => x.checked_add(y), 203 _ => None, 204 }; 205 206 (lower, upper) 207 } 208 Chain { a: Some(a), b: None } => a.size_hint(), 209 Chain { a: None, b: Some(b) } => b.size_hint(), 210 Chain { a: None, b: None } => (0, Some(0)), 211 } 212 } 213 } 214 215 #[stable(feature = "rust1", since = "1.0.0")] 216 impl<A, B> DoubleEndedIterator for Chain<A, B> 217 where 218 A: DoubleEndedIterator, 219 B: DoubleEndedIterator<Item = A::Item>, 220 { 221 #[inline] next_back(&mut self) -> Option<A::Item>222 fn next_back(&mut self) -> Option<A::Item> { 223 match fuse!(self.b.next_back()) { 224 None => maybe!(self.a.next_back()), 225 item => item, 226 } 227 } 228 229 #[inline] advance_back_by(&mut self, n: usize) -> Result<(), usize>230 fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { 231 let mut rem = n; 232 233 if let Some(ref mut b) = self.b { 234 match b.advance_back_by(rem) { 235 Ok(()) => return Ok(()), 236 Err(k) => rem -= k, 237 } 238 self.b = None; 239 } 240 241 if let Some(ref mut a) = self.a { 242 match a.advance_back_by(rem) { 243 Ok(()) => return Ok(()), 244 Err(k) => rem -= k, 245 } 246 // we don't fuse the second iterator 247 } 248 249 if rem == 0 { Ok(()) } else { Err(n - rem) } 250 } 251 252 #[inline] nth_back(&mut self, mut n: usize) -> Option<Self::Item>253 fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> { 254 if let Some(ref mut b) = self.b { 255 match b.advance_back_by(n) { 256 Ok(()) => match b.next_back() { 257 None => n = 0, 258 x => return x, 259 }, 260 Err(k) => n -= k, 261 } 262 263 self.b = None; 264 } 265 266 maybe!(self.a.nth_back(n)) 267 } 268 269 #[inline] rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item> where P: FnMut(&Self::Item) -> bool,270 fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item> 271 where 272 P: FnMut(&Self::Item) -> bool, 273 { 274 match fuse!(self.b.rfind(&mut predicate)) { 275 None => maybe!(self.a.rfind(predicate)), 276 item => item, 277 } 278 } 279 try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R where Self: Sized, F: FnMut(Acc, Self::Item) -> R, R: Try<Output = Acc>,280 fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R 281 where 282 Self: Sized, 283 F: FnMut(Acc, Self::Item) -> R, 284 R: Try<Output = Acc>, 285 { 286 if let Some(ref mut b) = self.b { 287 acc = b.try_rfold(acc, &mut f)?; 288 self.b = None; 289 } 290 if let Some(ref mut a) = self.a { 291 acc = a.try_rfold(acc, f)?; 292 // we don't fuse the second iterator 293 } 294 try { acc } 295 } 296 rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc where F: FnMut(Acc, Self::Item) -> Acc,297 fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc 298 where 299 F: FnMut(Acc, Self::Item) -> Acc, 300 { 301 if let Some(b) = self.b { 302 acc = b.rfold(acc, &mut f); 303 } 304 if let Some(a) = self.a { 305 acc = a.rfold(acc, f); 306 } 307 acc 308 } 309 } 310 311 // Note: *both* must be fused to handle double-ended iterators. 312 #[stable(feature = "fused", since = "1.26.0")] 313 impl<A, B> FusedIterator for Chain<A, B> 314 where 315 A: FusedIterator, 316 B: FusedIterator<Item = A::Item>, 317 { 318 } 319 320 #[unstable(feature = "trusted_len", issue = "37572")] 321 unsafe impl<A, B> TrustedLen for Chain<A, B> 322 where 323 A: TrustedLen, 324 B: TrustedLen<Item = A::Item>, 325 { 326 } 327