mtb-entity-slab 0.2.4

use crate::{
    EntityAlloc, IndexedID, PtrID, bitalloc::IBitAlloc, chunk::Chunk, gen_index::GenIndex,
    policy::IAllocPolicy,
};
use std::{
    cell::{BorrowError, Ref},
    marker::PhantomData,
};

pub(super) struct IterPos<E, P: IAllocPolicy> {
    index: u64,
    nelems_last: usize,
    nunits_last: usize,
    _mark: PhantomData<(E, P)>,
}
impl<E, P: IAllocPolicy> IterPos<E, P> {
    pub fn new(nelems_last: usize) -> Self {
        Self {
            index: 0,
            nelems_last,
            nunits_last: 0,
            _mark: PhantomData,
        }
    }

    pub fn at_end(&self) -> bool {
        self.nelems_last == 0
    }

    fn chunk_id(&self) -> u32 {
        P::chunkof_real(self.index as usize)
    }
    fn unit_id(&self) -> u16 {
        P::unitof_real(self.index as usize)
    }
    fn find_valid_chunk<'alloc>(
        &mut self,
        chunks: &'alloc [Chunk<E, P>],
    ) -> Option<&'alloc Chunk<E, P>> {
        while self.nunits_last == 0 {
            self.index = self.index.next_multiple_of(P::CHUNK_SIZE as u64);
            let chunk_id = self.chunk_id() as usize;
            if chunk_id >= chunks.len() {
                return None;
            }
            let chunk = &chunks[chunk_id];
            self.nunits_last = chunk.num_allocated as usize;
            debug_assert_eq!(self.nunits_last, chunk.allocated.count_allocated());
        }
        Some(&chunks[self.chunk_id() as usize])
    }
    pub fn find_valid_unit(&mut self, chunks: &[Chunk<E, P>]) {
        // 先定位到包含下一个元素的有效 chunk，再基于最新的 indexed_id 计算 unit_id
        let Some(chunk) = self.find_valid_chunk(chunks) else {
            return;
        };
        let unit_id = self.unit_id();
        let Some(next) = chunk.next_allocated(unit_id) else {
            unreachable!("find_valid_chunk ensures there is a valid unit");
        };
        self.index = {
            let chunk = chunk.chunk_id as u64;
            let unit = next as u64;
            chunk << P::CHUNK_SIZE_LOG2 | unit
        };
    }
    pub fn advance(&mut self, chunks: &[Chunk<E, P>]) {
        self.index += 1;
        self.nunits_last -= 1;
        self.nelems_last -= 1;
        self.find_valid_unit(chunks);
    }

    pub fn get_ids(&self, chunks: &[Chunk<E, P>]) -> Option<(IndexedID<E, P>, PtrID<E, P>)> {
        if self.at_end() {
            return None;
        }
        let chunk_id = self.chunk_id() as usize;
        let unit_id = self.unit_id() as usize;
        let chunk = chunks.get(chunk_id)?;
        let indexed_id = chunk.indexed_id_of_unit(unit_id as u16);
        let ptr = PtrID::from(chunk.unit_ptr_mut(unit_id as u16));
        Some((IndexedID::from(indexed_id), ptr))
    }
}

