mtb_entity_slab/
iter.rs

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

mtb_entity_slab/
iter.rs
