use crate::cells::cell_len;
use crate::console::{Console, ConsoleOptions, Renderable};
use crate::measure::Measurement;
use crate::segment::Segment;
use crate::style::Style;
use crate::text::Text;
const SPACE: usize = 0;
const CONTINUE: usize = 1;
const FORK: usize = 2;
const END: usize = 3;
const ASCII_GUIDES: [&str; 4] = [" ", "| ", "+-- ", "`-- "];
const TREE_GUIDES: [[&str; 4]; 3] = [
[
" ",
"\u{2502} ",
"\u{251c}\u{2500}\u{2500} ",
"\u{2514}\u{2500}\u{2500} ",
], [
" ",
"\u{2503} ",
"\u{2523}\u{2501}\u{2501} ",
"\u{2517}\u{2501}\u{2501} ",
], [
" ",
"\u{2551} ",
"\u{2560}\u{2550}\u{2550} ",
"\u{255a}\u{2550}\u{2550} ",
], ];
fn make_guide(index: usize, style: &Style, ascii_only: bool) -> Segment {
if ascii_only {
Segment::styled(ASCII_GUIDES[index], style.clone())
} else {
let guide_set = if style.bold() == Some(true) {
1
} else if style.underline2() == Some(true) {
2
} else {
0
};
Segment::styled(TREE_GUIDES[guide_set][index], style.clone())
}
}
pub struct Tree {
pub label: Text,
pub style: Style,
pub guide_style: Style,
pub children: Vec<Tree>,
pub expanded: bool,
pub hide_root: bool,
}
impl Tree {
pub fn new(label: Text) -> Self {
Tree {
label,
style: Style::null(),
guide_style: Style::null(),
children: Vec::new(),
expanded: true,
hide_root: false,
}
}
pub fn add(&mut self, label: Text) -> &mut Tree {
self.children.push(Tree {
label,
style: self.style.clone(),
guide_style: self.guide_style.clone(),
children: Vec::new(),
expanded: true,
hide_root: false,
});
self.children.last_mut().unwrap()
}
#[must_use]
pub fn style(mut self, style: Style) -> Self {
self.style = style;
self
}
#[must_use]
pub fn guide_style(mut self, style: Style) -> Self {
self.guide_style = style;
self
}
#[must_use]
pub fn expanded(mut self, expanded: bool) -> Self {
self.expanded = expanded;
self
}
#[must_use]
pub fn hide_root(mut self, hide_root: bool) -> Self {
self.hide_root = hide_root;
self
}
pub fn measure(&self, _console: &Console, _options: &ConsoleOptions) -> Measurement {
let mut minimum: usize = 0;
let mut maximum: usize = 0;
fn measure_recursive(
tree: &Tree,
level: usize,
min: &mut usize,
max: &mut usize,
hide_root: bool,
) {
let effective_level = if hide_root {
level.saturating_sub(1)
} else {
level
};
let indent = effective_level * 4;
let label_width = tree.label.cell_len();
let total = label_width + indent;
if !(level == 0 && hide_root) {
*min = (*min).max(total);
*max = (*max).max(total);
}
if tree.expanded {
for child in &tree.children {
measure_recursive(child, level + 1, min, max, hide_root);
}
}
}
measure_recursive(self, 0, &mut minimum, &mut maximum, self.hide_root);
Measurement::new(minimum, maximum)
}
}
struct StackFrame<'a> {
index: usize,
children: &'a [Tree],
}
impl Renderable for Tree {
fn rich_console(&self, console: &Console, options: &ConsoleOptions) -> Vec<Segment> {
let mut segments: Vec<Segment> = Vec::new();
let ascii_only = options.ascii_only();
let newline = Segment::line();
let mut levels: Vec<Segment> = vec![make_guide(CONTINUE, &self.guide_style, ascii_only)];
let mut stack: Vec<StackFrame> = Vec::new();
let root_slice = std::slice::from_ref(self);
stack.push(StackFrame {
index: 0,
children: root_slice,
});
let mut depth: usize = 0;
while let Some(frame) = stack.last_mut() {
if frame.index >= frame.children.len() {
stack.pop();
levels.pop();
if !levels.is_empty() {
let last_idx = levels.len() - 1;
let guide_style = levels[last_idx].style.clone().unwrap_or_else(Style::null);
levels[last_idx] = make_guide(FORK, &guide_style, ascii_only);
}
depth = depth.saturating_sub(1);
continue;
}
let child_idx = frame.index;
let total = frame.children.len();
let last = child_idx == total - 1;
let node = &frame.children[child_idx];
frame.index += 1;
if last {
let last_level = levels.len() - 1;
let guide_style = levels[last_level].style.clone().unwrap_or_else(Style::null);
levels[last_level] = make_guide(END, &guide_style, ascii_only);
}
let skip = if self.hide_root { 2 } else { 1 };
let prefix: Vec<Segment> = if levels.len() > skip {
levels[skip..].to_vec()
} else {
Vec::new()
};
let prefix_width: usize = prefix.iter().map(|s| cell_len(&s.text)).sum();
let child_width = options.max_width.saturating_sub(prefix_width);
let child_opts = options.update_width(child_width);
let rendered_lines =
console.render_lines(&node.label, Some(&child_opts), None, false, false);
let skip_node = depth == 0 && self.hide_root;
if !skip_node {
let mut current_prefix = prefix.clone();
for (i, line) in rendered_lines.iter().enumerate() {
for seg in ¤t_prefix {
segments.push(seg.clone());
}
segments.extend(line.iter().cloned());
segments.push(newline.clone());
if i == 0 && !current_prefix.is_empty() {
let last_idx = current_prefix.len() - 1;
let pstyle = current_prefix[last_idx]
.style
.clone()
.unwrap_or_else(Style::null);
current_prefix[last_idx] =
make_guide(if last { SPACE } else { CONTINUE }, &pstyle, ascii_only);
}
}
}
if node.expanded && !node.children.is_empty() {
let last_level = levels.len() - 1;
let guide_style = levels[last_level].style.clone().unwrap_or_else(Style::null);
levels[last_level] = make_guide(
if last { SPACE } else { CONTINUE },
&guide_style,
ascii_only,
);
let child_guide_style = &node.guide_style;
let child_count = node.children.len();
let guide_type = if child_count == 1 { END } else { FORK };
levels.push(make_guide(guide_type, child_guide_style, ascii_only));
stack.push(StackFrame {
index: 0,
children: &node.children,
});
depth += 1;
}
}
segments
}
}
impl std::fmt::Display for Tree {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut console = Console::builder()
.width(f.width().unwrap_or(80))
.force_terminal(true)
.no_color(true)
.build();
console.begin_capture();
console.print(self);
let output = console.end_capture();
write!(f, "{}", output.trim_end_matches('\n'))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn test_console(width: usize) -> Console {
Console::builder()
.width(width)
.markup(false)
.highlight(false)
.no_color(true)
.build()
}
fn render_tree(tree: &Tree, width: usize) -> String {
let console = test_console(width);
let opts = console.options();
let segments = tree.rich_console(&console, &opts);
segments
.iter()
.filter(|s| !s.is_control())
.map(|s| s.text.as_str())
.collect()
}
#[test]
fn test_single_node() {
let tree = Tree::new(Text::new("root", Style::null()));
let output = render_tree(&tree, 80);
assert!(output.contains("root"));
let lines: Vec<&str> = output.lines().collect();
assert_eq!(lines.len(), 1);
assert!(!output.contains("\u{251c}"));
assert!(!output.contains("\u{2514}"));
}
#[test]
fn test_one_child() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.add(Text::new("child", Style::null()));
let output = render_tree(&tree, 80);
let lines: Vec<&str> = output.lines().collect();
assert_eq!(lines.len(), 2);
assert!(lines[0].contains("root"));
assert!(lines[1].contains("child"));
assert!(output.contains("\u{2514}\u{2500}\u{2500}"));
}
#[test]
fn test_multiple_children() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.add(Text::new("child1", Style::null()));
tree.add(Text::new("child2", Style::null()));
tree.add(Text::new("child3", Style::null()));
let output = render_tree(&tree, 80);
let lines: Vec<&str> = output.lines().collect();
assert_eq!(lines.len(), 4);
assert!(lines[0].contains("root"));
assert!(lines[1].contains("\u{251c}\u{2500}\u{2500}"));
assert!(lines[2].contains("\u{251c}\u{2500}\u{2500}"));
assert!(lines[3].contains("\u{2514}\u{2500}\u{2500}"));
}
#[test]
fn test_nested_children() {
let mut tree = Tree::new(Text::new("root", Style::null()));
let child = tree.add(Text::new("child", Style::null()));
child
.children
.push(Tree::new(Text::new("grandchild", Style::null())));
let output = render_tree(&tree, 80);
let lines: Vec<&str> = output.lines().collect();
assert_eq!(lines.len(), 3);
assert!(lines[0].contains("root"));
assert!(lines[1].contains("child"));
assert!(lines[2].contains("grandchild"));
}
#[test]
fn test_hide_root() {
let mut tree = Tree::new(Text::new("root", Style::null())).hide_root(true);
tree.add(Text::new("child1", Style::null()));
tree.add(Text::new("child2", Style::null()));
let output = render_tree(&tree, 80);
assert!(!output.contains("root"));
let lines: Vec<&str> = output.lines().collect();
assert_eq!(lines.len(), 2);
assert!(lines[0].contains("child1"));
assert!(lines[1].contains("child2"));
}
#[test]
fn test_collapsed_node() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.children
.push(Tree::new(Text::new("branch", Style::null())).expanded(false));
tree.children[0]
.children
.push(Tree::new(Text::new("hidden", Style::null())));
let output = render_tree(&tree, 80);
assert!(output.contains("branch"));
assert!(!output.contains("hidden"));
}
#[test]
fn test_ascii_mode() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.add(Text::new("child1", Style::null()));
tree.add(Text::new("child2", Style::null()));
let console = Console::builder()
.width(80)
.markup(false)
.highlight(false)
.no_color(true)
.build();
let mut opts = console.options();
opts.encoding = "ascii".to_string();
let segments = tree.rich_console(&console, &opts);
let output: String = segments
.iter()
.filter(|s| !s.is_control())
.map(|s| s.text.as_str())
.collect();
assert!(output.contains("+-- "));
assert!(output.contains("`-- "));
assert!(!output.contains("\u{251c}"));
assert!(!output.contains("\u{2514}"));
}
#[test]
fn test_guide_characters() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.add(Text::new("a", Style::null()));
tree.add(Text::new("b", Style::null()));
let output = render_tree(&tree, 80);
let lines: Vec<&str> = output.lines().collect();
assert!(lines[1].starts_with("\u{251c}\u{2500}\u{2500} "));
assert!(lines[2].starts_with("\u{2514}\u{2500}\u{2500} "));
}
#[test]
fn test_multiline_label() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.add(Text::new(
"This is a very long label that should wrap",
Style::null(),
));
let output = render_tree(&tree, 20);
let lines: Vec<&str> = output.lines().collect();
assert!(lines.len() > 2);
}
#[test]
fn test_measure() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.add(Text::new("child", Style::null()));
let console = test_console(80);
let opts = console.options();
let measurement = tree.measure(&console, &opts);
assert_eq!(measurement.minimum, 9);
assert_eq!(measurement.maximum, 9);
}
#[test]
fn test_builder_pattern() {
let style = Style::parse("bold").unwrap();
let guide_style = Style::parse("red").unwrap();
let tree = Tree::new(Text::new("root", Style::null()))
.style(style.clone())
.guide_style(guide_style.clone())
.expanded(false)
.hide_root(true);
assert_eq!(tree.style, style);
assert_eq!(tree.guide_style, guide_style);
assert!(!tree.expanded);
assert!(tree.hide_root);
}
#[test]
fn test_add_returns_mut_ref() {
let mut tree = Tree::new(Text::new("root", Style::null()));
let child = tree.add(Text::new("child", Style::null()));
child
.children
.push(Tree::new(Text::new("grandchild", Style::null())));
assert_eq!(tree.children.len(), 1);
assert_eq!(tree.children[0].children.len(), 1);
assert_eq!(tree.children[0].children[0].label.plain(), "grandchild");
}
#[test]
fn test_deep_nesting() {
let mut tree = Tree::new(Text::new("L0", Style::null()));
let l1 = tree.add(Text::new("L1", Style::null()));
l1.children.push(Tree::new(Text::new("L2", Style::null())));
l1.children[0]
.children
.push(Tree::new(Text::new("L3", Style::null())));
let output = render_tree(&tree, 80);
let lines: Vec<&str> = output.lines().collect();
assert_eq!(lines.len(), 4);
assert!(lines[0].contains("L0"));
assert!(lines[1].contains("L1"));
assert!(lines[2].contains("L2"));
assert!(lines[3].contains("L3"));
for i in 1..lines.len() {
let current_content_start = lines[i].find('L').unwrap_or(0);
let prev_content_start = lines[i - 1].find('L').unwrap_or(0);
assert!(
current_content_start > prev_content_start,
"L{} should be indented more than L{}",
i,
i - 1
);
}
}
#[test]
fn test_add_inherits_styles() {
let style = Style::parse("bold").unwrap();
let guide_style = Style::parse("red").unwrap();
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.style = style.clone();
tree.guide_style = guide_style.clone();
tree.add(Text::new("child", Style::null()));
assert_eq!(tree.children[0].style, style);
assert_eq!(tree.children[0].guide_style, guide_style);
}
#[test]
fn test_empty_tree_no_guides() {
let tree = Tree::new(Text::new("alone", Style::null()));
let output = render_tree(&tree, 80);
let lines: Vec<&str> = output.lines().collect();
assert_eq!(lines.len(), 1);
assert_eq!(lines[0].trim(), "alone");
}
#[test]
fn test_continuation_guides() {
let mut tree = Tree::new(Text::new("root", Style::null()));
{
let child1 = tree.add(Text::new("child1", Style::null()));
child1
.children
.push(Tree::new(Text::new("grandchild1", Style::null())));
}
tree.add(Text::new("child2", Style::null()));
let output = render_tree(&tree, 80);
let lines: Vec<&str> = output.lines().collect();
assert_eq!(lines.len(), 4);
assert!(
lines[2].contains("\u{2502}"),
"grandchild1 line should contain continue guide: {:?}",
lines[2]
);
}
#[test]
fn test_hide_root_deep() {
let mut tree = Tree::new(Text::new("ROOT", Style::null())).hide_root(true);
let child = tree.add(Text::new("child", Style::null()));
child
.children
.push(Tree::new(Text::new("grandchild", Style::null())));
let output = render_tree(&tree, 80);
assert!(!output.contains("ROOT"));
assert!(output.contains("child"));
assert!(output.contains("grandchild"));
}
#[test]
fn test_measure_hide_root() {
let mut tree = Tree::new(Text::new("LONG_ROOT_NAME", Style::null())).hide_root(true);
tree.add(Text::new("short", Style::null()));
let console = test_console(80);
let opts = console.options();
let measurement = tree.measure(&console, &opts);
assert_eq!(measurement.minimum, 5);
}
#[test]
fn test_measure_deep() {
let mut tree = Tree::new(Text::new("r", Style::null()));
let c = tree.add(Text::new("cc", Style::null()));
c.children.push(Tree::new(Text::new("ggg", Style::null())));
let console = test_console(80);
let opts = console.options();
let measurement = tree.measure(&console, &opts);
assert_eq!(measurement.maximum, 11);
}
#[test]
fn test_measure_collapsed() {
let mut tree = Tree::new(Text::new("r", Style::null()));
let mut branch = Tree::new(Text::new("branch", Style::null())).expanded(false);
branch.children.push(Tree::new(Text::new(
"very_very_very_long_hidden_label",
Style::null(),
)));
tree.children.push(branch);
let console = test_console(80);
let opts = console.options();
let measurement = tree.measure(&console, &opts);
assert_eq!(measurement.maximum, 10);
}
#[test]
fn test_bold_guide_style() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.guide_style = Style::parse("bold").unwrap();
tree.add(Text::new("child1", Style::null()));
tree.add(Text::new("child2", Style::null()));
let output = render_tree(&tree, 80);
assert!(output.contains("\u{2523}\u{2501}\u{2501}"));
assert!(output.contains("\u{2517}\u{2501}\u{2501}"));
}
#[test]
fn test_double_guide_style() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.guide_style = Style::parse("underline2").unwrap();
tree.add(Text::new("child1", Style::null()));
tree.add(Text::new("child2", Style::null()));
let output = render_tree(&tree, 80);
assert!(output.contains("\u{2560}\u{2550}\u{2550}"));
assert!(output.contains("\u{255a}\u{2550}\u{2550}"));
}
#[test]
fn test_guide_width() {
for guide in &ASCII_GUIDES {
assert_eq!(cell_len(guide), 4, "ASCII guide {:?} is not 4 cells", guide);
}
for set in &TREE_GUIDES {
for guide in set {
assert_eq!(
cell_len(guide),
4,
"Unicode guide {:?} is not 4 cells",
guide
);
}
}
}
#[test]
fn test_segments_contain_newlines() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.add(Text::new("child", Style::null()));
let console = test_console(80);
let opts = console.options();
let segments = tree.rich_console(&console, &opts);
let newline_count = segments.iter().filter(|s| s.text == "\n").count();
assert_eq!(newline_count, 2, "Expected 2 newlines (one per line)");
}
#[test]
fn test_hide_root_no_children() {
let tree = Tree::new(Text::new("hidden", Style::null())).hide_root(true);
let output = render_tree(&tree, 80);
assert!(
output.trim().is_empty(),
"hide_root with no children should produce empty output"
);
}
#[test]
fn test_display_trait() {
let mut tree = Tree::new(Text::new("root", Style::null()));
tree.add(Text::new("child1", Style::null()));
tree.add(Text::new("child2", Style::null()));
let s = format!("{}", tree);
assert!(!s.is_empty());
assert!(s.contains("root"));
assert!(s.contains("child1"));
assert!(s.contains("child2"));
}
}