reingold-tilford 1.0.0

/// A vector backed tree implementation used to essentially cache the user's tree.
use crate::{MediumVec, NodeInfo, SmallVec};

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Node<D> {
	/// Data needed by the actual algorithm.
	pub data: D,

	/// The position of this node among it's siblings.
	///
	/// Can also be thought of as the number of left-siblings this node has.
	pub order: usize,
	/// The depth of this node.
	///
	/// Can be thought of as the number of edges between this node and the root node.
	pub depth: usize,

	/// The index into `Tree` of this node's parent.
	pub parent: Option<usize>,
	/// The indices into `Tree` of this node's children.
	pub children: SmallVec<usize>,
}

impl<D> Node<D> {
	pub fn is_leaf(&self) -> bool {
		self.children.is_empty()
	}

	pub fn is_root(&self) -> bool {
		self.parent.is_none()
	}
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Tree<D>(pub MediumVec<Node<D>>);

impl<D> Tree<D> {
	pub fn new<F, N, T>(user_tree: &T, root: N, data: F) -> Self
	where
		F: Fn(&T, N) -> D,
		N: Copy,
		T: NodeInfo<N>,
	{
		let mut tree = MediumVec::new();
		tree.push(Node {
			data: data(user_tree, root),

			order: 0,
			depth: 0,

			parent: None,
			children: SmallVec::new(),
		});

		let mut queue = std::collections::VecDeque::new();
		queue.push_back((0, root));

		while let Some((parent, node)) = queue.pop_front() {
			let index = tree.len();

			for (i, child) in user_tree.children(node).into_iter().enumerate() {
				let index = index + i;
				let depth = tree[parent].depth + 1;

				tree[parent].children.push(index);

				tree.push(Node {
					data: data(user_tree, child),

					order: i,
					depth,

					parent: Some(parent),
					children: SmallVec::new(),
				});

				queue.push_back((index, child));
			}
		}

		Tree(tree)
	}

	pub fn root(&self) -> Option<usize> {
		if self.0.is_empty() {
			None
		} else {
			Some(0)
		}
	}

	pub fn breadth_first(&self, node: usize) -> MediumVec<usize> {
		let mut breadth_first = MediumVec::from_elem(node, 1);
		let mut index = 0;

		while index < breadth_first.len() {
			let node = breadth_first[index];
			breadth_first.extend_from_slice(&self[node].children);
			index += 1;
		}

		breadth_first
	}

	pub fn post_order(&self, node: usize) -> MediumVec<usize> {
		let mut breadth_first = MediumVec::from_elem(node, 1);
		let mut post_order = MediumVec::new();

		while let Some(node) = breadth_first.pop() {
			breadth_first.extend_from_slice(&self[node].children);
			post_order.push(node);
		}

		post_order.reverse();
		post_order
	}

	pub fn left_siblings(&self, node: usize) -> SmallVec<usize> {
		let order = self[node].order;

		if let Some(parent) = self[node].parent {
			self[parent].children[0..order].into()
		} else {
			SmallVec::new()
		}
	}

	pub fn siblings_between(&self, left: usize, right: usize) -> SmallVec<usize> {
		let left_order = self[left].order;
		let right_order = self[right].order;

		if self[left].is_root() || self[right].is_root() {
			assert!(
				self[left].is_root(),
				"If one node is the root then both nodes must be."
			);
			assert!(
				self[right].is_root(),
				"If one node is the root then both nodes must be."
			);

			return SmallVec::new();
		}

		let left_parent = self[left]
			.parent
			.expect("`is_none` has already been checked.");

		let right_parent = self[right]
			.parent
			.expect("`is_none` has already been checked.");

		assert!(
			left_parent == right_parent,
			"Nodes must actually be siblings."
		);

		let parent = left_parent;

		self[parent].children[left_order + 1..right_order].into()
	}

	pub fn previous_sibling(&self, node: usize) -> Option<usize> {
		let order = self[node].order;
		if order == 0 {
			return None;
		}

		let parent = self[node]
			.parent
			.expect("Nodes where `order != 0` always have parents.");

		Some(self[parent].children[order - 1])
	}
}

impl<D> std::ops::Index<usize> for Tree<D> {
	type Output = Node<D>;

	fn index(&self, index: usize) -> &Self::Output {
		&self.0[index]
	}
}

impl<D> std::ops::IndexMut<usize> for Tree<D> {
	fn index_mut(&mut self, index: usize) -> &mut Self::Output {
		&mut self.0[index]
	}
}

#[cfg(test)]
mod test {
	use super::*;

	struct UserTree;

	struct UserNode {
		id: usize,
		children: Vec<UserNode>,
	}

	impl<'n> NodeInfo<&'n UserNode> for UserTree {
		type Key = usize;

		fn key(&self, node: &'n UserNode) -> Self::Key {
			node.id
		}

		fn children(&self, node: &'n UserNode) -> SmallVec<&'n UserNode> {
			node.children.iter().collect()
		}
	}

	fn user_tree() -> UserNode {
		UserNode {
			id: 5,
			children: vec![
				UserNode {
					id: 1,
					children: vec![
						UserNode {
							id: 0,
							children: vec![],
						},
						UserNode {
							id: 3,
							children: vec![
								UserNode {
									id: 2,
									children: vec![],
								},
								UserNode {
									id: 4,
									children: vec![],
								},
							],
						},
					],
				},
				UserNode {
					id: 6,
					children: vec![UserNode {
						id: 8,
						children: vec![UserNode {
							id: 7,
							children: vec![],
						}],
					}],
				},
			],
		}
	}

	fn tree() -> Tree<usize> {
		Tree::new(&UserTree, &user_tree(), |_, n| n.id)
	}

	#[test]
	fn new() {
		let tree = tree();

		assert_eq!(tree.0.len(), 9);
	}

	#[test]
	fn root() {
		let tree = tree();

		assert!(tree.root().is_some());
		assert_eq!(tree[tree.root().unwrap()].data, 5);
	}

	#[test]
	fn breadth_first() {
		let tree = tree();
		let order = tree
			.breadth_first(tree.root().unwrap())
			.iter()
			.map(|&n| tree[n].data)
			.collect::<Vec<_>>();

		assert_eq!(order, vec![5, 1, 6, 0, 3, 8, 2, 4, 7]);
	}

	#[test]
	fn post_order() {
		let tree = tree();
		let order = tree
			.post_order(tree.root().unwrap())
			.iter()
			.map(|&n| tree[n].data)
			.collect::<Vec<_>>();

		assert_eq!(order, vec![0, 2, 4, 3, 1, 7, 8, 6, 5]);
	}

	#[test]
	fn is_root_true() {
		let tree = tree();

		assert!(tree[tree.root().unwrap()].is_root());
	}

	#[test]
	fn is_root_false() {
		let tree = tree();
		let child = tree[tree.root().unwrap()].children[0];

		assert!(!tree[child].is_root());
	}

	#[test]
	fn is_leaf_true() {
		let tree = tree();
		let leaf = tree[tree[tree[tree.root().unwrap()].children[1]].children[0]].children[0];

		assert!(tree[leaf].is_leaf());
	}

	#[test]
	fn is_leaf_false() {
		let tree = tree();

		assert!(!tree[tree.root().unwrap()].is_leaf());
	}
}