ECS based tree data struture

packages/native-core/Cargo.toml

@@ -23,6 +23,8 @@ slab = "0.4"
 parking_lot = "0.12.1"
 crossbeam-deque = "0.8.2"
 dashmap = "5.4.0"
+hashbrown = { version = "0.13.2", features = ["raw"] }
+parking_lot = "0.12.1"
 
 # for parsing attributes
 lightningcss = "1.0.0-alpha.39"

packages/native-core/src/tree.rs

@@ -1,100 +1,271 @@
+use parking_lot::{
+    MappedRwLockReadGuard, MappedRwLockWriteGuard, RwLock, RwLockReadGuard, RwLockWriteGuard,
+};
+use rustc_hash::FxHasher;
+use std::any::{Any, TypeId};
 use std::collections::VecDeque;
-use std::marker::PhantomData;
-
-#[derive(Hash, PartialEq, Eq, Clone, Copy, Debug, PartialOrd, Ord)]
-pub struct NodeId(pub usize);
+use std::fmt::Debug;
+use std::hash::BuildHasherDefault;
 
 #[derive(PartialEq, Eq, Clone, Debug)]
-pub struct Node<T> {
-    value: T,
-    parent: Option<NodeId>,
-    children: Vec<NodeId>,
+pub struct Node {
+    parent: Option<usize>,
+    children: Vec<usize>,
     height: u16,
 }
 
+pub struct NodeView<'a> {
+    tree: &'a Tree,
+    id: usize,
+}
+
+impl NodeView<'_> {
+    pub fn insert<T>(&self, data: T)
+    where
+        T: Any,
+    {
+        self.tree.nodes.get_slab_mut::<T>().insert(self.id, data);
+    }
+
+    pub fn get<T>(&self) -> MappedRwLockReadGuard<'_, T>
+    where
+        T: Any,
+    {
+        MappedRwLockReadGuard::map(self.tree.nodes.get_slab::<T>(), |slab| {
+            slab.get(self.id).unwrap()
+        })
+    }
+
+    pub fn get_mut<T>(&self) -> MappedRwLockWriteGuard<'_, T>
+    where
+        T: Any,
+    {
+        MappedRwLockWriteGuard::map(self.tree.nodes.get_slab_mut::<T>(), |slab| {
+            slab.get_mut(self.id).unwrap()
+        })
+    }
+
+    pub fn height(&self) -> u16 {
+        self.get::<Node>().height
+    }
+
+    pub fn children_ids(&self) -> Vec<usize> {
+        self.get::<Node>().children.clone()
+    }
+
+    pub fn parent_id(&self) -> Option<usize> {
+        self.get::<Node>().parent
+    }
+
+    pub fn id(&self) -> usize {
+        self.id
+    }
+}
+
 #[derive(Debug)]
