pars_core/util/tree/
print.rs

1use std::collections::VecDeque;
2
3use anyhow::Result;
4use colored::Colorize;
5use log::debug;
6
7use super::{DirTree, NodeType, TreePrintConfig};
8
9impl DirTree<'_> {
10    pub fn print_tree(&self, config: &TreePrintConfig) -> Result<String> {
11        debug!("Start to print tree:\n{:?}", self.map);
12        let mut tree_builder = String::new(); //TODO(Vollate): we should use other structure for building string(huge dir case)
13        let mut node_stack = VecDeque::<(usize, usize)>::new();
14        node_stack.push_back((self.root, 0));
15        let mut level_stack = VecDeque::<bool>::new();
16        while let Some((parent_idx, vec_idx)) = node_stack.pop_back() {
17            level_stack.pop_back();
18            if vec_idx >= self.map[parent_idx].children.len() {
19                continue;
20            }
21            let child_idx = self.map[parent_idx].children[vec_idx];
22            let child = &self.map[child_idx];
23
24            let is_local_last = vec_idx + 1 == self.map[parent_idx].children.len();
25            level_stack.push_back(is_local_last);
26            if !child.visible {
27                node_stack.push_back((parent_idx, vec_idx + 1));
28                continue;
29            }
30            let is_local_last = vec_idx + 1 == self.map[parent_idx].children.len();
31            if level_stack.len() > 1 {
32                for level_status in level_stack.iter().take(level_stack.len() - 1) {
33                    if *level_status {
34                        tree_builder.push_str("    ");
35                    } else {
36                        tree_builder.push_str("│   ");
37                    }
38                }
39            }
40
41            if is_local_last {
42                tree_builder.push_str("└── ");
43            } else {
44                tree_builder.push_str("├── ");
45            }
46
47            node_stack.push_back((parent_idx, vec_idx + 1));
48
49            match child.node_type {
50                NodeType::Symlink => {
51                    let mut line = if let Some(color) = config.symbol_color {
52                        format!(
53                            "{} -> {}",
54                            child.name.color(color),
55                            child.symlink_target.as_ref().unwrap()
56                        )
57                    } else {
58                        format!("{} -> {}", child.name, child.symlink_target.as_ref().unwrap())
59                    };
60                    if child.is_recursive {
61                        line = format!("{line} [recursive, not followed]\n");
62                    } else {
63                        line.push('\n');
64                    }
65                    tree_builder.push_str(line.as_str());
66                    node_stack.push_back((child_idx, 0));
67                    level_stack.push_back(false);
68                }
69                NodeType::Dir => {
70                    let line = if let Some(color) = config.dir_color {
71                        format!("{}\n", child.name.color(color))
72                    } else {
73                        format!("{}\n", child.name)
74                    };
75                    tree_builder.push_str(line.as_str());
76                    node_stack.push_back((child_idx, 0));
77                    level_stack.push_back(false);
78                }
79                NodeType::File => {
80                    let line = if let Some(color) = config.file_color {
81                        format!("{}\n", child.name.color(color))
82                    } else {
83                        format!("{}\n", child.name)
84                    };
85                    tree_builder.push_str(line.as_str());
86                }
87                NodeType::Other => {
88                    let line = if let Some(color) = config.tree_color {
89                        format!("{}\n", child.name.color(color))
90                    } else {
91                        format!("{}\n", child.name)
92                    };
93                    tree_builder.push_str(line.as_str());
94                }
95                NodeType::Invalid => {
96                    panic!("Should not have any invalid type");
97                }
98            }
99        }
100        tree_builder.pop();
101        Ok(tree_builder)
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use bumpalo::Bump;
108    use pretty_assertions::assert_eq;
109    use regex::Regex;
110
111    use super::*;
112    use crate::util::defer::cleanup;
113    use crate::util::fs_util::create_symlink;
114    use crate::util::test_util::{create_dir_structure, gen_unique_temp_dir, log_test};
115    use crate::util::tree::{FilterType, TreeConfig};
116
117    #[test]
118    fn tree_normal_case() {
119        let no_color_print = TreePrintConfig {
120            dir_color: None,
121            file_color: None,
122            symbol_color: None,
123            tree_color: None,
124        };
125        // Create directory structure as below:
126        // ├── dir1
127        // │   ├── file1
128        // │   └── file2
129        // ├── dir2
130        // │   ├── file3
131        // │   └── file4
132        // └── dir3
133        //     ├── dir4
134        //     │   └── dir5
135        //     └── file5
136        let (_tmp_dir, root) = gen_unique_temp_dir();
137        let structure: &[(Option<&str>, &[&str])] = &[
138            (Some("dir1"), &["file1", "file2"][..]),
139            (Some("dir2"), &["file3", "file4"][..]),
140            (Some("dir3"), &["file5"][..]),
141            (Some("dir3/dir4"), &[][..]),
142            (Some("dir3/dir4/dir5"), &[][..]),
143        ];
144        create_dir_structure(&root, structure);
145        cleanup!(
146            {
147                let config = TreeConfig {
148                    root: &root,
149                    target: "",
150                    filter_type: FilterType::Disable,
151                    filters: Vec::new(),
152                };
153                let bump = Bump::new();
154                let tree = DirTree::new(&config, &bump).unwrap();
155                let result = tree.print_tree(&no_color_print).unwrap();
156                assert_eq!(
157                    result,
158                    r#"├── dir1
159│   ├── file1
160│   └── file2
161├── dir2
162│   ├── file3
163│   └── file4
164└── dir3
165    ├── dir4
166    │   └── dir5
167    └── file5"#
168                );
169            },
170            {}
171        )
172    }
173
174    #[test]
175    fn test_filtered_case() {
176        let no_color_print = TreePrintConfig {
177            dir_color: None,
178            file_color: None,
179            symbol_color: None,
180            tree_color: None,
181        };
182        // Create directory structure as below:
183        // ├── dir1
184        // │   └── file1
185        // ├── dir2
186        // │   └── file2
187        // └── dir3
188        let (_tmp_dir, root) = gen_unique_temp_dir();
189        let structure: &[(Option<&str>, &[&str])] = &[
190            (Some("dir1"), &["file1"][..]),
191            (Some("dir2"), &["file2"][..]),
192            (Some("dir3"), &[][..]),
193        ];
194        create_dir_structure(&root, structure);
195
196        // This case, only dir2 and file1 should be filtered
197        cleanup!(
198            {
199                let mut config = TreeConfig {
200                    root: &root,
201                    target: "",
202                    filter_type: FilterType::Exclude,
203                    filters: vec![Regex::new(r"dir2").unwrap(), Regex::new(r"file1").unwrap()],
204                };
205
206                let bump = Bump::new();
207                let tree = DirTree::new(&config, &bump).unwrap();
208                let result = tree.print_tree(&no_color_print).unwrap();
209                assert_eq!(
210                    result,
211                    r#"├── dir1
212└── dir3"#
213                );
214
215                // This case, nothing should be filtered
216                config.filters = Vec::new();
217                config.filter_type = FilterType::Disable;
218                let tree = DirTree::new(&config, &bump).unwrap();
219                let result = tree.print_tree(&no_color_print).unwrap();
220                assert_eq!(
221                    result,
222                    r#"├── dir1
223│   └── file1
224├── dir2
225│   └── file2
226└── dir3"#
227                );
228
229                // This case, everything should be filtered
230                config.filter_type = FilterType::Exclude;
231                config.filters = vec![
232                    Regex::new(r"dir1").unwrap(),
233                    Regex::new(r"dir2").unwrap(),
234                    Regex::new(r"dir3").unwrap(),
235                ];
236                let tree = DirTree::new(&config, &bump).unwrap();
237                let result = tree.print_tree(&no_color_print).unwrap();
238                assert_eq!(result, "");
239
240                // Now, test white-list mode
241                config.filter_type = FilterType::Include;
242                config.filters = vec![Regex::new(r"dir1").unwrap(), Regex::new(r"file2").unwrap()];
243                let tree = DirTree::new(&config, &bump).unwrap();
244                let result = tree.print_tree(&no_color_print).unwrap();
245                assert_eq!(
246                    result,
247                    r#"├── dir1
248│   └── file1
249└── dir2
250    └── file2"#
251                );
252            },
253            {}
254        )
255    }
256
257    #[test]
258    fn test_symbolic_link() {
259        let no_color_print = TreePrintConfig {
260            dir_color: None,
261            file_color: None,
262            symbol_color: None,
263            tree_color: None,
264        };
265        // Create directory structure as below:
266        // root1
267        // ├── dir1
268        // │   └── file1
269        // └── dir2 -> root2
270        //
271        // --------------------------------
272        //
273        // root2
274        // ├── dir3
275        // │   └── file2
276        // └── dir4
277
278        let (_tmp_dir, root1) = gen_unique_temp_dir();
279        let (_tmp_dir, root2) = gen_unique_temp_dir();
280        let structure1: &[(Option<&str>, &[&str])] = &[(Some("dir1"), &["file1"][..])];
281        let structure2: &[(Option<&str>, &[&str])] =
282            &[(Some("dir3"), &["file2"][..]), (Some("dir4"), &[][..])];
283        create_dir_structure(&root1, structure1);
284        create_dir_structure(&root2, structure2);
285        create_symlink(&root2, &root1.join("dir2")).unwrap();
286
287        cleanup!(
288            {
289                let config = TreeConfig {
290                    root: &root1,
291                    target: "",
292                    filter_type: FilterType::Exclude,
293                    filters: Vec::new(),
294                };
295                let bump = Bump::new();
296                let tree = DirTree::new(&config, &bump).unwrap();
297                let result = tree.print_tree(&no_color_print).unwrap();
298                assert_eq!(
299                    result,
300                    format!(
301                        r#"├── dir1
302│   └── file1
303└── dir2 -> {}
304    ├── dir3
305    │   └── file2
306    └── dir4"#,
307                        root2.to_str().unwrap()
308                    )
309                );
310            },
311            {}
312        );
313
314        // Create directory structure as below:
315        // root1
316        // ├── dir1
317        // │   └── file1
318        // ├── dir2 -> root2
319        // ├── dir3
320        // ├── dir4
321        // └── file114514
322        //
323        // --------------------------------
324        //
325        // root2
326        // ├── dir3
327        // │   └── file2
328        // └── dir4
329        let (_tmp_dir, root1) = gen_unique_temp_dir();
330        let (_tmp_dir, root2) = gen_unique_temp_dir();
331        let structure1: &[(Option<&str>, &[&str])] = &[
332            (Some("dir1"), &["file1"][..]),
333            (Some("dir3"), &[][..]),
334            (Some("dir4"), &[][..]),
335            (None, &["file114514"][..]),
336        ];
337        let structure2: &[(Option<&str>, &[&str])] =
338            &[(Some("dir3"), &["file2"][..]), (Some("dir4"), &[][..])];
339        create_dir_structure(&root1, structure1);
340        create_dir_structure(&root2, structure2);
341        create_symlink(&root2, &root1.join("dir2")).unwrap();
342        cleanup!(
343            {
344                let config = TreeConfig {
345                    root: &root1,
346                    target: "",
347                    filter_type: FilterType::Exclude,
348                    filters: Vec::new(),
349                };
350                let bump = Bump::new();
351                let tree = DirTree::new(&config, &bump).unwrap();
352                let result = tree.print_tree(&no_color_print).unwrap();
353                assert_eq!(
354                    result,
355                    format!(
356                        r#"├── dir1
357│   └── file1
358├── dir2 -> {}
359│   ├── dir3
360│   │   └── file2
361│   └── dir4
362├── dir3
363├── dir4
364└── file114514"#,
365                        root2.to_str().unwrap()
366                    )
367                );
368            },
369            {}
370        )
371    }
372
373    #[test]
374    fn test_recursive_detection() {
375        let no_color_print = TreePrintConfig {
376            dir_color: None,
377            file_color: None,
378            symbol_color: None,
379            tree_color: None,
380        };
381
382        // Test direct recursion
383        {
384            let (_tmp_dir, root) = gen_unique_temp_dir();
385            let structure: &[(Option<&str>, &[&str])] =
386                &[(Some("dir1"), &["file1.txt"][..]), (None, &["root_file.txt"][..])];
387            create_dir_structure(&root, structure);
388
389            cleanup!(
390                {
391                    // Create self-referential symlink
392                    create_symlink(&root.join("dir1"), &root.join("dir1/loop")).unwrap();
393
394                    let config = TreeConfig {
395                        root: &root,
396                        target: "",
397                        filter_type: FilterType::Exclude,
398                        filters: Vec::new(),
399                    };
400                    let bump = Bump::new();
401                    let tree = DirTree::new(&config, &bump).unwrap();
402                    let result = tree.print_tree(&no_color_print).unwrap();
403                    log_test!("{}", result);
404
405                    assert_eq!(
406                        result,
407                        format!(
408                            r#"├── dir1
409│   ├── file1.txt
410│   └── loop -> {} [recursive, not followed]
411└── root_file.txt"#,
412                            root.join("dir1").to_str().unwrap()
413                        )
414                    );
415                },
416                {}
417            );
418        }
419
420        // Test indirect recursion
421        {
422            let (_tmp_dir, root) = gen_unique_temp_dir();
423            let structure: &[(Option<&str>, &[&str])] = &[
424                (Some("dirA"), &["fileA.txt"][..]),
425                (Some("dirB"), &["fileB.txt"][..]),
426                (Some("dirC"), &["fileC.txt"][..]),
427            ];
428            create_dir_structure(&root, structure);
429
430            cleanup!(
431                {
432                    // Create circular reference: A -> B -> C -> A
433                    create_symlink(&root.join("dirB"), &root.join("dirA/linkB")).unwrap();
434                    create_symlink(&root.join("dirC"), &root.join("dirB/linkC")).unwrap();
435                    create_symlink(&root.join("dirA"), &root.join("dirC/linkA")).unwrap();
436
437                    let config = TreeConfig {
438                        root: &root,
439                        target: "",
440                        filter_type: FilterType::Exclude,
441                        filters: Vec::new(),
442                    };
443                    let bump = Bump::new();
444                    let tree = DirTree::new(&config, &bump).unwrap();
445                    let result = tree.print_tree(&no_color_print).unwrap();
446                    log_test!("{}", result);
447                    assert_eq!(
448                        result,
449                        format!(
450                            r#"├── dirA
451│   ├── fileA.txt
452│   └── linkB -> {}
453│       ├── fileB.txt
454│       └── linkC -> {}
455│           ├── fileC.txt
456│           └── linkA -> {} [recursive, not followed]
457├── dirB
458│   ├── fileB.txt
459│   └── linkC -> {} [recursive, not followed]
460└── dirC
461    ├── fileC.txt
462    └── linkA -> {} [recursive, not followed]"#,
463                            root.join("dirB").to_str().unwrap(),
464                            root.join("dirC").to_str().unwrap(),
465                            root.join("dirA").to_str().unwrap(),
466                            root.join("dirC").to_str().unwrap(),
467                            root.join("dirA").to_str().unwrap(),
468                        )
469                    );
470                },
471                {}
472            );
473        }
474    }
475}
pars_core/util/tree/print.rs

pars_core/util/tree/
print.rs
