artree 0.1.0

//! An arena based tree structure

use std::{
    collections::VecDeque,
    fmt::Debug,
    hash::Hash,
    marker::PhantomData,
    ops::{Deref, Index, IndexMut, RangeBounds},
};

/// Key of layout object.
pub struct ArenaKey<Item>(usize, PhantomData<Item>);

impl<Item> Debug for ArenaKey<Item> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("ArenaKey")
            .field(&self.0)
            .finish_non_exhaustive()
    }
}

impl<Item> PartialEq for ArenaKey<Item> {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.0 == other.0
    }
}

impl<Item> Eq for ArenaKey<Item> {}

impl<Item> PartialOrd for ArenaKey<Item> {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        self.0.partial_cmp(&other.0)
    }
}

impl<Item> Ord for ArenaKey<Item> {
    #[inline]
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.0.cmp(&other.0)
    }
}

impl<Item> Hash for ArenaKey<Item> {
    #[inline]
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

impl<Item> Clone for ArenaKey<Item> {
    #[inline]
    fn clone(&self) -> Self {
        Self(self.0, PhantomData)
    }
}

impl<Item> Copy for ArenaKey<Item> {}

impl<Item> ArenaKey<Item> {
    #[inline]
    pub fn new(idx: usize) -> Self {
        Self(idx, PhantomData)
    }
}

impl<Item> Deref for ArenaKey<Item> {
    type Target = usize;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Arena node.
#[derive(Default)]
struct Node<Item> {
    /// key of parent node.
    parent: Option<ArenaKey<Item>>,
    /// keys of children nodes.
    children: Vec<ArenaKey<Item>>,
    /// item
    item: Option<Item>,
}

/// Arena tree
pub struct Arena<Item> {
    /// root elements.
    roots: Vec<ArenaKey<Item>>,
    /// the slots of layout nodes.
    slots: Vec<Node<Item>>,
    /// unused layout slot keys.
    unused: Vec<ArenaKey<Item>>,
}

impl<Item> Default for Arena<Item> {
    fn default() -> Self {
        Self {
            slots: Default::default(),
            unused: Default::default(),
            roots: Default::default(),
        }
    }
}

impl<Item> Arena<Item> {
    /// Returns the key iterator of root nodes.
    #[inline]
    pub fn roots(&self) -> impl Iterator<Item = ArenaKey<Item>> {
        self.roots.iter().copied()
    }

    /// Returns the number of nodes int this tree.
    #[inline]
    pub fn len(&self) -> usize {
        self.slots.len() - self.unused.len()
    }

    /// Add a new layout node.
    #[inline]
    pub fn add(&mut self, parent: Option<ArenaKey<Item>>, item: Item) -> ArenaKey<Item> {
        let key = if let Some(key) = self.unused.pop() {
            self.slots[*key] = Node {
                parent,
                children: Default::default(),
                item: Some(item),
            };

            key
        } else {
            let key = ArenaKey::new(self.slots.len());

            self.slots.push(Node {
                parent,
                children: Default::default(),
                item: Some(item),
            });

            key
        };

        // add child to parent children list.
        if let Some(parent) = parent {
            self.slots[*parent].children.push(key);
        } else {
            self.roots.push(key);
        }

        key
    }

    /// Remove item and all its descendants.
    #[inline]
    #[allow(unused)]
    pub fn remove(&mut self, key: ArenaKey<Item>) {
        let mut fifo = VecDeque::from([key]);
        self.remove_prv(fifo);
    }

    #[inline]
    fn remove_root_key(&mut self, key: ArenaKey<Item>) {
        let Some((idx, _)) = self.roots.iter().enumerate().find(|(_, lhs)| **lhs == key) else {
            return;
        };

        assert_eq!(self.roots.swap_remove(idx), key);
    }

    fn remove_prv(&mut self, mut fifo: VecDeque<ArenaKey<Item>>) {
        while let Some(key) = fifo.pop_front() {
            if self.set_parent_prv(key, None, true) == None {
                self.remove_root_key(key);
            }

            let slot = &mut self.slots[*key];

            assert_eq!(slot.item.take().is_some(), true);

            // add children nodes into fifo.
            fifo.extend(slot.children.drain(..));

            self.unused.push(key);
        }
    }

    /// Try get immutable reference of item by key.
    #[inline]
    pub fn get(&self, key: ArenaKey<Item>) -> Option<&Item> {
        self.slots[*key].item.as_ref()
    }

    /// Try get mutable reference of item by key.
    #[inline]
    pub fn get_mut(&mut self, key: ArenaKey<Item>) -> Option<&mut Item> {
        self.slots[*key].item.as_mut()
    }

    /// Returns parent node key
    #[inline]
    pub fn get_parent(&self, key: ArenaKey<Item>) -> Option<ArenaKey<Item>> {
        self.slots[*key].parent
    }

    /// Update node's parent.
    #[inline]
    pub fn set_parent(
        &mut self,
        key: ArenaKey<Item>,
        parent: Option<ArenaKey<Item>>,
    ) -> Option<ArenaKey<Item>> {
        self.set_parent_prv(key, parent, false)
    }

