use std::error::Error;
use std::fmt;
use std::slice::Iter;
#[derive(Debug, Eq, PartialEq)]
pub struct TreeError {
message: String,
}
impl TreeError {
fn new(message: &str) -> Self {
Self {
message: message.to_string(),
}
}
}
impl fmt::Display for TreeError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.to_string())
}
}
impl Error for TreeError {
fn cause(&self) -> Option<&dyn Error> {
None
}
}
type Result<T> = std::result::Result<T, TreeError>;
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub struct NodeRef {
id: usize,
}
#[derive(Debug, Clone)]
struct Node<T> {
content: T,
parent: Option<NodeRef>,
children: Vec<NodeRef>,
}
#[derive(Debug, Clone)]
pub struct Tree<T> {
nodes: Vec<Option<Node<T>>>,
root: Option<NodeRef>,
len: usize,
}
impl<T> Tree<T> {
pub fn new() -> Self {
Self {
nodes: Vec::new(),
root: None,
len: 0,
}
}
pub fn root(&mut self, content: T) -> Result<NodeRef> {
if self.root.is_some() {
return Err(TreeError::new("Another root node already exists."));
}
let node_ref = self.node(content);
self.root = Some(node_ref);
Ok(node_ref)
}
pub fn node(&mut self, content: T) -> NodeRef {
let id = self.nodes.len();
self.nodes.push(Some(Node {
content,
parent: None,
children: Vec::new(),
}));
self.len += 1;
NodeRef { id }
}
pub fn child_node(&mut self, parent: NodeRef, content: T) -> Result<NodeRef> {
let child = self.node(content);
self.append_child(parent, child)?;
Ok(child)
}
pub fn remove(&mut self, node_ref: NodeRef) -> Result<()> {
match self.nodes.get(node_ref.id) {
None => return Err(TreeError::new("Invalid node reference.")),
Some(node) => match node {
None => return Err(TreeError::new("Node already removed.")),
Some(_) => self.nodes[node_ref.id] = None,
},
}
self.len -= 1;
Ok(())
}
pub fn len(&self) -> usize {
self.len
}
fn get_node(&self, node_ref: NodeRef) -> Option<&Node<T>> {
match self.nodes.get(node_ref.id) {
None => None,
Some(node) => node.as_ref(),
}
}
fn get_node_mut(&mut self, node_ref: NodeRef) -> Option<&mut Node<T>> {
match self.nodes.get_mut(node_ref.id) {
None => None,
Some(node) => node.as_mut(),
}
}
pub fn get(&self, node_ref: NodeRef) -> Option<&T> {
match self.get_node(node_ref) {
None => None,
Some(node) => Some(&node.content),
}
}
pub fn get_mut(&mut self, node_ref: NodeRef) -> Option<&mut T> {
match self.get_node_mut(node_ref) {
None => None,
Some(node) => Some(&mut node.content),
}
}
pub fn append_child(&mut self, parent_ref: NodeRef, child_ref: NodeRef) -> Result<()> {
if self.get_node_mut(parent_ref).is_none() {
return Err(TreeError::new("Parent node does not exist."));
}
if self.get_node_mut(child_ref).is_none() {
return Err(TreeError::new("Child node does not exist."));
}
let parent_node = self.get_node_mut(parent_ref).unwrap();
parent_node.children.push(child_ref);
let child_node = self.get_node_mut(child_ref).unwrap();
child_node.parent = Some(parent_ref);
Ok(())
}
pub fn append_children(
&mut self,
parent_ref: NodeRef,
children_refs: &[NodeRef],
) -> Result<()> {
for child_ref in children_refs.iter() {
self.append_child(parent_ref, *child_ref)?;
}
Ok(())
}
pub fn get_children(&self, parent_ref: NodeRef) -> Result<Iter<NodeRef>> {
match self.get_node(parent_ref) {
None => Err(TreeError::new("Parent node does not exist.")),
Some(parent_node) => Ok(parent_node.children.iter()),
}
}
pub fn get_parent(&self, child_ref: NodeRef) -> Result<Option<NodeRef>> {
match self.get_node(child_ref) {
None => Err(TreeError::new("Child node does not exist.")),
Some(child_node) => Ok(child_node.parent),
}
}
pub fn depth_first_of(
&self,
node_ref: NodeRef,
include_start: bool,
) -> Result<DepthFirstIterator<T>> {
let mut iterator = DepthFirstIterator::new(&self, node_ref)?;
if !include_start {
iterator.next();
}
Ok(iterator)
}
pub fn depth_first(&self, include_root: bool) -> Result<DepthFirstIterator<T>> {
match self.root {
None => Err(TreeError::new("Cannot iterate a tree without a root node.")),
Some(root_ref) => self.depth_first_of(root_ref, include_root),
}
}
}
#[doc(hidden)]
pub struct DepthFirstIterator<'a, T> {
tree: &'a Tree<T>,
current: NodeRef,
child_iterator: Box<dyn Iterator<Item = &'a NodeRef> + 'a>,
current_iterator: Option<Box<dyn Iterator<Item = NodeRef> + 'a>>,
finished: bool,
}
impl<'a, T> DepthFirstIterator<'a, T> {
fn new(tree: &'a Tree<T>, current: NodeRef) -> Result<Self> {
Ok(Self {
tree,
current,
child_iterator: Box::new(tree.get_children(current)?),
current_iterator: None,
finished: false,
})
}
fn next_child(&mut self) -> bool {
match self.child_iterator.next() {
None => {
self.finished = true;
false
}
Some(next_child) => {
let child_iterator = self.tree.depth_first_of(*next_child, true).unwrap();
self.current_iterator = Some(Box::new(child_iterator));
true
}
}
}
}
impl<'a, T> Iterator for DepthFirstIterator<'a, T> {
type Item = NodeRef;
fn next(&mut self) -> Option<Self::Item> {
if self.finished {
return None;
}
match &mut self.current_iterator {
None => {
self.next_child();
Some(self.current)
}
Some(iterator) => match iterator.next() {
None => {
if self.next_child() {
self.next()
} else {
None
}
}
Some(value) => Some(value),
},
}
}
}
#[cfg(test)]
mod tests;