atomic_arena/
iter.rs

1//! [`Arena`] iterators.
2
3use std::marker::PhantomData;
4
5use crate::{
6	slot::{ArenaSlot, ArenaSlotState},
7	Arena, Key,
8};
9
10/// Iterates over shared references to the items in
11/// the [`Arena`].
12///
13/// The most recently added items will be visited first.
14pub struct Iter<'a, T> {
15	next_occupied_slot_index: Option<usize>,
16	arena: &'a Arena<T>,
17}
18
19impl<'a, T> Iter<'a, T> {
20	pub(super) fn new(arena: &'a Arena<T>) -> Self {
21		Self {
22			next_occupied_slot_index: arena.first_occupied_slot_index,
23			arena,
24		}
25	}
26}
27
28impl<'a, T> Iterator for Iter<'a, T> {
29	type Item = (Key, &'a T);
30
31	fn next(&mut self) -> Option<Self::Item> {
32		if let Some(index) = self.next_occupied_slot_index {
33			let slot = &self.arena.slots[index];
34			if let ArenaSlotState::Occupied {
35				data,
36				next_occupied_slot_index,
37				..
38			} = &slot.state
39			{
40				self.next_occupied_slot_index = *next_occupied_slot_index;
41				Some((
42					Key {
43						index,
44						generation: slot.generation,
45					},
46					data,
47				))
48			} else {
49				panic!("the iterator should not encounter a free slot");
50			}
51		} else {
52			None
53		}
54	}
55}
56
57/// Iterates over mutable references to the items in
58/// the [`Arena`].
59///
60/// The most recently added items will be visited first.
61pub struct IterMut<'a, T> {
62	next_occupied_slot_index: Option<usize>,
63	slots: *mut [ArenaSlot<T>],
64	marker: PhantomData<&'a mut Arena<T>>,
65}
66
67impl<'a, T> IterMut<'a, T> {
68	pub(super) fn new(arena: &'a mut Arena<T>) -> Self {
69		Self {
70			next_occupied_slot_index: arena.first_occupied_slot_index,
71			slots: arena.slots.as_mut_slice(),
72			marker: PhantomData,
73		}
74	}
75}
76
77impl<'a, T> Iterator for IterMut<'a, T> {
78	type Item = (Key, &'a mut T);
79
80	fn next(&mut self) -> Option<Self::Item> {
81		if let Some(index) = self.next_occupied_slot_index {
82			let slot = {
83				// as_mut_ptr and get_unchecked_mut on *mut [T] are unstable :(
84				let start_ptr = self.slots.cast::<ArenaSlot<T>>();
85				// SAFETY: This is always in bounds.
86				let slot_ptr = unsafe { start_ptr.add(index) };
87				// SAFETY:
88				// * This relies on the invariant that `next_occupied_slot_index` never repeats. If
89				//   it did repeat, we could create aliasing mutable references here.
90				// * Lifetime is the same that we mutably borrow the Arena for.
91				unsafe { slot_ptr.as_mut::<'a>() }.unwrap()
92			};
93
94			if let ArenaSlotState::Occupied {
95				data,
96				next_occupied_slot_index,
97				..
98			} = &mut slot.state
99			{
100				self.next_occupied_slot_index = *next_occupied_slot_index;
101				Some((
102					Key {
103						index,
104						generation: slot.generation,
105					},
106					data,
107				))
108			} else {
109				panic!("the iterator should not encounter a free slot");
110			}
111		} else {
112			None
113		}
114	}
115}
116
117/// An iterator that removes and yields elements from an
118/// [`Arena`] according to a filter function.
119pub struct DrainFilter<'a, T, F: FnMut(&T) -> bool> {
120	arena: &'a mut Arena<T>,
121	filter: F,
122	next_occupied_slot_index: Option<usize>,
123}
124
125impl<'a, T, F: FnMut(&T) -> bool> DrainFilter<'a, T, F> {
126	pub(super) fn new(arena: &'a mut Arena<T>, filter: F) -> Self {
127		Self {
128			next_occupied_slot_index: arena.first_occupied_slot_index,
129			arena,
130			filter,
131		}
132	}
133}
134
135impl<T, F: FnMut(&T) -> bool> Iterator for DrainFilter<'_, T, F> {
136	type Item = (Key, T);
137
138	fn next(&mut self) -> Option<Self::Item> {
139		while let Some(index) = self.next_occupied_slot_index {
140			let slot = &mut self.arena.slots[index];
141			if let ArenaSlotState::Occupied {
142				data,
143				next_occupied_slot_index,
144				..
145			} = &mut slot.state
146			{
147				self.next_occupied_slot_index = *next_occupied_slot_index;
148				if (self.filter)(data) {
149					let key = Key {
150						index,
151						generation: slot.generation,
152					};
153					return self
154						.arena
155						.remove_from_slot(index)
156						.map(|element| (key, element));
157				}
158			} else {
159				panic!("the iterator should not encounter a free slot");
160			}
161		}
162		None
163	}
164}
atomic_arena/iter.rs

atomic_arena/
iter.rs