/// Immutable iterator over allocated entities in an `EntityAlloc`.
///
/// # Yields
///
/// Yields tuples of `(IndexedID<E, P>, PtrID<E, P>, &E)`, where:
/// - `IndexedID<E, P>` is the indexed ID of the entity.
/// - `PtrID<E, P>` is the pointer ID of the entity.
/// - `&E` is a reference to the allocated entity itself.
///
/// # Borrowing
///
/// This iterator borrows the `EntityAlloc` immutably for its entire lifetime.
/// Attempting to mutate the allocator while this iterator is active will result
/// in a runtime borrow error.
///
/// 不可变迭代器：遍历 `EntityAlloc` 中已分配的实体。
pub struct EntityAllocReadIter<'alloc, E, P: IAllocPolicy> {
    chunks: Ref<'alloc, [Chunk<E, P>]>,
    pos: IterPos<E, P>,
}
impl<'alloc, E, P: IAllocPolicy> Iterator for EntityAllocReadIter<'alloc, E, P> {
    type Item = (IndexedID<E, P>, PtrID<E, P>, &'alloc E);

    fn next(&mut self) -> Option<Self::Item> {
        if self.pos.at_end() {
            return None;
        }
        let Some((index, ptr)) = self.pos.get_ids(&self.chunks) else {
            panic!("IterPos ensures valid pointer when not at end")
        };

        // 准备下一个
        self.pos.advance(&self.chunks);
        let e = unsafe { ptr.direct_deref().unwrap() };
        Some((index, ptr, e))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.pos.nelems_last;
        (len, Some(len))
    }
}
impl<'alloc, E, P: IAllocPolicy> ExactSizeIterator for EntityAllocReadIter<'alloc, E, P> {
    fn len(&self) -> usize {
        self.pos.nelems_last
    }
}
impl<'alloc, E, P: IAllocPolicy> EntityAllocReadIter<'alloc, E, P> {
    /// Create a new iterator over the given `EntityAlloc`.
    ///
    /// # Errors
    ///
    /// Returns a `BorrowError` if the `EntityAlloc` is already mutably borrowed.
    pub fn try_new(alloc: &'alloc EntityAlloc<E, P>) -> Result<Self, BorrowError> {
        let chunks = alloc.chunks.try_borrow()?;
        let nelems_last = chunks.iter().map(|c| c.num_allocated as usize).sum();
        let mut pos = IterPos::new(nelems_last);
        pos.find_valid_unit(&chunks);
        Ok(Self {
            chunks: Ref::map(chunks, |x| x.as_slice()),
            pos,
        })
    }
    /// Create a new iterator over the given `EntityAlloc`.
    ///
    /// # Panics
    ///
    /// Panics if the `EntityAlloc` is already mutably borrowed.
    pub fn new(alloc: &'alloc EntityAlloc<E, P>) -> Self {
        Self::try_new(alloc).expect(
            "Cannot borrow EntityAlloc for iteration: are you allocating elements while iterating?",
        )
    }
    /// Get the current `IndexedID` at the iterator's position, if not at end.
    pub fn current_indexed_id(&self) -> Option<IndexedID<E, P>> {
        self.pos
            .get_ids(&self.chunks)
            .map(|(indexed_id, _)| indexed_id)
    }
    /// Get the current chunk ID at the iterator's position, if not at end.
    pub fn current_chunk_id(&self) -> Option<u32> {
        if self.pos.at_end() {
            None
        } else {
            Some(self.pos.chunk_id())
        }
    }
    /// Get the current unit ID at the iterator's position, if not at end.
    pub fn current_unit_id(&self) -> Option<u16> {
        if self.pos.at_end() {
            None
        } else {
            Some(self.pos.unit_id())
        }
    }
}

/// Mutable iterator over allocated entities in an `EntityAlloc`.
///
/// # Yields
///
/// Yields tuples of `(IndexedID<E, P>, PtrID<E, P>, &mut E)`, where:
/// - `IndexedID<E, P>` is the indexed ID of the entity.
/// - `PtrID<E, P>` is the pointer ID of the entity.
/// - `&mut E` is a mutable reference to the allocated entity itself.
///
/// # Borrowing
///
/// This iterator borrows the `EntityAlloc` mutably for its entire lifetime.
/// The entity allocator cannot be accessed in any other way while this
/// iterator is active. (guarenteed by Rust's mutable borrow rules)
///
/// 可变迭代器：遍历 `EntityAlloc` 中已分配的实体。
pub struct EntityAllocEditIter<'alloc, E, P: IAllocPolicy> {
    chunks: &'alloc mut [Chunk<E, P>],
    pos: IterPos<E, P>,
}
impl<'alloc, E, P: IAllocPolicy> Iterator for EntityAllocEditIter<'alloc, E, P> {
    type Item = (IndexedID<E, P>, PtrID<E, P>, &'alloc mut E);

    fn next(&mut self) -> Option<Self::Item> {
        if self.pos.at_end() {
            return None;
        }
        let Some((indexed_id, ptr)) = self.pos.get_ids(self.chunks) else {
            panic!("IterPos ensures valid pointer when not at end")
        };

        // 准备下一个
        self.pos.advance(self.chunks);
        let e = unsafe { ptr.direct_deref_mut().unwrap() };
        Some((indexed_id, ptr, e))
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.pos.nelems_last;
        (len, Some(len))
    }
}
impl<'alloc, E, P: IAllocPolicy> ExactSizeIterator for EntityAllocEditIter<'alloc, E, P> {
    fn len(&self) -> usize {
        self.pos.nelems_last
    }
}
impl<'alloc, E, P: IAllocPolicy> EntityAllocEditIter<'alloc, E, P> {
    /// Create a new iterator over the given `EntityAlloc`.
    pub fn new(alloc: &'alloc mut EntityAlloc<E, P>) -> Self {
        let chunks = alloc.chunks.get_mut();
        let nelems_last = chunks.iter().map(|c| c.num_allocated as usize).sum();
        let mut pos = IterPos::new(nelems_last);
        pos.find_valid_unit(chunks);
        Self { chunks, pos }
    }
}

/// Consuming iterator over allocated entities in an `EntityAlloc`.
///
/// # Yields
///
/// Yields tuples of `(GenIndex, E)`, where:
/// - `GenIndex` is the original generation-indexed ID of the entity.
/// - `E` is the owned value of the allocated entity itself.
///
/// The pointer ID is not yielded since the entity is being consumed.
///
/// # Ownership
///
/// This iterator takes ownership of the `EntityAlloc` and consumes its entities.
/// After the iterator is dropped or fully consumed, the allocator can be
/// retrieved back using the `release()` method.
pub struct EntityAllocConsumeIter<E, P: IAllocPolicy> {
    alloc: EntityAlloc<E, P>,
    pos: IterPos<E, P>,
}

impl<E, P: IAllocPolicy> Iterator for EntityAllocConsumeIter<E, P> {
    type Item = (GenIndex, E);

