1 use crate::tree_config::{tree_config, TreeConfig};
2 use std::cmp::max;
3 use std::sync::{Arc, Mutex};
4 
5 /// Tree that holds `text` for the current leaf and a list of `children` that are the branches.
6 #[derive(Debug)]
7 pub struct Tree {
8     pub text: Option<String>,
9     pub children: Vec<Tree>,
10 }
11 
12 /// Position of the element relative to its siblings
13 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
14 pub enum Position {
15     Inside,
16     First,
17     Last,
18     Only,
19 }
20 
21 impl Tree {
22     /// Create a new tree with some optional text.
new(text: Option<&str>) -> Tree23     pub fn new(text: Option<&str>) -> Tree {
24         Tree {
25             text: text.map(|x| x.to_string()),
26             children: Vec::new(),
27         }
28     }
29 
30     /// Navigate to the branch at the given `path` relative to this tree.
31     /// If a valid branch is found by following the path, it is returned.
at_mut(&mut self, path: &[usize]) -> Option<&mut Tree>32     pub fn at_mut(&mut self, path: &[usize]) -> Option<&mut Tree> {
33         match path.first() {
34             Some(&i) => match self.children.get_mut(i) {
35                 Some(x) => x.at_mut(&path[1..]),
36                 _ => None,
37             },
38             _ => Some(self),
39         }
40     }
41 
42     /// "Render" this tree as a list of `String`s.
43     /// Each string represents a line in the tree.
44     /// `does_continue` is a bool for each column indicating whether the tree continues.
lines( &self, does_continue: &Vec<bool>, index: usize, pool_size: usize, config: &TreeConfig, ) -> Vec<String>45     pub fn lines(
46         &self,
47         does_continue: &Vec<bool>,
48         index: usize,
49         pool_size: usize,
50         config: &TreeConfig,
51     ) -> Vec<String> {
52         let does_continue = if config.show_first_level && does_continue.is_empty() {
53             vec![true]
54         } else {
55             does_continue.clone()
56         };
57         let position = match index {
58             _ if pool_size == 1 => Position::Only,
59             _ if (index + 1) == pool_size => Position::Last,
60             0 => Position::First,
61             _ => Position::Inside,
62         };
63         let mut next_continue = does_continue.clone();
64         next_continue.push(match position {
65             Position::Inside | Position::First => true,
66             Position::Last | Position::Only => false,
67         });
68 
69         let mut txt = String::new();
70         let pad: String;
71         if does_continue.len() > 1 {
72             for &i in &does_continue[2..] {
73                 txt.push_str(&format!(
74                     "{}{:indent$}",
75                     if i { config.symbols.continued } else { " " },
76                     "",
77                     indent = max(config.indent, 1) - 1
78                 ));
79             }
80             pad = txt.clone();
81             let branch_size = max(config.indent, 2usize) - 2;
82             let branch = match config.symbols.branch.len() {
83                 0 => "-".repeat(branch_size),
84                 1 => config.symbols.branch.repeat(branch_size),
85                 _n => config
86                     .symbols
87                     .branch
88                     .repeat(branch_size)
89                     .chars()
90                     .take(branch_size)
91                     .collect::<String>(),
92             };
93 
94             let is_multiline = self
95                 .text
96                 .as_ref()
97                 .map(|x| x.contains("\n"))
98                 .unwrap_or(false);
99 
100             let first_leaf = match (is_multiline, config.symbols.multiline_first) {
101                 (true, Some(x)) => x,
102                 _ => config.symbols.leaf,
103             };
104             txt.push_str(&format!(
105                 "{}{}{}",
106                 match position {
107                     Position::Only => config.symbols.join_only,
108                     Position::First => config.symbols.join_first,
109                     Position::Last => config.symbols.join_last,
110                     Position::Inside => config.symbols.join_inner,
111                 },
112                 branch,
113                 first_leaf,
114             ));
115 
116             let s = match &self.text {
117                 Some(x) => match is_multiline {
118                     true => format!(
119                         "{}",
120                         x.replace(
121                             "\n",
122                             &format!(
123                                 "\n{}{}{}{}",
124                                 &pad,
125                                 match position {
126                                     Position::Only | Position::Last =>
127                                         " ".repeat(config.symbols.continued.chars().count()),
128                                     _ => config.symbols.continued.to_string(),
129                                 },
130                                 " ".repeat(branch_size),
131                                 match &config.symbols.multiline_continued {
132                                     Some(multi) => multi.to_string(),
133                                     _ => " ".repeat(first_leaf.chars().count()),
134                                 }
135                             ),
136                         )
137                     ),
138                     false => x.clone(),
139                 },
140                 _ => String::new(),
141             };
142             txt.push_str(&s);
143         } else {
144             if let Some(x) = &self.text {
145                 txt.push_str(&x);
146             }
147         }
148         let mut ret = vec![txt];
149         for (index, x) in self.children.iter().enumerate() {
150             for line in x.lines(&next_continue, index, self.children.len(), config) {
151                 ret.push(line);
152             }
153         }
154         ret
155     }
156 }
157 
158 /// Holds the current state of the tree, including the path to the branch.
159 /// Multiple trees may point to the same data.
160 #[derive(Debug, Clone)]
161 pub(crate) struct TreeBuilderBase {
162     data: Arc<Mutex<Tree>>,
163     path: Vec<usize>,
164     dive_count: usize,
165     config: Option<TreeConfig>,
166     is_enabled: bool,
167 }
168 
169 impl TreeBuilderBase {
170     /// Create a new state
new() -> TreeBuilderBase171     pub fn new() -> TreeBuilderBase {
172         TreeBuilderBase {
173             data: Arc::new(Mutex::new(Tree::new(None))),
174             path: vec![],
175             dive_count: 1,
176             config: None,
177             is_enabled: true,
178         }
179     }
180 
set_enabled(&mut self, enabled: bool)181     pub fn set_enabled(&mut self, enabled: bool) {
182         self.is_enabled = enabled;
183     }
is_enabled(&self) -> bool184     pub fn is_enabled(&self) -> bool {
185         self.is_enabled
186     }
187 
add_leaf(&mut self, text: &str)188     pub fn add_leaf(&mut self, text: &str) {
189         let &dive_count = &self.dive_count;
190         if dive_count > 0 {
191             for i in 0..dive_count {
192                 let mut n = 0;
193                 if let Some(x) = self.data.lock().unwrap().at_mut(&self.path) {
194                     x.children.push(Tree::new(if i == max(1, dive_count) - 1 {
195                         Some(&text)
196                     } else {
197                         None
198                     }));
199                     n = x.children.len() - 1;
200                 }
201                 self.path.push(n);
202             }
203             self.dive_count = 0;
204         } else {
205             if let Some(x) = self
206                 .data
207                 .lock()
208                 .unwrap()
209                 .at_mut(&self.path[..max(1, self.path.len()) - 1])
210             {
211                 x.children.push(Tree::new(Some(&text)));
212                 let n = match self.path.last() {
213                     Some(&x) => x + 1,
214                     _ => 0,
215                 };
216                 self.path.last_mut().map(|x| *x = n);
217             }
218         }
219     }
220 
set_config_override(&mut self, config: Option<TreeConfig>)221     pub fn set_config_override(&mut self, config: Option<TreeConfig>) {
222         self.config = config;
223     }
224 
config_override(&self) -> &Option<TreeConfig>225     pub fn config_override(&self) -> &Option<TreeConfig> {
226         &self.config
227     }
config_override_mut(&mut self) -> &mut Option<TreeConfig>228     pub fn config_override_mut(&mut self) -> &mut Option<TreeConfig> {
229         &mut self.config
230     }
231 
enter(&mut self)232     pub fn enter(&mut self) {
233         self.dive_count += 1;
234     }
235 
236     /// Try stepping up to the parent tree branch.
237     /// Returns false if already at the top branch.
exit(&mut self) -> bool238     pub fn exit(&mut self) -> bool {
239         if self.dive_count > 0 {
240             self.dive_count -= 1;
241             true
242         } else {
243             if self.path.len() > 1 {
244                 self.path.pop();
245                 true
246             } else {
247                 false
248             }
249         }
250     }
251 
depth(&self) -> usize252     pub fn depth(&self) -> usize {
253         max(1, self.path.len() + self.dive_count) - 1
254     }
255 
peek_print(&self)256     pub fn peek_print(&self) {
257         println!("{}", self.peek_string());
258     }
259 
print(&mut self)260     pub fn print(&mut self) {
261         self.peek_print();
262         self.clear();
263     }
clear(&mut self)264     pub fn clear(&mut self) {
265         *self = Self::new();
266     }
267 
string(&mut self) -> String268     pub fn string(&mut self) -> String {
269         let s = self.peek_string();
270         self.clear();
271         s
272     }
273 
peek_string(&self) -> String274     pub fn peek_string(&self) -> String {
275         let config = self
276             .config_override()
277             .clone()
278             .unwrap_or_else(|| tree_config().clone());
279         (&self.data.lock().unwrap().lines(&vec![], 0, 1, &config)[1..]).join("\n")
280     }
281 }
282