collect_with/collect_index.rs
1use core::hash::Hash;
2
3pub use indexmap::{IndexMap, IndexSet};
4
5use crate::collect::CollectWith;
6
7impl<I: Iterator> CollectIndex for I {}
8
9/// Trait for collecting items into IndexMap or IndexSet with specified capacity
10/// while preserving insertion order.
11pub trait CollectIndex: Iterator {
12 /// Collects items into an `IndexMap` with a specified capacity.
13 ///
14 /// Preserves insertion order and allows lookups by key.
15 ///
16 /// ## Example
17 ///
18 /// ```
19 /// use indexmap::IndexMap;
20 /// use collect_with::CollectIndex;
21 ///
22 /// let map = ('a'..='i')
23 /// .zip(100..=109)
24 /// .collect_indexmap_with(|u| u + 1); // u + 1 => 9 + 1 = 10
25 ///
26 /// assert_eq!(map.get(&'a'), Some(&100));
27 /// assert_eq!(map.get_index(0), Some((&'a', &100)));
28 /// assert_eq!(map.get_index(2), Some((&'c', &102)));
29 /// assert_eq!(map.capacity(), 10);
30 /// ```
31 #[cfg(not(feature = "ahash"))]
32 fn collect_indexmap_with<K, V>(
33 self,
34 capacity: impl FnOnce(usize) -> usize,
35 ) -> IndexMap<K, V>
36 where
37 Self: Sized + Iterator<Item = (K, V)>,
38 K: Hash + Eq,
39 {
40 self.collect_with(capacity)
41 }
42
43 /// Collects items into an `IndexMap<K, V, ahash::RandomState>` with a
44 /// specified capacity.
45 ///
46 /// ## Example
47 ///
48 /// ```
49 /// use indexmap::IndexMap;
50 /// use collect_with::CollectIndex;
51 ///
52 /// let map = ('a'..='i')
53 /// .zip(100..=109)
54 /// .collect_indexmap_with(|u| u + 1); // u + 1 => 9 + 1 = 10
55 ///
56 /// assert_eq!(map.get(&'a'), Some(&100));
57 /// assert_eq!(map.get_index(0), Some((&'a', &100)));
58 /// assert_eq!(map.get_index(2), Some((&'c', &102)));
59 /// assert_eq!(map.capacity(), 10);
60 /// ```
61 #[cfg(feature = "ahash")]
62 fn collect_indexmap_with<K, V>(
63 self,
64 capacity: impl FnOnce(usize) -> usize,
65 ) -> IndexMap<K, V, ahash::RandomState>
66 where
67 Self: Sized + Iterator<Item = (K, V)>,
68 K: Hash + Eq,
69 {
70 self.collect_with(capacity)
71 }
72
73 /// Collects items into an `IndexMap` with exact specified capacity.
74 ///
75 /// Preserves insertion order and strictly uses calculated capacity.
76 ///
77 /// # Example
78 ///
79 /// ```
80 /// use indexmap::IndexMap;
81 /// use collect_with::CollectIndex;
82 ///
83 /// let map = [(1, "a"), (2, "b"), (3, "c")]
84 /// .into_iter()
85 /// .collect_indexmap_with_exact(|size_hint| size_hint);
86 /// assert_eq!(map.len(), 3);
87 /// assert_eq!(map.capacity(), 3);
88 /// ```
89 fn collect_indexmap_with_exact<K, V>(
90 self,
91 capacity: impl FnOnce(usize) -> usize,
92 ) -> IndexMap<K, V>
93 where
94 Self: Sized + Iterator<Item = (K, V)>,
95 K: Hash + Eq,
96 {
97 self.collect_with_exact(capacity)
98 }
99
100 /// Collects items into an `IndexSet` with specified capacity.
101 ///
102 /// Preserves insertion order and maintains unique elements.
103 ///
104 /// # Example
105 ///
106 /// ```
107 /// use indexmap::IndexSet;
108 /// use collect_with::CollectIndex;
109 ///
110 /// let set = (0..3)
111 /// .collect_indexset_with(|size_hint| size_hint + 2);
112 /// assert_eq!(set.len(), 3);
113 /// assert!(set.capacity() >= 3);
114 /// ```
115 #[cfg(not(feature = "ahash"))]
116 fn collect_indexset_with<K>(
117 self,
118 capacity: impl FnOnce(usize) -> usize,
119 ) -> IndexSet<K>
120 where
121 Self: Sized + Iterator<Item = K>,
122 K: Hash + Eq,
123 {
124 self.collect_with(capacity)
125 }
126
127 /// Collects items into an `IndexSet<K, ahash::RandomState>` with specified
128 /// capacity.
129 ///
130 /// Preserves insertion order and maintains unique elements.
131 ///
132 /// # Example
133 ///
134 /// ```
135 /// use indexmap::IndexSet;
136 /// use collect_with::CollectIndex;
137 ///
138 /// let set = (0..3)
139 /// .collect_indexset_with(|size_hint| size_hint + 2);
140 /// assert_eq!(set.len(), 3);
141 /// assert!(set.capacity() >= 3);
142 /// ```
143 #[cfg(feature = "ahash")]
144 fn collect_indexset_with<K>(
145 self,
146 capacity: impl FnOnce(usize) -> usize,
147 ) -> IndexSet<K, ahash::RandomState>
148 where
149 Self: Sized + Iterator<Item = K>,
150 K: Hash + Eq,
151 {
152 self.collect_with(capacity)
153 }
154
155 /// Collects items into an `IndexSet` with exact specified capacity.
156 ///
157 /// Preserves insertion order and strictly uses calculated capacity.
158 ///
159 /// # Example
160 ///
161 /// ```
162 /// use indexmap::IndexSet;
163 /// use collect_with::CollectIndex;
164 ///
165 /// let set = (0..3)
166 /// .into_iter()
167 /// .collect_indexset_with_exact(|size_hint| size_hint);
168 /// assert_eq!(set.len(), 3);
169 /// assert_eq!(set.capacity(), 3);
170 /// ```
171 fn collect_indexset_with_exact<K>(
172 self,
173 capacity: impl FnOnce(usize) -> usize,
174 ) -> IndexSet<K>
175 where
176 Self: Sized + Iterator<Item = K>,
177 K: Hash + Eq,
178 {
179 self.collect_with_exact(capacity)
180 }
181}
182
183#[cfg(test)]
184mod tests {
185 use super::*;
186 use crate::TryCollectWith;
187
188 #[ignore]
189 #[test]
190 fn test_index_map() {
191 let map = (1..=9)
192 .zip('a'..='i')
193 .collect_indexmap_with(|u| u + 1);
194 assert_eq!(map.get(&1), Some(&'a'));
195 assert_eq!(map.capacity(), 10);
196
197 assert_eq!(map.get_index(2), Some((&3, &'c')));
198 }
199
200 #[ignore]
201 #[test]
202 fn test_indexmap2() {
203 let result = ["42", "76", "03"]
204 .into_iter()
205 // &str -> Result<i32>
206 .map(|x| x.parse::<i32>())
207 .try_collect_with::<Vec<_>, _, _>(|u| u + 3); // Result<Vec<i32>, ParseIntError>
208 assert_eq!(
209 result
210 .expect("Invalid IndexMap")
211 .capacity(),
212 6
213 );
214 // dbg!(result);
215 }
216}