diff --git a/src/display.rs b/src/display.rs new file mode 100644 index 0000000..f552aff --- /dev/null +++ b/src/display.rs @@ -0,0 +1,86 @@ +use std::fmt::Display; + +/// Indentation token +#[derive(Debug)] +struct Token { + /// Is followed by a brother + siblings: bool, + /// Is intermediate while printing children + children: bool, +} + +impl ToString for Token { + fn to_string(&self) -> String { + let Token { siblings, children } = self; + + match (siblings, children) { + (true, true) => "│ ", + (true, false) => "├── ", + (false, true) => " ", + (false, false) => "└── ", + } + .to_string() + } +} + +impl Token { + /// Create a new indentation token + fn new(siblings: bool) -> Self { + Token { + siblings, + children: false, + } + } + + /// Set children flag before starting displaying children + fn set_children(&mut self) { + self.children = true; + } +} + +/// Manages the state during the display operation +#[derive(Debug)] +pub struct Indentation { + tokens: Vec, + ignore_root: bool, +} + +impl Display for Indentation { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + let first: usize = if self.ignore_root { 1 } else { 0 }; + + for token in &self.tokens[first..] { + write!(f, "{}", token.to_string())?; + } + + Ok(()) + } +} + +impl Indentation { + /// Creates a new indentation handler + pub fn new(ignore_root: bool) -> Self { + Indentation { + tokens: Vec::new(), + ignore_root, + } + } + + /// Adds a new layer of indentation + pub fn indent(&mut self, siblings: bool) -> &mut Self { + // Setup children mode for previous tokens + let len = self.tokens.len(); + if len > 0 { + self.tokens[len - 1].set_children(); + } + + self.tokens.push(Token::new(siblings)); + self + } + + /// Removes the last layer of indentation + pub fn deindent(&mut self) -> &mut Self { + self.tokens.pop(); + self + } +} diff --git a/src/lib.rs b/src/lib.rs index 99b3137..b7ce9d2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -719,36 +719,34 @@ impl Debug for Tree { } } +// Handles display +mod display; + impl Display for Tree { fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> { + use display::Indentation; use iter::Edge; - fn indent(f: &mut Formatter, level: usize) -> Result<(), fmt::Error> { - for _ in 0..level { - write!(f, " ")?; - } - Ok(()) - } - - let mut level: usize = 0; + let mut indent: Indentation = Indentation::new(true); for edge in self.root().traverse() { match edge { Edge::Open(node) if node.has_children() => { - indent(f, level)?; - write!(f, "{}\n", node.value())?; - level += 1; + indent.indent(node.next_sibling().is_some()); + write!(f, "{}{}\n", indent.to_string(), node.value())?; } Edge::Open(node) if node.next_sibling().is_some() => { - indent(f, level)?; - write!(f, "{}\n", node.value())?; + indent.indent(node.next_sibling().is_some()); + write!(f, "{}{}\n", indent.to_string(), node.value())?; + indent.deindent(); } Edge::Open(node) => { - indent(f, level)?; - write!(f, "{}\n", node.value())?; + indent.indent(node.next_sibling().is_some()); + write!(f, "{}{}\n", indent.to_string(), node.value())?; + indent.deindent(); } Edge::Close(node) if node.has_children() => { - level -= 1; + indent.deindent(); } _ => {} } diff --git a/tests/tree.rs b/tests/tree.rs index dc2683c..78eee6a 100644 --- a/tests/tree.rs +++ b/tests/tree.rs @@ -171,7 +171,7 @@ fn test_display() { }; let repr = format!("{tree}"); - let expected = "root\n a\n child 1\n b\n child 2\n"; + let expected = "root\n├── a\n│ └── child 1\n└── b\n └── child 2\n"; assert_eq!(repr, expected); @@ -184,7 +184,7 @@ fn test_display() { }; let repr = format!("{tree}"); - let expected = "root\n a\n b\n x\n y\n c\n"; + let expected = "root\n├── a\n├── b\n│ ├── x\n│ └── y\n└── c\n"; assert_eq!(repr, expected); }