summaryrefslogtreecommitdiffstats
path: root/third_party/rust/itertools/benches/extra/zipslices.rs
blob: 633be59068691674320b2d3d171b3df06c01ef7c (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
use std::cmp;

// Note: There are different ways to implement ZipSlices.
// This version performed the best in benchmarks.
//
// I also implemented a version with three pointers (tptr, tend, uptr),
// that mimiced slice::Iter and only checked bounds by using tptr == tend,
// but that was inferior to this solution.

/// An iterator which iterates two slices simultaneously.
///
/// `ZipSlices` acts like a double-ended `.zip()` iterator.
///
/// It was intended to be more efficient than `.zip()`, and it was, then
/// rustc changed how it optimizes so it can not promise improved performance
/// at this time.
///
/// Note that elements past the end of the shortest of the two slices are ignored.
///
/// Iterator element type for `ZipSlices<T, U>` is `(T::Item, U::Item)`. For example,
/// for a `ZipSlices<&'a [A], &'b mut [B]>`, the element type is `(&'a A, &'b mut B)`.
#[derive(Clone)]
pub struct ZipSlices<T, U> {
    t: T,
    u: U,
    len: usize,
    index: usize,
}

impl<'a, 'b, A, B> ZipSlices<&'a [A], &'b [B]> {
    /// Create a new `ZipSlices` from slices `a` and `b`.
    ///
    /// Act like a double-ended `.zip()` iterator, but more efficiently.
    ///
    /// Note that elements past the end of the shortest of the two slices are ignored.
    #[inline(always)]
    pub fn new(a: &'a [A], b: &'b [B]) -> Self {
        let minl = cmp::min(a.len(), b.len());
        ZipSlices {
            t: a,
            u: b,
            len: minl,
            index: 0,
        }
    }
}

impl<T, U> ZipSlices<T, U>
    where T: Slice,
          U: Slice
{
    /// Create a new `ZipSlices` from slices `a` and `b`.
    ///
    /// Act like a double-ended `.zip()` iterator, but more efficiently.
    ///
    /// Note that elements past the end of the shortest of the two slices are ignored.
    #[inline(always)]
    pub fn from_slices(a: T, b: U) -> Self {
        let minl = cmp::min(a.len(), b.len());
        ZipSlices {
            t: a,
            u: b,
            len: minl,
            index: 0,
        }
    }
}

impl<T, U> Iterator for ZipSlices<T, U>
    where T: Slice,
          U: Slice
{
    type Item = (T::Item, U::Item);

    #[inline(always)]
    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            if self.index >= self.len {
                None
            } else {
                let i = self.index;
                self.index += 1;
                Some((
                    self.t.get_unchecked(i),
                    self.u.get_unchecked(i)))
            }
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.len - self.index;
        (len, Some(len))
    }
}

impl<T, U> DoubleEndedIterator for ZipSlices<T, U>
    where T: Slice,
          U: Slice
{
    #[inline(always)]
    fn next_back(&mut self) -> Option<Self::Item> {
        unsafe {
            if self.index >= self.len {
                None
            } else {
                self.len -= 1;
                let i = self.len;
                Some((
                    self.t.get_unchecked(i),
                    self.u.get_unchecked(i)))
            }
        }
    }
}

impl<T, U> ExactSizeIterator for ZipSlices<T, U>
    where T: Slice,
          U: Slice
{}

unsafe impl<T, U> Slice for ZipSlices<T, U>
    where T: Slice,
          U: Slice
{
    type Item = (T::Item, U::Item);

    fn len(&self) -> usize {
        self.len - self.index
    }

    unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item {
        (self.t.get_unchecked(i),
         self.u.get_unchecked(i))
    }
}

/// A helper trait to let `ZipSlices` accept both `&[T]` and `&mut [T]`.
///
/// Unsafe trait because:
///
/// - Implementors must guarantee that `get_unchecked` is valid for all indices `0..len()`.
pub unsafe trait Slice {
    /// The type of a reference to the slice's elements
    type Item;
    #[doc(hidden)]
    fn len(&self) -> usize;
    #[doc(hidden)]
    unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item;
}

unsafe impl<'a, T> Slice for &'a [T] {
    type Item = &'a T;
    #[inline(always)]
    fn len(&self) -> usize { (**self).len() }
    #[inline(always)]
    unsafe fn get_unchecked(&mut self, i: usize) -> &'a T {
        debug_assert!(i < self.len());
        (**self).get_unchecked(i)
    }
}

unsafe impl<'a, T> Slice for &'a mut [T] {
    type Item = &'a mut T;
    #[inline(always)]
    fn len(&self) -> usize { (**self).len() }
    #[inline(always)]
    unsafe fn get_unchecked(&mut self, i: usize) -> &'a mut T {
        debug_assert!(i < self.len());
        // override the lifetime constraints of &mut &'a mut [T]
        (*(*self as *mut [T])).get_unchecked_mut(i)
    }
}

#[test]
fn zipslices() {

    let xs = [1, 2, 3, 4, 5, 6];
    let ys = [1, 2, 3, 7];
    ::itertools::assert_equal(ZipSlices::new(&xs, &ys), xs.iter().zip(&ys));

    let xs = [1, 2, 3, 4, 5, 6];
    let mut ys = [0; 6];
    for (x, y) in ZipSlices::from_slices(&xs[..], &mut ys[..]) {
        *y = *x;
    }
    ::itertools::assert_equal(&xs, &ys);
}