-pub struct Tree<T> {
-    nodes: Slab<Node<T>>,
-    root: NodeId,
+pub struct Tree {
+    nodes: AnySlab,
+    root: usize,
 }
 
-impl<T> Tree<T> {
-    pub fn new(root: T) -> Self {
-        let mut nodes = Slab::default();
-        let root = NodeId(nodes.insert(Node {
-            value: root,
+impl Tree {
+    pub fn new() -> Self {
+        let mut nodes = AnySlab::default();
+        let mut node = nodes.insert();
+        node.insert(Node {
             parent: None,
             children: Vec::new(),
             height: 0,
-        }));
+        });
+        let root = node.id();
         Self { nodes, root }
     }
 
-    fn try_remove(&mut self, id: NodeId) -> Option<Node<T>> {
-        self.nodes.try_remove(id.0).map(|node| {
-            if let Some(parent) = node.parent {
-                self.nodes
-                    .get_mut(parent.0)
-                    .unwrap()
-                    .children
-                    .retain(|child| child != &id);
+    fn node_slab(&self) -> MappedRwLockReadGuard<'_, SlabStorage<Node>> {
+        self.nodes.get_slab()
+    }
+
+    pub fn get_node_data(&self, id: usize) -> MappedRwLockReadGuard<'_, Node> {
+        MappedRwLockReadGuard::map(self.node_slab(), |slab| slab.get(id).unwrap())
+    }
+
+    fn node_slab_mut(&self) -> MappedRwLockWriteGuard<'_, SlabStorage<Node>> {
+        self.nodes.get_slab_mut()
+    }
+
+    pub fn get_node_data_mut(&self, id: usize) -> MappedRwLockWriteGuard<'_, Node> {
+        MappedRwLockWriteGuard::map(self.node_slab_mut(), |slab| slab.get_mut(id).unwrap())
+    }
+
+    pub fn remove(&mut self, id: usize) {
+        fn recurse(tree: &mut Tree, id: usize) {
+            let children = { tree.get_node_data(id).children.clone() };
+            for child in children {
+                recurse(tree, child);
             }
-            for child in &node.children {
-                self.remove_recursive(*child);
+
+            tree.nodes.remove(id);
+        }
+        {
+            let mut node_data_mut = self.node_slab_mut();
+            if let Some(parent) = node_data_mut.get(id).unwrap().parent {
+                let parent = node_data_mut.get_mut(parent).unwrap();
+                parent.children.retain(|&child| child != id);
             }
-            node
-        })
+        }
+
+        recurse(self, id);
     }
 
-    fn remove_recursive(&mut self, node: NodeId) {
-        let node = self.nodes.remove(node.0);
-        for child in node.children {
-            self.remove_recursive(child);
+    fn set_height(&self, node: usize, height: u16) {
+        let children = {
+            let mut node = self.get_node_data_mut(node);
+            node.height = height;
+            node.children.clone()
+        };
+        for child in children {
+            self.set_height(child, height + 1);
         }
     }
 
-    fn set_height(&mut self, node: NodeId, height: u16) {
-        let self_mut = self as *mut Self;
-        let node = self.nodes.get_mut(node.0).unwrap();
-        node.height = height;
-        unsafe {
-            // Safety: No node has itself as a child
-            for child in &node.children {
-                (*self_mut).set_height(*child, height + 1);
+    pub fn create_node(&mut self) -> Entry<'_> {
+        let mut node = self.nodes.insert();
+        node.insert(Node {
+            parent: None,
+            children: Vec::new(),
+            height: 0,
+        });
+        node
+    }
+
+    pub fn add_child(&mut self, parent: usize, new: usize) {
+        let mut node_state = self.node_slab_mut();
+        node_state.get_mut(new).unwrap().parent = Some(parent);
+        let parent = node_state.get_mut(parent).unwrap();
+        parent.children.push(new);
+        let height = parent.height + 1;
+        drop(node_state);
+        self.set_height(new, height);
+    }
+
+    pub fn replace(&mut self, old_id: usize, new_id: usize) {
+        {
+            let mut node_state = self.node_slab_mut();
+            // update the parent's link to the child
+            if let Some(parent_id) = node_state.get(old_id).unwrap().parent {
+                let parent = node_state.get_mut(parent_id).unwrap();
+                for id in &mut parent.children {
+                    if *id == old_id {
+                        *id = new_id;
+                        break;
+                    }
+                }
+                let height = parent.height + 1;
+                drop(node_state);
+                self.set_height(new_id, height);
             }
         }
+        // remove the old node
+        self.remove(old_id);
     }
-}
 
-pub trait TreeView<T>: Sized {
-    type Iterator<'a>: Iterator<Item = &'a T>
-    where
-        T: 'a,
-        Self: 'a;
+    pub fn insert_before(&mut self, old_id: usize, new_id: usize) {
+        let mut node_state = self.node_slab_mut();
+        let old_node = node_state.get(old_id).unwrap();
+        let parent_id = old_node.parent.expect("tried to insert before root");
+        node_state.get_mut(new_id).unwrap().parent = Some(parent_id);
+        let parent = node_state.get_mut(parent_id).unwrap();
+        let index = parent
+            .children
+            .iter()
+            .position(|child| *child == old_id)
+            .unwrap();
+        parent.children.insert(index, new_id);
+        let height = parent.height + 1;
+        drop(node_state);
+        self.set_height(new_id, height);
+    }
 
-    fn root(&self) -> NodeId;
+    pub fn insert_after(&mut self, old_id: usize, new_id: usize) {
+        let mut node_state = self.node_slab_mut();
+        let old_node = node_state.get(old_id).unwrap();
+        let parent_id = old_node.parent.expect("tried to insert before root");
+        node_state.get_mut(new_id).unwrap().parent = Some(parent_id);
+        let parent = node_state.get_mut(parent_id).unwrap();
+        let index = parent
+            .children
+            .iter()
+            .position(|child| *child == old_id)
+            .unwrap();
+        parent.children.insert(index + 1, new_id);
+        let height = parent.height + 1;
+        drop(node_state);
+        self.set_height(new_id, height);
+    }
+
+    pub fn insert<T: Any>(&mut self, id: usize, value: T) {
+        self.nodes.add(id, value);
+    }
 
-    fn contains(&self, id: NodeId) -> bool {
-        self.get(id).is_some()
+    pub fn size(&self) -> usize {
+        self.nodes.len()
     }
 
-    fn get(&self, id: NodeId) -> Option<&T>;
+    pub fn state_view<T: Any>(&self) -> TreeStateView<'_, T> {
+        TreeStateView {
+            nodes_data: self.node_slab(),
+            nodes: self.nodes.get_slab(),
+            root: self.root,
+        }
+    }
 
-    fn get_unchecked(&self, id: NodeId) -> &T {
-        unsafe { self.get(id).unwrap_unchecked() }
+    pub fn state_view_mut<T: Any>(&mut self) -> TreeStateViewMut<'_, T> {
+        TreeStateViewMut {
+            nodes_data: self.node_slab(),
+            nodes: self.nodes.get_slab_mut(),
+            root: self.root,
+        }
+    }
+}
+
+impl NodeDataView for Tree {
+    fn root(&self) -> usize {
+        self.root
+    }
+
+    fn parent_id(&self, id: usize) -> Option<usize> {
+        self.get_node_data(id).parent
+    }
+
+    fn children_ids(&self, id: usize) -> Option<Vec<usize>> {
+        Some(self.get_node_data(id).children.clone())
+    }
+
+    fn height(&self, id: usize) -> Option<u16> {
+        Some(self.get_node_data(id).height)
     }
+}
 
-    fn children(&self, id: NodeId) -> Option<Self::Iterator<'_>>;
+pub trait NodeDataView {
+    fn root(&self) -> usize;
 
-    fn children_ids(&self, id: NodeId) -> Option<&[NodeId]>;
+    fn parent_id(&self, id: usize) -> Option<usize>;
 
-    fn parent(&self, id: NodeId) -> Option<&T>;
+    fn children_ids(&self, id: usize) -> Option<Vec<usize>>;
 
-    fn parent_id(&self, id: NodeId) -> Option<NodeId>;
+    fn height(&self, id: usize) -> Option<u16>;
+}
+
+pub trait TreeView<T>: Sized + NodeDataView {
+    fn get(&self, id: usize) -> Option<&T>;
+
+    fn get_unchecked(&self, id: usize) -> &T {
+        unsafe { self.get(id).unwrap_unchecked() }
+    }
 
-    fn height(&self, id: NodeId) -> Option<u16>;
+    fn children(&self, id: usize) -> Option<Vec<&T>>;
 
-    fn size(&self) -> usize;
+    fn parent(&self, id: usize) -> Option<&T>;
 
     fn traverse_depth_first(&self, mut f: impl FnMut(&T)) {
         let mut stack = vec![self.root()];
@@ -116,7 +287,7 @@ pub trait TreeView<T>: Sized {
                 f(node);
                 if let Some(children) = self.children_ids(id) {
                     for id in children {
-                        queue.push_back(*id);
+                        queue.push_back(id);
                     }
                 }
             }
@@ -125,23 +296,15 @@ pub trait TreeView<T>: Sized {
 }
 
 pub trait TreeViewMut<T>: Sized + TreeView<T> {
-    type IteratorMut<'a>: Iterator<Item = &'a mut T>
-    where
-        T: 'a,
-        Self: 'a;
+    fn get_mut(&mut self, id: usize) -> Option<&mut T>;
 
-    fn get_mut(&mut self, id: NodeId) -> Option<&mut T>;
-
-    fn get_mut_unchecked(&mut self, id: NodeId) -> &mut T {
+    unsafe fn get_mut_unchecked(&mut self, id: usize) -> &mut T {
         unsafe { self.get_mut(id).unwrap_unchecked() }
     }
 
-    fn children_mut(&mut self, id: NodeId) -> Option<<Self as TreeViewMut<T>>::IteratorMut<'_>>;
+    fn children_mut(&mut self, id: usize) -> Option<Vec<&mut T>>;
 
-    fn parent_child_mut(
-        &mut self,
-        id: NodeId,
-    ) -> Option<(&mut T, <Self as TreeViewMut<T>>::IteratorMut<'_>)> {
+    fn parent_child_mut(&mut self, id: usize) -> Option<(&mut T, Vec<&mut T>)> {
         let mut_ptr: *mut Self = self;
         unsafe {
             // Safety: No node has itself as a child.
@@ -153,19 +316,15 @@ pub trait TreeViewMut<T>: Sized + TreeView<T> {
         }
     }
 
-    fn parent_mut(&mut self, id: NodeId) -> Option<&mut T>;
+    fn parent_mut(&mut self, id: usize) -> Option<&mut T>;
 
-    fn node_parent_mut(&mut self, id: NodeId) -> Option<(&mut T, Option<&mut T>)>;
+    fn node_parent_mut(&mut self, id: usize) -> Option<(&mut T, Option<&mut T>)>;
 
     #[allow(clippy::type_complexity)]
     fn node_parent_children_mut(
         &mut self,
-        id: NodeId,
-    ) -> Option<(
-        &mut T,
-        Option<&mut T>,
-        Option<<Self as TreeViewMut<T>>::IteratorMut<'_>>,
-    )>;
+        id: usize,
+    ) -> Option<(&mut T, Option<&mut T>, Option<Vec<&mut T>>)>;
 
     fn traverse_depth_first_mut(&mut self, mut f: impl FnMut(&mut T)) {
         let mut stack = vec![self.root()];
@@ -187,7 +346,7 @@ pub trait TreeViewMut<T>: Sized + TreeView<T> {
                 f(node);
                 if let Some(children) = self.children_ids(id) {
                     for id in children {
-                        queue.push_back(*id);
+                        queue.push_back(id);
                     }
                 }
             }
@@ -195,160 +354,150 @@ pub trait TreeViewMut<T>: Sized + TreeView<T> {
     }
 }
 
-pub trait TreeLike<T> {
-    fn create_node(&mut self, value: T) -> NodeId;
+pub trait TreeLike {
+    fn create_node(&mut self) -> Entry;
 
-    fn add_child(&mut self, parent: NodeId, child: NodeId);
+    fn add_child(&mut self, parent: usize, child: usize);
 
-    fn remove(&mut self, id: NodeId) -> Option<T>;
+    fn remove(&mut self, id: usize);
 
-    fn remove_all_children(&mut self, id: NodeId) -> Vec<T>;
+    fn replace(&mut self, old: usize, new: usize);
 
-    fn replace(&mut self, old: NodeId, new: NodeId);
+    fn insert_before(&mut self, id: usize, new: usize);
 
-    fn insert_before(&mut self, id: NodeId, new: NodeId);
-
-    fn insert_after(&mut self, id: NodeId, new: NodeId);
+    fn insert_after(&mut self, id: usize, new: usize);
 }
 
-pub struct ChildNodeIterator<'a, T> {
-    nodes: &'a Slab<Node<T>>,
-    children_ids: Vec<NodeId>,
-    index: usize,
-    node_type: PhantomData<T>,
+pub struct TreeStateView<'a, T> {
+    nodes_data: MappedRwLockReadGuard<'a, SlabStorage<Node>>,
+    nodes: MappedRwLockReadGuard<'a, SlabStorage<T>>,
+    root: usize,
 }
 
-impl<'a, T: 'a> Iterator for ChildNodeIterator<'a, T> {
-    type Item = &'a T;
+pub struct TreeStateViewMut<'a, T> {
+    nodes_data: MappedRwLockReadGuard<'a, SlabStorage<Node>>,
+    nodes: MappedRwLockWriteGuard<'a, SlabStorage<T>>,
+    root: usize,
+}
 
-    fn next(&mut self) -> Option<Self::Item> {
-        self.children_ids.get(self.index).map(|id| {
-            self.index += 1;
-            &self.nodes.get(id.0).unwrap().value
-        })
+impl<'a, T> NodeDataView for TreeStateView<'a, T> {
+    fn root(&self) -> usize {
+        self.root
     }
-}
 
-pub struct ChildNodeIteratorMut<'a, T> {
-    children: Vec<&'a mut Node<T>>,
-}
+    fn children_ids(&self, id: usize) -> Option<Vec<usize>> {
+        self.nodes_data.get(id).map(|node| node.children.clone())
+    }
 
-impl<'a, T> Iterator for ChildNodeIteratorMut<'a, T> {
-    type Item = &'a mut T;
+    fn parent_id(&self, id: usize) -> Option<usize> {
+        self.nodes_data.get(id).and_then(|node| node.parent)
+    }
 
-    fn next(&mut self) -> Option<Self::Item> {
-        self.children.pop().map(|node| &mut node.value)
+    fn height(&self, id: usize) -> Option<u16> {
+        self.nodes_data.get(id).map(|n| n.height)
     }
 }
 
-impl<T> TreeView<T> for Tree<T> {
-    type Iterator<'a> = ChildNodeIterator<'a, T> where T: 'a;
+impl<'a, T> TreeView<T> for TreeStateView<'a, T> {
+    fn get(&self, id: usize) -> Option<&T> {
+        self.nodes.get(id)
+    }
 
-    fn root(&self) -> NodeId {
-        self.root
+    fn children(&self, id: usize) -> Option<Vec<&T>> {
+        let ids = self.children_ids(id);
+        ids.map(|ids| ids.iter().map(|id| self.get_unchecked(*id)).collect())
     }
 
-    fn get(&self, id: NodeId) -> Option<&T> {
-        self.nodes.get(id.0).map(|node| &node.value)
+    fn parent(&self, id: usize) -> Option<&T> {
+        self.nodes_data
+            .get(id)
+            .and_then(|node| node.parent.map(|id| self.nodes.get(id).unwrap()))
     }
 
-    fn children(&self, id: NodeId) -> Option<Self::Iterator<'_>> {
-        self.children_ids(id).map(|children_ids| ChildNodeIterator {
-            nodes: &self.nodes,
-            children_ids: children_ids.to_vec(),
-            index: 0,
-            node_type: PhantomData,
-        })
+    fn get_unchecked(&self, id: usize) -> &T {
+        unsafe { &self.nodes.get_unchecked(id) }
     }
+}
 
-    fn children_ids(&self, id: NodeId) -> Option<&[NodeId]> {
-        self.nodes.get(id.0).map(|node| node.children.as_slice())
+impl<'a, T> NodeDataView for TreeStateViewMut<'a, T> {
+    fn root(&self) -> usize {
+        self.root
     }
 
-    fn parent(&self, id: NodeId) -> Option<&T> {
-        self.nodes
-            .get(id.0)
-            .and_then(|node| node.parent.map(|id| self.nodes.get(id.0).unwrap()))
-            .map(|node| &node.value)
+    fn children_ids(&self, id: usize) -> Option<Vec<usize>> {
+        self.nodes_data.get(id).map(|node| node.children.clone())
     }
 
-    fn parent_id(&self, id: NodeId) -> Option<NodeId> {
-        self.nodes.get(id.0).and_then(|node| node.parent)
+    fn parent_id(&self, id: usize) -> Option<usize> {
+        self.nodes_data.get(id).and_then(|node| node.parent)
     }
 
-    fn height(&self, id: NodeId) -> Option<u16> {
-        self.nodes.get(id.0).map(|n| n.height)
+    fn height(&self, id: usize) -> Option<u16> {
+        self.nodes_data.get(id).map(|n| n.height)
     }
+}
 
-    fn get_unchecked(&self, id: NodeId) -> &T {
-        unsafe { &self.nodes.get_unchecked(id.0).value }
+impl<'a, T> TreeView<T> for TreeStateViewMut<'a, T> {
+    fn get(&self, id: usize) -> Option<&T> {
+        self.nodes.get(id)
     }
 
-    fn size(&self) -> usize {
-        self.nodes.len()
+    fn children(&self, id: usize) -> Option<Vec<&T>> {
+        let ids = self.children_ids(id);
+        ids.map(|ids| ids.iter().map(|id| self.get_unchecked(*id)).collect())
+    }
+    fn parent(&self, id: usize) -> Option<&T> {
+        self.nodes_data
+            .get(id)
+            .and_then(|node| node.parent.map(|id| self.nodes.get(id).unwrap()))
     }
-}
 
-impl<T> TreeViewMut<T> for Tree<T> {
-    type IteratorMut<'a> = ChildNodeIteratorMut<'a, T> where T: 'a;
+    fn get_unchecked(&self, id: usize) -> &T {
+        unsafe { &self.nodes.get_unchecked(id) }
+    }
+}
 
-    fn get_mut(&mut self, id: NodeId) -> Option<&mut T> {
-        self.nodes.get_mut(id.0).map(|node| &mut node.value)
+impl<'a, T> TreeViewMut<T> for TreeStateViewMut<'a, T> {
+    fn get_mut(&mut self, id: usize) -> Option<&mut T> {
+        self.nodes.get_mut(id)
     }
 
-    fn parent_mut(&mut self, id: NodeId) -> Option<&mut T> {
-        let self_ptr = self as *mut Self;
-        unsafe {
-            // Safety: No node has itself as a parent.
-            self.nodes
-                .get_mut(id.0)
-                .and_then(move |node| {
-                    node.parent
-                        .map(move |id| (*self_ptr).nodes.get_mut(id.0).unwrap())
-                })
-                .map(|node| &mut node.value)
-        }
+    fn parent_mut(&mut self, id: usize) -> Option<&mut T> {
+        let parent_id = self.parent_id(id)?;
+        unsafe { Some(self.get_mut_unchecked(parent_id)) }
     }
 
-    fn children_mut(&mut self, id: NodeId) -> Option<Self::IteratorMut<'_>> {
+    fn children_mut(&mut self, id: usize) -> Option<Vec<&mut T>> {
         // Safety: No node has itself as a parent.
         if let Some(children_ids) = self.children_ids(id) {
             let children_ids = children_ids.to_vec();
-            let children_ids = children_ids.to_vec();
-            Some(ChildNodeIteratorMut {
-                children: unsafe {
-                    self.nodes
-                        .get_many_mut_unchecked(children_ids.into_iter().rev().map(|id| id.0))
-                        .unwrap()
-                },
-            })
+            unsafe {
+                self.nodes
+                    .get_many_mut_unchecked(children_ids.into_iter().rev().map(|id| id))
+            }
         } else {
             None
         }
     }
 
-    fn get_mut_unchecked(&mut self, id: NodeId) -> &mut T {
-        unsafe { &mut self.nodes.get_unchecked_mut(id.0).value }
+    unsafe fn get_mut_unchecked(&mut self, id: usize) -> &mut T {
+        unsafe { self.nodes.get_unchecked_mut(id) }
     }
 
-    fn node_parent_mut(&mut self, id: NodeId) -> Option<(&mut T, Option<&mut T>)> {
+    fn node_parent_mut(&mut self, id: usize) -> Option<(&mut T, Option<&mut T>)> {
         if let Some(parent_id) = self.parent_id(id) {
             self.nodes
-                .get2_mut(id.0, parent_id.0)
-                .map(|(node, parent)| (&mut node.value, Some(&mut parent.value)))
+                .get2_mut(id, parent_id)
+                .map(|(node, parent)| (node, Some(parent)))
         } else {
-            self.nodes.get_mut(id.0).map(|node| (&mut node.value, None))
+            self.nodes.get_mut(id).map(|node| (node, None))
         }
     }
 
     fn node_parent_children_mut(
         &mut self,
-        id: NodeId,
-    ) -> Option<(
-        &mut T,
-        Option<&mut T>,
-        Option<<Self as TreeViewMut<T>>::IteratorMut<'_>>,
-    )> {
+        id: usize,
+    ) -> Option<(&mut T, Option<&mut T>, Option<Vec<&mut T>>)> {
         // Safety: No node has itself as a parent.
         let unbounded_self = unsafe { &mut *(self as *mut Self) };
         self.node_parent_mut(id).map(move |(node, parent)| {
@@ -357,7 +506,7 @@ impl<T> TreeViewMut<T> for Tree<T> {
         })
     }
 
-    fn parent_child_mut(&mut self, id: NodeId) -> Option<(&mut T, Self::IteratorMut<'_>)> {
+    fn parent_child_mut(&mut self, id: usize) -> Option<(&mut T, Vec<&mut T>)> {
         // Safety: No node will appear as a child twice
         if let Some(children_ids) = self.children_ids(id) {
             debug_assert!(!children_ids.iter().any(|child_id| *child_id == id));
@@ -365,131 +514,57 @@ impl<T> TreeViewMut<T> for Tree<T> {
                 let as_vec = children_ids.to_vec();
                 self.nodes
                     .get_many_mut_unchecked(
-                        as_vec
-                            .into_iter()
-                            .rev()
-                            .map(|id| id.0)
-                            .chain(std::iter::once(id.0)),
+                        as_vec.into_iter().map(|id| id).chain(std::iter::once(id)),
                     )
                     .unwrap()
             };
-            let node = &mut borrowed.pop().unwrap().value;
-            Some((node, ChildNodeIteratorMut { children: borrowed }))
+            let node = borrowed.pop().unwrap();
+            Some((node, borrowed))
         } else {
             None
         }
     }
 }
 
-impl<T> TreeLike<T> for Tree<T> {
-    fn create_node(&mut self, value: T) -> NodeId {
-        NodeId(self.nodes.insert(Node {
-            value,
-            parent: None,
-            children: Vec::new(),
-            height: 0,
-        }))
-    }
-
-    fn add_child(&mut self, parent: NodeId, new: NodeId) {
-        self.nodes.get_mut(new.0).unwrap().parent = Some(parent);
-        let parent = self.nodes.get_mut(parent.0).unwrap();
-        parent.children.push(new);
-        let height = parent.height + 1;
-        self.set_height(new, height);
-    }
-
-    fn remove(&mut self, id: NodeId) -> Option<T> {
-        self.try_remove(id).map(|node| node.value)
-    }
-
-    fn remove_all_children(&mut self, id: NodeId) -> Vec<T> {
-        let mut children = Vec::new();
-        let self_mut = self as *mut Self;
-        for child in self.children_ids(id).unwrap() {
-            unsafe {
-                // Safety: No node has itself as a child
-                children.push((*self_mut).remove(*child).unwrap());
-            }
-        }
-        children
-    }
-
-    fn replace(&mut self, old_id: NodeId, new_id: NodeId) {
-        // remove the old node
-        let old = self
-            .try_remove(old_id)
-            .expect("tried to replace a node that doesn't exist");
-        // update the parent's link to the child
-        if let Some(parent_id) = old.parent {
-            let parent = self.nodes.get_mut(parent_id.0).unwrap();
-            for id in &mut parent.children {
-                if *id == old_id {
-                    *id = new_id;
-                }
-            }
-            let height = parent.height + 1;
-            self.set_height(new_id, height);
-        }
-    }
-
-    fn insert_before(&mut self, id: NodeId, new: NodeId) {
-        let node = self.nodes.get(id.0).unwrap();
-        let parent_id = node.parent.expect("tried to insert before root");
-        self.nodes.get_mut(new.0).unwrap().parent = Some(parent_id);
-        let parent = self.nodes.get_mut(parent_id.0).unwrap();
-        let index = parent
-            .children
-            .iter()
-            .position(|child| child == &id)
-            .unwrap();
-        parent.children.insert(index, new);
-        let height = parent.height + 1;
-        self.set_height(new, height);
-    }
-
-    fn insert_after(&mut self, id: NodeId, new: NodeId) {
-        let node = self.nodes.get(id.0).unwrap();
-        let parent_id = node.parent.expect("tried to insert before root");
-        self.nodes.get_mut(new.0).unwrap().parent = Some(parent_id);
-        let parent = self.nodes.get_mut(parent_id.0).unwrap();
-        let index = parent
-            .children
-            .iter()
-            .position(|child| child == &id)
-            .unwrap();
-        parent.children.insert(index + 1, new);
-        let height = parent.height + 1;
-        self.set_height(new, height);
-    }
-}
 #[test]
 fn creation() {
-    let mut tree = Tree::new(1);
-    let parent = tree.root();
-    let child = tree.create_node(0);
-    tree.add_child(parent, child);
+    let mut tree = Tree::new();
+    let parent_id = tree.root;
+    tree.insert(parent_id, 1i32);
+    let mut child = tree.create_node();
+    child.insert(0i32);
+    let child_id = child.id();
+
+    tree.add_child(parent_id, child_id);
 
     println!("Tree: {:#?}", tree);
     assert_eq!(tree.size(), 2);
-    assert_eq!(tree.height(parent), Some(0));
-    assert_eq!(tree.height(child), Some(1));
-    assert_eq!(*tree.get(parent).unwrap(), 1);
-    assert_eq!(*tree.get(child).unwrap(), 0);
-    assert_eq!(tree.parent_id(parent), None);
-    assert_eq!(tree.parent_id(child).unwrap(), parent);
-    assert_eq!(tree.children_ids(parent).unwrap(), &[child]);
+    assert_eq!(tree.height(parent_id), Some(0));
+    assert_eq!(tree.height(child_id), Some(1));
+    assert_eq!(tree.parent_id(parent_id), None);
+    assert_eq!(tree.parent_id(child_id).unwrap(), parent_id);
+    assert_eq!(tree.children_ids(parent_id).unwrap(), &[child_id]);
+    let view = tree.state_view::<i32>();
+    assert_eq!(*view.get(parent_id).unwrap(), 1);
+    assert_eq!(*view.get(child_id).unwrap(), 0);
 }
 
 #[test]
 fn insertion() {
-    let mut tree = Tree::new(0);
+    let mut tree = Tree::new();
     let parent = tree.root();
-    let child = tree.create_node(2);
+    tree.insert(parent, 0);
+    let mut child = tree.create_node();
+    child.insert(2);
+    let child = child.id();
     tree.add_child(parent, child);
-    let before = tree.create_node(1);
+    let mut before = tree.create_node();
+    before.insert(1);
+    let before = before.id();
     tree.insert_before(child, before);
-    let after = tree.create_node(3);
+    let mut after = tree.create_node();
+    after.insert(3);
+    let after = after.id();
     tree.insert_after(child, after);
 
     println!("Tree: {:#?}", tree);
@@ -498,25 +573,33 @@ fn insertion() {
     assert_eq!(tree.height(child), Some(1));
     assert_eq!(tree.height(before), Some(1));
     assert_eq!(tree.height(after), Some(1));
-    assert_eq!(*tree.get(parent).unwrap(), 0);
-    assert_eq!(*tree.get(before).unwrap(), 1);
-    assert_eq!(*tree.get(child).unwrap(), 2);
-    assert_eq!(*tree.get(after).unwrap(), 3);
     assert_eq!(tree.parent_id(before).unwrap(), parent);
     assert_eq!(tree.parent_id(child).unwrap(), parent);
     assert_eq!(tree.parent_id(after).unwrap(), parent);
     assert_eq!(tree.children_ids(parent).unwrap(), &[before, child, after]);
+    let view = tree.state_view::<i32>();
+    assert_eq!(*view.get(parent).unwrap(), 0);
+    assert_eq!(*view.get(before).unwrap(), 1);
+    assert_eq!(*view.get(child).unwrap(), 2);
+    assert_eq!(*view.get(after).unwrap(), 3);
 }
 
 #[test]
 fn deletion() {
-    let mut tree = Tree::new(0);
+    let mut tree = Tree::new();
     let parent = tree.root();
-    let child = tree.create_node(2);
+    tree.insert(parent, 0);
+    let mut child = tree.create_node();
+    child.insert(2);
+    let child = child.id();
     tree.add_child(parent, child);
-    let before = tree.create_node(1);
+    let mut before = tree.create_node();
+    before.insert(1);
+    let before = before.id();
     tree.insert_before(child, before);
-    let after = tree.create_node(3);
+    let mut after = tree.create_node();
+    after.insert(3);
+    let after = after.id();
     tree.insert_after(child, after);
 
     println!("Tree: {:#?}", tree);
@@ -525,14 +608,17 @@ fn deletion() {
     assert_eq!(tree.height(child), Some(1));
     assert_eq!(tree.height(before), Some(1));
     assert_eq!(tree.height(after), Some(1));
-    assert_eq!(*tree.get(parent).unwrap(), 0);
-    assert_eq!(*tree.get(before).unwrap(), 1);
-    assert_eq!(*tree.get(child).unwrap(), 2);
-    assert_eq!(*tree.get(after).unwrap(), 3);
     assert_eq!(tree.parent_id(before).unwrap(), parent);
     assert_eq!(tree.parent_id(child).unwrap(), parent);
     assert_eq!(tree.parent_id(after).unwrap(), parent);
     assert_eq!(tree.children_ids(parent).unwrap(), &[before, child, after]);
+    {
+        let view = tree.state_view::<i32>();
+        assert_eq!(*view.get(parent).unwrap(), 0);
+        assert_eq!(*view.get(before).unwrap(), 1);
+        assert_eq!(*view.get(child).unwrap(), 2);
+        assert_eq!(*view.get(after).unwrap(), 3);
+    }
 
     tree.remove(child);
 
@@ -541,13 +627,16 @@ fn deletion() {
     assert_eq!(tree.height(parent), Some(0));
     assert_eq!(tree.height(before), Some(1));
     assert_eq!(tree.height(after), Some(1));
-    assert_eq!(*tree.get(parent).unwrap(), 0);
-    assert_eq!(*tree.get(before).unwrap(), 1);
-    assert_eq!(tree.get(child), None);
-    assert_eq!(*tree.get(after).unwrap(), 3);
     assert_eq!(tree.parent_id(before).unwrap(), parent);
     assert_eq!(tree.parent_id(after).unwrap(), parent);
     assert_eq!(tree.children_ids(parent).unwrap(), &[before, after]);
+    {
+        let view = tree.state_view::<i32>();
+        assert_eq!(*view.get(parent).unwrap(), 0);
+        assert_eq!(*view.get(before).unwrap(), 1);
+        assert_eq!(view.get(child), None);
+        assert_eq!(*view.get(after).unwrap(), 3);
+    }
 
     tree.remove(before);
 
@@ -555,37 +644,53 @@ fn deletion() {
     assert_eq!(tree.size(), 2);
     assert_eq!(tree.height(parent), Some(0));
     assert_eq!(tree.height(after), Some(1));
-    assert_eq!(*tree.get(parent).unwrap(), 0);
-    assert_eq!(tree.get(before), None);
-    assert_eq!(*tree.get(after).unwrap(), 3);
     assert_eq!(tree.parent_id(after).unwrap(), parent);
     assert_eq!(tree.children_ids(parent).unwrap(), &[after]);
+    {
+        let view = tree.state_view::<i32>();
+        assert_eq!(*view.get(parent).unwrap(), 0);
+        assert_eq!(view.get(before), None);
+        assert_eq!(*view.get(after).unwrap(), 3);
+    }
 
     tree.remove(after);
 
     println!("Tree: {:#?}", tree);
     assert_eq!(tree.size(), 1);
     assert_eq!(tree.height(parent), Some(0));
-    assert_eq!(*tree.get(parent).unwrap(), 0);
-    assert_eq!(tree.get(after), None);
     assert_eq!(tree.children_ids(parent).unwrap(), &[]);
+    {
+        let view = tree.state_view::<i32>();
+        assert_eq!(*view.get(parent).unwrap(), 0);
+        assert_eq!(view.get(after), None);
+    }
 }
 
 #[test]
 fn traverse_depth_first() {
-    let mut tree = Tree::new(0);
+    let mut tree = Tree::new();
     let parent = tree.root();
-    let child1 = tree.create_node(1);
+    tree.insert(parent, 0);
+    let mut child1 = tree.create_node();
+    child1.insert(1);
+    let child1 = child1.id();
     tree.add_child(parent, child1);
-    let grandchild1 = tree.create_node(2);
+    let mut grandchild1 = tree.create_node();
+    grandchild1.insert(2);
+    let grandchild1 = grandchild1.id();
     tree.add_child(child1, grandchild1);
-    let child2 = tree.create_node(3);
+    let mut child2 = tree.create_node();
+    child2.insert(3);
+    let child2 = child2.id();
     tree.add_child(parent, child2);
-    let grandchild2 = tree.create_node(4);
+    let mut grandchild2 = tree.create_node();
+    grandchild2.insert(4);
+    let grandchild2 = grandchild2.id();
     tree.add_child(child2, grandchild2);
 
+    let view = tree.state_view::<i32>();
     let mut node_count = 0;
-    tree.traverse_depth_first(move |node| {
+    view.traverse_depth_first(move |node| {
         assert_eq!(*node, node_count);
         node_count += 1;
     });
@@ -593,60 +698,55 @@ fn traverse_depth_first() {
 
 #[test]
 fn get_node_children_mut() {
-    let mut tree = Tree::new(0);
+    let mut tree = Tree::new();
     let parent = tree.root();
-    let child1 = tree.create_node(1);
+    tree.insert(parent, 0);
+    let mut child1 = tree.create_node();
+    child1.insert(1);
+    let child1 = child1.id();
     tree.add_child(parent, child1);
-    let child2 = tree.create_node(2);
+    let mut child2 = tree.create_node();
+    child2.insert(2);
+    let child2 = child2.id();
     tree.add_child(parent, child2);
-    let child3 = tree.create_node(3);
+    let mut child3 = tree.create_node();
+    child3.insert(3);
+    let child3 = child3.id();
     tree.add_child(parent, child3);
 
-    let (parent, children) = tree.parent_child_mut(parent).unwrap();
-    for (i, child) in children.enumerate() {
-        assert_eq!(*child, i + 1);
-    }
+    let mut view = tree.state_view_mut::<i32>();
+    let (parent, children) = view.parent_child_mut(parent).unwrap();
     println!("Parent: {:#?}", parent);
-}
-
-#[test]
-fn get_many_mut_unchecked() {
-    let mut slab = Slab::new();
-    let parent = slab.insert(0);
-    let child = slab.insert(1);
-    let grandchild = slab.insert(2);
-
-    let all =
-        unsafe { slab.get_many_mut_unchecked([parent, child, grandchild].into_iter()) }.unwrap();
-    println!("All: {:#?}", all);
+    println!("Children: {:#?}", children);
+    for (i, child) in children.into_iter().enumerate() {
+        assert_eq!(*child, i as i32 + 1);
+    }
 }
 
 #[derive(Debug)]
-struct Slab<T> {
+struct SlabStorage<T> {
     data: Vec<Option<T>>,
-    free: VecDeque<usize>,
 }
 
-impl<T> Default for Slab<T> {
+impl<T> Default for SlabStorage<T> {
     fn default() -> Self {
-        Self::new()
+        Self { data: Vec::new() }
     }
 }
 
-impl<T> Slab<T> {
-    fn new() -> Self {
-        Self {
-            data: Vec::new(),
-            free: VecDeque::new(),
-        }
-    }
+trait AnySlabStorageImpl: Any {
+    fn remove(&mut self, id: usize);
+    fn as_any(&self) -> &dyn Any;
+    fn as_any_mut(&mut self) -> &mut dyn Any;
+}
 
+impl<T> SlabStorage<T> {
     fn get(&self, id: usize) -> Option<&T> {
         self.data.get(id).and_then(|x| x.as_ref())
     }
 
     unsafe fn get_unchecked(&self, id: usize) -> &T {
-        self.data.get_unchecked(id).as_ref().unwrap()
+        self.data.get_unchecked(id).as_ref().unwrap_unchecked()
     }
 
     fn get_mut(&mut self, id: usize) -> Option<&mut T> {
@@ -654,7 +754,14 @@ impl<T> Slab<T> {
     }
 
     unsafe fn get_unchecked_mut(&mut self, id: usize) -> &mut T {
-        self.data.get_unchecked_mut(id).as_mut().unwrap()
+        self.data.get_unchecked_mut(id).as_mut().unwrap_unchecked()
+    }
+
+    fn insert(&mut self, id: usize, value: T) {
+        if id >= self.data.len() {
+            self.data.resize_with(id + 1, || None);
+        }
+        self.data[id] = Some(value);
     }
 
     fn get2_mut(&mut self, id1: usize, id2: usize) -> Option<(&mut T, &mut T)> {
@@ -676,6 +783,9 @@ impl<T> Slab<T> {
         let ptr = self.data.as_mut_ptr();
         let mut result = Vec::new();
         for id in ids {
+            if id >= self.data.len() {
+                return None;
+            }
             let item = unsafe { &mut *ptr.add(id) };
             if let Some(item) = item {
                 result.push(item);
@@ -685,29 +795,222 @@ impl<T> Slab<T> {
         }
         Some(result)
     }
+}
+
+impl<T: 'static> AnySlabStorageImpl for SlabStorage<T> {
+    fn remove(&mut self, id: usize) {
+        self.data[id].take();
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        self
+    }
+}
+
+struct AnySlab {
+    data: hashbrown::HashMap<
+        TypeId,
+        RwLock<Box<dyn AnySlabStorageImpl>>,
+        BuildHasherDefault<FxHasher>,
+    >,
+    free: VecDeque<usize>,
+    len: usize,
+}
+
+impl Debug for AnySlab {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("AnySlab")
+            .field("data", &self.data.keys().collect::<Vec<_>>())
+            .field("free", &self.free)
+            .field("len", &self.len)
+            .finish()
+    }
+}
+
+impl Default for AnySlab {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl AnySlab {
+    fn new() -> Self {
+        Self {
+            data: Default::default(),
+            free: VecDeque::new(),
+            len: 0,
+        }
+    }
+
+    fn try_get_slab<T: Any>(&self) -> Option<MappedRwLockReadGuard<'_, SlabStorage<T>>> {
+        self.data
+            .get(&TypeId::of::<T>())
+            .map(|x| RwLockReadGuard::map(x.read(), |x| x.as_any().downcast_ref().unwrap()))
+    }
+
+    fn get_slab<T: Any>(&self) -> MappedRwLockReadGuard<'_, SlabStorage<T>> {
+        self.try_get_slab().unwrap()
+    }
+
+    fn try_get_slab_mut<T: Any>(&self) -> Option<MappedRwLockWriteGuard<'_, SlabStorage<T>>> {
+        self.data
+            .get(&TypeId::of::<T>())
+            .map(|x| RwLockWriteGuard::map(x.write(), |x| x.as_any_mut().downcast_mut().unwrap()))
+    }
 
-    fn insert(&mut self, value: T) -> usize {
-        if let Some(id) = self.free.pop_front() {
-            self.data[id] = Some(value);
+    fn get_slab_mut<T: Any>(&self) -> MappedRwLockWriteGuard<'_, SlabStorage<T>> {
+        self.try_get_slab_mut().unwrap()
+    }
+
+    fn get_or_insert_slab_mut<T: Any>(&mut self) -> MappedRwLockWriteGuard<'_, SlabStorage<T>> {
+        RwLockWriteGuard::map(
+            self.data
+                .entry(TypeId::of::<T>())
+                .or_insert_with(|| RwLock::new(Box::new(SlabStorage::<T>::default())))
+                .write(),
+            |x| x.as_any_mut().downcast_mut().unwrap(),
+        )
+    }
+
+    fn insert(&mut self) -> Entry<'_> {
+        let id = if let Some(id) = self.free.pop_front() {
             id
         } else {
-            self.data.push(Some(value));
-            self.data.len() - 1
-        }
+            let id = self.len;
+            self.len += 1;
+            id
+        };
+        Entry { id, inner: self }
     }
 
-    fn try_remove(&mut self, id: usize) -> Option<T> {
-        self.data.get_mut(id).and_then(|x| {
-            self.free.push_back(id);
-            x.take()
-        })
+    fn add<T: Any>(&mut self, id: usize, value: T) {
+        self.get_or_insert_slab_mut().insert(id, value);
     }
 
-    fn remove(&mut self, id: usize) -> T {
-        self.try_remove(id).unwrap()
+    fn remove(&mut self, id: usize) {
+        for slab in self.data.values_mut() {
+            slab.write().remove(id);
+        }
+        self.free.push_back(id);
     }
 
     fn len(&self) -> usize {
-        self.data.len() - self.free.len()
+        self.len - self.free.len()
     }
 }
+
+pub struct Entry<'a> {
+    id: usize,
+    inner: &'a mut AnySlab,
+}
+
+impl Entry<'_> {
+    pub fn insert<T: Any>(&mut self, value: T) {
+        self.inner
+            .get_or_insert_slab_mut::<T>()
+            .insert(self.id, value);
+    }
+
+    pub fn remove(self) {
+        self.inner.remove(self.id);
+    }
+
+    pub fn id(&self) -> usize {
+        self.id
+    }
+}
+
+#[test]
+fn remove() {
+    let mut slab = AnySlab::new();
+    let mut node1 = slab.insert();
+    node1.insert(0i32);
+    let node1_id = node1.id();
+    let mut node2 = slab.insert();
+    node2.insert(0i32);
+
+    assert_eq!(slab.len(), 2);
+
+    slab.remove(node1_id);
+
+    assert!(slab.get_slab::<i32>().get(node1_id).is_none());
+    assert_eq!(slab.len(), 1);
+}
+
+#[test]
+fn get_many_mut_unchecked() {
+    let mut slab = AnySlab::new();
+    let mut parent = slab.insert();
+    parent.insert(0i32);
+    let parent = parent.id();
+    let mut child = slab.insert();
+    child.insert(1i32);
+    let child = child.id();
+    let mut grandchild = slab.insert();
+    grandchild.insert(2i32);
+    let grandchild = grandchild.id();
+    assert_eq!(slab.len(), 3);
+    println!("ids: {:#?}", [parent, child, grandchild]);
+
+    {
+        let mut i32_slab = slab.get_slab_mut::<i32>();
+        let all =
+            unsafe { i32_slab.get_many_mut_unchecked([parent, child, grandchild].into_iter()) }
+                .unwrap();
+        assert_eq!(all, [&mut 0, &mut 1, &mut 2]);
+    }
+
+    {
+        let mut i32_slab = slab.get_slab_mut::<i32>();
+        assert!(unsafe { i32_slab.get_many_mut_unchecked([3, 100].into_iter()) }.is_none())
+    }
+}
+
+#[test]
+fn get_many_many_mut_unchecked() {
+    let mut slab = AnySlab::new();
+    let mut parent = slab.insert();
+    parent.insert(0i32);
+    parent.insert("0");
+    let parent = parent.id();
+    let mut child = slab.insert();
+    child.insert(1i32);
+    child.insert("1");
+    let child = child.id();
+    let mut grandchild = slab.insert();
+    grandchild.insert(2i32);
+    grandchild.insert("2");
+    let grandchild = grandchild.id();
+    println!("ids: {:#?}", [parent, child, grandchild]);
+
+    println!("slab: {:#?}", slab);
+
+    let mut num_entries = slab.get_slab_mut::<i32>();
+    let mut str_entries = slab.get_slab_mut::<&str>();
+
+    let all_num = unsafe {
+        num_entries
+            .as_any_mut()
+            .downcast_mut::<SlabStorage<i32>>()
+            .unwrap()
+            .get_many_mut_unchecked([parent, child, grandchild].into_iter())
+    }
+    .unwrap();
+
+    assert_eq!(all_num, [&mut 0, &mut 1, &mut 2]);
+
+    let all_str = unsafe {
+        str_entries
+            .as_any_mut()
+            .downcast_mut::<SlabStorage<&str>>()
+            .unwrap()
+            .get_many_mut_unchecked([parent, child, grandchild].into_iter())
+    }
+    .unwrap();
+
+    assert_eq!(all_str, [&mut "0", &mut "1", &mut "2"]);
+}