    fn next(&mut self) -> Option<Self::Item> {
        let Self { alloc, pos } = self;
        if pos.at_end() {
            return None;
        }
        // 准备下一个
        let chunks = alloc.chunks.get_mut();
        pos.advance(chunks);
        let indexed = match self.pos.get_ids(chunks) {
            Some((indexed, _)) => indexed.indexed,
            None => panic!("IterPos ensures valid pointer when not at end"),
        };
        let e = alloc
            .free_gen_index(indexed)
            .expect("UAF detected: pointer should be valid during consume iteration");
        Some((indexed, e))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.pos.nelems_last;
        (len, Some(len))
    }
}
impl<E, P: IAllocPolicy> ExactSizeIterator for EntityAllocConsumeIter<E, P> {
    fn len(&self) -> usize {
        self.pos.nelems_last
    }
}
impl<E, P: IAllocPolicy> EntityAllocConsumeIter<E, P> {
    /// Create a new iterator over the given `EntityAlloc`.
    pub fn new(mut alloc: EntityAlloc<E, P>) -> Self {
        let nelems_last = alloc.num_allocated.get();
        let mut pos = IterPos::new(nelems_last);
        pos.find_valid_unit(alloc.chunks.get_mut());
        Self { alloc, pos }
    }

    /// Release the owned `EntityAlloc` after iteration.
    pub fn release(self) -> EntityAlloc<E, P> {
        self.alloc
    }
}