1use crate::git_diff::{FileDiff, FileStatus, StageState};
2
3pub enum FileTreeNode {
4 Directory { name: String, children: Vec<FileTreeNode>, expanded: bool },
5 File { file_index: usize, name: String, status: FileStatus, staged: StageState, 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, staged: StageState, 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 self.selected_visible = self.selected_visible.saturating_add_signed(delta).min(self.cached_entries.len() - 1);
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 staged: file.staged,
139 additions: file.additions(),
140 deletions: file.deletions(),
141 });
142 return;
143 }
144
145 let dir_name = parts[0];
146 let existing = nodes.iter_mut().find(|n| matches!(n, FileTreeNode::Directory { name, .. } if name == dir_name));
147
148 if let Some(FileTreeNode::Directory { children, .. }) = existing {
149 insert_into_tree(children, &parts[1..], file_index, file);
150 } else {
151 let mut children = Vec::new();
152 insert_into_tree(&mut children, &parts[1..], file_index, file);
153 nodes.push(FileTreeNode::Directory { name: dir_name.to_string(), children, expanded: true });
154 }
155}
156
157fn sort_tree(nodes: &mut [FileTreeNode]) {
158 nodes.sort_by(|a, b| {
159 let a_is_dir = matches!(a, FileTreeNode::Directory { .. });
160 let b_is_dir = matches!(b, FileTreeNode::Directory { .. });
161 b_is_dir.cmp(&a_is_dir).then_with(|| node_name(a).cmp(node_name(b)))
162 });
163 for node in nodes.iter_mut() {
164 if let FileTreeNode::Directory { children, .. } = node {
165 sort_tree(children);
166 }
167 }
168}
169
170fn compress_paths(nodes: &mut [FileTreeNode]) {
171 for node in nodes.iter_mut() {
172 loop {
173 let should_compress = matches!(
174 node,
175 FileTreeNode::Directory { children, .. }
176 if children.len() == 1 && matches!(children[0], FileTreeNode::Directory { .. })
177 );
178 if !should_compress {
179 break;
180 }
181 if let FileTreeNode::Directory { name, children, .. } = node {
182 let child = children.remove(0);
183 if let FileTreeNode::Directory {
184 name: child_name,
185 children: child_children,
186 expanded: child_expanded,
187 } = child
188 {
189 *name = format!("{name}/{child_name}");
190 *children = child_children;
191 if let FileTreeNode::Directory { expanded, .. } = node {
192 *expanded = child_expanded;
193 }
194 }
195 }
196 }
197 if let FileTreeNode::Directory { children, .. } = node {
198 compress_paths(children);
199 }
200 }
201}
202
203fn node_name(node: &FileTreeNode) -> &str {
204 match node {
205 FileTreeNode::Directory { name, .. } | FileTreeNode::File { name, .. } => name,
206 }
207}
208
209fn collect_visible(node: &FileTreeNode, depth: usize, entries: &mut Vec<FileTreeEntry>) {
210 match node {
211 FileTreeNode::Directory { name, children, expanded } => {
212 entries.push(FileTreeEntry {
213 depth,
214 kind: FileTreeEntryKind::Directory { name: name.clone(), expanded: *expanded },
215 });
216 if *expanded {
217 for child in children {
218 collect_visible(child, depth + 1, entries);
219 }
220 }
221 }
222 FileTreeNode::File { file_index, name, status, staged, additions, deletions } => {
223 entries.push(FileTreeEntry {
224 depth,
225 kind: FileTreeEntryKind::File {
226 file_index: *file_index,
227 name: name.clone(),
228 status: *status,
229 staged: *staged,
230 additions: *additions,
231 deletions: *deletions,
232 },
233 });
234 }
235 }
236}
237
238fn toggle_at(nodes: &mut [FileTreeNode], target_visible_idx: usize) {
239 let mut counter = 0;
240 toggle_at_inner(nodes, target_visible_idx, &mut counter);
241}
242
243fn toggle_at_inner(nodes: &mut [FileTreeNode], target: usize, counter: &mut usize) -> bool {
244 for node in nodes.iter_mut() {
245 if *counter == target {
246 if let FileTreeNode::Directory { expanded, .. } = node {
247 *expanded = !*expanded;
248 }
249 return true;
250 }
251 *counter += 1;
252 if let FileTreeNode::Directory { children, expanded: true, .. } = node
253 && toggle_at_inner(children, target, counter)
254 {
255 return true;
256 }
257 }
258 false
259}
260
261fn find_parent_dir(entries: &[FileTreeEntry], idx: usize) -> Option<usize> {
262 let current_depth = entries.get(idx)?.depth;
263 if current_depth == 0 {
264 return None;
265 }
266 (0..idx)
267 .rev()
268 .find(|&i| entries[i].depth < current_depth && matches!(entries[i].kind, FileTreeEntryKind::Directory { .. }))
269}
270
271#[cfg(test)]
272mod tests {
273 use super::*;
274 use crate::git_diff::{FileDiff, FileStatus, Hunk, PatchLine, PatchLineKind, StageState};
275
276 fn file(path: &str, status: FileStatus, additions: usize, deletions: usize) -> FileDiff {
277 let mut lines = Vec::new();
278 for i in 0..additions {
279 lines.push(PatchLine {
280 kind: PatchLineKind::Added,
281 text: format!("added {i}"),
282 old_line_no: None,
283 new_line_no: Some(i + 1),
284 });
285 }
286 for i in 0..deletions {
287 lines.push(PatchLine {
288 kind: PatchLineKind::Removed,
289 text: format!("removed {i}"),
290 old_line_no: Some(i + 1),
291 new_line_no: None,
292 });
293 }
294 FileDiff {
295 old_path: None,
296 path: path.to_string(),
297 status,
298 staged: StageState::Unstaged,
299 hunks: if lines.is_empty() {
300 vec![]
301 } else {
302 vec![Hunk {
303 header: "@@ -1 +1 @@".to_string(),
304 old_start: 1,
305 old_count: deletions,
306 new_start: 1,
307 new_count: additions,
308 lines,
309 }]
310 },
311 binary: false,
312 }
313 }
314
315 fn modified(path: &str) -> FileDiff {
316 file(path, FileStatus::Modified, 1, 1)
317 }
318
319 fn added(path: &str) -> FileDiff {
320 file(path, FileStatus::Added, 2, 0)
321 }
322
323 #[test]
324 fn from_files_groups_by_directory() {
325 let files = vec![modified("src/a.rs"), modified("src/b.rs"), modified("lib/c.rs")];
326 let tree = FileTree::from_files(&files);
327 let entries = tree.visible_entries();
328 assert_eq!(entries.len(), 5);
329 assert!(matches!(&entries[0].kind, FileTreeEntryKind::Directory { name, .. } if name == "lib"));
330 assert!(matches!(&entries[1].kind, FileTreeEntryKind::File { name, .. } if name == "c.rs"));
331 assert!(matches!(&entries[2].kind, FileTreeEntryKind::Directory { name, .. } if name == "src"));
332 assert!(matches!(&entries[3].kind, FileTreeEntryKind::File { name, .. } if name == "a.rs"));
333 assert!(matches!(&entries[4].kind, FileTreeEntryKind::File { name, .. } if name == "b.rs"));
334 }
335
336 #[test]
337 fn visible_entries_respects_collapse() {
338 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
339 let mut tree = FileTree::from_files(&files);
340 assert_eq!(tree.visible_entries().len(), 3);
341
342 tree.collapse_or_parent();
343 assert_eq!(tree.visible_entries().len(), 1);
344 }
345
346 #[test]
347 fn navigate_clamps_at_bounds() {
348 let files = vec![modified("a.rs"), modified("b.rs")];
349 let mut tree = FileTree::from_files(&files);
350 assert_eq!(tree.selected_visible, 0);
351 tree.navigate(1);
352 assert_eq!(tree.selected_visible, 1);
353 tree.navigate(1);
354 assert_eq!(tree.selected_visible, 1, "navigating past the last entry should stay on it");
355 tree.navigate(-1);
356 assert_eq!(tree.selected_visible, 0);
357 tree.navigate(-1);
358 assert_eq!(tree.selected_visible, 0, "navigating before the first entry should stay on it");
359 }
360
361 #[test]
362 fn collapse_or_parent_collapses_dir() {
363 let files = vec![modified("src/a.rs")];
364 let mut tree = FileTree::from_files(&files);
365 tree.selected_visible = 0;
366 assert_eq!(tree.visible_entries().len(), 2);
367 tree.collapse_or_parent();
368 assert_eq!(tree.visible_entries().len(), 1);
369 }
370
371 #[test]
372 fn collapse_or_parent_moves_to_parent_from_file() {
373 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
374 let mut tree = FileTree::from_files(&files);
375 tree.selected_visible = 1;
376 tree.collapse_or_parent();
377 assert_eq!(tree.selected_visible, 0);
378 }
379
380 #[test]
381 fn expand_or_enter_returns_true_for_file() {
382 let files = vec![modified("a.rs")];
383 let mut tree = FileTree::from_files(&files);
384 assert!(tree.expand_or_enter());
385 }
386
387 #[test]
388 fn expand_or_enter_expands_collapsed_dir() {
389 let files = vec![modified("src/a.rs")];
390 let mut tree = FileTree::from_files(&files);
391 tree.collapse_or_parent();
392 assert_eq!(tree.visible_entries().len(), 1);
393 let result = tree.expand_or_enter();
394 assert!(!result);
395 assert_eq!(tree.visible_entries().len(), 2);
396 }
397
398 #[test]
399 fn path_compression_for_single_child_dirs() {
400 let files = vec![modified("src/deep/nested/file.rs")];
401 let tree = FileTree::from_files(&files);
402 let entries = tree.visible_entries();
403 assert_eq!(entries.len(), 2);
404 match &entries[0].kind {
405 FileTreeEntryKind::Directory { name, .. } => {
406 assert_eq!(name, "src/deep/nested");
407 }
408 FileTreeEntryKind::File { .. } => panic!("expected directory"),
409 }
410 }
411
412 #[test]
413 fn selected_file_index_returns_none_for_dir() {
414 let files = vec![modified("src/a.rs")];
415 let tree = FileTree::from_files(&files);
416 assert!(tree.selected_file_index().is_none());
417 }
418
419 #[test]
420 fn selected_file_index_returns_index_for_file() {
421 let files = vec![modified("a.rs")];
422 let tree = FileTree::from_files(&files);
423 assert_eq!(tree.selected_file_index(), Some(0));
424 }
425
426 #[test]
427 fn flat_files_no_grouping() {
428 let files = vec![modified("a.rs"), added("b.rs")];
429 let tree = FileTree::from_files(&files);
430 let entries = tree.visible_entries();
431 assert_eq!(entries.len(), 2);
432 assert!(matches!(&entries[0].kind, FileTreeEntryKind::File { name, .. } if name == "a.rs"));
433 assert!(matches!(&entries[1].kind, FileTreeEntryKind::File { name, .. } if name == "b.rs"));
434 }
435
436 #[test]
437 fn selected_visible_clamped_after_collapse() {
438 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
439 let mut tree = FileTree::from_files(&files);
440 tree.selected_visible = 2;
441 tree.collapse_or_parent();
442 assert!(tree.selected_visible < tree.visible_entries().len());
443 }
444
445 #[test]
446 fn expand_already_expanded_dir_moves_to_first_child() {
447 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
448 let mut tree = FileTree::from_files(&files);
449 tree.selected_visible = 0;
450 let result = tree.expand_or_enter();
451 assert!(!result);
452 assert_eq!(tree.selected_visible, 1);
453 }
454}