blob: cc9e5529f576ac8ecf7c8a4675ca94f23eaa2646 [file] [log] [blame] [edit]
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! A Vec-based container for a tree structure.
use std::num::NonZeroUsize;
use std::ops::{Add, Sub};
#[derive(Debug, Eq, PartialEq, Copy, Clone)]
pub enum TreePointer {
Nil,
Valid(TreeIndex),
}
impl TreePointer {
pub fn unwrap(self) -> TreeIndex {
match self {
TreePointer::Nil => panic!("Called unwrap on a Nil value"),
TreePointer::Valid(ix) => ix,
}
}
}
#[derive(Debug, Eq, PartialEq, Copy, Clone, PartialOrd)]
pub struct TreeIndex(NonZeroUsize);
impl TreeIndex {
fn new(i: usize) -> Self {
TreeIndex(NonZeroUsize::new(i).unwrap())
}
pub fn get(self) -> usize {
self.0.get()
}
}
impl Add<usize> for TreeIndex {
type Output = TreeIndex;
fn add(self, rhs: usize) -> Self {
let inner = self.0.get() + rhs;
TreeIndex::new(inner)
}
}
impl Sub<usize> for TreeIndex {
type Output = TreeIndex;
fn sub(self, rhs: usize) -> Self {
let inner = self.0.get().checked_sub(rhs).unwrap();
TreeIndex::new(inner)
}
}
#[derive(Debug, Clone, Copy)]
pub struct Node<T> {
pub child: TreePointer,
pub next: TreePointer,
pub item: T,
}
/// A tree abstraction, intended for fast building as a preorder traversal.
#[derive(Clone)]
pub struct Tree<T> {
nodes: Vec<Node<T>>,
spine: Vec<TreeIndex>, // indices of nodes on path to current node
cur: TreePointer,
}
impl<T: Default> Tree<T> {
// Indices start at one, so we place a dummy value at index zero.
// The alternative would be subtracting one from every TreeIndex
// every time we convert it to usize to index our nodes.
pub fn with_capacity(cap: usize) -> Tree<T> {
let mut nodes = Vec::with_capacity(cap);
nodes.push(Node {
child: TreePointer::Nil,
next: TreePointer::Nil,
item: <T as Default>::default(),
});
Tree {
nodes,
spine: Vec::new(),
cur: TreePointer::Nil,
}
}
/// Returns the index of the element currently in focus.
pub fn cur(&self) -> TreePointer {
self.cur
}
/// Append one item to the current position in the tree.
pub fn append(&mut self, item: T) -> TreeIndex {
let ix = self.create_node(item);
let this = TreePointer::Valid(ix);
if let TreePointer::Valid(ix) = self.cur {
self[ix].next = this;
} else if let Some(&parent) = self.spine.last() {
self[parent].child = this;
}
self.cur = this;
ix
}
/// Create an isolated node.
pub fn create_node(&mut self, item: T) -> TreeIndex {
let this = self.nodes.len();
self.nodes.push(Node {
child: TreePointer::Nil,
next: TreePointer::Nil,
item,
});
TreeIndex::new(this)
}
/// Push down one level, so that new items become children of the current node.
/// The new focus index is returned.
pub fn push(&mut self) -> TreeIndex {
let cur_ix = self.cur.unwrap();
self.spine.push(cur_ix);
self.cur = self[cur_ix].child;
cur_ix
}
/// Pop back up a level.
pub fn pop(&mut self) -> Option<TreeIndex> {
let ix = self.spine.pop()?;
self.cur = TreePointer::Valid(ix);
Some(ix)
}
/// Look at the parent node.
pub fn peek_up(&self) -> Option<TreeIndex> {
self.spine.last().cloned()
}
/// Look at grandparent node.
pub fn peek_grandparent(&self) -> Option<TreeIndex> {
if self.spine.len() >= 2 {
Some(self.spine[self.spine.len() - 2])
} else {
None
}
}
/// Returns true when there are no nodes in the tree, false otherwise.
pub fn is_empty(&self) -> bool {
self.nodes.len() <= 1
}
/// Returns the length of the spine.
pub fn spine_len(&self) -> usize {
self.spine.len()
}
/// Resets the focus to the first node added to the tree, if it exists.
pub fn reset(&mut self) {
self.cur = if self.is_empty() {
TreePointer::Nil
} else {
TreePointer::Valid(TreeIndex::new(1))
};
self.spine.truncate(0);
}
/// Walks the spine from a root node up to, but not including, the current node.
pub fn walk_spine(&self) -> impl std::iter::DoubleEndedIterator<Item = &TreeIndex> {
self.spine.iter()
}
/// Moves focus to the next sibling of the given node.
pub fn next_sibling(&mut self, cur_ix: TreeIndex) -> TreePointer {
self.cur = self[cur_ix].next;
self.cur
}
}
impl<T> std::fmt::Debug for Tree<T>
where
T: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
fn debug_tree<T>(
tree: &Tree<T>,
cur: TreeIndex,
indent: usize,
f: &mut std::fmt::Formatter,
) -> std::fmt::Result
where
T: std::fmt::Debug,
{
for _ in 0..indent {
write!(f, " ")?;
}
writeln!(f, "{:?}", &tree[cur].item)?;
if let TreePointer::Valid(child_ix) = tree[cur].child {
debug_tree(tree, child_ix, indent + 1, f)?;
}
if let TreePointer::Valid(next_ix) = tree[cur].next {
debug_tree(tree, next_ix, indent, f)?;
}
Ok(())
}
if self.nodes.len() > 1 {
let cur = TreeIndex(NonZeroUsize::new(1).unwrap());
debug_tree(self, cur, 0, f)
} else {
write!(f, "Empty tree")
}
}
}
impl<T> std::ops::Index<TreeIndex> for Tree<T> {
type Output = Node<T>;
fn index(&self, ix: TreeIndex) -> &Self::Output {
self.nodes.index(ix.get())
}
}
impl<T> std::ops::IndexMut<TreeIndex> for Tree<T> {
fn index_mut(&mut self, ix: TreeIndex) -> &mut Node<T> {
self.nodes.index_mut(ix.get())
}
}