mz_ore/
iter.rs

Help
1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Iterator utilities.
17
18use std::fmt::Debug;
19use std::iter::{self, Chain, Once, Peekable};
20
21/// Extension methods for iterators.
22pub trait IteratorExt
23where
24    Self: Iterator + Sized,
25{
26    /// Chains a single `item` onto the end of this iterator.
27    ///
28    /// Equivalent to `self.chain(iter::once(item))`.
29    fn chain_one(self, item: Self::Item) -> Chain<Self, Once<Self::Item>> {
30        self.chain(iter::once(item))
31    }
32
33    /// Reports whether all the elements of the iterator are the same.
34    ///
35    /// This condition is trivially true for iterators with zero or one elements.
36    fn all_equal(mut self) -> bool
37    where
38        Self::Item: PartialEq,
39    {
40        match self.next() {
41            None => true,
42            Some(v1) => self.all(|v2| v1 == v2),
43        }
44    }
45
46    /// Converts the the iterator into an `ExactSizeIterator` reporting the given size.
47    ///
48    /// The caller is responsible for providing the correct size of the iterator! Providing an
49    /// incorrect size value will lead to panics and/or incorrect responses to size queries.
50    ///
51    /// # Panics
52    ///
53    /// Panics if the given length is not consistent with this iterator's `size_hint`.
54    fn exact_size(self, len: usize) -> ExactSize<Self> {
55        let (lower, upper) = self.size_hint();
56        assert!(
57            lower <= len && upper.map_or(true, |upper| upper >= len),
58            "provided length {len} inconsistent with `size_hint`: {:?}",
59            (lower, upper)
60        );
61
62        ExactSize { inner: self, len }
63    }
64
65    /// Wrap this iterator with one that yields a tuple of the iterator element and the extra
66    /// value on each iteration. The extra value is cloned for each but the last `Some` element
67    /// returned.
68    ///
69    /// This is useful to provide an owned extra value to each iteration, but only clone it
70    /// when necessary.
71    ///
72    /// NOTE: Once the iterator starts producing `None` values, the extra value will be consumed
73    /// and no longer be available. This should not be used for iterators that may produce
74    /// `Some` values after producing `None`.
75    fn repeat_clone<A: Clone>(self, extra_val: A) -> RepeatClone<Self, A> {
76        RepeatClone {
77            iter: self.peekable(),
78            extra_val: Some(extra_val),
79        }
80    }
81}
82
83impl<I> IteratorExt for I where I: Iterator {}
84
85/// Iterator type returned by [`IteratorExt::exact_size`].
86#[derive(Debug)]
87pub struct ExactSize<I> {
88    inner: I,
89    len: usize,
90}
91
92impl<I: Iterator> Iterator for ExactSize<I> {
93    type Item = I::Item;
94
95    fn next(&mut self) -> Option<Self::Item> {
96        self.len = self.len.saturating_sub(1);
97        self.inner.next()
98    }
99
100    fn size_hint(&self) -> (usize, Option<usize>) {
101        (self.len, Some(self.len))
102    }
103}
104
105impl<I: Iterator> ExactSizeIterator for ExactSize<I> {}
106
107/// Iterator type returned by [`IteratorExt::repeat_clone`].
108pub struct RepeatClone<I: Iterator, A> {
109    iter: Peekable<I>,
110    extra_val: Option<A>,
111}
112
113impl<I: Iterator, A: Clone> Iterator for RepeatClone<I, A> {
114    type Item = (I::Item, A);
115
116    fn next(&mut self) -> Option<Self::Item> {
117        let next = self.iter.next()?;
118
119        // Clone the extra_val only if there is an item to return on the next call to `next`.
120        let val = match self.iter.peek() {
121            Some(_) => self.extra_val.clone(),
122            None => self.extra_val.take(),
123        };
124
125        // We should always return a value if there is a current element.
126        Some((next, val.expect("RepeatClone invariant violated")))
127    }
128}
129
130impl<I: Iterator<Item: Debug> + Debug, A: Debug> Debug for RepeatClone<I, A> {
131    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
132        f.debug_struct("RepeatClone")
133            .field("iter", &self.iter)
134            .field("extra_val", &self.extra_val)
135            .finish()
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use crate::iter::IteratorExt;
142
143    #[crate::test]
144    fn test_all_equal() {
145        let empty: [i64; 0] = [];
146        assert!(empty.iter().all_equal());
147        assert!([1].iter().all_equal());
148        assert!([1, 1].iter().all_equal());
149        assert!(![1, 2].iter().all_equal());
150    }
151}
mz_ore/iter.rs

mz_ore/
iter.rs
