#![forbid(unsafe_code)]
extern crate generational_arena;
use generational_arena::Arena;
pub use generational_arena::Index;
use core::ops;
use std::{fmt, mem};
#[derive(Debug)]
pub struct VecTree<T> {
nodes: Arena<Node<T>>,
root_index: Option<Index>,
}
#[derive(Clone, Debug)]
struct Node<T> {
parent: Option<Index>,
previous_sibling: Option<Index>,
next_sibling: Option<Index>,
first_child: Option<Index>,
last_child: Option<Index>,
data: T,
}
const DEFAULT_CAPACITY: usize = 4;
impl<T> Default for VecTree<T> {
fn default() -> Self {
VecTree::with_capacity(DEFAULT_CAPACITY)
}
}
impl<T> VecTree<T> {
pub fn new() -> VecTree<T> {
VecTree::with_capacity(DEFAULT_CAPACITY)
}
pub fn with_capacity(n: usize) -> VecTree<T> {
VecTree {
nodes: Arena::with_capacity(n),
root_index: None,
}
}
#[inline]
pub fn reserve(&mut self, additional_capacity: usize) {
self.nodes.reserve(additional_capacity);
}
#[inline]
pub fn try_insert(&mut self, data: T, parent_id: Index) -> Result<Index, T> {
let node_result = self.try_create_node(data);
match node_result {
Ok(node) => {
self.append_child(parent_id, node);
node_result
}
Err(err) => Err(err),
}
}
#[inline]
pub fn insert(&mut self, data: T, parent_id: Index) -> Index {
let node = self.create_node(data);
self.append_child(parent_id, node);
node
}
#[inline]
pub fn try_insert_root(&mut self, data: T) -> Result<Index, T> {
if self.root_index.is_some() {
panic!("A root node already exists");
}
match self.try_create_node(data) {
Ok(node_id) => {
self.root_index = Some(node_id);
Ok(node_id)
}
Err(error) => Err(error),
}
}
#[inline]
pub fn insert_root(&mut self, data: T) -> Index {
if self.root_index.is_some() {
panic!("A root node already exists");
}
let node_id = self.create_node(data);
self.root_index = Some(node_id);
node_id
}
#[inline]
fn try_create_node(&mut self, data: T) -> Result<Index, T> {
let new_node = Node {
parent: None,
first_child: None,
last_child: None,
previous_sibling: None,
next_sibling: None,
data,
};
match self.nodes.try_insert(new_node) {
Ok(index) => Ok(index),
Err(Node { data, .. }) => Err(data),
}
}
#[inline]
fn create_node(&mut self, data: T) -> Index {
let new_node = Node {
parent: None,
first_child: None,
last_child: None,
previous_sibling: None,
next_sibling: None,
data,
};
self.nodes.insert(new_node)
}
pub fn remove(&mut self, node_id: Index) -> Option<T> {
if !self.contains(node_id) {
return None;
}
let descendants = self.descendants(node_id).skip(1).collect::<Vec<Index>>();
let node = self.nodes.remove(node_id).unwrap();
let previous_sibling_opt = node.previous_sibling;
let next_sibling_opt = node.next_sibling;
if let Some(previous_sibling_idx) = previous_sibling_opt {
if let Some(next_sibling_idx) = next_sibling_opt {
let (previous_sibling, next_sibling) =
self.nodes.get2_mut(previous_sibling_idx, next_sibling_idx);
previous_sibling.unwrap().next_sibling = Some(next_sibling_idx);
next_sibling.unwrap().previous_sibling = Some(previous_sibling_idx);
} else if let Some(parent_idx) = node.parent {
let previous_sibling = &mut self.nodes[previous_sibling_idx];
previous_sibling.next_sibling = None;
let parent = &mut self.nodes[parent_idx];
parent.last_child = Some(previous_sibling_idx);
}
} else if let Some(next_sibling_idx) = next_sibling_opt {
let next_sibling = &mut self.nodes[next_sibling_idx];
next_sibling.previous_sibling = None;
if let Some(parent_idx) = node.parent {
let parent = &mut self.nodes[parent_idx];
parent.first_child = Some(next_sibling_idx);
}
} else if let Some(parent_idx) = node.parent {
let parent = &mut self.nodes[parent_idx];
parent.first_child = None;
parent.last_child = None;
}
for node_id in descendants {
self.nodes.remove(node_id);
}
if let Some(root_index) = self.root_index {
if root_index == node_id {
self.root_index = None;
}
}
Some(node.data)
}
pub fn contains(&self, i: Index) -> bool {
self.nodes.get(i).is_some()
}
#[inline]
pub fn append_child(&mut self, node_id: Index, new_child_id: Index) {
self.detach(new_child_id);
let last_child_opt;
{
let (node_opt, new_child_node_opt) = self.nodes.get2_mut(node_id, new_child_id);
if node_opt.is_none() {
panic!("The node you are trying to append to is invalid");
}
if new_child_node_opt.is_none() {
panic!("The node you are trying to append is invalid");
}
let node = node_opt.unwrap();
let new_child_node = new_child_node_opt.unwrap();
new_child_node.parent = Some(node_id);
last_child_opt = mem::replace(&mut node.last_child, Some(new_child_id));
if let Some(last_child) = last_child_opt {
new_child_node.previous_sibling = Some(last_child);
} else {
debug_assert!(node.first_child.is_none());
node.first_child = Some(new_child_id);
}
}
if let Some(last_child) = last_child_opt {
debug_assert!(self.nodes[last_child].next_sibling.is_none());
self.nodes[last_child].next_sibling = Some(new_child_id);
}
}
#[inline]
pub fn detach(&mut self, node_id: Index) {
let (parent, previous_sibling, next_sibling) = {
let node = &mut self.nodes[node_id];
(
node.parent.take(),
node.previous_sibling.take(),
node.next_sibling.take(),
)
};
if let Some(next_sibling) = next_sibling {
self.nodes[next_sibling].previous_sibling = previous_sibling;
} else if let Some(parent) = parent {
self.nodes[parent].last_child = previous_sibling;
}
if let Some(previous_sibling) = previous_sibling {
self.nodes[previous_sibling].next_sibling = next_sibling;
} else if let Some(parent) = parent {
self.nodes[parent].first_child = next_sibling;
}
}
pub fn get(&self, node_id: Index) -> Option<&T> {
match self.nodes.get(node_id) {
Some(Node { ref data, .. }) => Some(data),
_ => None,
}
}
pub fn get_mut(&mut self, node_id: Index) -> Option<&mut T> {
match self.nodes.get_mut(node_id) {
Some(Node { ref mut data, .. }) => Some(data),
_ => None,
}
}
pub fn get_root_index(&self) -> Option<Index> {
self.root_index
}
pub fn capacity(&self) -> usize {
self.nodes.capacity()
}
pub fn clear(&mut self) {
self.nodes.clear();
self.root_index = None;
}
pub fn parent(&self, node_id: Index) -> Option<Index> {
match self.nodes.get(node_id) {
Some(node) => node.parent,
_ => None,
}
}
pub fn children(&self, node_id: Index) -> ChildrenIter<T> {
ChildrenIter {
tree: self,
node_id: self.nodes[node_id].first_child,
}
}
pub fn preceding_siblings(&self, node_id: Index) -> PrecedingSiblingsIter<T> {
PrecedingSiblingsIter {
tree: self,
node_id: Some(node_id),
}
}
pub fn following_siblings(&self, node_id: Index) -> FollowingSiblingsIter<T> {
FollowingSiblingsIter {
tree: self,
node_id: Some(node_id),
}
}
pub fn ancestors(&self, node_id: Index) -> AncestorsIter<T> {
AncestorsIter {
tree: self,
node_id: Some(node_id),
}
}
fn traverse(&self, node_id: Index) -> TraverseIter<T> {
TraverseIter {
tree: self,
root: node_id,
next: Some(NodeEdge::Start(node_id)),
}
}
fn traverse_with_depth(&self, node_id: Index) -> TraverseWithDepthIter<T> {
TraverseWithDepthIter {
tree: self,
root: node_id,
next: Some(NodeEdgeWithDepth::Start(node_id, 0)),
}
}
pub fn descendants(&self, node_id: Index) -> DescendantsIter<T> {
DescendantsIter(self.traverse(node_id))
}
pub fn descendants_with_depth(&self, node_id: Index) -> DescendantsWithDepthIter<T> {
DescendantsWithDepthIter(self.traverse_with_depth(node_id))
}
}
impl<T> fmt::Display for Node<T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "Parent: {:?}, ", self.parent)?;
write!(f, "Previous sibling: {:?}, ", self.previous_sibling)?;
write!(f, "Next sibling: {:?}, ", self.next_sibling)?;
write!(f, "First child: {:?}, ", self.first_child)?;
write!(f, "Last child: {:?}", self.last_child)
}
}
impl<T> ops::Index<Index> for VecTree<T> {
type Output = T;
fn index(&self, index: Index) -> &Self::Output {
self.get(index).unwrap()
}
}
impl<T> ops::IndexMut<Index> for VecTree<T> {
fn index_mut(&mut self, index: Index) -> &mut Self::Output {
self.get_mut(index).unwrap()
}
}
macro_rules! impl_node_iterator {
($name:ident, $next:expr) => {
impl<'a, T> Iterator for $name<'a, T> {
type Item = Index;
fn next(&mut self) -> Option<Index> {
match self.node_id.take() {
Some(node_id) => {
self.node_id = $next(&self.tree.nodes[node_id]);
Some(node_id)
}
None => None,
}
}
}
};
}
pub struct ChildrenIter<'a, T: 'a> {
tree: &'a VecTree<T>,
node_id: Option<Index>,
}
impl_node_iterator!(ChildrenIter, |node: &Node<T>| node.next_sibling);
pub struct PrecedingSiblingsIter<'a, T: 'a> {
tree: &'a VecTree<T>,
node_id: Option<Index>,
}
impl_node_iterator!(PrecedingSiblingsIter, |node: &Node<T>| node
.previous_sibling);
pub struct FollowingSiblingsIter<'a, T: 'a> {
tree: &'a VecTree<T>,
node_id: Option<Index>,
}
impl_node_iterator!(FollowingSiblingsIter, |node: &Node<T>| node.next_sibling);
pub struct AncestorsIter<'a, T: 'a> {
tree: &'a VecTree<T>,
node_id: Option<Index>,
}
impl_node_iterator!(AncestorsIter, |node: &Node<T>| node.parent);
#[derive(Debug, Clone)]
pub enum NodeEdge<T> {
Start(T),
End(T),
}
pub struct TraverseIter<'a, T: 'a> {
tree: &'a VecTree<T>,
root: Index,
next: Option<NodeEdge<Index>>,
}
impl<'a, T> Iterator for TraverseIter<'a, T> {
type Item = NodeEdge<Index>;
fn next(&mut self) -> Option<NodeEdge<Index>> {
match self.next.take() {
Some(item) => {
self.next = match item {
NodeEdge::Start(node_id) => match self.tree.nodes[node_id].first_child {
Some(first_child) => Some(NodeEdge::Start(first_child)),
None => Some(NodeEdge::End(node_id)),
},
NodeEdge::End(node_id) => {
if node_id == self.root {
None
} else {
match self.tree.nodes[node_id].next_sibling {
Some(next_sibling) => Some(NodeEdge::Start(next_sibling)),
None => {
match self.tree.nodes[node_id].parent {
Some(parent) => Some(NodeEdge::End(parent)),
None => None,
}
}
}
}
}
};
Some(item)
}
None => None,
}
}
}
pub struct DescendantsIter<'a, T: 'a>(pub TraverseIter<'a, T>);
impl<'a, T> Iterator for DescendantsIter<'a, T> {
type Item = Index;
fn next(&mut self) -> Option<Index> {
loop {
match self.0.next() {
Some(NodeEdge::Start(node_id)) => return Some(node_id),
Some(NodeEdge::End(_)) => {}
None => return None,
}
}
}
}
#[derive(Debug, Clone)]
pub enum NodeEdgeWithDepth<T> {
Start(T, u32),
End(T, u32),
}
pub struct TraverseWithDepthIter<'a, T: 'a> {
tree: &'a VecTree<T>,
root: Index,
next: Option<NodeEdgeWithDepth<Index>>,
}
impl<'a, T> Iterator for TraverseWithDepthIter<'a, T> {
type Item = NodeEdgeWithDepth<Index>;
fn next(&mut self) -> Option<NodeEdgeWithDepth<Index>> {
match self.next.take() {
Some(item) => {
self.next = match item {
NodeEdgeWithDepth::Start(node_id, depth) => {
match self.tree.nodes[node_id].first_child {
Some(first_child) => {
Some(NodeEdgeWithDepth::Start(first_child, depth + 1))
}
None => Some(NodeEdgeWithDepth::End(node_id, depth)),
}
}
NodeEdgeWithDepth::End(node_id, depth) => {
if node_id == self.root {
None
} else {
match self.tree.nodes[node_id].next_sibling {
Some(next_sibling) => {
Some(NodeEdgeWithDepth::Start(next_sibling, depth))
}
None => {
match self.tree.nodes[node_id].parent {
Some(parent) => {
Some(NodeEdgeWithDepth::End(parent, depth - 1))
}
None => None,
}
}
}
}
}
};
Some(item)
}
None => None,
}
}
}
pub struct DescendantsWithDepthIter<'a, T: 'a>(pub TraverseWithDepthIter<'a, T>);
impl<'a, T> Iterator for DescendantsWithDepthIter<'a, T> {
type Item = (Index, u32);
fn next(&mut self) -> Option<(Index, u32)> {
loop {
match self.0.next() {
Some(NodeEdgeWithDepth::Start(node_id, depth)) => return Some((node_id, depth)),
Some(NodeEdgeWithDepth::End(_, _)) => {}
None => return None,
}
}
}
}