daggrs 0.1.0

use alloc::vec::Vec;

/// A slot tracking next/prev indices for the circular doubly-linked list.
#[derive(Clone, Copy)]
struct Slot {
    next: u32,
    prev: u32,
}

/// Manages a circular doubly-linked list of free positions in the double-array.
/// Kept parallel to the states array - same size, same indices.
pub struct Allocator {
    slots: Vec<Slot>,
    head: Option<u32>,
}

impl Allocator {
    /// Create a new Allocator with initial capacity.
    /// Positions 0 and 1 (dead and root) are not added to the free list.
    pub fn new(initial_size: usize) -> Self {
        let mut allocator = Allocator {
            slots: Vec::new(),
            head: None,
        };
        if initial_size > 2 {
            allocator.extend(2, initial_size as u32);
        } else {
            // Just allocate slots for dead and root, but don't add to list
            allocator.slots.resize(initial_size, Slot { next: 0, prev: 0 });
        }
        allocator
    }

    /// Extend the free list to match a new states array size.
    /// New positions from `old_size` to `new_size` are added to the free list.
    pub fn extend(&mut self, old_size: u32, new_size: u32) {
        if new_size <= old_size {
            return;
        }

        // Ensure slots vec is big enough
        self.slots.resize(new_size as usize, Slot { next: 0, prev: 0 });

        let start = old_size.max(2); // Never add 0 or 1 to free list
        if start >= new_size {
            return;
        }

        // Link new indices to each other
        for idx in start..new_size {
            self.slots[idx as usize].next = if idx + 1 < new_size { idx + 1 } else { start };
            self.slots[idx as usize].prev = if idx > start { idx - 1 } else { new_size - 1 };
        }

        // Splice into existing list or set as head
        match self.head {
            None => {
                self.head = Some(start);
            }
            Some(h) => {
                let old_tail = self.slots[h as usize].prev;
                let new_tail = new_size - 1;

                self.slots[old_tail as usize].next = start;
                self.slots[start as usize].prev = old_tail;
                self.slots[new_tail as usize].next = h;
                self.slots[h as usize].prev = new_tail;
            }
        }
    }

    /// Remove an index from the free list (when the position is used).
    pub fn delete(&mut self, idx: u32) {
        let prev = self.slots[idx as usize].prev;
        let next = self.slots[idx as usize].next;

        if prev == idx {
            // Only element in list
            self.head = None;
        } else {
            self.slots[prev as usize].next = next;
            self.slots[next as usize].prev = prev;

            // Update head if we deleted it
            if self.head == Some(idx) {
                self.head = Some(next);
            }
        }
    }

    /// Iterate through all free indices.
    pub fn iter(&self) -> AllocatorIter<'_> {
        AllocatorIter {
            slots: &self.slots,
            head: self.head,
            current: self.head,
        }
    }
}

pub struct AllocatorIter<'a> {
    slots: &'a [Slot],
    head: Option<u32>,
    current: Option<u32>,
}

impl<'a> Iterator for AllocatorIter<'a> {
    type Item = u32;

    fn next(&mut self) -> Option<u32> {
        let idx = self.current?;
        let next = self.slots[idx as usize].next;

        // Stop when we've wrapped back to head
        self.current = if Some(next) == self.head {
            None
        } else {
            Some(next)
        };

        Some(idx)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_new_with_initial_size() {
        let alloc = Allocator::new(10);

        // Should have indices 2..10 in the free list (0 and 1 excluded)
        let indices: Vec<u32> = alloc.iter().collect();
        assert_eq!(indices, vec![2, 3, 4, 5, 6, 7, 8, 9]);
    }

    #[test]
    fn test_new_small_size() {
        let alloc = Allocator::new(2);

        // Only dead and root, no free positions
        assert_eq!(alloc.iter().count(), 0);
    }

    #[test]
    fn test_extend() {
        let mut alloc = Allocator::new(4); // indices 2, 3
        alloc.extend(4, 7); // add indices 4, 5, 6

        let indices: Vec<u32> = alloc.iter().collect();
        assert_eq!(indices, vec![2, 3, 4, 5, 6]);
    }

    #[test]
    fn test_extend_from_empty() {
        let mut alloc = Allocator::new(2); // no free positions
        alloc.extend(2, 5); // add indices 2, 3, 4

        let indices: Vec<u32> = alloc.iter().collect();
        assert_eq!(indices, vec![2, 3, 4]);
    }

    #[test]
    fn test_delete_middle() {
        let mut alloc = Allocator::new(6); // indices 2, 3, 4, 5

        alloc.delete(3);

        let indices: Vec<u32> = alloc.iter().collect();
        assert_eq!(indices, vec![2, 4, 5]);
    }

    #[test]
    fn test_delete_head() {
        let mut alloc = Allocator::new(6); // indices 2, 3, 4, 5

        alloc.delete(2); // delete head

        let indices: Vec<u32> = alloc.iter().collect();
        assert_eq!(indices, vec![3, 4, 5]);
    }

    #[test]
    fn test_delete_tail() {
        let mut alloc = Allocator::new(6); // indices 2, 3, 4, 5

        alloc.delete(5); // delete tail

        let indices: Vec<u32> = alloc.iter().collect();
        assert_eq!(indices, vec![2, 3, 4]);
    }

    #[test]
    fn test_delete_all() {
        let mut alloc = Allocator::new(4); // indices 2, 3

        alloc.delete(2);
        alloc.delete(3);

        assert_eq!(alloc.iter().count(), 0);
    }

    #[test]
    fn test_circular_structure() {
        let alloc = Allocator::new(5); // indices 2, 3, 4

        // Check circular links
        assert_eq!(alloc.slots[2].next, 3);
        assert_eq!(alloc.slots[2].prev, 4);
        assert_eq!(alloc.slots[4].next, 2);
        assert_eq!(alloc.slots[4].prev, 3);
    }
}