termtree/
lib.rs

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}
termtree/lib.rs

termtree/
lib.rs
