1 use super::plumbing::*; 2 use super::*; 3 use std::cmp; 4 use std::iter::Fuse; 5 6 /// `Interleave` is an iterator that interleaves elements of iterators 7 /// `i` and `j` in one continuous iterator. This struct is created by 8 /// the [`interleave()`] method on [`IndexedParallelIterator`] 9 /// 10 /// [`interleave()`]: trait.IndexedParallelIterator.html#method.interleave 11 /// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html 12 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] 13 #[derive(Debug, Clone)] 14 pub struct Interleave<I, J> 15 where 16 I: IndexedParallelIterator, 17 J: IndexedParallelIterator<Item = I::Item>, 18 { 19 i: I, 20 j: J, 21 } 22 23 impl<I, J> Interleave<I, J> 24 where 25 I: IndexedParallelIterator, 26 J: IndexedParallelIterator<Item = I::Item>, 27 { 28 /// Creates a new `Interleave` iterator new(i: I, j: J) -> Self29 pub(super) fn new(i: I, j: J) -> Self { 30 Interleave { i, j } 31 } 32 } 33 34 impl<I, J> ParallelIterator for Interleave<I, J> 35 where 36 I: IndexedParallelIterator, 37 J: IndexedParallelIterator<Item = I::Item>, 38 { 39 type Item = I::Item; 40 drive_unindexed<C>(self, consumer: C) -> C::Result where C: Consumer<I::Item>,41 fn drive_unindexed<C>(self, consumer: C) -> C::Result 42 where 43 C: Consumer<I::Item>, 44 { 45 bridge(self, consumer) 46 } 47 opt_len(&self) -> Option<usize>48 fn opt_len(&self) -> Option<usize> { 49 Some(self.len()) 50 } 51 } 52 53 impl<I, J> IndexedParallelIterator for Interleave<I, J> 54 where 55 I: IndexedParallelIterator, 56 J: IndexedParallelIterator<Item = I::Item>, 57 { drive<C>(self, consumer: C) -> C::Result where C: Consumer<Self::Item>,58 fn drive<C>(self, consumer: C) -> C::Result 59 where 60 C: Consumer<Self::Item>, 61 { 62 bridge(self, consumer) 63 } 64 len(&self) -> usize65 fn len(&self) -> usize { 66 self.i.len().checked_add(self.j.len()).expect("overflow") 67 } 68 with_producer<CB>(self, callback: CB) -> CB::Output where CB: ProducerCallback<Self::Item>,69 fn with_producer<CB>(self, callback: CB) -> CB::Output 70 where 71 CB: ProducerCallback<Self::Item>, 72 { 73 let (i_len, j_len) = (self.i.len(), self.j.len()); 74 return self.i.with_producer(CallbackI { 75 callback, 76 i_len, 77 j_len, 78 i_next: false, 79 j: self.j, 80 }); 81 82 struct CallbackI<CB, J> { 83 callback: CB, 84 i_len: usize, 85 j_len: usize, 86 i_next: bool, 87 j: J, 88 } 89 90 impl<CB, J> ProducerCallback<J::Item> for CallbackI<CB, J> 91 where 92 J: IndexedParallelIterator, 93 CB: ProducerCallback<J::Item>, 94 { 95 type Output = CB::Output; 96 97 fn callback<I>(self, i_producer: I) -> Self::Output 98 where 99 I: Producer<Item = J::Item>, 100 { 101 self.j.with_producer(CallbackJ { 102 i_producer, 103 i_len: self.i_len, 104 j_len: self.j_len, 105 i_next: self.i_next, 106 callback: self.callback, 107 }) 108 } 109 } 110 111 struct CallbackJ<CB, I> { 112 callback: CB, 113 i_len: usize, 114 j_len: usize, 115 i_next: bool, 116 i_producer: I, 117 } 118 119 impl<CB, I> ProducerCallback<I::Item> for CallbackJ<CB, I> 120 where 121 I: Producer, 122 CB: ProducerCallback<I::Item>, 123 { 124 type Output = CB::Output; 125 126 fn callback<J>(self, j_producer: J) -> Self::Output 127 where 128 J: Producer<Item = I::Item>, 129 { 130 let producer = InterleaveProducer::new( 131 self.i_producer, 132 j_producer, 133 self.i_len, 134 self.j_len, 135 self.i_next, 136 ); 137 self.callback.callback(producer) 138 } 139 } 140 } 141 } 142 143 struct InterleaveProducer<I, J> 144 where 145 I: Producer, 146 J: Producer<Item = I::Item>, 147 { 148 i: I, 149 j: J, 150 i_len: usize, 151 j_len: usize, 152 i_next: bool, 153 } 154 155 impl<I, J> InterleaveProducer<I, J> 156 where 157 I: Producer, 158 J: Producer<Item = I::Item>, 159 { new(i: I, j: J, i_len: usize, j_len: usize, i_next: bool) -> InterleaveProducer<I, J>160 fn new(i: I, j: J, i_len: usize, j_len: usize, i_next: bool) -> InterleaveProducer<I, J> { 161 InterleaveProducer { 162 i, 163 j, 164 i_len, 165 j_len, 166 i_next, 167 } 168 } 169 } 170 171 impl<I, J> Producer for InterleaveProducer<I, J> 172 where 173 I: Producer, 174 J: Producer<Item = I::Item>, 175 { 176 type Item = I::Item; 177 type IntoIter = InterleaveSeq<I::IntoIter, J::IntoIter>; 178 into_iter(self) -> Self::IntoIter179 fn into_iter(self) -> Self::IntoIter { 180 InterleaveSeq { 181 i: self.i.into_iter().fuse(), 182 j: self.j.into_iter().fuse(), 183 i_next: self.i_next, 184 } 185 } 186 min_len(&self) -> usize187 fn min_len(&self) -> usize { 188 cmp::max(self.i.min_len(), self.j.min_len()) 189 } 190 max_len(&self) -> usize191 fn max_len(&self) -> usize { 192 cmp::min(self.i.max_len(), self.j.max_len()) 193 } 194 195 /// We know 0 < index <= self.i_len + self.j_len 196 /// 197 /// Find a, b satisfying: 198 /// 199 /// (1) 0 < a <= self.i_len 200 /// (2) 0 < b <= self.j_len 201 /// (3) a + b == index 202 /// 203 /// For even splits, set a = b = index/2. 204 /// For odd splits, set a = (index/2)+1, b = index/2, if `i` 205 /// should yield the next element, otherwise, if `j` should yield 206 /// the next element, set a = index/2 and b = (index/2)+1 split_at(self, index: usize) -> (Self, Self)207 fn split_at(self, index: usize) -> (Self, Self) { 208 #[inline] 209 fn odd_offset(flag: bool) -> usize { 210 (!flag) as usize 211 } 212 213 let even = index % 2 == 0; 214 let idx = index >> 1; 215 216 // desired split 217 let (i_idx, j_idx) = ( 218 idx + odd_offset(even || self.i_next), 219 idx + odd_offset(even || !self.i_next), 220 ); 221 222 let (i_split, j_split) = if self.i_len >= i_idx && self.j_len >= j_idx { 223 (i_idx, j_idx) 224 } else if self.i_len >= i_idx { 225 // j too short 226 (index - self.j_len, self.j_len) 227 } else { 228 // i too short 229 (self.i_len, index - self.i_len) 230 }; 231 232 let trailing_i_next = even == self.i_next; 233 let (i_left, i_right) = self.i.split_at(i_split); 234 let (j_left, j_right) = self.j.split_at(j_split); 235 236 ( 237 InterleaveProducer::new(i_left, j_left, i_split, j_split, self.i_next), 238 InterleaveProducer::new( 239 i_right, 240 j_right, 241 self.i_len - i_split, 242 self.j_len - j_split, 243 trailing_i_next, 244 ), 245 ) 246 } 247 } 248 249 /// Wrapper for Interleave to implement DoubleEndedIterator and 250 /// ExactSizeIterator. 251 /// 252 /// This iterator is fused. 253 struct InterleaveSeq<I, J> { 254 i: Fuse<I>, 255 j: Fuse<J>, 256 257 /// Flag to control which iterator should provide the next element. When 258 /// `false` then `i` produces the next element, otherwise `j` produces the 259 /// next element. 260 i_next: bool, 261 } 262 263 /// Iterator implementation for InterleaveSeq. This implementation is 264 /// taken more or less verbatim from itertools. It is replicated here 265 /// (instead of calling itertools directly), because we also need to 266 /// implement `DoubledEndedIterator` and `ExactSizeIterator`. 267 impl<I, J> Iterator for InterleaveSeq<I, J> 268 where 269 I: Iterator, 270 J: Iterator<Item = I::Item>, 271 { 272 type Item = I::Item; 273 274 #[inline] next(&mut self) -> Option<Self::Item>275 fn next(&mut self) -> Option<Self::Item> { 276 self.i_next = !self.i_next; 277 if self.i_next { 278 match self.i.next() { 279 None => self.j.next(), 280 r => r, 281 } 282 } else { 283 match self.j.next() { 284 None => self.i.next(), 285 r => r, 286 } 287 } 288 } 289 size_hint(&self) -> (usize, Option<usize>)290 fn size_hint(&self) -> (usize, Option<usize>) { 291 let (ih, jh) = (self.i.size_hint(), self.j.size_hint()); 292 let min = ih.0.saturating_add(jh.0); 293 let max = match (ih.1, jh.1) { 294 (Some(x), Some(y)) => x.checked_add(y), 295 _ => None, 296 }; 297 (min, max) 298 } 299 } 300 301 // The implementation for DoubleEndedIterator requires 302 // ExactSizeIterator to provide `next_back()`. The last element will 303 // come from the iterator that runs out last (ie has the most elements 304 // in it). If the iterators have the same number of elements, then the 305 // last iterator will provide the last element. 306 impl<I, J> DoubleEndedIterator for InterleaveSeq<I, J> 307 where 308 I: DoubleEndedIterator + ExactSizeIterator, 309 J: DoubleEndedIterator<Item = I::Item> + ExactSizeIterator<Item = I::Item>, 310 { 311 #[inline] next_back(&mut self) -> Option<I::Item>312 fn next_back(&mut self) -> Option<I::Item> { 313 if self.i.len() == self.j.len() { 314 if self.i_next { 315 self.i.next_back() 316 } else { 317 self.j.next_back() 318 } 319 } else if self.i.len() < self.j.len() { 320 self.j.next_back() 321 } else { 322 self.i.next_back() 323 } 324 } 325 } 326 327 impl<I, J> ExactSizeIterator for InterleaveSeq<I, J> 328 where 329 I: ExactSizeIterator, 330 J: ExactSizeIterator<Item = I::Item>, 331 { 332 #[inline] len(&self) -> usize333 fn len(&self) -> usize { 334 self.i.len() + self.j.len() 335 } 336 } 337