Skip to main content

rusticity_term/ui/
tree.rs

1use ratatui::prelude::*;
2use ratatui::widgets::Cell;
3use std::collections::{HashMap, HashSet};
4
5pub const CURSOR_EXPANDED: &str = "▼";
6pub const CURSOR_COLLAPSED: &str = "▶";
7
8/// Trait for items that can be displayed in a tree structure
9pub trait TreeItem {
10    /// Unique identifier for this item
11    fn id(&self) -> &str;
12
13    /// Display name for this item
14    fn display_name(&self) -> &str;
15
16    /// Whether this item can have children (is expandable)
17    fn is_expandable(&self) -> bool;
18
19    /// Icon to display for this item
20    fn icon(&self) -> &str {
21        if self.is_expandable() {
22            "📁"
23        } else {
24            "📄"
25        }
26    }
27}
28
29/// Generic tree renderer that handles hierarchical display with expand/collapse
30pub struct TreeRenderer<'a, T: TreeItem> {
31    /// All items at the current level
32    pub items: &'a [T],
33
34    /// Set of expanded item IDs
35    pub expanded_ids: &'a HashSet<String>,
36
37    /// Map of parent ID to children
38    pub children_map: &'a HashMap<String, Vec<T>>,
39
40    /// Currently selected row index
41    pub selected_row: usize,
42
43    /// Starting row index for this render
44    pub start_row: usize,
45}
46
47impl<'a, T: TreeItem> TreeRenderer<'a, T> {
48    pub fn new(
49        items: &'a [T],
50        expanded_ids: &'a HashSet<String>,
51        children_map: &'a HashMap<String, Vec<T>>,
52        selected_row: usize,
53        start_row: usize,
54    ) -> Self {
55        Self {
56            items,
57            expanded_ids,
58            children_map,
59            selected_row,
60            start_row,
61        }
62    }
63
64    /// Render tree items recursively, returning rows with tree structure
65    pub fn render<F>(&self, mut render_cell: F) -> Vec<(Vec<Cell<'a>>, Style)>
66    where
67        F: FnMut(&T, &str) -> Vec<Cell<'a>>,
68    {
69        let mut result = Vec::new();
70        let mut current_row = self.start_row;
71
72        self.render_recursive(
73            self.items,
74            &mut current_row,
75            &mut result,
76            "",
77            &[],
78            &mut render_cell,
79        );
80
81        result
82    }
83
84    fn render_recursive<F>(
85        &self,
86        items: &[T],
87        current_row: &mut usize,
88        result: &mut Vec<(Vec<Cell<'a>>, Style)>,
89        _parent_id: &str,
90        is_last: &[bool],
91        render_cell: &mut F,
92    ) where
93        F: FnMut(&T, &str) -> Vec<Cell<'a>>,
94    {
95        for (idx, item) in items.iter().enumerate() {
96            let is_last_item = idx == items.len() - 1;
97            let is_expanded = self.expanded_ids.contains(item.id());
98
99            // Build prefix with tree characters
100            let mut prefix = String::new();
101            for &last in is_last.iter() {
102                prefix.push_str(if last { "  " } else { "│ " });
103            }
104
105            let tree_char = if is_last_item { "╰─" } else { "├─" };
106            let expand_char = if item.is_expandable() {
107                if is_expanded {
108                    CURSOR_EXPANDED
109                } else {
110                    CURSOR_COLLAPSED
111                }
112            } else {
113                ""
114            };
115
116            let icon = item.icon();
117            let tree_prefix = format!("{}{}{} {} ", prefix, tree_char, expand_char, icon);
118
119            // Determine style based on selection
120            let style = if *current_row == self.selected_row {
121                Style::default().bg(Color::DarkGray)
122            } else {
123                Style::default()
124            };
125
126            // Render cells for this item
127            let cells = render_cell(item, &tree_prefix);
128            result.push((cells, style));
129            *current_row += 1;
130
131            // Recursively render children if expanded
132            if item.is_expandable() && is_expanded {
133                if let Some(children) = self.children_map.get(item.id()) {
134                    let mut new_is_last = is_last.to_vec();
135                    new_is_last.push(is_last_item);
136                    self.render_recursive(
137                        children,
138                        current_row,
139                        result,
140                        item.id(),
141                        &new_is_last,
142                        render_cell,
143                    );
144                }
145            }
146        }
147    }
148
149    /// Count total visible rows including expanded children
150    pub fn count_visible_rows(
151        items: &[T],
152        expanded_ids: &HashSet<String>,
153        children_map: &HashMap<String, Vec<T>>,
154    ) -> usize {
155        let mut count = 0;
156        for item in items {
157            count += 1;
158            if item.is_expandable() && expanded_ids.contains(item.id()) {
159                if let Some(children) = children_map.get(item.id()) {
160                    count += Self::count_visible_rows(children, expanded_ids, children_map);
161                }
162            }
163        }
164        count
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    #[derive(Debug, Clone)]
173    struct TestItem {
174        id: String,
175        name: String,
176        is_folder: bool,
177    }
178
179    impl TreeItem for TestItem {
180        fn id(&self) -> &str {
181            &self.id
182        }
183
184        fn display_name(&self) -> &str {
185            &self.name
186        }
187
188        fn is_expandable(&self) -> bool {
189            self.is_folder
190        }
191    }
192
193    #[test]
194    fn test_count_visible_rows_no_expansion() {
195        let items = vec![
196            TestItem {
197                id: "1".to_string(),
198                name: "folder1".to_string(),
199                is_folder: true,
200            },
201            TestItem {
202                id: "2".to_string(),
203                name: "file1".to_string(),
204                is_folder: false,
205            },
206        ];
207
208        let expanded = HashSet::new();
209        let children = HashMap::new();
210
211        let count = TreeRenderer::count_visible_rows(&items, &expanded, &children);
212        assert_eq!(count, 2);
213    }
214
215    #[test]
216    fn test_count_visible_rows_with_expansion() {
217        let items = vec![TestItem {
218            id: "1".to_string(),
219            name: "folder1".to_string(),
220            is_folder: true,
221        }];
222
223        let mut expanded = HashSet::new();
224        expanded.insert("1".to_string());
225
226        let mut children = HashMap::new();
227        children.insert(
228            "1".to_string(),
229            vec![
230                TestItem {
231                    id: "1/a".to_string(),
232                    name: "file_a".to_string(),
233                    is_folder: false,
234                },
235                TestItem {
236                    id: "1/b".to_string(),
237                    name: "file_b".to_string(),
238                    is_folder: false,
239                },
240            ],
241        );
242
243        let count = TreeRenderer::count_visible_rows(&items, &expanded, &children);
244        assert_eq!(count, 3); // 1 folder + 2 children
245    }
246
247    #[test]
248    fn test_count_visible_rows_nested_expansion() {
249        let items = vec![TestItem {
250            id: "1".to_string(),
251            name: "folder1".to_string(),
252            is_folder: true,
253        }];
254
255        let mut expanded = HashSet::new();
256        expanded.insert("1".to_string());
257        expanded.insert("1/a".to_string());
258
259        let mut children = HashMap::new();
260        children.insert(
261            "1".to_string(),
262            vec![TestItem {
263                id: "1/a".to_string(),
264                name: "folder_a".to_string(),
265                is_folder: true,
266            }],
267        );
268        children.insert(
269            "1/a".to_string(),
270            vec![TestItem {
271                id: "1/a/x".to_string(),
272                name: "file_x".to_string(),
273                is_folder: false,
274            }],
275        );
276
277        let count = TreeRenderer::count_visible_rows(&items, &expanded, &children);
278        assert_eq!(count, 3); // 1 folder + 1 subfolder + 1 file
279    }
280
281    #[test]
282    fn test_tree_item_default_icons() {
283        let folder = TestItem {
284            id: "1".to_string(),
285            name: "folder".to_string(),
286            is_folder: true,
287        };
288        let file = TestItem {
289            id: "2".to_string(),
290            name: "file".to_string(),
291            is_folder: false,
292        };
293
294        assert_eq!(folder.icon(), "📁");
295        assert_eq!(file.icon(), "📄");
296    }
297}