use crate::behaviors::*;
use crate::core_tree::CoreTree;
use crate::node::*;
use crate::NodeId;
pub struct TreeBuilder<T> {
root: Option<T>,
capacity: Option<usize>,
}
impl<T> Default for TreeBuilder<T> {
fn default() -> Self {
TreeBuilder::new()
}
}
impl<T> TreeBuilder<T> {
pub fn new() -> TreeBuilder<T> {
TreeBuilder {
root: None,
capacity: None,
}
}
pub fn with_root(self, root: T) -> TreeBuilder<T> {
TreeBuilder {
root: Some(root),
capacity: self.capacity,
}
}
pub fn with_capacity(self, capacity: usize) -> TreeBuilder<T> {
TreeBuilder {
root: self.root,
capacity: Some(capacity),
}
}
pub fn build(self) -> Tree<T> {
let capacity = self.capacity.unwrap_or(0);
let mut core_tree: CoreTree<T> = CoreTree::new(capacity);
let root_id = self.root.map(|val| core_tree.insert(val));
Tree { root_id, core_tree }
}
}
#[derive(Debug, PartialEq)]
pub struct Tree<T> {
pub(crate) root_id: Option<NodeId>,
pub(crate) core_tree: CoreTree<T>,
}
impl<T> Tree<T> {
pub fn new() -> Tree<T> {
TreeBuilder::new().build()
}
pub fn set_root(&mut self, root: T) -> NodeId {
let old_root_id = self.root_id.take();
let new_root_id = self.core_tree.insert(root);
self.root_id = Some(new_root_id);
self.set_first_child(new_root_id, old_root_id);
if let Some(node_id) = old_root_id {
self.set_parent(node_id, self.root_id);
}
new_root_id
}
pub fn capacity(&self) -> usize {
self.core_tree.capacity()
}
pub fn root_id(&self) -> Option<NodeId> {
self.root_id
}
pub fn root(&self) -> Option<NodeRef<T>> {
self.root_id.map(|id| self.new_node_ref(id))
}
pub fn root_mut(&mut self) -> Option<NodeMut<T>> {
self.root_id.map(move |id| self.new_node_mut(id))
}
pub fn get(&self, node_id: NodeId) -> Option<NodeRef<T>> {
let _ = self.core_tree.get(node_id)?;
Some(self.new_node_ref(node_id))
}
pub fn get_mut(&mut self, node_id: NodeId) -> Option<NodeMut<T>> {
let _ = self.core_tree.get_mut(node_id)?;
Some(self.new_node_mut(node_id))
}
pub fn remove(&mut self, node_id: NodeId, behavior: RemoveBehavior) -> Option<T> {
if let Some(node) = self.get_node(node_id) {
let Relatives {
parent,
prev_sibling,
next_sibling,
..
} = node.relatives;
let (is_first_child, is_last_child) = self.is_node_first_last_child(node_id);
if is_first_child {
self.set_first_child(parent.expect("parent must exist"), next_sibling);
}
if is_last_child {
self.set_last_child(parent.expect("parent must exist"), prev_sibling);
}
if let Some(prev) = prev_sibling {
self.set_next_sibling(prev, next_sibling);
}
if let Some(next) = next_sibling {
self.set_prev_sibling(next, prev_sibling);
}
match behavior {
RemoveBehavior::DropChildren => self.drop_children(node_id),
RemoveBehavior::OrphanChildren => self.orphan_children(node_id),
};
if self.root_id == Some(node_id) {
self.root_id = None;
}
self.core_tree.remove(node_id)
} else {
None
}
}
pub(crate) fn get_node(&self, node_id: NodeId) -> Option<&Node<T>> {
self.core_tree.get(node_id)
}
pub(crate) fn get_node_mut(&mut self, node_id: NodeId) -> Option<&mut Node<T>> {
self.core_tree.get_mut(node_id)
}
pub(crate) fn set_prev_siblings_next_sibling(
&mut self,
current_id: NodeId,
next_sibling: Option<NodeId>,
) {
if let Some(prev_sibling_id) = self.get_node_prev_sibling_id(current_id) {
self.set_next_sibling(prev_sibling_id, next_sibling);
}
}
pub(crate) fn set_next_siblings_prev_sibling(
&mut self,
current_id: NodeId,
prev_sibling: Option<NodeId>,
) {
if let Some(next_sibling_id) = self.get_node_next_sibling_id(current_id) {
self.set_prev_sibling(next_sibling_id, prev_sibling);
}
}
pub(crate) fn set_parent(&mut self, node_id: NodeId, parent_id: Option<NodeId>) {
if let Some(node) = self.get_node_mut(node_id) {
node.relatives.parent = parent_id;
} else {
unreachable!()
}
}
pub(crate) fn set_prev_sibling(&mut self, node_id: NodeId, prev_sibling: Option<NodeId>) {
if let Some(node) = self.get_node_mut(node_id) {
node.relatives.prev_sibling = prev_sibling;
} else {
unreachable!()
}
}
pub(crate) fn set_next_sibling(&mut self, node_id: NodeId, next_sibling: Option<NodeId>) {
if let Some(node) = self.get_node_mut(node_id) {
node.relatives.next_sibling = next_sibling;
} else {
unreachable!()
}
}
pub(crate) fn set_first_child(&mut self, node_id: NodeId, first_child: Option<NodeId>) {
if let Some(node) = self.get_node_mut(node_id) {
node.relatives.first_child = first_child;
} else {
unreachable!()
}
}
pub(crate) fn set_last_child(&mut self, node_id: NodeId, last_child: Option<NodeId>) {
if let Some(node) = self.get_node_mut(node_id) {
node.relatives.last_child = last_child;
} else {
unreachable!()
}
}
pub(crate) fn get_node_prev_sibling_id(&self, node_id: NodeId) -> Option<NodeId> {
if let Some(node) = self.get_node(node_id) {
node.relatives.prev_sibling
} else {
unreachable!()
}
}
pub(crate) fn get_node_next_sibling_id(&self, node_id: NodeId) -> Option<NodeId> {
if let Some(node) = self.get_node(node_id) {
node.relatives.next_sibling
} else {
unreachable!()
}
}
pub(crate) fn get_node_relatives(&self, node_id: NodeId) -> Relatives {
if let Some(node) = self.get_node(node_id) {
node.relatives
} else {
unreachable!()
}
}
fn drop_children(&mut self, node_id: NodeId) {
let sub_tree_ids: Vec<NodeId> = self
.get(node_id)
.expect("node must exist")
.traverse_level_order()
.skip(1)
.map(|node_ref| node_ref.node_id())
.collect();
for id in sub_tree_ids {
self.core_tree.remove(id);
}
}
fn orphan_children(&mut self, node_id: NodeId) {
let child_ids: Vec<NodeId> = self
.get(node_id)
.expect("node must exist")
.children()
.map(|node_ref| node_ref.node_id())
.collect();
for id in child_ids {
self.set_parent(id, None);
}
}
fn new_node_ref(&self, node_id: NodeId) -> NodeRef<T> {
NodeRef::new(node_id, self)
}
fn new_node_mut(&mut self, node_id: NodeId) -> NodeMut<T> {
NodeMut::new(node_id, self)
}
fn is_node_first_last_child(&self, node_id: NodeId) -> (bool, bool) {
if let Some(node) = self.get_node(node_id) {
node.relatives
.parent
.and_then(|parent_id| self.get_node(parent_id))
.map(|parent| {
let Relatives {
first_child: first,
last_child: last,
..
} = parent.relatives;
(
first.map(|child_id| child_id == node_id).unwrap_or(false),
last.map(|child_id| child_id == node_id).unwrap_or(false),
)
})
.unwrap_or((false, false))
} else {
(false, false)
}
}
}
impl<T> Default for Tree<T> {
fn default() -> Self {
TreeBuilder::new().build()
}
}
impl<T: std::fmt::Debug> Tree<T> {
pub fn write_formatted<W: std::fmt::Write>(&self, w: &mut W) -> std::fmt::Result {
if let Some(root) = self.root() {
let node_id = root.node_id();
let childn = 0;
let level = 0;
let last = vec![];
let mut stack = vec![(node_id, childn, level, last)];
while let Some((node_id, childn, level, last)) = stack.pop() {
debug_assert_eq!(
last.len(),
level,
"each previous level should indicate whether it has reached the last node"
);
let node = self
.get(node_id)
.expect("getting node of existing node ref id");
if childn == 0 {
for i in 1..level {
if last[i - 1] {
write!(w, " ")?;
} else {
write!(w, "│ ")?;
}
}
if level > 0 {
if last[level - 1] {
write!(w, "└── ")?;
} else {
write!(w, "├── ")?;
}
}
writeln!(w, "{:?}", node.data())?;
}
let mut children = node.children().skip(childn);
if let Some(child) = children.next() {
let mut next_last = last.clone();
if children.next().is_some() {
stack.push((node_id, childn + 1, level, last));
next_last.push(false);
} else {
next_last.push(true);
}
stack.push((child.node_id(), 0, level + 1, next_last));
}
}
}
Ok(())
}
}
#[cfg_attr(tarpaulin, skip)]
#[cfg(test)]
mod tree_tests {
use super::*;
use crate::behaviors::RemoveBehavior::{DropChildren, OrphanChildren};
#[test]
fn capacity() {
let tree = TreeBuilder::new().with_root(1).with_capacity(5).build();
assert_eq!(tree.capacity(), 5);
}
#[test]
fn root_id() {
let tree = TreeBuilder::new().with_root(1).build();
let root_id = tree.root_id().expect("root doesn't exist?");
let root = tree.get(root_id).unwrap();
assert_eq!(root.data(), &1);
}
#[test]
fn remove_root_drop() {
let mut tree = TreeBuilder::new().with_root(1).build();
let root_id = tree.root_id().expect("root doesn't exist?");
tree.remove(root_id, RemoveBehavior::DropChildren);
assert!(tree.root().is_none());
}
#[test]
fn remove_root_orphan() {
let mut tree = TreeBuilder::new().with_root(1).build();
let root_id = tree.root_id().expect("root doesn't exist?");
tree.remove(root_id, RemoveBehavior::OrphanChildren);
assert!(tree.root().is_none());
}
#[test]
fn root() {
let tree = TreeBuilder::new().with_root(1).build();
let root = tree.root().expect("root doesn't exist?");
assert_eq!(root.data(), &1);
}
#[test]
fn root_mut() {
let mut tree = TreeBuilder::new().with_root(1).build();
let mut root = tree.root_mut().expect("root doesn't exist?");
assert_eq!(root.data(), &mut 1);
*root.data() = 2;
assert_eq!(root.data(), &mut 2);
}
#[test]
fn get() {
let tree = TreeBuilder::new().with_root(1).build();
let root_id = tree.root_id().expect("root doesn't exist?");
let root = tree.get(root_id);
assert!(root.is_some());
let root = root.unwrap();
assert_eq!(root.data(), &1);
}
#[test]
fn get_mut() {
let mut tree = TreeBuilder::new().with_root(1).build();
let root_id = tree.root_id().expect("root doesn't exist?");
let root = tree.get_mut(root_id);
assert!(root.is_some());
let mut root = root.unwrap();
assert_eq!(root.data(), &mut 1);
*root.data() = 2;
assert_eq!(root.data(), &mut 2);
}
#[test]
fn get_node() {
let tree = TreeBuilder::new().with_root(1).build();
let root_id = tree.root_id().expect("root doesn't exist?");
let root = tree.get_node(root_id);
assert!(root.is_some());
let root = root.unwrap();
assert_eq!(root.data, 1);
}
#[test]
fn get_node_mut() {
let mut tree = TreeBuilder::new().with_root(1).build();
let root_id = tree.root_id().expect("root doesn't exist?");
let root = tree.get_node_mut(root_id);
assert!(root.is_some());
let root = root.unwrap();
assert_eq!(root.data, 1);
root.data = 2;
assert_eq!(root.data, 2);
}
#[test]
fn remove_drop() {
let mut tree = TreeBuilder::new().with_root(1).build();
let two_id;
let three_id;
let four_id;
let five_id;
{
let mut root = tree.root_mut().expect("root doesn't exist?");
two_id = root.append(2).node_id();
three_id = root.append(3).node_id();
four_id = root.append(4).node_id();
}
{
five_id = tree
.get_mut(three_id)
.expect("three doesn't exist?")
.append(5)
.node_id();
}
tree.remove(three_id, DropChildren);
let root = tree
.get_node(tree.root_id().expect("tree doesn't exist?"))
.unwrap();
assert!(root.relatives.first_child.is_some());
assert!(root.relatives.last_child.is_some());
assert_eq!(root.relatives.first_child.unwrap(), two_id);
assert_eq!(root.relatives.last_child.unwrap(), four_id);
let two = tree.get_node(two_id);
assert!(two.is_some());
let two = two.unwrap();
assert_eq!(two.relatives.next_sibling, Some(four_id));
let four = tree.get_node(four_id);
assert!(four.is_some());
let four = four.unwrap();
assert_eq!(four.relatives.prev_sibling, Some(two_id));
let five = tree.get_node(five_id);
assert!(five.is_none());
}
#[test]
fn address_dropped() {
let mut tree = TreeBuilder::new().with_root(1).build();
let two_id = tree.root_mut().expect("root doesn't exist").node_id();
tree.remove(two_id, DropChildren);
tree.remove(two_id, DropChildren);
}
#[test]
fn remove_orphan() {
let mut tree = TreeBuilder::new().with_root(1).build();
let two_id;
let three_id;
let four_id;
let five_id;
{
let mut root = tree.root_mut().expect("root doesn't exist?");
two_id = root.append(2).node_id();
three_id = root.append(3).node_id();
four_id = root.append(4).node_id();
}
{
five_id = tree
.get_mut(three_id)
.expect("three doesn't exist?")
.append(5)
.node_id();
}
tree.remove(three_id, OrphanChildren);
let root = tree
.get_node(tree.root_id().expect("tree doesn't exist?"))
.unwrap();
assert!(root.relatives.first_child.is_some());
assert!(root.relatives.last_child.is_some());
assert_eq!(root.relatives.first_child.unwrap(), two_id);
assert_eq!(root.relatives.last_child.unwrap(), four_id);
let two = tree.get_node(two_id);
assert!(two.is_some());
let two = two.unwrap();
assert_eq!(two.relatives.next_sibling, Some(four_id));
let four = tree.get_node(four_id);
assert!(four.is_some());
let four = four.unwrap();
assert_eq!(four.relatives.prev_sibling, Some(two_id));
let five = tree.get_node(five_id);
assert!(five.is_some());
let five = five.unwrap();
assert_eq!(five.relatives.parent, None);
}
}