    #[inline]
    fn set_parent_prv(
        &mut self,
        key: ArenaKey<Item>,
        parent: Option<ArenaKey<Item>>,
        removing: bool,
    ) -> Option<ArenaKey<Item>> {
        let old = self.slots[*key].parent;

        if old != parent {
            if let Some(old) = old {
                let parent = &mut self.slots[*old];

                if let Some((idx, _)) = parent
                    .children
                    .iter()
                    .enumerate()
                    .find(|(_, lhs)| **lhs == key)
                {
                    parent.children.swap_remove(idx);
                }
            }

            if let Some(parent) = parent {
                self.slots[*parent].children.push(key);
            } else if !removing {
                self.roots.push(key);
            }
        }

        self.slots[*key].parent = parent;

        old
    }

    /// Returns the child item of parent in the given offset.
    #[inline]
    pub fn get_child(&mut self, key: ArenaKey<Item>, offset: usize) -> Option<ArenaKey<Item>> {
        self.slots[*key].children.get(offset).copied()
    }

    /// Returns the children slice.
    #[inline]
    pub fn get_children_slice(&self, key: ArenaKey<Item>) -> &[ArenaKey<Item>] {
        // safety: remove fn will clear the children list.
        &self.slots[*key].children
    }

    /// swap two children's position.
    #[inline]
    pub fn swap_children(&mut self, parent: ArenaKey<Item>, c1: usize, c2: usize) {
        if c1 == c2 {
            return;
        }

        let parent = &mut self.slots[*parent];
        let value = parent.children[c2];
        parent.children[c2] = parent.children[c1];
        parent.children[c1] = value;
    }

    /// Removes the subslice indicated by the given range from the children vector.
    #[inline]
    pub fn remove_children<R>(&mut self, parent: ArenaKey<Item>, range: R)
    where
        R: RangeBounds<usize>,
    {
        let removed = self.slots[*parent].children.drain(range).collect();

        self.remove_prv(removed);
    }
}

impl<Item> Index<ArenaKey<Item>> for Arena<Item> {
    type Output = Item;

    #[inline]
    fn index(&self, index: ArenaKey<Item>) -> &Self::Output {
        self.get(index).expect("Item does not exists")
    }
}

impl<Item> IndexMut<ArenaKey<Item>> for Arena<Item> {
    #[inline]
    fn index_mut(&mut self, index: ArenaKey<Item>) -> &mut Self::Output {
        self.get_mut(index).expect("Item does not exists")
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn root_item() {
        let mut arena = Arena::default();
        assert_eq!(arena.add(None, ()), ArenaKey::new(0));
        assert_eq!(arena.get_parent(ArenaKey::new(0)), None);
    }

    #[test]
    fn remove() {
        let mut arena = Arena::default();

        let mut parent = None;

        for i in 0..100 {
            parent = Some(arena.add(parent, i));
        }

        for i in 0..100 {
            assert_eq!(arena[ArenaKey::new(i)], i, "{}", i);
        }

        assert_eq!(arena.len(), 100);

        arena.remove(ArenaKey::new(0));

        assert_eq!(arena.len(), 0);

        assert_eq!(arena.unused.len(), 100);
    }

    #[test]
    fn remove_children() {
        let mut arena = Arena::default();
        let parent = Some(arena.add(None, ()));

        for _ in 0..10 {
            let parent = Some(arena.add(parent, ()));
            arena.add(parent, ());
        }

        assert_eq!(arena.len(), 21);

        arena.remove_children(parent.unwrap(), 3..10);

        assert_eq!(arena.len(), 7);
        assert_eq!(arena.unused.len(), 14);
    }

    #[test]
    fn set_parent() {
        let mut arena = Arena::default();

        let mut parent = None;

        for i in 0..100 {
            parent = Some(arena.add(parent, i));
        }

        assert_eq!(arena.roots().collect::<Vec<_>>(), &[ArenaKey::new(0)]);

        assert_eq!(arena.get_parent(ArenaKey::new(90)), Some(ArenaKey::new(89)));

        assert_eq!(
            arena.get_children_slice(ArenaKey::new(89)),
            &[ArenaKey::new(90)]
        );

        arena.set_parent(ArenaKey::new(90), None);

        assert_eq!(
            arena.roots().collect::<Vec<_>>(),
            &[ArenaKey::new(0), ArenaKey::new(90)]
        );

        assert_eq!(arena.get_parent(ArenaKey::new(90)), None);

        assert_eq!(arena.get_children_slice(ArenaKey::new(89)), &[]);

        arena.set_parent(ArenaKey::new(90), Some(ArenaKey::new(89)));

        assert_eq!(
            arena.get_children_slice(ArenaKey::new(89)),
            &[ArenaKey::new(90)]
        );
    }

    #[test]
    fn swap_children() {
        let mut arena = Arena::default();
        let parent = Some(arena.add(None, ()));

        for _ in 0..10 {
            let parent = Some(arena.add(parent, ()));
            arena.add(parent, ());
        }

        assert_eq!(arena.get_child(ArenaKey::new(0), 0), Some(ArenaKey::new(1)));
        assert_eq!(
            arena.get_child(ArenaKey::new(0), 9),
            Some(ArenaKey::new(19))
        );

        assert_eq!(arena.get_child(ArenaKey::new(0), 10), None, "Out of range");

        arena.swap_children(ArenaKey::new(0), 0, 9);

        assert_eq!(
            arena.get_child(ArenaKey::new(0), 0),
            Some(ArenaKey::new(19))
        );
        assert_eq!(arena.get_child(ArenaKey::new(0), 9), Some(ArenaKey::new(1)));
    }
}