mtb_entity_slab/
iter.rs

1use crate::{
2    EntityAlloc, IDProxy, IEntityAllocID, IEntityAllocPolicy, PtrID, alloc::IEntityAllocatable,
3    chunk::Chunk,
4};
5use std::{
6    cell::{BorrowError, Ref},
7    marker::PhantomData,
8    ptr::NonNull,
9};
10
11struct IterPos<E: IEntityAllocatable> {
12    indexed_id: usize,
13    nelems_last: usize,
14    nunits_last: usize,
15    _mark: PhantomData<E>,
16}
17impl<E: IEntityAllocatable> IterPos<E> {
18    fn new(nelems_last: usize) -> Self {
19        Self {
20            indexed_id: 0,
21            nelems_last,
22            nunits_last: 0,
23            _mark: PhantomData::<E>,
24        }
25    }
26}
27
28impl<E: IEntityAllocatable> IterPos<E> {
29    fn chunk_id(&self) -> u32 {
30        (self.indexed_id / E::AllocatePolicyT::CHUNK_SIZE) as u32
31    }
32    fn unit_id(&self) -> u16 {
33        (self.indexed_id % E::AllocatePolicyT::CHUNK_SIZE) as u16
34    }
35
36    fn at_end(&self) -> bool {
37        self.nelems_last == 0
38    }
39    fn find_valid_chunk<'alloc>(
40        &mut self,
41        chunks: &'alloc [Chunk<E, E::AllocatePolicyT>],
42    ) -> Option<&'alloc Chunk<E, E::AllocatePolicyT>> {
43        while self.nunits_last == 0 {
44            self.indexed_id = self
45                .indexed_id
46                .next_multiple_of(E::AllocatePolicyT::CHUNK_SIZE);
47            let chunk_id = self.chunk_id() as usize;
48            if chunk_id >= chunks.len() {
49                return None;
50            }
51            self.nunits_last = chunks[chunk_id].num_allocated as usize;
52        }
53        Some(&chunks[self.chunk_id() as usize])
54    }
55    fn find_valid_unit<'alloc>(&mut self, chunks: &'alloc [Chunk<E, E::AllocatePolicyT>]) {
56        let unit_id = self.unit_id();
57        let Some(chunk) = self.find_valid_chunk(chunks) else {
58            return;
59        };
60        let Some(next) = chunk.next_allocated(unit_id) else {
61            unreachable!("find_valid_chunk ensures there is a valid unit");
62        };
63        self.indexed_id = chunk.indexed_id_of_unit(next);
64    }
65    fn advance<'alloc>(&mut self, chunks: &'alloc [Chunk<E, E::AllocatePolicyT>]) {
66        self.indexed_id += 1;
67        self.nunits_last -= 1;
68        self.nelems_last -= 1;
69        self.find_valid_unit(chunks);
70    }
71
72    fn get_ptr(&self, chunks: &[Chunk<E, E::AllocatePolicyT>]) -> Option<PtrID<E>> {
73        if self.at_end() {
74            return None;
75        }
76        let chunk_id = self.chunk_id() as usize;
77        let unit_id = self.unit_id() as usize;
78        let chunk = chunks.get(chunk_id)?;
79        let ptr = PtrID(NonNull::from_ref(chunk.unit(unit_id as u16)));
80        Some(ptr)
81    }
82}
83
84pub struct EntityAllocReadIter<'alloc, E: IEntityAllocatable> {
85    chunks: Ref<'alloc, [Chunk<E, E::AllocatePolicyT>]>,
86    pos: IterPos<E>,
87}
88
89impl<'alloc, E: IEntityAllocatable> Iterator for EntityAllocReadIter<'alloc, E> {
90    type Item = (usize, E::PtrID, &'alloc E);
91
92    fn next(&mut self) -> Option<Self::Item> {
93        if self.pos.at_end() {
94            return None;
95        }
96        let ptr = self
97            .pos
98            .get_ptr(&self.chunks)
99            .expect("IterPos ensures valid pointer when not at end");
100        let indexed_id = self.pos.indexed_id;
101
102        // 准备下一个
103        self.pos.advance(&self.chunks);
104        let e = unsafe { ptr.direct_deref() };
105        Some((indexed_id, E::id_of_ptr(ptr), e))
106    }
107
108    fn size_hint(&self) -> (usize, Option<usize>) {
109        let len = self.pos.nelems_last;
110        (len, Some(len))
111    }
112}
113
114impl<'alloc, E: IEntityAllocatable> ExactSizeIterator for EntityAllocReadIter<'alloc, E> {
115    fn len(&self) -> usize {
116        self.pos.nelems_last
117    }
118}
119
120impl<'alloc, E: IEntityAllocatable> EntityAllocReadIter<'alloc, E> {
121    pub fn try_new(alloc: &'alloc EntityAlloc<E>) -> Result<Self, BorrowError> {
122        let chunks = alloc.chunks.try_borrow()?;
123        let nelems_last = chunks.iter().map(|c| c.num_allocated as usize).sum();
124        let mut pos = IterPos::new(nelems_last);
125        pos.find_valid_unit(&chunks);
126        Ok(Self {
127            chunks: Ref::map(chunks, |x| x.as_slice()),
128            pos,
129        })
130    }
131
132    pub fn new(alloc: &'alloc EntityAlloc<E>) -> Self {
133        Self::try_new(alloc).expect(
134            "Cannot borrow EntityAlloc for iteration: are you allocating elements while iterating?",
135        )
136    }
137
138    pub fn current_indexed_id(&self) -> Option<usize> {
139        if self.pos.at_end() {
140            None
141        } else {
142            Some(self.pos.indexed_id)
143        }
144    }
145    pub fn current_chunk_id(&self) -> Option<u32> {
146        if self.pos.at_end() {
147            None
148        } else {
149            Some(self.pos.chunk_id())
150        }
151    }
152    pub fn current_unit_id(&self) -> Option<u16> {
153        if self.pos.at_end() {
154            None
155        } else {
156            Some(self.pos.unit_id())
157        }
158    }
159}
160
161pub struct EntityAllocEditIter<'alloc, E: IEntityAllocatable> {
162    chunks: &'alloc mut [Chunk<E, E::AllocatePolicyT>],
163    pos: IterPos<E>,
164}
165
166impl<'alloc, E: IEntityAllocatable> Iterator for EntityAllocEditIter<'alloc, E> {
167    type Item = (usize, E::PtrID, &'alloc mut E);
168
169    fn next(&mut self) -> Option<Self::Item> {
170        if self.pos.at_end() {
171            return None;
172        }
173        let ptr = self
174            .pos
175            .get_ptr(&self.chunks)
176            .expect("IterPos ensures valid pointer when not at end");
177        let indexed_id = self.pos.indexed_id;
178        // 准备下一个
179        self.pos.advance(&self.chunks);
180        let e = unsafe { ptr.direct_deref_mut() };
181        Some((indexed_id, E::id_of_ptr(ptr), e))
182    }
183
184    fn size_hint(&self) -> (usize, Option<usize>) {
185        let len = self.pos.nelems_last;
186        (len, Some(len))
187    }
188}
189impl<'alloc, E: IEntityAllocatable> ExactSizeIterator for EntityAllocEditIter<'alloc, E> {
190    fn len(&self) -> usize {
191        self.pos.nelems_last
192    }
193}
194impl<'alloc, E: IEntityAllocatable> EntityAllocEditIter<'alloc, E> {
195    pub fn new(alloc: &'alloc mut EntityAlloc<E>) -> Self {
196        let chunks = alloc.chunks.get_mut();
197        let nelems_last = chunks.iter().map(|c| c.num_allocated as usize).sum();
198        let mut pos = IterPos::new(nelems_last);
199        pos.find_valid_unit(&chunks);
200        Self { chunks, pos }
201    }
202}
203
204pub struct EntityAllocConsumeIter<E: IEntityAllocatable> {
205    alloc: EntityAlloc<E>,
206    pos: IterPos<E>,
207}
208
209impl<E: IEntityAllocatable> Iterator for EntityAllocConsumeIter<E> {
210    type Item = (usize, E);
211
212    fn next(&mut self) -> Option<Self::Item> {
213        let Self { alloc, pos } = self;
214        let chunks = alloc.chunks.get_mut();
215        if pos.at_end() {
216            return None;
217        }
218        let ptr = pos
219            .get_ptr(chunks)
220            .expect("IterPos ensures valid pointer when not at end");
221        let indexed_id = pos.indexed_id;
222
223        // 准备下一个
224        pos.advance(alloc.chunks.get_mut());
225        let e = ptr
226            .free(alloc)
227            .expect("UAF detected: pointer should be valid during consume iteration");
228        Some((indexed_id, e))
229    }
230
231    fn size_hint(&self) -> (usize, Option<usize>) {
232        let len = self.pos.nelems_last;
233        (len, Some(len))
234    }
235}
236impl<E: IEntityAllocatable> ExactSizeIterator for EntityAllocConsumeIter<E> {
237    fn len(&self) -> usize {
238        self.pos.nelems_last
239    }
240}
241impl<E: IEntityAllocatable> EntityAllocConsumeIter<E> {
242    pub fn new(mut alloc: EntityAlloc<E>) -> Self {
243        let nelems_last = alloc.num_allocated.get();
244        let mut pos = IterPos::new(nelems_last);
245        pos.find_valid_unit(alloc.chunks.get_mut());
246        Self { alloc, pos }
247    }
248
249    pub fn into_alloc(self) -> EntityAlloc<E> {
250        self.alloc
251    }
252}
253
254/// Iterator that yields a freeable slot IDProxy
255pub(crate) struct EntitySlotIter<'alloc, E: IEntityAllocatable> {
256    chunks: &'alloc mut [Chunk<E, E::AllocatePolicyT>],
257    pos: IterPos<E>,
258}
259impl<'alloc, E: IEntityAllocatable> Iterator for EntitySlotIter<'alloc, E> {
260    type Item = IDProxy<'alloc, E>;
261
262    fn next(&mut self) -> Option<Self::Item> {
263        if self.pos.at_end() {
264            return None;
265        }
266        let chunk_id = self.pos.chunk_id() as usize;
267        let unit_id = self.pos.unit_id() as usize;
268        // 准备下一个
269        self.pos.advance(&self.chunks);
270        let chunk = &mut self.chunks[chunk_id];
271
272        let unit = chunk.unit_mut(unit_id as u16);
273        let unit = unsafe { &mut *(unit as *mut _) };
274        Some(IDProxy { unit })
275    }
276}
277impl<'alloc, E: IEntityAllocatable> EntitySlotIter<'alloc, E> {
278    pub fn new(alloc: &'alloc mut EntityAlloc<E>) -> Self {
279        let chunks = alloc.chunks.get_mut();
280        let nelems_last = chunks.iter().map(|c| c.num_allocated as usize).sum();
281        let mut pos = IterPos::new(nelems_last);
282        pos.find_valid_unit(&chunks);
283        Self { chunks, pos }
284    }
285}
mtb_entity_slab/iter.rs

mtb_entity_slab/
iter.rs
