#![warn(
// ---------- Stylistic
future_incompatible,
nonstandard_style,
rust_2018_idioms,
trivial_casts,
trivial_numeric_casts,
// ---------- Public
missing_debug_implementations,
missing_docs,
unreachable_pub,
// ---------- Unsafe
unsafe_code,
// ---------- Unused
unused_extern_crates,
unused_import_braces,
unused_qualifications,
unused_results,
)]
use std::fmt::{Display, Formatter};
use std::io::Result;
use std::io::Write;
#[derive(Clone, Debug, PartialEq)]
pub enum TreeOrientation {
TopDown,
}
#[derive(Clone, Debug, PartialEq)]
pub enum AnchorPosition {
Below,
Left,
}
#[derive(Clone, Debug)]
pub struct TreeFormatting {
pub prefix_str: Option<String>,
pub orientation: TreeOrientation,
pub anchor: AnchorPosition,
pub chars: FormatCharacters,
}
#[derive(Clone, Debug)]
pub struct FormatCharacters {
pub down_facing_angle: char,
pub down_facing_tee: char,
pub vertical_line: char,
pub horizontal_line: char,
pub horizontal_space: char,
pub horizontal_line_count: usize,
pub right_facing_tee: char,
pub right_facing_angle: char,
pub label_space_char: char,
pub label_space_count: usize,
}
#[derive(Clone, Debug)]
pub struct TreeNode<T>
where
T: Display,
{
data: T,
children: Vec<TreeNode<T>>,
}
pub type StringTreeNode = TreeNode<String>;
impl Default for TreeFormatting {
fn default() -> Self {
Self::dir_tree(Default::default())
}
}
impl TreeFormatting {
pub fn dir_tree(chars: FormatCharacters) -> Self {
Self {
prefix_str: None,
orientation: TreeOrientation::TopDown,
anchor: AnchorPosition::Below,
chars,
}
}
pub fn dir_tree_with_prefix(chars: FormatCharacters, prefix_str: String) -> Self {
Self {
prefix_str: Some(prefix_str),
orientation: TreeOrientation::TopDown,
anchor: AnchorPosition::Below,
chars,
}
}
pub fn dir_tree_left(chars: FormatCharacters) -> Self {
Self {
prefix_str: None,
orientation: TreeOrientation::TopDown,
anchor: AnchorPosition::Left,
chars,
}
}
pub fn dir_tree_left_with_prefix(chars: FormatCharacters, prefix_str: String) -> Self {
Self {
prefix_str: Some(prefix_str),
orientation: TreeOrientation::TopDown,
anchor: AnchorPosition::Left,
chars,
}
}
#[inline]
pub(crate) fn just_space(&self) -> String {
format!(
"{}{}",
self.chars.just_space(),
if self.anchor == AnchorPosition::Below {
self.chars.horizontal_space.to_string()
} else {
String::new()
}
)
}
#[inline]
pub(crate) fn bar_and_space(&self) -> String {
format!(
"{}{}",
self.chars.bar_and_space(),
if self.anchor == AnchorPosition::Below {
self.chars.horizontal_space.to_string()
} else {
String::new()
}
)
}
#[inline]
pub(crate) fn tee(&self, has_children: bool) -> String {
format!(
"{}{}{}{}",
self.chars.right_facing_tee,
self.chars.horizontal_line(),
if self.anchor == AnchorPosition::Below {
String::new()
} else if has_children {
self.chars.down_facing_tee.to_string()
} else {
self.chars.horizontal_line.to_string()
},
self.chars.label_space()
)
}
#[inline]
pub(crate) fn angle(&self, has_children: bool) -> String {
format!(
"{}{}{}{}",
self.chars.right_facing_angle,
self.chars.horizontal_line(),
if self.anchor == AnchorPosition::Below {
String::new()
} else if has_children {
self.chars.down_facing_tee.to_string()
} else {
self.chars.horizontal_line.to_string()
},
self.chars.label_space(),
)
}
}
impl Default for FormatCharacters {
fn default() -> Self {
Self::ascii()
}
}
impl FormatCharacters {
pub fn ascii() -> Self {
Self {
down_facing_angle: '+',
down_facing_tee: ',',
vertical_line: '|',
horizontal_line: '-',
horizontal_space: ' ',
horizontal_line_count: 2,
right_facing_tee: '+',
right_facing_angle: '\'',
label_space_char: ' ',
label_space_count: 1,
}
}
pub fn box_chars() -> Self {
Self {
down_facing_angle: '┌',
down_facing_tee: '┬',
vertical_line: '│',
horizontal_line: '─',
horizontal_space: ' ',
horizontal_line_count: 2,
right_facing_tee: '├',
right_facing_angle: '└',
label_space_char: ' ',
label_space_count: 1,
}
}
#[inline]
pub(crate) fn just_space(&self) -> String {
format!("{}{}", self.horizontal_space, self.horizontal_space(),)
}
#[inline]
pub(crate) fn bar_and_space(&self) -> String {
format!("{}{}", self.vertical_line, self.horizontal_space(),)
}
#[inline]
pub(crate) fn horizontal_line(&self) -> String {
char_repeat(self.horizontal_line, self.horizontal_line_count)
}
#[inline]
pub(crate) fn horizontal_space(&self) -> String {
char_repeat(self.horizontal_space, self.horizontal_line_count)
}
#[inline]
pub(crate) fn label_space(&self) -> String {
char_repeat(self.label_space_char, self.label_space_count)
}
}
impl<T> Display for TreeNode<T>
where
T: Display,
{
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}",
self.to_string_with_format(&Default::default()).unwrap()
)
}
}
impl<T> TreeNode<T>
where
T: Display,
{
pub fn new(data: T) -> Self {
Self {
data,
children: Default::default(),
}
}
pub fn with_children(data: T, children: impl Iterator<Item = T>) -> Self
where
T: Sized,
{
Self::with_child_nodes(data, children.map(TreeNode::new))
}
pub fn with_child_nodes(data: T, children: impl Iterator<Item = TreeNode<T>>) -> Self
where
T: Sized,
{
Self {
data,
children: children.collect(),
}
}
pub fn data(&self) -> &T {
&self.data
}
pub fn label(&self) -> String {
self.data.to_string()
}
pub fn has_children(&self) -> bool {
!self.children.is_empty()
}
pub fn children(&self) -> impl Iterator<Item = &TreeNode<T>> {
self.children.iter()
}
pub fn push(&mut self, data: T) {
self.push_node(TreeNode {
data,
children: Default::default(),
})
}
pub fn push_node(&mut self, child: TreeNode<T>) {
self.children.push(child)
}
pub fn extend<V>(&mut self, children: impl Iterator<Item = T>) {
self.children.extend(children.map(TreeNode::new))
}
pub fn to_string_with_format(&self, format: &TreeFormatting) -> Result<String> {
use std::io::Cursor;
let mut buffer = Cursor::new(Vec::new());
self.write_with_format(&mut buffer, format)?;
Ok(String::from_utf8(buffer.into_inner()).unwrap())
}
pub fn write(&self, to_writer: &mut impl Write) -> Result<()>
where
T: Display,
{
self.write_with_format(
to_writer,
&TreeFormatting::dir_tree(FormatCharacters::ascii()),
)
}
pub fn write_with_format(
&self,
to_writer: &mut impl Write,
format: &TreeFormatting,
) -> Result<()>
where
T: Display,
{
write_tree_inner(self, to_writer, format, Default::default())
}
}
impl<T> From<T> for TreeNode<T>
where
T: Display,
{
fn from(v: T) -> Self {
Self {
data: v,
children: Default::default(),
}
}
}
impl<T> From<&T> for TreeNode<T>
where
T: Display + Clone,
{
fn from(v: &T) -> Self {
Self::from(v.clone())
}
}
impl<T> PartialEq for TreeNode<T>
where
T: Display + PartialEq,
{
fn eq(&self, other: &Self) -> bool {
self.data == other.data && self.children == other.children
}
}
impl From<&str> for TreeNode<String> {
fn from(v: &str) -> Self {
Self {
data: v.to_string(),
children: Default::default(),
}
}
}
fn write_tree_inner<T>(
node: &TreeNode<T>,
w: &mut impl Write,
format: &TreeFormatting,
remaining_children_stack: Vec<usize>,
) -> Result<()>
where
T: Display,
{
if let Some(prefix_str) = &format.prefix_str {
write!(w, "{}", prefix_str)?;
}
if !(format.anchor == AnchorPosition::Below) && remaining_children_stack.is_empty() {
write!(
w,
"{}{}",
format.chars.down_facing_angle,
char_repeat(
format.chars.label_space_char,
format.chars.label_space_count
)
)?;
}
let stack_depth = remaining_children_stack.len();
for (row, remaining_children) in remaining_children_stack.iter().enumerate() {
write!(
w,
"{}",
match (*remaining_children, row == (stack_depth - 1)) {
(1, true) => format.angle(node.has_children()),
(1, false) => format.just_space(),
(_, true) => format.tee(node.has_children()),
(_, false) => format.bar_and_space(),
}
)?;
}
if node.has_children() {
writeln!(w, "{}", node.label())?;
let mut d = node.children.len();
for child in &node.children {
let mut new_child_stack = remaining_children_stack.clone();
new_child_stack.push(d);
d -= 1;
write_tree_inner(child, w, format, new_child_stack)?;
}
} else {
writeln!(w, "{}", node.label())?;
}
Ok(())
}
#[inline]
fn char_repeat(c: char, n: usize) -> String {
c.to_string().as_str().repeat(n)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_node_new() {
let node = TreeNode::new(String::from("hello"));
assert_eq!(
node,
TreeNode {
data: "hello".to_string(),
children: vec![]
}
);
}
#[test]
fn test_node_with_children() {
let node = TreeNode::with_children(String::from("hello"), vec!["world".into()].into_iter());
assert_eq!(
node,
TreeNode {
data: "hello".to_string(),
children: vec![TreeNode {
data: "world".to_string(),
children: vec![]
}]
}
);
}
#[test]
fn test_node_from_string() {
let node: TreeNode<String> = String::from("hello").into();
assert_eq!(
node,
TreeNode {
data: "hello".to_string(),
children: vec![]
}
);
}
}