collect_with/
collect.rs

1use crate::{ExtendWithCapacity, common::collect_iterator};
2
3// Implement CollectWithCapacity trait for Iterator
4impl<I: Iterator> CollectWithCapacity for I {}
5/// Trait for collecting iterator elements into a collection with specified
6/// capacity
7pub trait CollectWithCapacity: Iterator {
8  /// Collect iterator elements into a collection with pre-allocated capacity.
9  ///
10  /// - `capacity`
11  ///   - Initial capacity to allocate for the collection
12  /// - `T`
13  ///   - Collection type that implements ExtendWithCapacity (e.g., `Vec<X>`,
14  ///     String, `HashMap<K, V>`)
15  ///
16  /// > See also: [collect_with()](crate::CollectWith::collect_with)
17  fn collect_with_capacity<T>(self, capacity: usize) -> T
18  where
19    T: ExtendWithCapacity<Self::Item>,
20    Self: Sized,
21  {
22    collect_iterator(self, false, |_| capacity)
23  }
24}
25/// Implement CollectWith trait for Iterator
26impl<I: Iterator> CollectWith for I {}
27
28/// Trait for collecting iterator elements with flexible capacity calculation
29pub trait CollectWith: Iterator {
30  /// Collect elements using a capacity calculated from a closure
31  ///
32  /// - `capacity`
33  ///   - Closure that calculates capacity based on iterator size hints
34  ///
35  /// ## About the Final Capacity Size
36  ///
37  /// For example, `(0..10).collect_with::<Vec<_>>(|_size_bound| 2)`
38  ///
39  /// 1. `(0..10).size_hint()` returns `(10, Some(10))`.
40  /// 2. _size_bound is 10.
41  /// 3. The closure returns 2, the final capacity is `max(10, 2)` = 10.
42  /// 4. The vector is created with Vec::with_capacity(10).
43  ///
44  /// ## Example
45  ///
46  /// ```
47  /// use collect_with::CollectWith;
48  ///
49  /// let s = [vec!["a"], vec!["b", "c", "d"]]
50  ///   .into_iter()
51  ///   .flatten()
52  ///   .collect_with::<String>(|size| match size {
53  ///     0 => 8,
54  ///     n => n,
55  ///   });
56  /// assert_eq!(s.len(), 4);
57  /// assert_eq!(s.capacity(), 8);
58  /// ```
59  ///
60  /// The collection may allocate more capacity than calculated if needed.
61  /// If you need an exact capacity size, please use
62  /// [collect_with_exact()](crate::CollectWith::collect_with_exact)
63  fn collect_with<T>(self, capacity: impl FnOnce(usize) -> usize) -> T
64  where
65    T: ExtendWithCapacity<Self::Item>,
66    Self: Sized,
67  {
68    collect_iterator(self, false, capacity)
69  }
70
71  /// Collect elements using exact capacity calculated from a closure.
72  ///
73  /// - `capacity`
74  ///   - Closure that calculates exact capacity requirement
75  ///
76  /// The collection will strictly use the calculated capacity without
77  /// overallocation.
78  fn collect_with_exact<T>(self, capacity: impl FnOnce(usize) -> usize) -> T
79  where
80    T: ExtendWithCapacity<Self::Item>,
81    Self: Sized,
82  {
83    collect_iterator(self, true, capacity)
84  }
85}
86
87#[cfg(test)]
88mod tests {
89  use alloc::{string::String, vec};
90
91  use super::*;
92
93  #[ignore]
94  #[test]
95  fn test_collect_with() {
96    let s = [vec!["a"], vec!["b", "c", "d"]]
97      .into_iter()
98      .flatten()
99      .collect_with::<String>(|size| match size {
100        0 => 8,
101        n => n,
102      });
103    assert_eq!(s.len(), 4);
104    assert_eq!(s.capacity(), 8);
105  }
106}
collect_with/collect.rs

collect_with/
collect.rs
