artree 0.1.1

//! An arena based tree structure

#![cfg_attr(docsrs, feature(doc_cfg))]

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

/// A lightweight handle pointing to actual node data.
pub struct Node<Item>(usize, PhantomData<Item>);

impl<Item> Debug for Node<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 Node<Item> {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.0 == other.0
    }
}

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

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

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

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

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

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

impl<Item> Node<Item> {
    /// Create a new arena `token`.
    #[inline]
    pub const fn new(idx: usize) -> Self {
        Self(idx, PhantomData)
    }
}

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

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

/// Inner node slot.
struct Slot<Item> {
    /// key of parent node.
    parent: Option<Node<Item>>,
    /// keys of children nodes.
    children: Vec<Node<Item>>,
    /// item
    item: Option<Item>,
}

impl<Item> Default for Slot<Item> {
    fn default() -> Self {
        Self {
            parent: Default::default(),
            children: Default::default(),
            item: None,
        }
    }
}

/// Arena allocator.
pub struct Arena<Item> {
    /// the slots of layout nodes.
    slots: Vec<Slot<Item>>,
    /// unused layout slot keys.
    unused: Vec<Node<Item>>,
}

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

impl<Item> Arena<Item> {
    /// Create an empty `Arena`.
    pub fn new() -> Self {
        Self {
            slots: vec![Slot::default()],
            unused: Default::default(),
        }
    }

    /// Returns an iterator over the root nodes.
    #[inline]
    pub fn roots(&self) -> impl Iterator<Item = Node<Item>> {
        self.slots[0].children.iter().copied()
    }

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

    /// Returns true if the arena contains no elements.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Creates a new node with the given data and append to the given `parent`.
    ///
    /// If the parent is `None`, append the new node to root node list.
    #[inline]
    pub fn append(&mut self, parent: Option<Node<Item>>, value: Item) -> Node<Item> {
        let (parent, token) = self.create_node(parent, value);

        self.slots[*parent].children.push(token);

        token
    }

    #[inline]
    fn create_node(&mut self, parent: Option<Node<Item>>, value: Item) -> (Node<Item>, Node<Item>) {
        let parent = parent.or_else(|| Some(self.virtual_root()));

        let token = if let Some(token) = self.unused.pop() {
            self.slots[*token] = Slot {
                parent,
                item: Some(value),
                ..Default::default()
            };

            token
        } else {
            let token = Node::new(self.slots.len());

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

            token
        };

        (parent.unwrap(), token)
    }

    /// Creates a new node with the given data and sets as the previous sibling of the node.
    #[inline]
    pub fn insert_before(&mut self, token: Node<Item>, value: Item) -> Node<Item> {
        let parent = self.slots[*token].parent;

        let (parent, token) = self.create_node(parent, value);

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

        let offset = slot
            .children
            .iter()
            .position(|item| *item == token)
            .expect("offset in parent children list.");

        slot.children.insert(offset, token);

        token
    }

    /// Creates a new node with the given data and sets as the next sibling of the node.
    #[inline]
    pub fn insert_after(&mut self, token: Node<Item>, value: Item) -> Node<Item> {
        let parent = self.slots[*token].parent;

        let (parent, token) = self.create_node(parent, value);

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

        let offset = slot
            .children
            .iter()
            .position(|item| *item == token)
            .expect("offset in parent children list.");

        slot.children.insert(offset + 1, token);

        token
    }

    /// Detaches the given node and its descendants into its own tree while keeping it in the same arena.
    #[inline]
    pub fn detach(&mut self, token: Node<Item>) {
        self.set_parent(token, None);
    }

    /// Move the subtree under a new parent node
    #[inline]
    pub fn set_parent(&mut self, token: Node<Item>, parent: Option<Node<Item>>) {
        let old = self.slots[*token].parent.unwrap();
        let parent = parent.unwrap_or_else(|| self.virtual_root());

        if parent == old {
            return;
        }

        self.remove_from_parent(old, token);

        self.slots[*token].parent = Some(parent);

        self.slots[*parent].children.push(token);
    }

    /// Get the parent node token.
    ///
    /// Returns `None` if this node is a root element.
    #[inline]
    pub fn parent(&self, token: Node<Item>) -> Option<Node<Item>> {
        self.slots[*token]
            .parent
            .filter(|t| *t != self.virtual_root())
    }

