1#![allow(clippy::branches_sharing_code)]
2
3#[cfg(test)]
4mod tests;
5
6use std::collections::VecDeque;
7use std::fmt::{self, Display};
8use std::rc::Rc;
9
10/// a simple recursive type which is able to render its
11/// components in a tree-like format
12#[derive(Debug, Clone)]
13pub struct Tree<D: Display> {
14 pub root: D,
15 pub leaves: Vec<Tree<D>>,
16 multiline: bool,
17 glyphs: GlyphPalette,
18}
19
20impl<D: Display> Tree<D> {
21 pub fn new(root: D) -> Self {
22 Tree {
23 root,
24 leaves: Vec::new(),
25 multiline: false,
26 glyphs: GlyphPalette::new(),
27 }
28 }
29
30 pub fn with_leaves(mut self, leaves: impl IntoIterator<Item = impl Into<Tree<D>>>) -> Self {
31 self.leaves = leaves.into_iter().map(Into::into).collect();
32 self
33 }
34
35 /// Ensure all lines for `root` are indented
36 pub fn with_multiline(mut self, yes: bool) -> Self {
37 self.multiline = yes;
38 self
39 }
40
41 /// Customize the rendering of this node
42 pub fn with_glyphs(mut self, glyphs: GlyphPalette) -> Self {
43 self.glyphs = glyphs;
44 self
45 }
46}
47
48impl<D: Display> Tree<D> {
49 /// Ensure all lines for `root` are indented
50 pub fn set_multiline(&mut self, yes: bool) -> &mut Self {
51 self.multiline = yes;
52 self
53 }
54
55 /// Customize the rendering of this node
56 pub fn set_glyphs(&mut self, glyphs: GlyphPalette) -> &mut Self {
57 self.glyphs = glyphs;
58 self
59 }
60}
61
62impl<D: Display> Tree<D> {
63 pub fn push(&mut self, leaf: impl Into<Tree<D>>) -> &mut Self {
64 self.leaves.push(leaf.into());
65 self
66 }
67}
68
69impl<D: Display> From<D> for Tree<D> {
70 fn from(inner: D) -> Self {
71 Self::new(inner)
72 }
73}
74
75impl<D: Display> Extend<D> for Tree<D> {
76 fn extend<T: IntoIterator<Item = D>>(&mut self, iter: T) {
77 self.leaves.extend(iter.into_iter().map(Into::into))
78 }
79}
80
81impl<D: Display> Extend<Tree<D>> for Tree<D> {
82 fn extend<T: IntoIterator<Item = Tree<D>>>(&mut self, iter: T) {
83 self.leaves.extend(iter)
84 }
85}
86
87impl<D: Display> Display for Tree<D> {
88 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
89 self.root.fmt(f)?; // Pass along `f.alternate()`
90 writeln!(f)?;
91 let mut queue = DisplauQueue::new();
92 let no_space = Rc::new(Vec::new());
93 enqueue_leaves(&mut queue, self, no_space);
94 while let Some((last, leaf, spaces)) = queue.pop_front() {
95 let mut prefix = (
96 if last {
97 leaf.glyphs.last_item
98 } else {
99 leaf.glyphs.middle_item
100 },
101 leaf.glyphs.item_indent,
102 );
103
104 if leaf.multiline {
105 let rest_prefix = (
106 if last {
107 leaf.glyphs.last_skip
108 } else {
109 leaf.glyphs.middle_skip
110 },
111 leaf.glyphs.skip_indent,
112 );
113 debug_assert_eq!(prefix.0.chars().count(), rest_prefix.0.chars().count());
114 debug_assert_eq!(prefix.1.chars().count(), rest_prefix.1.chars().count());
115
116 let root = if f.alternate() {
117 format!("{:#}", leaf.root)
118 } else {
119 format!("{:}", leaf.root)
120 };
121 for line in root.lines() {
122 // print single line
123 for s in spaces.as_slice() {
124 if *s {
125 self.glyphs.last_skip.fmt(f)?;
126 self.glyphs.skip_indent.fmt(f)?;
127 } else {
128 self.glyphs.middle_skip.fmt(f)?;
129 self.glyphs.skip_indent.fmt(f)?;
130 }
131 }
132 prefix.0.fmt(f)?;
133 prefix.1.fmt(f)?;
134 line.fmt(f)?;
135 writeln!(f)?;
136 prefix = rest_prefix;
137 }
138 } else {
139 // print single line
140 for s in spaces.as_slice() {
141 if *s {
142 self.glyphs.last_skip.fmt(f)?;
143 self.glyphs.skip_indent.fmt(f)?;
144 } else {
145 self.glyphs.middle_skip.fmt(f)?;
146 self.glyphs.skip_indent.fmt(f)?;
147 }
148 }
149 prefix.0.fmt(f)?;
150 prefix.1.fmt(f)?;
151 leaf.root.fmt(f)?; // Pass along `f.alternate()`
152 writeln!(f)?;
153 }
154
155 // recurse
156 if !leaf.leaves.is_empty() {
157 let s: &Vec<bool> = &spaces;
158 let mut child_spaces = s.clone();
159 child_spaces.push(last);
160 let child_spaces = Rc::new(child_spaces);
161 enqueue_leaves(&mut queue, leaf, child_spaces);
162 }
163 }
164 Ok(())
165 }
166}
167
168type DisplauQueue<'t, D> = VecDeque<(bool, &'t Tree<D>, Rc<Vec<bool>>)>;
169
170fn enqueue_leaves<'t, D: Display>(
171 queue: &mut DisplauQueue<'t, D>,
172 parent: &'t Tree<D>,
173 spaces: Rc<Vec<bool>>,
174) {
175 for (i, leaf) in parent.leaves.iter().rev().enumerate() {
176 let last = i == 0;
177 queue.push_front((last, leaf, spaces.clone()));
178 }
179}
180
181#[derive(Copy, Clone, Debug, PartialEq, Eq)]
182pub struct GlyphPalette {
183 pub middle_item: &'static str,
184 pub last_item: &'static str,
185 pub item_indent: &'static str,
186
187 pub middle_skip: &'static str,
188 pub last_skip: &'static str,
189 pub skip_indent: &'static str,
190}
191
192impl GlyphPalette {
193 pub const fn new() -> Self {
194 Self {
195 middle_item: "├",
196 last_item: "└",
197 item_indent: "── ",
198
199 middle_skip: "│",
200 last_skip: " ",
201 skip_indent: " ",
202 }
203 }
204}
205
206impl Default for GlyphPalette {
207 fn default() -> Self {
208 Self::new()
209 }
210}
211