parentree 0.0.1

use serde::{Deserialize, Serialize};
use slotmap::{DefaultKey, DenseSlotMap};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Node<T> {
    parent_key: Option<DefaultKey>,
    children: Vec<DefaultKey>,
    val: T,
}

pub struct Tree<T>(DenseSlotMap<DefaultKey, Node<T>>);

impl<T> Tree<T> {
    pub fn root(val: T) -> (Self, DefaultKey) {
        let mut this = Self(DenseSlotMap::new());
        let root = this.0.insert(Node {
            parent_key: None,
            children: vec![],
            val,
        });
        (this, root)
    }

    pub fn add_child(&mut self, parent_key: DefaultKey, val: T) -> DefaultKey {
        let new = self.0.insert(Node {
            parent_key: Some(parent_key),
            children: vec![],
            val,
        });
        let parent = self.0.get_mut(parent_key).expect("key does not exist");
        parent.children.push(new);
        new
    }

    pub fn modify_recursive<F, Arg>(&mut self, key: DefaultKey, func: F, arg: Arg) -> Arg
    where
        F: Fn(&mut T, Arg) -> Arg,
    {
        let my_value = self.0.get_mut(key).expect("key does not exist");
        let arg = func(&mut my_value.val, arg);

        if let Some(parent_key) = my_value.parent_key {
            return self.modify_recursive(parent_key, func, arg);
        }
        return arg;
    }
}
impl<T> std::ops::Deref for Tree<T> {
    type Target = DenseSlotMap<DefaultKey, Node<T>>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

#[cfg(test)]
mod test {
    use crate::Tree;

    #[test]
    fn test_tree() {
        fn reducer(val: &mut i32, arg: i32) -> i32 {
            *val += arg;
            arg
        }

        let (mut tree, root) = Tree::root(2);

        let child1 = tree.add_child(root, 2);
        let child2 = tree.add_child(child1, 2);
        tree.modify_recursive(child2, reducer, 1);

        assert_eq!(tree[root].val, 3);
        assert_eq!(tree[child1].val, 3);
        assert_eq!(tree[child2].val, 3);
    }
}