    /// Returns an iterator over one's children nodes.
    #[inline]
    pub fn children(&self, token: Node<Item>) -> impl DoubleEndedIterator<Item = Node<Item>> {
        self.slots[*token].children.iter().copied()
    }

    /// Returns child at position.
    #[inline]
    pub fn child(&self, token: Node<Item>, offset: usize) -> Option<Node<Item>> {
        self.slots[*token].children.get(offset).copied()
    }

    /// Create an iterator over all ancestor nodes.
    #[inline]
    pub fn ancestor(&self, token: Node<Item>) -> Ancestor<'_, Item> {
        Ancestor {
            current: token,
            arena: self,
        }
    }

    /// Create an iterator over all descendant nodes.
    #[inline]
    pub fn descendants(&self, token: Node<Item>) -> Descendant<'_, Item> {
        Descendant {
            fifo: VecDeque::from_iter(self.slots[*token].children.iter().copied()),
            arena: self,
        }
    }

    /// Create an iterator that performs the depth-first search algorithm over all descendant nodes .
    #[inline]
    pub fn descendants_dfs(&self, token: Node<Item>) -> DescendantDFS<'_, Item> {
        DescendantDFS {
            stack: self.slots[*token].children.iter().rev().copied().collect(),
            arena: self,
        }
    }

    /// Swap two children nodes.
    #[inline]
    pub fn swap(&mut self, token: Node<Item>, x: usize, y: usize) {
        let slot = &mut self.slots[*token];

        let token = slot.children[x];
        slot.children[x] = slot.children[y];
        slot.children[y] = token;
    }

    /// Returns a reference to an element
    #[inline]
    pub fn get(&self, token: Node<Item>) -> Option<&Item> {
        self.slots[*token].item.as_ref()
    }

    /// Returns a mutable reference to an element
    #[inline]
    pub fn get_mut(&mut self, token: Node<Item>) -> Option<&mut Item> {
        self.slots[*token].item.as_mut()
    }

    /// Remove a subtree.
    ///
    /// Returns the `number` of removed elements.
    #[inline]
    pub fn remove(&mut self, token: Node<Item>) -> usize {
        self.remove_prv(VecDeque::from_iter([token]))
    }

    /// Removes the sub-children slice indicated by the given range from the element.
    #[inline]
    pub fn remove_children<R>(&mut self, token: Node<Item>, range: R)
    where
        R: RangeBounds<usize>,
    {
        self.unused.extend(self.slots[*token].children.drain(range));
    }

    #[inline]
    fn remove_prv(&mut self, mut tokens: VecDeque<Node<Item>>) -> usize {
        let mut count = 0;

        while let Some(token) = tokens.pop_front() {
            count += 1;

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

            assert!(slot.item.take().is_some(), "Remove twice. {:?}", token);

            tokens.extend(slot.children.drain(..));

            let parent = slot.parent.unwrap();

            self.remove_from_parent(parent, token);

            self.unused.push(token);
        }

        count
    }

    #[inline]
    fn remove_from_parent(&mut self, parent: Node<Item>, token: Node<Item>) {
        let slot = &mut self.slots[*parent];

        let Some(idx) = slot.children.iter().position(|item| *item == token) else {
            return;
        };

        assert_eq!(slot.children.swap_remove(idx), token);
    }

    /// Return the token for the virtual root node
    #[inline]
    const fn virtual_root(&self) -> Node<Item> {
        Node::new(0)
    }
}

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

    #[inline]
    fn index(&self, index: Node<Item>) -> &Self::Output {
        self.get(index).unwrap()
    }
}

impl<Item> IndexMut<Node<Item>> for Arena<Item> {
    #[inline]
    fn index_mut(&mut self, index: Node<Item>) -> &mut Self::Output {
        self.get_mut(index).unwrap()
    }
}

/// An iterator over ancestor nodes.
pub struct Ancestor<'a, Item> {
    current: Node<Item>,
    arena: &'a Arena<Item>,
}

impl<'a, Item> Iterator for Ancestor<'a, Item> {
    type Item = Node<Item>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let Some(next) = self.arena.parent(self.current) else {
            return None;
        };

        self.current = next;

        Some(next)
    }
}

/// An iterator over descendant nodes.
pub struct Descendant<'a, Item> {
    fifo: VecDeque<Node<Item>>,
    arena: &'a Arena<Item>,
}

impl<'a, Item> Iterator for Descendant<'a, Item> {
    type Item = Node<Item>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let Some(next) = self.fifo.pop_front() else {
            return None;
        };

