estimate_size/
lib.rs

1/// An iterator wrapper that estimates the size of the underlying iterator.
2pub struct SizeEstimate<I> {
3    iter: I,
4    lower: usize,
5    upper: Option<usize>,
6}
7
8impl<I: Iterator + Sized> Iterator for SizeEstimate<I> {
9    type Item = I::Item;
10
11    #[inline]
12    fn next(&mut self) -> Option<Self::Item> {
13        self.lower = self.lower.saturating_sub(1);
14        self.upper = self.upper.map(|u| u.saturating_sub(1));
15
16        self.iter.next()
17    }
18
19    #[inline]
20    fn size_hint(&self) -> (usize, Option<usize>) {
21        (self.lower, self.upper)
22    }
23}
24
25/// An extension trait for iterators that allows suggesting custom size hints.
26/// Particularly useful for efficient pre-allocation of collections into Vec.
27pub trait EstimateSize: Iterator + Sized {
28    fn estimate_size(self, lower: usize, upper: Option<usize>) -> SizeEstimate<Self> {
29        SizeEstimate {
30            iter: self,
31            lower,
32            upper,
33        }
34    }
35
36    fn estimate_exact_size(self, size: usize) -> SizeEstimate<Self> {
37        self.estimate_size(size, Some(size))
38    }
39
40    fn estimate_min_size(self, lower: usize) -> SizeEstimate<Self> {
41        let (_, prev_upper) = self.size_hint();
42        let upper = prev_upper.map(|u| u.max(lower));
43        self.estimate_size(lower, upper)
44    }
45
46    fn estimate_max_size(self, upper: Option<usize>) -> SizeEstimate<Self> {
47        let (prev_lower, _) = self.size_hint();
48        let lower = if let Some(u) = upper {
49            prev_lower.min(u)
50        } else {
51            prev_lower
52        };
53        self.estimate_size(lower, upper)
54    }
55}
56
57// Implement the trait for all iterators
58impl<I: Iterator> EstimateSize for I {}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[test]
65    fn test_underflow() {
66        let bad_hint_iter = std::iter::successors(Some(0), |n| Some(n + 1).filter(|&x| x < 5))
67            .estimate_size(3, Some(3));
68
69        let vec: Vec<i32> = bad_hint_iter.collect();
70
71        assert!(vec.capacity() >= 5);
72        assert!(vec.len() >= 5);
73    }
74
75    #[test]
76    fn test_estimate_exact_size() {
77        let iter = (0..10).estimate_exact_size(10);
78        assert_eq!(iter.size_hint(), (10, Some(10)));
79
80        let collected: Vec<_> = iter.collect();
81        assert_eq!(collected.len(), 10);
82    }
83
84    #[test]
85    fn test_estimate_min_size() {
86        let iter = (0..7).estimate_min_size(10);
87        assert_eq!(iter.size_hint(), (10, Some(10)));
88
89        let collected: Vec<_> = iter.collect();
90        assert_eq!(collected.len(), 7); // Actual length is 7
91    }
92
93    #[test]
94    fn test_estimate_max_size() {
95        let iter = (0..15).estimate_max_size(Some(10));
96        assert_eq!(iter.size_hint(), (10, Some(10))); // Lower bound is min of actual and specified
97
98        let collected: Vec<_> = iter.collect();
99        assert_eq!(collected.len(), 15); // Actual length is 15
100    }
101
102    #[test]
103    fn test_iterator_behavior() {
104        let mut iter = (0..5).estimate_size(10, Some(20));
105
106        assert_eq!(iter.size_hint(), (10, Some(20)));
107        assert_eq!(iter.next(), Some(0));
108        assert_eq!(iter.size_hint(), (9, Some(19)));
109        assert_eq!(iter.next(), Some(1));
110        assert_eq!(iter.size_hint(), (8, Some(18)));
111
112        // Collect the rest to consume
113        let remaining: Vec<_> = iter.collect();
114        assert_eq!(remaining, vec![2, 3, 4]);
115        // Size hint isn't checked after collection as the iterator is consumed
116    }
117
118    #[test]
119    fn test_with_empty_iterator() {
120        let empty_iter: std::vec::IntoIter<i32> = Vec::new().into_iter();
121        let mut sized_iter = empty_iter.estimate_size(5, Some(10));
122
123        assert_eq!(sized_iter.size_hint(), (5, Some(10)));
124        assert_eq!(sized_iter.next(), None); // Iterator is empty
125        assert_eq!(sized_iter.size_hint(), (4, Some(9))); // But size hint is still decremented
126    }
127
128    #[test]
129    fn test_with_filter() {
130        let iter = (0..20)
131            .filter(|x| x % 2 == 0) // Only even numbers
132            .estimate_size(8, Some(12));
133
134        let collected: Vec<_> = iter.collect();
135        assert_eq!(collected.len(), 10); // Actually 10 even numbers
136    }
137
138    #[test]
139    fn test_saturating_behavior() {
140        let mut iter = (0..3).estimate_size(2, Some(2));
141
142        assert_eq!(iter.next(), Some(0));
143        assert_eq!(iter.size_hint(), (1, Some(1)));
144
145        assert_eq!(iter.next(), Some(1));
146        assert_eq!(iter.size_hint(), (0, Some(0)));
147
148        assert_eq!(iter.next(), Some(2));
149        // Lower bound saturates at 0, doesn't underflow
150        assert_eq!(iter.size_hint(), (0, None));
151    }
152
153    #[test]
154    fn test_chained_iterators() {
155        let first = (0..3).estimate_size(5, Some(5));
156        let second = (3..6).estimate_size(5, Some(5));
157        let chained = first.chain(second);
158
159        assert_eq!(chained.count(), 6);
160    }
161
162    #[test]
163    fn test_capacity_optimization() {
164        // Test if vector pre-allocation works with our size hint
165        let iter = (0..100).estimate_exact_size(200);
166        let vec: Vec<_> = iter.collect();
167
168        assert_eq!(vec.len(), 100); // Actual elements
169        assert!(vec.capacity() >= 100); // Capacity should be at least 100
170    }
171
172    #[test]
173    fn test_none_upper_bound() {
174        let iter = (0..10).estimate_size(5, None);
175        assert_eq!(iter.size_hint(), (5, None));
176
177        let collected: Vec<_> = iter.collect();
178        assert_eq!(collected.len(), 10);
179    }
180}
estimate_size/lib.rs

estimate_size/
lib.rs
