atomic_arena/
controller.rs

1use std::sync::{
2	atomic::{AtomicBool, AtomicUsize, Ordering},
3	Arc,
4};
5
6use crate::{ArenaFull, Key};
7
8/// Represents that a [`ControllerSlot`] does not have a free slot
9/// after it.
10///
11/// This is used because the next free slot variable is an
12/// [`AtomicUsize`], but we still need some way to represent the
13/// absence of a next free slot.
14const NO_NEXT_FREE_SLOT: usize = usize::MAX;
15
16#[derive(Debug)]
17struct ControllerSlot {
18	free: AtomicBool,
19	generation: AtomicUsize,
20	next_free_slot_index: AtomicUsize,
21}
22
23/// The shared state for all [`Controller`]s for an [`Arena`](super::Arena).
24#[derive(Debug)]
25struct ControllerInner {
26	slots: Vec<ControllerSlot>,
27	first_free_slot_index: AtomicUsize,
28}
29
30impl ControllerInner {
31	fn new(capacity: usize) -> Self {
32		Self {
33			slots: (0..capacity)
34				.map(|i| ControllerSlot {
35					free: AtomicBool::new(true),
36					generation: AtomicUsize::new(0),
37					next_free_slot_index: AtomicUsize::new(if i < capacity - 1 {
38						i + 1
39					} else {
40						NO_NEXT_FREE_SLOT
41					}),
42				})
43				.collect(),
44			first_free_slot_index: AtomicUsize::new(0),
45		}
46	}
47
48	fn capacity(&self) -> usize {
49		self.slots.len()
50	}
51
52	fn len(&self) -> usize {
53		self.slots
54			.iter()
55			.filter(|slot| !slot.free.load(Ordering::SeqCst))
56			.count()
57	}
58
59	fn try_reserve(&self) -> Result<Key, ArenaFull> {
60		loop {
61			let first_free_slot_index = self.first_free_slot_index.load(Ordering::SeqCst);
62			if first_free_slot_index == NO_NEXT_FREE_SLOT {
63				return Err(ArenaFull);
64			}
65			let slot = &self.slots[first_free_slot_index];
66			if self
67				.first_free_slot_index
68				.compare_exchange_weak(
69					first_free_slot_index,
70					slot.next_free_slot_index.load(Ordering::SeqCst),
71					Ordering::SeqCst,
72					Ordering::SeqCst,
73				)
74				.is_ok()
75			{
76				slot.free.store(false, Ordering::SeqCst);
77				return Ok(Key {
78					index: first_free_slot_index,
79					generation: slot.generation.load(Ordering::SeqCst),
80				});
81			}
82		}
83	}
84
85	fn free(&self, index: usize) {
86		let slot = &self.slots[index];
87		slot.free.store(true, Ordering::SeqCst);
88		slot.generation.fetch_add(1, Ordering::SeqCst);
89		loop {
90			let first_free_slot_index = self.first_free_slot_index.load(Ordering::SeqCst);
91			slot.next_free_slot_index
92				.store(first_free_slot_index, Ordering::SeqCst);
93			if self
94				.first_free_slot_index
95				.compare_exchange_weak(
96					first_free_slot_index,
97					index,
98					Ordering::SeqCst,
99					Ordering::SeqCst,
100				)
101				.is_ok()
102			{
103				break;
104			}
105		}
106	}
107}
108
109/// Manages [`Key`] reservations for an [`Arena`](super::Arena).
110#[derive(Debug, Clone)]
111pub struct Controller(Arc<ControllerInner>);
112
113impl Controller {
114	pub(crate) fn new(capacity: usize) -> Self {
115		Self(Arc::new(ControllerInner::new(capacity)))
116	}
117
118	/// Returns the total capacity of the arena.
119	pub fn capacity(&self) -> usize {
120		self.0.capacity()
121	}
122
123	/// Returns the number of items in the arena.
124	pub fn len(&self) -> usize {
125		self.0.len()
126	}
127
128	/// Returns `true` if the arena is empty.
129	pub fn is_empty(&self) -> bool {
130		self.len() == 0
131	}
132
133	/// Tries to reserve a key for the [`Arena`](super::Arena).
134	pub fn try_reserve(&self) -> Result<Key, ArenaFull> {
135		self.0.try_reserve()
136	}
137
138	pub(crate) fn free(&self, index: usize) {
139		self.0.free(index);
140	}
141}