        self.fifo.extend(self.arena.children(next));

        Some(next)
    }
}

/// An iterator over descendant nodes.
pub struct DescendantDFS<'a, Item> {
    stack: Vec<Node<Item>>,
    arena: &'a Arena<Item>,
}

impl<'a, Item> Iterator for DescendantDFS<'a, Item> {
    type Item = Node<Item>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let Some(next) = self.stack.pop() else {
            return None;
        };

        self.stack.extend(self.arena.children(next).rev());

        Some(next)
    }
}

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

    #[test]
    fn append_as_roots() {
        let mut arena = Arena::new();

        for i in 0..100 {
            arena.append(None, i);
        }

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

        for (idx, token) in arena.roots().enumerate() {
            assert_eq!(arena.parent(token), None);
            assert_eq!(arena[token], idx);
        }
    }

    #[test]
    fn children() {
        let mut arena = Arena::new();
        let root = arena.append(None, -1);

        for i in 0..100i32 {
            arena.append(Some(root), i);
        }

        for (idx, token) in arena.children(root).enumerate() {
            assert_eq!(arena.parent(token), Some(root));
            assert_eq!(arena[token], idx as i32);
        }
    }

    #[test]
    fn detach() {
        let mut arena = Arena::new();
        let root = arena.append(None, ());
        let child = arena.append(Some(root), ());

        assert_eq!(arena.parent(child), Some(root));
        assert_eq!(arena.children(root).next(), Some(child));

        // do nothing
        arena.detach(root);

        assert_eq!(arena.roots().next(), Some(root));

        assert_eq!(arena.parent(child), Some(root));
        assert_eq!(arena.children(root).next(), Some(child));

        // split off the subtree,
        arena.detach(child);

        assert_eq!(arena.parent(child), None);
        assert_eq!(arena.children(root).next(), None);

        assert_eq!(arena.roots().collect::<Vec<_>>(), [root, child]);
    }

    #[test]
    fn set_parent() {
        let mut arena = Arena::new();
        let root1 = arena.append(None, ());
        let root2 = arena.append(None, ());
        let child = arena.append(Some(root1), ());

        assert_eq!(arena.parent(child), Some(root1));
        assert_eq!(arena.children(root1).next(), Some(child));

        arena.set_parent(child, Some(root2));

        assert_eq!(arena.parent(child), Some(root2));
        assert_eq!(arena.children(root1).next(), None);
        assert_eq!(arena.children(root2).next(), Some(child));
    }

    #[test]
    fn ancestors_descendants() {
        let mut arena = Arena::new();
        let root = arena.append(None, 0);
        let mut current = root;

        for i in 1..100 {
            current = arena.append(Some(current), i);
        }

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

        for (idx, node) in arena.ancestor(current).enumerate() {
            assert_eq!(arena[node], 98 - idx);
        }

        for (idx, node) in arena.descendants(root).enumerate() {
            assert_eq!(arena[node], idx + 1);
        }
    }

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

        let mut current = arena.append(None, 0);

        for i in 1..100 {
            current = arena.append(Some(current), i);
        }

        let root = arena.roots().next().unwrap();

        let child = arena.children(root).next().unwrap();

        arena.remove(child);

        assert_eq!(arena.len(), 1);
    }

    #[test]
    fn remove2() {
        let mut arena = Arena::new();

        let root = arena.append(None, 0);

        for i in 1..100 {
            arena.append(Some(root), i);
        }

        arena.remove(root);

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

    #[test]
    fn remove_children() {
        let mut arena = Arena::new();

        let root = arena.append(None, 0);

        for i in 1..100 {
            arena.append(Some(root), i);
        }

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

        arena.remove_children(root, 0..99);

        assert_eq!(arena.len(), 1);
        assert_eq!(arena.roots().next(), Some(root));
    }

    #[test]
    fn swap() {
        let mut arena = Arena::new();

        let root = arena.append(None, 0);
        let c1 = arena.append(Some(root), 1);
        let c2 = arena.append(Some(root), 2);

        assert_eq!(arena.child(root, 0), Some(c1));
        assert_eq!(arena.child(root, 1), Some(c2));

        arena.swap(root, 0, 1);

        assert_eq!(arena.child(root, 0), Some(c2));
        assert_eq!(arena.child(root, 1), Some(c1));

        assert_eq!(arena[c1], 1);
        assert_eq!(arena[c2], 2);
    }
}