1use crate::git_diff::{FileDiff, FileStatus};
2
3pub enum FileTreeNode {
4 Directory { name: String, children: Vec<FileTreeNode>, expanded: bool },
5 File { file_index: usize, name: String, status: FileStatus, additions: usize, deletions: usize },
6}
7
8#[derive(Debug, Clone)]
9pub struct FileTreeEntry {
10 pub depth: usize,
11 pub kind: FileTreeEntryKind,
12}
13
14#[derive(Debug, Clone)]
15pub enum FileTreeEntryKind {
16 Directory { name: String, expanded: bool },
17 File { file_index: usize, name: String, status: FileStatus, additions: usize, deletions: usize },
18}
19
20pub struct FileTree {
21 roots: Vec<FileTreeNode>,
22 selected_visible: usize,
23 cached_entries: Vec<FileTreeEntry>,
24}
25
26impl FileTree {
27 pub fn empty() -> Self {
28 Self { roots: Vec::new(), selected_visible: 0, cached_entries: Vec::new() }
29 }
30
31 pub fn from_files(files: &[FileDiff]) -> Self {
32 let mut roots: Vec<FileTreeNode> = Vec::new();
33
34 for (idx, file) in files.iter().enumerate() {
35 let parts: Vec<&str> = file.path.split('/').collect();
36 insert_into_tree(&mut roots, &parts, idx, file);
37 }
38
39 sort_tree(&mut roots);
40 compress_paths(&mut roots);
41
42 let mut tree = Self { roots, selected_visible: 0, cached_entries: Vec::new() };
43 tree.rebuild_visible_entries();
44 tree
45 }
46
47 pub fn select_file_index(&mut self, file_index: usize) {
48 if let Some(pos) = self
49 .cached_entries
50 .iter()
51 .position(|e| matches!(&e.kind, FileTreeEntryKind::File { file_index: fi, .. } if *fi == file_index))
52 {
53 self.selected_visible = pos;
54 }
55 }
56
57 pub fn visible_entries(&self) -> &[FileTreeEntry] {
58 &self.cached_entries
59 }
60
61 pub fn selected_visible(&self) -> usize {
62 self.selected_visible
63 }
64
65 pub fn selected_file_index(&self) -> Option<usize> {
66 self.cached_entries.get(self.selected_visible).and_then(|e| match &e.kind {
67 FileTreeEntryKind::File { file_index, .. } => Some(*file_index),
68 FileTreeEntryKind::Directory { .. } => None,
69 })
70 }
71
72 pub fn navigate(&mut self, delta: isize) {
73 if self.cached_entries.is_empty() {
74 return;
75 }
76 crate::components::wrap_selection(&mut self.selected_visible, self.cached_entries.len(), delta);
77 }
78
79 pub fn collapse_or_parent(&mut self) {
80 let Some(entry) = self.cached_entries.get(self.selected_visible) else {
81 return;
82 };
83
84 match &entry.kind {
85 FileTreeEntryKind::Directory { expanded: true, .. } => {
86 toggle_at(&mut self.roots, self.selected_visible);
87 self.rebuild_visible_entries();
88 }
89 _ => {
90 if let Some(parent_idx) = find_parent_dir(&self.cached_entries, self.selected_visible) {
91 self.selected_visible = parent_idx;
92 }
93 }
94 }
95 }
96
97 pub fn expand_or_enter(&mut self) -> bool {
98 let Some(entry) = self.cached_entries.get(self.selected_visible) else {
99 return false;
100 };
101
102 match &entry.kind {
103 FileTreeEntryKind::Directory { expanded: false, .. } => {
104 toggle_at(&mut self.roots, self.selected_visible);
105 self.rebuild_visible_entries();
106 false
107 }
108 FileTreeEntryKind::Directory { expanded: true, .. } => {
109 if self.selected_visible + 1 < self.cached_entries.len() {
110 self.selected_visible += 1;
111 }
112 false
113 }
114 FileTreeEntryKind::File { .. } => true,
115 }
116 }
117
118 fn rebuild_visible_entries(&mut self) {
119 let mut entries = Vec::new();
120 for node in &self.roots {
121 collect_visible(node, 0, &mut entries);
122 }
123 self.cached_entries = entries;
124 if self.cached_entries.is_empty() {
125 self.selected_visible = 0;
126 } else {
127 self.selected_visible = self.selected_visible.min(self.cached_entries.len() - 1);
128 }
129 }
130}
131
132fn insert_into_tree(nodes: &mut Vec<FileTreeNode>, parts: &[&str], file_index: usize, file: &FileDiff) {
133 if parts.len() == 1 {
134 nodes.push(FileTreeNode::File {
135 file_index,
136 name: parts[0].to_string(),
137 status: file.status,
138 additions: file.additions(),
139 deletions: file.deletions(),
140 });
141 return;
142 }
143
144 let dir_name = parts[0];
145 let existing = nodes.iter_mut().find(|n| matches!(n, FileTreeNode::Directory { name, .. } if name == dir_name));
146
147 if let Some(FileTreeNode::Directory { children, .. }) = existing {
148 insert_into_tree(children, &parts[1..], file_index, file);
149 } else {
150 let mut children = Vec::new();
151 insert_into_tree(&mut children, &parts[1..], file_index, file);
152 nodes.push(FileTreeNode::Directory { name: dir_name.to_string(), children, expanded: true });
153 }
154}
155
156fn sort_tree(nodes: &mut [FileTreeNode]) {
157 nodes.sort_by(|a, b| {
158 let a_is_dir = matches!(a, FileTreeNode::Directory { .. });
159 let b_is_dir = matches!(b, FileTreeNode::Directory { .. });
160 b_is_dir.cmp(&a_is_dir).then_with(|| node_name(a).cmp(node_name(b)))
161 });
162 for node in nodes.iter_mut() {
163 if let FileTreeNode::Directory { children, .. } = node {
164 sort_tree(children);
165 }
166 }
167}
168
169fn compress_paths(nodes: &mut [FileTreeNode]) {
170 for node in nodes.iter_mut() {
171 loop {
172 let should_compress = matches!(
173 node,
174 FileTreeNode::Directory { children, .. }
175 if children.len() == 1 && matches!(children[0], FileTreeNode::Directory { .. })
176 );
177 if !should_compress {
178 break;
179 }
180 if let FileTreeNode::Directory { name, children, .. } = node {
181 let child = children.remove(0);
182 if let FileTreeNode::Directory {
183 name: child_name,
184 children: child_children,
185 expanded: child_expanded,
186 } = child
187 {
188 *name = format!("{name}/{child_name}");
189 *children = child_children;
190 if let FileTreeNode::Directory { expanded, .. } = node {
191 *expanded = child_expanded;
192 }
193 }
194 }
195 }
196 if let FileTreeNode::Directory { children, .. } = node {
197 compress_paths(children);
198 }
199 }
200}
201
202fn node_name(node: &FileTreeNode) -> &str {
203 match node {
204 FileTreeNode::Directory { name, .. } | FileTreeNode::File { name, .. } => name,
205 }
206}
207
208fn collect_visible(node: &FileTreeNode, depth: usize, entries: &mut Vec<FileTreeEntry>) {
209 match node {
210 FileTreeNode::Directory { name, children, expanded } => {
211 entries.push(FileTreeEntry {
212 depth,
213 kind: FileTreeEntryKind::Directory { name: name.clone(), expanded: *expanded },
214 });
215 if *expanded {
216 for child in children {
217 collect_visible(child, depth + 1, entries);
218 }
219 }
220 }
221 FileTreeNode::File { file_index, name, status, additions, deletions } => {
222 entries.push(FileTreeEntry {
223 depth,
224 kind: FileTreeEntryKind::File {
225 file_index: *file_index,
226 name: name.clone(),
227 status: *status,
228 additions: *additions,
229 deletions: *deletions,
230 },
231 });
232 }
233 }
234}
235
236fn toggle_at(nodes: &mut [FileTreeNode], target_visible_idx: usize) {
237 let mut counter = 0;
238 toggle_at_inner(nodes, target_visible_idx, &mut counter);
239}
240
241fn toggle_at_inner(nodes: &mut [FileTreeNode], target: usize, counter: &mut usize) -> bool {
242 for node in nodes.iter_mut() {
243 if *counter == target {
244 if let FileTreeNode::Directory { expanded, .. } = node {
245 *expanded = !*expanded;
246 }
247 return true;
248 }
249 *counter += 1;
250 if let FileTreeNode::Directory { children, expanded: true, .. } = node
251 && toggle_at_inner(children, target, counter)
252 {
253 return true;
254 }
255 }
256 false
257}
258
259fn find_parent_dir(entries: &[FileTreeEntry], idx: usize) -> Option<usize> {
260 let current_depth = entries.get(idx)?.depth;
261 if current_depth == 0 {
262 return None;
263 }
264 (0..idx)
265 .rev()
266 .find(|&i| entries[i].depth < current_depth && matches!(entries[i].kind, FileTreeEntryKind::Directory { .. }))
267}
268
269#[cfg(test)]
270mod tests {
271 use super::*;
272 use crate::git_diff::{FileDiff, FileStatus, Hunk, PatchLine, PatchLineKind};
273
274 fn file(path: &str, status: FileStatus, additions: usize, deletions: usize) -> FileDiff {
275 let mut lines = Vec::new();
276 for i in 0..additions {
277 lines.push(PatchLine {
278 kind: PatchLineKind::Added,
279 text: format!("added {i}"),
280 old_line_no: None,
281 new_line_no: Some(i + 1),
282 });
283 }
284 for i in 0..deletions {
285 lines.push(PatchLine {
286 kind: PatchLineKind::Removed,
287 text: format!("removed {i}"),
288 old_line_no: Some(i + 1),
289 new_line_no: None,
290 });
291 }
292 FileDiff {
293 old_path: None,
294 path: path.to_string(),
295 status,
296 hunks: if lines.is_empty() {
297 vec![]
298 } else {
299 vec![Hunk {
300 header: "@@ -1 +1 @@".to_string(),
301 old_start: 1,
302 old_count: deletions,
303 new_start: 1,
304 new_count: additions,
305 lines,
306 }]
307 },
308 binary: false,
309 }
310 }
311
312 fn modified(path: &str) -> FileDiff {
313 file(path, FileStatus::Modified, 1, 1)
314 }
315
316 fn added(path: &str) -> FileDiff {
317 file(path, FileStatus::Added, 2, 0)
318 }
319
320 #[test]
321 fn from_files_groups_by_directory() {
322 let files = vec![modified("src/a.rs"), modified("src/b.rs"), modified("lib/c.rs")];
323 let tree = FileTree::from_files(&files);
324 let entries = tree.visible_entries();
325 assert_eq!(entries.len(), 5);
326 assert!(matches!(&entries[0].kind, FileTreeEntryKind::Directory { name, .. } if name == "lib"));
327 assert!(matches!(&entries[1].kind, FileTreeEntryKind::File { name, .. } if name == "c.rs"));
328 assert!(matches!(&entries[2].kind, FileTreeEntryKind::Directory { name, .. } if name == "src"));
329 assert!(matches!(&entries[3].kind, FileTreeEntryKind::File { name, .. } if name == "a.rs"));
330 assert!(matches!(&entries[4].kind, FileTreeEntryKind::File { name, .. } if name == "b.rs"));
331 }
332
333 #[test]
334 fn visible_entries_respects_collapse() {
335 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
336 let mut tree = FileTree::from_files(&files);
337 assert_eq!(tree.visible_entries().len(), 3);
338
339 tree.collapse_or_parent();
340 assert_eq!(tree.visible_entries().len(), 1);
341 }
342
343 #[test]
344 fn navigate_wraps() {
345 let files = vec![modified("a.rs"), modified("b.rs")];
346 let mut tree = FileTree::from_files(&files);
347 assert_eq!(tree.selected_visible, 0);
348 tree.navigate(1);
349 assert_eq!(tree.selected_visible, 1);
350 tree.navigate(1);
351 assert_eq!(tree.selected_visible, 0);
352 }
353
354 #[test]
355 fn collapse_or_parent_collapses_dir() {
356 let files = vec![modified("src/a.rs")];
357 let mut tree = FileTree::from_files(&files);
358 tree.selected_visible = 0;
359 assert_eq!(tree.visible_entries().len(), 2);
360 tree.collapse_or_parent();
361 assert_eq!(tree.visible_entries().len(), 1);
362 }
363
364 #[test]
365 fn collapse_or_parent_moves_to_parent_from_file() {
366 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
367 let mut tree = FileTree::from_files(&files);
368 tree.selected_visible = 1;
369 tree.collapse_or_parent();
370 assert_eq!(tree.selected_visible, 0);
371 }
372
373 #[test]
374 fn expand_or_enter_returns_true_for_file() {
375 let files = vec![modified("a.rs")];
376 let mut tree = FileTree::from_files(&files);
377 assert!(tree.expand_or_enter());
378 }
379
380 #[test]
381 fn expand_or_enter_expands_collapsed_dir() {
382 let files = vec![modified("src/a.rs")];
383 let mut tree = FileTree::from_files(&files);
384 tree.collapse_or_parent();
385 assert_eq!(tree.visible_entries().len(), 1);
386 let result = tree.expand_or_enter();
387 assert!(!result);
388 assert_eq!(tree.visible_entries().len(), 2);
389 }
390
391 #[test]
392 fn path_compression_for_single_child_dirs() {
393 let files = vec![modified("src/deep/nested/file.rs")];
394 let tree = FileTree::from_files(&files);
395 let entries = tree.visible_entries();
396 assert_eq!(entries.len(), 2);
397 match &entries[0].kind {
398 FileTreeEntryKind::Directory { name, .. } => {
399 assert_eq!(name, "src/deep/nested");
400 }
401 FileTreeEntryKind::File { .. } => panic!("expected directory"),
402 }
403 }
404
405 #[test]
406 fn selected_file_index_returns_none_for_dir() {
407 let files = vec![modified("src/a.rs")];
408 let tree = FileTree::from_files(&files);
409 assert!(tree.selected_file_index().is_none());
410 }
411
412 #[test]
413 fn selected_file_index_returns_index_for_file() {
414 let files = vec![modified("a.rs")];
415 let tree = FileTree::from_files(&files);
416 assert_eq!(tree.selected_file_index(), Some(0));
417 }
418
419 #[test]
420 fn flat_files_no_grouping() {
421 let files = vec![modified("a.rs"), added("b.rs")];
422 let tree = FileTree::from_files(&files);
423 let entries = tree.visible_entries();
424 assert_eq!(entries.len(), 2);
425 assert!(matches!(&entries[0].kind, FileTreeEntryKind::File { name, .. } if name == "a.rs"));
426 assert!(matches!(&entries[1].kind, FileTreeEntryKind::File { name, .. } if name == "b.rs"));
427 }
428
429 #[test]
430 fn selected_visible_clamped_after_collapse() {
431 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
432 let mut tree = FileTree::from_files(&files);
433 tree.selected_visible = 2;
434 tree.collapse_or_parent();
435 assert!(tree.selected_visible < tree.visible_entries().len());
436 }
437
438 #[test]
439 fn expand_already_expanded_dir_moves_to_first_child() {
440 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
441 let mut tree = FileTree::from_files(&files);
442 tree.selected_visible = 0;
443 let result = tree.expand_or_enter();
444 assert!(!result);
445 assert_eq!(tree.selected_visible, 1);
446 }
447}