eza/output/
tree.rs

1// SPDX-FileCopyrightText: 2024 Christina Sørensen
2// SPDX-License-Identifier: EUPL-1.2
3//
4// SPDX-FileCopyrightText: 2023-2024 Christina Sørensen, eza contributors
5// SPDX-FileCopyrightText: 2014 Benjamin Sago
6// SPDX-License-Identifier: MIT
7//! Tree structures, such as `├──` or `└──`, used in a tree view.
8//!
9//! ## Constructing Tree Views
10//!
11//! When using the `--tree` argument, instead of a vector of cells, each row
12//! has a `depth` field that indicates how far deep in the tree it is: the top
13//! level has depth 0, its children have depth 1, and *their* children have
14//! depth 2, and so on.
15//!
16//! On top of this, it also has a `last` field that specifies whether this is
17//! the last row of this particular consecutive set of rows. This doesn’t
18//! affect the file’s information; it’s just used to display a different set of
19//! Unicode tree characters! The resulting table looks like this:
20//!
21//! ```text
22//!     ┌───────┬───────┬───────────────────────┐
23//!     │ Depth │ Last  │ Output                │
24//!     ├───────┼───────┼───────────────────────┤
25//!     │     0 │       │ documents             │
26//!     │     1 │ false │ ├── this_file.txt     │
27//!     │     1 │ false │ ├── that_file.txt     │
28//!     │     1 │ false │ ├── features          │
29//!     │     2 │ false │ │   ├── feature_1.rs  │
30//!     │     2 │ false │ │   ├── feature_2.rs  │
31//!     │     2 │ true  │ │   └── feature_3.rs  │
32//!     │     1 │ true  │ └── pictures          │
33//!     │     2 │ false │     ├── garden.jpg    │
34//!     │     2 │ false │     ├── flowers.jpg   │
35//!     │     2 │ false │     ├── library.png   │
36//!     │     2 │ true  │     └── space.tiff    │
37//!     └───────┴───────┴───────────────────────┘
38//! ```
39//!
40//! Creating the table like this means that each file has to be tested to see
41//! if it’s the last one in the group. This is usually done by putting all the
42//! files in a vector beforehand, getting its length, then comparing the index
43//! of each file to see if it’s the last one. (As some files may not be
44//! successfully `stat`ted, we don’t know how many files are going to exist in
45//! each directory)
46
47#[derive(PartialEq, Eq, Debug, Copy, Clone)]
48pub enum TreePart {
49    /// Rightmost column, *not* the last in the directory.
50    Edge,
51
52    /// Not the rightmost column, and the directory has not finished yet.
53    Line,
54
55    /// Rightmost column, and the last in the directory.
56    Corner,
57
58    /// Not the rightmost column, and the directory *has* finished.
59    Blank,
60}
61
62impl TreePart {
63    /// Turn this tree part into ASCII-licious box drawing characters!
64    /// (Warning: not actually ASCII)
65    pub fn ascii_art(self) -> &'static str {
66        #[rustfmt::skip]
67        return match self {
68            Self::Edge    => "├── ",
69            Self::Line    => "│   ",
70            Self::Corner  => "└── ",
71            Self::Blank   => "    ",
72        };
73    }
74}
75
76/// A **tree trunk** builds up arrays of tree parts over multiple depths.
77#[derive(Debug, Default)]
78pub struct TreeTrunk {
79    /// A stack tracks which tree characters should be printed. It’s
80    /// necessary to maintain information about the previously-printed
81    /// lines, as the output will change based on any previous entries.
82    stack: Vec<TreePart>,
83
84    /// A tuple for the last ‘depth’ and ‘last’ parameters that are passed in.
85    last_params: Option<TreeParams>,
86}
87
88#[derive(Debug, Copy, Clone)]
89pub struct TreeParams {
90    /// How many directories deep into the tree structure this is. Directories
91    /// on top have depth 0.
92    depth: TreeDepth,
93
94    /// Whether this is the last entry in the directory.
95    last: bool,
96}
97
98#[derive(Debug, Copy, Clone)]
99pub struct TreeDepth(pub usize);
100
101impl TreeTrunk {
102    /// Calculates the tree parts for an entry at the given depth and
103    /// last-ness. The depth is used to determine where in the stack the tree
104    /// part should be inserted, and the last-ness is used to determine which
105    /// type of tree part to insert.
106    ///
107    /// This takes a `&mut self` because the results of each file are stored
108    /// and used in future rows.
109    pub fn new_row(&mut self, params: TreeParams) -> &[TreePart] {
110        // If this isn’t our first iteration, then update the tree parts thus
111        // far to account for there being another row after it.
112        if let Some(last) = self.last_params {
113            self.stack[last.depth.0] = if last.last {
114                TreePart::Blank
115            } else {
116                TreePart::Line
117            };
118        }
119
120        // Make sure the stack has enough space, then add or modify another
121        // part into it.
122        self.stack.resize(params.depth.0 + 1, TreePart::Edge);
123        self.stack[params.depth.0] = if params.last {
124            TreePart::Corner
125        } else {
126            TreePart::Edge
127        };
128
129        self.last_params = Some(params);
130
131        // Return the tree parts as a slice of the stack.
132        //
133        // Ignore the first element here to prevent a ‘zeroth level’ from
134        // appearing before the very first directory. This level would
135        // join unrelated directories without connecting to anything:
136        //
137        //     with [0..]        with [1..]
138        //     ==========        ==========
139        //      ├── folder        folder
140        //      │   └── file       └── file
141        //      └── folder        folder
142        //          └── file       └──file
143        //
144        &self.stack[1..]
145    }
146}
147
148impl TreeParams {
149    pub fn new(depth: TreeDepth, last: bool) -> Self {
150        Self { depth, last }
151    }
152}
153
154impl TreeDepth {
155    pub fn root() -> Self {
156        Self(0)
157    }
158
159    pub fn deeper(self) -> Self {
160        Self(self.0 + 1)
161    }
162
163    /// Creates an iterator that, as well as yielding each value, yields a
164    /// `TreeParams` with the current depth and last flag filled in.
165    pub fn iterate_over<I, T>(self, inner: I) -> Iter<I>
166    where
167        I: ExactSizeIterator + Iterator<Item = T>,
168    {
169        Iter {
170            current_depth: self,
171            inner,
172        }
173    }
174}
175
176pub struct Iter<I> {
177    current_depth: TreeDepth,
178    inner: I,
179}
180
181impl<I, T> Iterator for Iter<I>
182where
183    I: ExactSizeIterator + Iterator<Item = T>,
184{
185    type Item = (TreeParams, T);
186
187    fn next(&mut self) -> Option<Self::Item> {
188        let t = self.inner.next()?;
189
190        // TODO: use exact_size_is_empty API soon
191        let params = TreeParams::new(self.current_depth, self.inner.len() == 0);
192        Some((params, t))
193    }
194}
195
196#[cfg(test)]
197mod trunk_test {
198    use super::*;
199
200    fn params(depth: usize, last: bool) -> TreeParams {
201        TreeParams::new(TreeDepth(depth), last)
202    }
203
204    #[rustfmt::skip]
205    #[test]
206    fn empty_at_first() {
207        let mut tt = TreeTrunk::default();
208        assert_eq!(tt.new_row(params(0, true)),  &[ ]);
209    }
210
211    #[rustfmt::skip]
212    #[test]
213    fn one_child() {
214        let mut tt = TreeTrunk::default();
215        assert_eq!(tt.new_row(params(0, true)),  &[ ]);
216        assert_eq!(tt.new_row(params(1, true)),  &[ TreePart::Corner ]);
217    }
218
219    #[rustfmt::skip]
220    #[test]
221    fn two_children() {
222        let mut tt = TreeTrunk::default();
223        assert_eq!(tt.new_row(params(0, true)),  &[ ]);
224        assert_eq!(tt.new_row(params(1, false)), &[ TreePart::Edge ]);
225        assert_eq!(tt.new_row(params(1, true)),  &[ TreePart::Corner ]);
226    }
227
228    #[rustfmt::skip]
229    #[test]
230    fn two_times_two_children() {
231        let mut tt = TreeTrunk::default();
232        assert_eq!(tt.new_row(params(0, false)), &[ ]);
233        assert_eq!(tt.new_row(params(1, false)), &[ TreePart::Edge ]);
234        assert_eq!(tt.new_row(params(1, true)),  &[ TreePart::Corner ]);
235
236        assert_eq!(tt.new_row(params(0, true)),  &[ ]);
237        assert_eq!(tt.new_row(params(1, false)), &[ TreePart::Edge ]);
238        assert_eq!(tt.new_row(params(1, true)),  &[ TreePart::Corner ]);
239    }
240
241    #[rustfmt::skip]
242    #[test]
243    fn two_times_two_nested_children() {
244        let mut tt = TreeTrunk::default();
245        assert_eq!(tt.new_row(params(0, true)),  &[ ]);
246
247        assert_eq!(tt.new_row(params(1, false)), &[ TreePart::Edge ]);
248        assert_eq!(tt.new_row(params(2, false)), &[ TreePart::Line, TreePart::Edge ]);
249        assert_eq!(tt.new_row(params(2, true)),  &[ TreePart::Line, TreePart::Corner ]);
250
251        assert_eq!(tt.new_row(params(1, true)),  &[ TreePart::Corner ]);
252        assert_eq!(tt.new_row(params(2, false)), &[ TreePart::Blank, TreePart::Edge ]);
253        assert_eq!(tt.new_row(params(2, true)),  &[ TreePart::Blank, TreePart::Corner ]);
254    }
255}
256
257#[cfg(test)]
258mod iter_test {
259    use super::*;
260
261    #[test]
262    fn test_iteration() {
263        let foos = &["first", "middle", "last"];
264        let mut iter = TreeDepth::root().iterate_over(foos.iter());
265
266        let next = iter.next().unwrap();
267        assert_eq!(&"first", next.1);
268        assert!(!next.0.last);
269
270        let next = iter.next().unwrap();
271        assert_eq!(&"middle", next.1);
272        assert!(!next.0.last);
273
274        let next = iter.next().unwrap();
275        assert_eq!(&"last", next.1);
276        assert!(next.0.last);
277
278        assert!(iter.next().is_none());
279    }
280
281    #[test]
282    fn test_empty() {
283        let nothing: &[usize] = &[];
284        let mut iter = TreeDepth::root().iterate_over(nothing.iter());
285        assert!(iter.next().is_none());
286    }
287}