mago_allocator/
iter.rs

1//! Collecting iterators directly into an arena.
2
3use core::hash::BuildHasher;
4use core::hash::Hash;
5
6use crate::Arena;
7use crate::collections::HashMap;
8use crate::collections::HashSet;
9use crate::vec::Vec;
10
11/// Construct a collection from an iterator, allocating its storage in an arena.
12///
13/// The arena-equivalent of [`FromIterator`]: the implementing collection decides
14/// what allocator handle it needs via [`Alloc`](FromIteratorIn::Alloc).
15pub trait FromIteratorIn<T> {
16    /// The allocator handle the collection is built into (e.g. `&'arena A`).
17    type Alloc;
18
19    /// Builds `Self` from `iter`, allocating into `alloc`.
20    fn from_iter_in<I>(iter: I, alloc: Self::Alloc) -> Self
21    where
22        I: IntoIterator<Item = T>;
23}
24
25impl<'arena, T, A: Arena> FromIteratorIn<T> for Vec<'arena, T, A> {
26    type Alloc = &'arena A;
27
28    #[inline]
29    fn from_iter_in<I>(iter: I, alloc: Self::Alloc) -> Self
30    where
31        I: IntoIterator<Item = T>,
32    {
33        let mut collection = Self::new_in(alloc);
34        collection.extend(iter);
35        collection
36    }
37}
38
39impl<'arena, K, V, A, S> FromIteratorIn<(K, V)> for HashMap<'arena, K, V, A, S>
40where
41    K: Eq + Hash,
42    S: BuildHasher + Default,
43    A: Arena,
44{
45    type Alloc = &'arena A;
46
47    #[inline]
48    fn from_iter_in<I>(iter: I, alloc: Self::Alloc) -> Self
49    where
50        I: IntoIterator<Item = (K, V)>,
51    {
52        let mut collection = Self::with_hasher_in(S::default(), alloc);
53        collection.extend(iter);
54        collection
55    }
56}
57
58impl<'arena, T, A, S> FromIteratorIn<T> for HashSet<'arena, T, A, S>
59where
60    T: Eq + Hash,
61    S: BuildHasher + Default,
62    A: Arena,
63{
64    type Alloc = &'arena A;
65
66    #[inline]
67    fn from_iter_in<I>(iter: I, alloc: Self::Alloc) -> Self
68    where
69        I: IntoIterator<Item = T>,
70    {
71        let mut collection = Self::with_hasher_in(S::default(), alloc);
72        collection.extend(iter);
73        collection
74    }
75}
76
77/// Collect an iterator into an arena-allocated collection.
78///
79/// The arena-equivalent of [`Iterator::collect`]:
80///
81/// ```ignore
82/// let evens: Vec<i32, A> = (0..10).filter(|n| n % 2 == 0).collect_in(arena);
83/// ```
84pub trait CollectIn: Iterator + Sized {
85    /// Collects this iterator into `C`, allocating its storage in `alloc`.
86    #[inline]
87    fn collect_in<C>(self, alloc: C::Alloc) -> C
88    where
89        C: FromIteratorIn<Self::Item>,
90    {
91        C::from_iter_in(self, alloc)
92    }
93}
94
95impl<I: Iterator> CollectIn for I {}
96
97#[cfg(feature = "rayon")]
98pub use parallel::ParallelCollectIn;
99
100#[cfg(feature = "rayon")]
101mod parallel {
102    use rayon::iter::IndexedParallelIterator;
103    use rayon::iter::IntoParallelRefMutIterator;
104    use rayon::iter::ParallelIterator;
105
106    use crate::Arena;
107    use crate::vec::Vec;
108
109    /// Collect an indexed parallel iterator into an arena-allocated [`Vec`], with
110    /// zero global allocations.
111    ///
112    /// The destination is allocated once, in the arena, at the iterator's exact
113    /// length; each worker writes its element directly into its uninitialized slot.
114    /// There is no intermediate buffer.
115    pub trait ParallelCollectIn: IndexedParallelIterator
116    where
117        Self::Item: Send,
118    {
119        /// Collects this parallel iterator into an arena-allocated [`Vec`].
120        fn collect_in<A>(self, arena: &A) -> Vec<'_, Self::Item, A>
121        where
122            A: Arena + Sync,
123        {
124            let length = self.len();
125            let mut collection = Vec::with_capacity_in(length, arena);
126
127            collection.spare_capacity_mut()[..length].par_iter_mut().zip(self).for_each(|(slot, item)| {
128                slot.write(item);
129            });
130
131            // SAFETY: every one of the first `length` slots was initialized above.
132            unsafe { collection.set_len(length) };
133
134            collection
135        }
136    }
137
138    impl<I: IndexedParallelIterator> ParallelCollectIn for I where I::Item: Send {}
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use crate::arena::LocalArena;
145
146    #[test]
147    fn collect_in_vec() {
148        let arena = LocalArena::new();
149        let collected: Vec<'_, i32, LocalArena> = (0..5).collect_in(&arena);
150
151        assert_eq!(collected.as_slice(), &[0, 1, 2, 3, 4]);
152    }
153
154    #[test]
155    fn collect_in_hashset_deduplicates() {
156        let arena = LocalArena::new();
157        let collected: HashSet<'_, i32, LocalArena> = [1, 2, 2, 3].into_iter().collect_in(&arena);
158
159        assert_eq!(collected.len(), 3);
160        assert!(collected.contains(&2));
161    }
162
163    #[test]
164    fn collect_in_hashmap() {
165        let arena = LocalArena::new();
166        let collected: HashMap<'_, i32, i32, LocalArena> = [(1, 10), (2, 20)].into_iter().collect_in(&arena);
167
168        assert_eq!(collected.get(&2), Some(&20));
169    }
170}
mago_allocator/iter.rs

mago_allocator/
iter.rs
