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}