sparsey/entity/
entity_allocator.rs

1use crate::entity::Entity;
2use alloc::collections::VecDeque;
3use core::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
4
5#[derive(Default, Debug)]
6pub(crate) struct EntityAllocator {
7    next_index: AtomicU64,
8    last_maintained_index: u64,
9    recycled: VecDeque<Entity>,
10    recycled_since_maintain: AtomicUsize,
11}
12
13impl EntityAllocator {
14    #[must_use]
15    pub fn allocate(&mut self) -> Option<Entity> {
16        debug_assert!(!self.should_maintain_recyled());
17
18        let next_index = self.next_index.get_mut();
19
20        if let Some(entity) = self.recycled.pop_front() {
21            Some(entity)
22        } else if let Ok(index) = u32::try_from(*next_index) {
23            *next_index += 1;
24            self.last_maintained_index = *next_index;
25            Some(Entity::with_index(index))
26        } else {
27            None
28        }
29    }
30
31    #[must_use]
32    pub fn allocate_atomic(&self) -> Option<Entity> {
33        if let Some(recycled_index) = self.increment_recycled_since_maintain() {
34            Some(self.recycled[recycled_index])
35        } else {
36            self.increment_next_index().map(Entity::with_index)
37        }
38    }
39
40    pub fn recycle(&mut self, entity: Entity) {
41        if let Some(next_version) = entity.version.next() {
42            self.recycled
43                .push_back(Entity::new(entity.index, next_version));
44        }
45    }
46
47    #[inline]
48    #[must_use]
49    pub fn should_maintain_recyled(&mut self) -> bool {
50        *self.recycled_since_maintain.get_mut() != 0
51    }
52
53    pub fn maintain_recycled(&mut self) -> impl Iterator<Item = Entity> + '_ {
54        let recyled_since_maintain = *self.recycled_since_maintain.get_mut();
55        *self.recycled_since_maintain.get_mut() = 0;
56        self.recycled.drain(..recyled_since_maintain)
57    }
58
59    pub fn maintain_new(&mut self) -> impl Iterator<Item = Entity> + '_ {
60        let next_index = *self.next_index.get_mut();
61        let last_maintained_index = self.last_maintained_index;
62        self.last_maintained_index = next_index;
63
64        (last_maintained_index..next_index).map(|i| Entity::with_index(i as u32))
65    }
66
67    pub fn reset(&mut self) {
68        *self.next_index.get_mut() = 0;
69        self.last_maintained_index = 0;
70        self.recycled.clear();
71        *self.recycled_since_maintain.get_mut() = 0;
72    }
73
74    fn increment_next_index(&self) -> Option<u32> {
75        let mut prev = self.next_index.load(Ordering::Relaxed);
76
77        while u32::try_from(prev).is_ok() {
78            match self.next_index.compare_exchange_weak(
79                prev,
80                prev + 1,
81                Ordering::Relaxed,
82                Ordering::Relaxed,
83            ) {
84                Ok(prev) => return Some(prev as u32),
85                Err(next_prev) => prev = next_prev,
86            }
87        }
88
89        None
90    }
91
92    fn increment_recycled_since_maintain(&self) -> Option<usize> {
93        let recycled_len = self.recycled.len();
94        let mut prev = self.recycled_since_maintain.load(Ordering::Relaxed);
95
96        while prev < recycled_len {
97            match self.recycled_since_maintain.compare_exchange_weak(
98                prev,
99                prev + 1,
100                Ordering::Relaxed,
101                Ordering::Relaxed,
102            ) {
103                Ok(prev) => return Some(prev),
104                Err(next_prev) => prev = next_prev,
105            }
106        }
107
108        None
109    }
110}
sparsey/entity/entity_allocator.rs

sparsey/entity/
entity_allocator.rs
