1 // Copyright (c) 2023 Huawei Device Co., Ltd.
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 // http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13
14 use std::iter;
15
16 use super::ParallelIterator;
17
zip<A, B>(mut a: A, mut b: B) -> Zip<A, B> where A: ParallelIterator, B: ParallelIterator,18 pub fn zip<A, B>(mut a: A, mut b: B) -> Zip<A, B>
19 where
20 A: ParallelIterator,
21 B: ParallelIterator,
22 {
23 let len_a = a.len();
24 let len_b = b.len();
25 if len_a > len_b {
26 a = a.reduce(len_b);
27 } else {
28 b = b.reduce(len_a);
29 }
30 Zip { a, b }
31 }
32
33 pub struct Zip<A, B> {
34 a: A,
35 b: B,
36 }
37
38 impl<A, B> ParallelIterator for Zip<A, B>
39 where
40 A: ParallelIterator,
41 B: ParallelIterator,
42 {
43 type Item = (A::Item, B::Item);
44 type Iter = iter::Zip<A::Iter, B::Iter>;
45
len(&self) -> usize46 fn len(&self) -> usize {
47 self.a.len()
48 }
49
reduce(self, len: usize) -> Self50 fn reduce(self, len: usize) -> Self {
51 Self {
52 a: self.a.reduce(len),
53 b: self.b.reduce(len),
54 }
55 }
56
split(self) -> (Self, Option<Self>)57 fn split(self) -> (Self, Option<Self>) {
58 let (left_a, right_a) = self.a.split();
59 let (left_b, right_b) = self.b.split();
60 let left = Zip {
61 a: left_a,
62 b: left_b,
63 };
64 let right = match (right_a, right_b) {
65 (Some(a), Some(b)) => Some(Zip { a, b }),
66 _ => None,
67 };
68 (left, right)
69 }
70
iter(self) -> Self::Iter71 fn iter(self) -> Self::Iter {
72 self.a.iter().zip(self.b.iter())
73 }
74 }
75