1 use super::plumbing::*; 2 use super::*; 3 4 /// `MultiZip` is an iterator that zips up a tuple of parallel iterators to 5 /// produce tuples of their items. 6 /// 7 /// It is created by calling `into_par_iter()` on a tuple of types that 8 /// implement `IntoParallelIterator`, or `par_iter()`/`par_iter_mut()` with 9 /// types that are iterable by reference. 10 /// 11 /// The implementation currently support tuples up to length 12. 12 /// 13 /// # Examples 14 /// 15 /// ``` 16 /// use rayon::prelude::*; 17 /// 18 /// // This will iterate `r` by mutable reference, like `par_iter_mut()`, while 19 /// // ranges are all iterated by value like `into_par_iter()`. 20 /// // Note that the zipped iterator is only as long as the shortest input. 21 /// let mut r = vec![0; 3]; 22 /// (&mut r, 1..10, 10..100, 100..1000).into_par_iter() 23 /// .for_each(|(r, x, y, z)| *r = x * y + z); 24 /// 25 /// assert_eq!(&r, &[1 * 10 + 100, 2 * 11 + 101, 3 * 12 + 102]); 26 /// ``` 27 /// 28 /// For a group that should all be iterated by reference, you can use a tuple reference. 29 /// 30 /// ``` 31 /// use rayon::prelude::*; 32 /// 33 /// let xs: Vec<_> = (1..10).collect(); 34 /// let ys: Vec<_> = (10..100).collect(); 35 /// let zs: Vec<_> = (100..1000).collect(); 36 /// 37 /// // Reference each input separately with `IntoParallelIterator`: 38 /// let r1: Vec<_> = (&xs, &ys, &zs).into_par_iter() 39 /// .map(|(x, y, z)| x * y + z) 40 /// .collect(); 41 /// 42 /// // Reference them all together with `IntoParallelRefIterator`: 43 /// let r2: Vec<_> = (xs, ys, zs).par_iter() 44 /// .map(|(x, y, z)| x * y + z) 45 /// .collect(); 46 /// 47 /// assert_eq!(r1, r2); 48 /// ``` 49 /// 50 /// Mutable references to a tuple will work similarly. 51 /// 52 /// ``` 53 /// use rayon::prelude::*; 54 /// 55 /// let mut xs: Vec<_> = (1..4).collect(); 56 /// let mut ys: Vec<_> = (-4..-1).collect(); 57 /// let mut zs = vec![0; 3]; 58 /// 59 /// // Mutably reference each input separately with `IntoParallelIterator`: 60 /// (&mut xs, &mut ys, &mut zs).into_par_iter().for_each(|(x, y, z)| { 61 /// *z += *x + *y; 62 /// std::mem::swap(x, y); 63 /// }); 64 /// 65 /// assert_eq!(xs, (vec![-4, -3, -2])); 66 /// assert_eq!(ys, (vec![1, 2, 3])); 67 /// assert_eq!(zs, (vec![-3, -1, 1])); 68 /// 69 /// // Mutably reference them all together with `IntoParallelRefMutIterator`: 70 /// let mut tuple = (xs, ys, zs); 71 /// tuple.par_iter_mut().for_each(|(x, y, z)| { 72 /// *z += *x + *y; 73 /// std::mem::swap(x, y); 74 /// }); 75 /// 76 /// assert_eq!(tuple, (vec![1, 2, 3], vec![-4, -3, -2], vec![-6, -2, 2])); 77 /// ``` 78 #[derive(Debug, Clone)] 79 pub struct MultiZip<T> { 80 tuple: T, 81 } 82 83 // These macros greedily consume 4 or 2 items first to achieve log2 nesting depth. 84 // For example, 5 => 4,1 => (2,2),1. 85 // 86 // The tuples go up to 12, so we might want to greedily consume 8 too, but 87 // the depth works out the same if we let that expand on the right: 88 // 9 => 4,5 => (2,2),(4,1) => (2,2),((2,2),1) 89 // 12 => 4,8 => (2,2),(4,4) => (2,2),((2,2),(2,2)) 90 // 91 // But if we ever increase to 13, we would want to split 8,5 rather than 4,9. 92 93 macro_rules! reduce { 94 ($a:expr, $b:expr, $c:expr, $d:expr, $( $x:expr ),+ => $fn:path) => { 95 reduce!(reduce!($a, $b, $c, $d => $fn), 96 reduce!($( $x ),+ => $fn) 97 => $fn) 98 }; 99 ($a:expr, $b:expr, $( $x:expr ),+ => $fn:path) => { 100 reduce!(reduce!($a, $b => $fn), 101 reduce!($( $x ),+ => $fn) 102 => $fn) 103 }; 104 ($a:expr, $b:expr => $fn:path) => { $fn($a, $b) }; 105 ($a:expr => $fn:path) => { $a }; 106 } 107 108 macro_rules! nest { 109 ($A:tt, $B:tt, $C:tt, $D:tt, $( $X:tt ),+) => { 110 (nest!($A, $B, $C, $D), nest!($( $X ),+)) 111 }; 112 ($A:tt, $B:tt, $( $X:tt ),+) => { 113 (($A, $B), nest!($( $X ),+)) 114 }; 115 ($A:tt, $B:tt) => { ($A, $B) }; 116 ($A:tt) => { $A }; 117 } 118 119 macro_rules! flatten { 120 ($( $T:ident ),+) => {{ 121 #[allow(non_snake_case)] 122 fn flatten<$( $T ),+>(nest!($( $T ),+) : nest!($( $T ),+)) -> ($( $T, )+) { 123 ($( $T, )+) 124 } 125 flatten 126 }}; 127 } 128 129 macro_rules! multizip_impls { 130 ($( 131 $Tuple:ident { 132 $(($idx:tt) -> $T:ident)+ 133 } 134 )+) => { 135 $( 136 impl<$( $T, )+> IntoParallelIterator for ($( $T, )+) 137 where 138 $( 139 $T: IntoParallelIterator, 140 $T::Iter: IndexedParallelIterator, 141 )+ 142 { 143 type Item = ($( $T::Item, )+); 144 type Iter = MultiZip<($( $T::Iter, )+)>; 145 146 fn into_par_iter(self) -> Self::Iter { 147 MultiZip { 148 tuple: ( $( self.$idx.into_par_iter(), )+ ), 149 } 150 } 151 } 152 153 impl<'a, $( $T, )+> IntoParallelIterator for &'a ($( $T, )+) 154 where 155 $( 156 $T: IntoParallelRefIterator<'a>, 157 $T::Iter: IndexedParallelIterator, 158 )+ 159 { 160 type Item = ($( $T::Item, )+); 161 type Iter = MultiZip<($( $T::Iter, )+)>; 162 163 fn into_par_iter(self) -> Self::Iter { 164 MultiZip { 165 tuple: ( $( self.$idx.par_iter(), )+ ), 166 } 167 } 168 } 169 170 impl<'a, $( $T, )+> IntoParallelIterator for &'a mut ($( $T, )+) 171 where 172 $( 173 $T: IntoParallelRefMutIterator<'a>, 174 $T::Iter: IndexedParallelIterator, 175 )+ 176 { 177 type Item = ($( $T::Item, )+); 178 type Iter = MultiZip<($( $T::Iter, )+)>; 179 180 fn into_par_iter(self) -> Self::Iter { 181 MultiZip { 182 tuple: ( $( self.$idx.par_iter_mut(), )+ ), 183 } 184 } 185 } 186 187 impl<$( $T, )+> ParallelIterator for MultiZip<($( $T, )+)> 188 where 189 $( $T: IndexedParallelIterator, )+ 190 { 191 type Item = ($( $T::Item, )+); 192 193 fn drive_unindexed<CONSUMER>(self, consumer: CONSUMER) -> CONSUMER::Result 194 where 195 CONSUMER: UnindexedConsumer<Self::Item>, 196 { 197 self.drive(consumer) 198 } 199 200 fn opt_len(&self) -> Option<usize> { 201 Some(self.len()) 202 } 203 } 204 205 impl<$( $T, )+> IndexedParallelIterator for MultiZip<($( $T, )+)> 206 where 207 $( $T: IndexedParallelIterator, )+ 208 { 209 fn drive<CONSUMER>(self, consumer: CONSUMER) -> CONSUMER::Result 210 where 211 CONSUMER: Consumer<Self::Item>, 212 { 213 reduce!($( self.tuple.$idx ),+ => IndexedParallelIterator::zip) 214 .map(flatten!($( $T ),+)) 215 .drive(consumer) 216 } 217 218 fn len(&self) -> usize { 219 reduce!($( self.tuple.$idx.len() ),+ => Ord::min) 220 } 221 222 fn with_producer<CB>(self, callback: CB) -> CB::Output 223 where 224 CB: ProducerCallback<Self::Item>, 225 { 226 reduce!($( self.tuple.$idx ),+ => IndexedParallelIterator::zip) 227 .map(flatten!($( $T ),+)) 228 .with_producer(callback) 229 } 230 } 231 )+ 232 } 233 } 234 235 multizip_impls! { 236 Tuple1 { 237 (0) -> A 238 } 239 Tuple2 { 240 (0) -> A 241 (1) -> B 242 } 243 Tuple3 { 244 (0) -> A 245 (1) -> B 246 (2) -> C 247 } 248 Tuple4 { 249 (0) -> A 250 (1) -> B 251 (2) -> C 252 (3) -> D 253 } 254 Tuple5 { 255 (0) -> A 256 (1) -> B 257 (2) -> C 258 (3) -> D 259 (4) -> E 260 } 261 Tuple6 { 262 (0) -> A 263 (1) -> B 264 (2) -> C 265 (3) -> D 266 (4) -> E 267 (5) -> F 268 } 269 Tuple7 { 270 (0) -> A 271 (1) -> B 272 (2) -> C 273 (3) -> D 274 (4) -> E 275 (5) -> F 276 (6) -> G 277 } 278 Tuple8 { 279 (0) -> A 280 (1) -> B 281 (2) -> C 282 (3) -> D 283 (4) -> E 284 (5) -> F 285 (6) -> G 286 (7) -> H 287 } 288 Tuple9 { 289 (0) -> A 290 (1) -> B 291 (2) -> C 292 (3) -> D 293 (4) -> E 294 (5) -> F 295 (6) -> G 296 (7) -> H 297 (8) -> I 298 } 299 Tuple10 { 300 (0) -> A 301 (1) -> B 302 (2) -> C 303 (3) -> D 304 (4) -> E 305 (5) -> F 306 (6) -> G 307 (7) -> H 308 (8) -> I 309 (9) -> J 310 } 311 Tuple11 { 312 (0) -> A 313 (1) -> B 314 (2) -> C 315 (3) -> D 316 (4) -> E 317 (5) -> F 318 (6) -> G 319 (7) -> H 320 (8) -> I 321 (9) -> J 322 (10) -> K 323 } 324 Tuple12 { 325 (0) -> A 326 (1) -> B 327 (2) -> C 328 (3) -> D 329 (4) -> E 330 (5) -> F 331 (6) -> G 332 (7) -> H 333 (8) -> I 334 (9) -> J 335 (10) -> K 336 (11) -> L 337 } 338 } 339