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

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License in the LICENSE file at the
// root of this repository, or online at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Iterator utilities.

use std::fmt::Debug;
use std::iter::{self, Chain, Once, Peekable};

/// Extension methods for iterators.
pub trait IteratorExt
where
    Self: Iterator + Sized,
{
    /// Chains a single `item` onto the end of this iterator.
    ///
    /// Equivalent to `self.chain(iter::once(item))`.
    fn chain_one(self, item: Self::Item) -> Chain<Self, Once<Self::Item>> {
        self.chain(iter::once(item))
    }

    /// Reports whether all the elements of the iterator are the same.
    ///
    /// This condition is trivially true for iterators with zero or one elements.
    fn all_equal(mut self) -> bool
    where
        Self::Item: PartialEq,
    {
        match self.next() {
            None => true,
            Some(v1) => self.all(|v2| v1 == v2),
        }
    }

    /// Converts the the iterator into an `ExactSizeIterator` reporting the given size.
    ///
    /// The caller is responsible for providing the correct size of the iterator! Providing an
    /// incorrect size value will lead to panics and/or incorrect responses to size queries.
    ///
    /// # Panics
    ///
    /// Panics if the given length is not consistent with this iterator's `size_hint`.
    fn exact_size(self, len: usize) -> ExactSize<Self> {
        let (lower, upper) = self.size_hint();
        assert!(
            lower <= len && upper.map_or(true, |upper| upper >= len),
            "provided length {len} inconsistent with `size_hint`: {:?}",
            (lower, upper)
        );

        ExactSize { inner: self, len }
    }

    /// Wrap this iterator with one that yields a tuple of the iterator element and the extra
    /// value on each iteration. The extra value is cloned for each but the last `Some` element
    /// returned.
    ///
    /// This is useful to provide an owned extra value to each iteration, but only clone it
    /// when necessary.
    ///
    /// NOTE: Once the iterator starts producing `None` values, the extra value will be consumed
    /// and no longer be available. This should not be used for iterators that may produce
    /// `Some` values after producing `None`.
    fn repeat_clone<A: Clone>(self, extra_val: A) -> RepeatClone<Self, A> {
        RepeatClone {
            iter: self.peekable(),
            extra_val: Some(extra_val),
        }
    }
}

impl<I> IteratorExt for I where I: Iterator {}

/// Iterator type returned by [`IteratorExt::exact_size`].
#[derive(Debug)]
pub struct ExactSize<I> {
    inner: I,
    len: usize,
}

impl<I: Iterator> Iterator for ExactSize<I> {
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        self.len = self.len.saturating_sub(1);
        self.inner.next()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.len, Some(self.len))
    }
}

impl<I: Iterator> ExactSizeIterator for ExactSize<I> {}

/// Iterator type returned by [`IteratorExt::repeat_clone`].
pub struct RepeatClone<I: Iterator, A> {
    iter: Peekable<I>,
    extra_val: Option<A>,
}

impl<I: Iterator, A: Clone> Iterator for RepeatClone<I, A> {
    type Item = (I::Item, A);

    fn next(&mut self) -> Option<Self::Item> {
        let next = self.iter.next()?;

        // Clone the extra_val only if there is an item to return on the next call to `next`.
        let val = match self.iter.peek() {
            Some(_) => self.extra_val.clone(),
            None => self.extra_val.take(),
        };

        // We should always return a value if there is a current element.
        Some((next, val.expect("RepeatClone invariant violated")))
    }
}

impl<I: Iterator<Item: Debug> + Debug, A: Debug> Debug for RepeatClone<I, A> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RepeatClone")
            .field("iter", &self.iter)
            .field("extra_val", &self.extra_val)
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use crate::iter::IteratorExt;

    #[crate::test]
    fn test_all_equal() {
        let empty: [i64; 0] = [];
        assert!(empty.iter().all_equal());
        assert!([1].iter().all_equal());
        assert!([1, 1].iter().all_equal());
        assert!(![1, 2].iter().all_equal());
    }
}