use crate::node::Node;
use crate::slab::{self, Slab};
use crate::NodeId;
use snowflake::ProcessUniqueId;
#[derive(Debug)]
pub(crate) struct CoreTree<T> {
id: ProcessUniqueId,
slab: Slab<Node<T>>,
}
impl<T> CoreTree<T> {
pub(crate) fn new(capacity: usize) -> CoreTree<T> {
CoreTree {
id: ProcessUniqueId::new(),
slab: Slab::new(capacity),
}
}
pub(crate) fn capacity(&self) -> usize {
self.slab.capacity()
}
pub(crate) fn insert(&mut self, data: T) -> NodeId {
let key = self.slab.insert(Node::new(data));
self.new_node_id(key)
}
pub(crate) fn remove(&mut self, node_id: NodeId) -> Option<T> {
self.filter_by_tree_id(node_id)
.and_then(|id| self.slab.try_remove(id.index))
.map(|node| node.data)
}
pub(crate) fn get(&self, node_id: NodeId) -> Option<&Node<T>> {
self.filter_by_tree_id(node_id)
.and_then(|id| self.slab.get(id.index))
}
pub(crate) fn get_mut(&mut self, node_id: NodeId) -> Option<&mut Node<T>> {
self.filter_by_tree_id(node_id)
.and_then(move |id| self.slab.get_mut(id.index))
}
#[cfg(feature = "experimental")]
pub(crate) fn compact(&mut self) -> usize {
use std::collections::HashMap;
let mut rekey_tuples = Vec::new();
self.slab.compact(|from, to| {
rekey_tuples.push((from, to));
});
let mut rekeys = HashMap::new();
for (from, to) in rekey_tuples.into_iter() {
rekeys.insert(from, self.new_node_id(to));
}
for (from, to) in &rekeys {
self.fix_rekeyed_node(*from, *to, &rekeys);
}
self.capacity()
}
#[cfg(feature = "experimental")]
fn fix_rekeyed_node(
&mut self,
from: usize,
to: NodeId,
rekeys: &std::collections::HashMap<usize, NodeId>,
) {
let relatives = { self.get(to).unwrap().relatives };
if let Some(parent) = relatives.parent {
let parent = *rekeys.get(&parent.index.index).unwrap_or(&parent);
let parent = self.get_mut(parent).unwrap();
if parent.relatives.first_child.unwrap().index.index == from {
parent.relatives.first_child = Some(to);
}
if parent.relatives.last_child.unwrap().index.index == from {
parent.relatives.last_child = Some(to);
}
}
if let Some(prev) = relatives.prev_sibling {
let prev = *rekeys.get(&prev.index.index).unwrap_or(&prev);
let prev = self.get_mut(prev).unwrap();
if prev.relatives.next_sibling.unwrap().index.index == from {
prev.relatives.next_sibling = Some(to);
}
}
if let Some(next) = relatives.next_sibling {
let next = *rekeys.get(&next.index.index).unwrap_or(&next);
let next = self.get_mut(next).unwrap();
if next.relatives.prev_sibling.unwrap().index.index == from {
next.relatives.prev_sibling = Some(to);
}
}
let mut next_child_id = relatives.first_child;
while let Some(child_id) = next_child_id {
let child = self.get_mut(child_id).unwrap();
child.relatives.parent = Some(to);
next_child_id = child.relatives.next_sibling;
}
}
fn new_node_id(&self, index: slab::Index) -> NodeId {
NodeId {
tree_id: self.id,
index,
}
}
fn filter_by_tree_id(&self, node_id: NodeId) -> Option<NodeId> {
if node_id.tree_id != self.id {
return None;
}
Some(node_id)
}
pub(crate) fn shrink_to_fit(&mut self) {
self.slab.shrink_to_fit();
}
}
#[cfg_attr(tarpaulin, skip)]
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn capacity() {
let capacity = 5;
let tree = CoreTree::<i32>::new(capacity);
assert_eq!(tree.capacity(), capacity);
}
#[test]
fn insert() {
let mut tree = CoreTree::new(0);
let id = tree.insert(1);
let id2 = tree.insert(3);
assert_eq!(tree.get(id).unwrap().data, 1);
assert_eq!(tree.get(id2).unwrap().data, 3);
}
#[test]
fn remove() {
let mut tree = CoreTree::new(0);
let id = tree.insert(1);
assert_eq!(tree.get(id).unwrap().data, 1);
let one = tree.remove(id);
assert!(one.is_some());
let one = one.unwrap();
assert_eq!(one, 1);
}
#[test]
fn get() {
let mut tree = CoreTree::new(0);
let id = tree.insert(1);
let id2 = tree.insert(3);
assert_eq!(tree.get(id).unwrap().data, 1);
assert_eq!(tree.get(id2).unwrap().data, 3);
}
#[test]
fn get_mut() {
let mut tree = CoreTree::new(0);
let id = tree.insert(1);
let id2 = tree.insert(3);
assert_eq!(tree.get_mut(id).unwrap().data, 1);
assert_eq!(tree.get_mut(id2).unwrap().data, 3);
}
#[test]
fn get_with_bad_id() {
let mut tree = CoreTree::new(0);
let tree2: CoreTree<i32> = CoreTree::new(0);
let mut id = tree.insert(1);
id.tree_id = tree2.id; let result = tree.get(id);
assert!(result.is_none());
}
}