1use crate::git_diff::{FileDiff, FileStatus, StageState};
2use std::collections::HashSet;
3
4#[derive(Debug, Clone)]
5pub struct FileTreeEntry {
6 pub depth: usize,
7 pub kind: FileTreeEntryKind,
8}
9
10#[derive(Debug, Clone)]
11pub enum FileTreeEntryKind {
12 Directory {
13 path: String,
14 name: String,
15 expanded: bool,
16 staged: StageState,
17 file_indices: Vec<usize>,
18 },
19 File {
20 path: String,
21 file_index: usize,
22 name: String,
23 status: FileStatus,
24 staged: StageState,
25 additions: usize,
26 deletions: usize,
27 },
28}
29
30pub struct FileTree {
31 entries: Vec<FileTreeEntry>,
32 visible: Vec<usize>,
33 selected_visible: usize,
34}
35
36impl FileTree {
37 pub fn empty() -> Self {
38 Self { entries: Vec::new(), visible: Vec::new(), selected_visible: 0 }
39 }
40
41 pub fn from_files(files: &[FileDiff]) -> Self {
42 let mut tree = Self::empty();
43 tree.rebuild_from_files(files);
44 tree
45 }
46
47 pub fn rebuild_from_files(&mut self, files: &[FileDiff]) {
48 let selected_path = self.selected_entry().map(|entry| entry_path(entry).to_owned());
49 let collapsed: HashSet<String> = self
50 .entries
51 .iter()
52 .filter_map(|entry| match &entry.kind {
53 FileTreeEntryKind::Directory { path, expanded: false, .. } => Some(path.clone()),
54 _ => None,
55 })
56 .collect();
57
58 self.entries = build_entries(files);
59 for entry in &mut self.entries {
60 if let FileTreeEntryKind::Directory { path, expanded, .. } = &mut entry.kind {
61 *expanded = !collapsed.contains(path);
62 }
63 }
64 self.rebuild_visible();
65 if let Some(path) = selected_path
66 && let Some(position) = self.visible.iter().position(|&index| entry_path(&self.entries[index]) == path)
67 {
68 self.selected_visible = position;
69 }
70 }
71
72 pub fn visible_entries(&self) -> Vec<&FileTreeEntry> {
73 self.visible.iter().map(|&index| &self.entries[index]).collect()
74 }
75
76 pub fn selected_visible(&self) -> usize {
77 self.selected_visible
78 }
79
80 pub fn selected_file_index(&self) -> Option<usize> {
81 self.selected_entry().and_then(|entry| match &entry.kind {
82 FileTreeEntryKind::File { file_index, .. } => Some(*file_index),
83 FileTreeEntryKind::Directory { .. } => None,
84 })
85 }
86
87 pub fn selected_file_indices(&self) -> Vec<usize> {
88 match self.selected_entry().map(|entry| &entry.kind) {
89 Some(FileTreeEntryKind::File { file_index, .. }) => vec![*file_index],
90 Some(FileTreeEntryKind::Directory { file_indices, .. }) => file_indices.clone(),
91 None => Vec::new(),
92 }
93 }
94
95 pub fn select_file_index(&mut self, file_index: usize) {
96 if let Some(position) = self.visible.iter().position(|&index| {
97 matches!(&self.entries[index].kind, FileTreeEntryKind::File { file_index: fi, .. } if *fi == file_index)
98 }) {
99 self.selected_visible = position;
100 }
101 }
102
103 pub fn navigate(&mut self, delta: isize) {
104 if self.visible.is_empty() {
105 return;
106 }
107 self.selected_visible = self.selected_visible.saturating_add_signed(delta).min(self.visible.len() - 1);
108 }
109
110 pub fn collapse_or_parent(&mut self) {
111 let Some(&index) = self.visible.get(self.selected_visible) else {
112 return;
113 };
114 if let FileTreeEntryKind::Directory { expanded, .. } = &mut self.entries[index].kind
115 && *expanded
116 {
117 *expanded = false;
118 self.rebuild_visible();
119 return;
120 }
121 if let Some(parent) = self.parent_position(self.selected_visible) {
122 self.selected_visible = parent;
123 }
124 }
125
126 pub fn expand_or_enter(&mut self) -> bool {
127 let Some(&index) = self.visible.get(self.selected_visible) else {
128 return false;
129 };
130 match &mut self.entries[index].kind {
131 FileTreeEntryKind::File { .. } => true,
132 FileTreeEntryKind::Directory { expanded, .. } => {
133 if *expanded {
134 if self.selected_visible + 1 < self.visible.len() {
135 self.selected_visible += 1;
136 }
137 } else {
138 *expanded = true;
139 self.rebuild_visible();
140 }
141 false
142 }
143 }
144 }
145
146 fn selected_entry(&self) -> Option<&FileTreeEntry> {
147 self.visible.get(self.selected_visible).map(|&index| &self.entries[index])
148 }
149
150 fn parent_position(&self, position: usize) -> Option<usize> {
151 let current_depth = self.entries[*self.visible.get(position)?].depth;
152 if current_depth == 0 {
153 return None;
154 }
155 (0..position).rev().find(|&candidate| {
156 let entry = &self.entries[self.visible[candidate]];
157 entry.depth < current_depth && matches!(entry.kind, FileTreeEntryKind::Directory { .. })
158 })
159 }
160
161 fn rebuild_visible(&mut self) {
162 self.visible.clear();
163 let mut hidden_below: Option<usize> = None;
164 for (index, entry) in self.entries.iter().enumerate() {
165 if let Some(depth) = hidden_below {
166 if entry.depth > depth {
167 continue;
168 }
169 hidden_below = None;
170 }
171 self.visible.push(index);
172 if matches!(entry.kind, FileTreeEntryKind::Directory { expanded: false, .. }) {
173 hidden_below = Some(entry.depth);
174 }
175 }
176 self.selected_visible = self.selected_visible.min(self.visible.len().saturating_sub(1));
177 }
178}
179
180enum BuildNode {
181 Directory { name: String, children: Vec<BuildNode> },
182 File { file_index: usize, name: String },
183}
184
185fn build_entries(files: &[FileDiff]) -> Vec<FileTreeEntry> {
186 let mut roots: Vec<BuildNode> = Vec::new();
187 for (index, file) in files.iter().enumerate() {
188 let parts: Vec<&str> = file.path.split('/').collect();
189 insert_into_tree(&mut roots, &parts, index);
190 }
191 sort_tree(&mut roots);
192 compress_paths(&mut roots);
193
194 let mut entries = Vec::new();
195 flatten_into(&roots, 0, "", files, &mut entries);
196 entries
197}
198
199fn insert_into_tree(nodes: &mut Vec<BuildNode>, parts: &[&str], file_index: usize) {
200 if parts.len() == 1 {
201 nodes.push(BuildNode::File { file_index, name: parts[0].to_string() });
202 return;
203 }
204
205 let dir_name = parts[0];
206 let existing = nodes.iter_mut().find(|node| matches!(node, BuildNode::Directory { name, .. } if name == dir_name));
207
208 if let Some(BuildNode::Directory { children, .. }) = existing {
209 insert_into_tree(children, &parts[1..], file_index);
210 } else {
211 let mut children = Vec::new();
212 insert_into_tree(&mut children, &parts[1..], file_index);
213 nodes.push(BuildNode::Directory { name: dir_name.to_string(), children });
214 }
215}
216
217fn sort_tree(nodes: &mut [BuildNode]) {
218 nodes.sort_by(|a, b| {
219 let a_is_dir = matches!(a, BuildNode::Directory { .. });
220 let b_is_dir = matches!(b, BuildNode::Directory { .. });
221 b_is_dir.cmp(&a_is_dir).then_with(|| node_name(a).cmp(node_name(b)))
222 });
223 for node in nodes.iter_mut() {
224 if let BuildNode::Directory { children, .. } = node {
225 sort_tree(children);
226 }
227 }
228}
229
230fn compress_paths(nodes: &mut [BuildNode]) {
231 for node in nodes.iter_mut() {
232 while let BuildNode::Directory { name, children } = node {
233 if children.len() != 1 || !matches!(children[0], BuildNode::Directory { .. }) {
234 break;
235 }
236 let BuildNode::Directory { name: child_name, children: grandchildren } = children.remove(0) else {
237 unreachable!("checked above that the only child is a directory");
238 };
239 *name = format!("{name}/{child_name}");
240 *children = grandchildren;
241 }
242 if let BuildNode::Directory { children, .. } = node {
243 compress_paths(children);
244 }
245 }
246}
247
248fn node_name(node: &BuildNode) -> &str {
249 match node {
250 BuildNode::Directory { name, .. } | BuildNode::File { name, .. } => name,
251 }
252}
253
254fn flatten_into(
255 nodes: &[BuildNode],
256 depth: usize,
257 parent_path: &str,
258 files: &[FileDiff],
259 entries: &mut Vec<FileTreeEntry>,
260) -> Vec<usize> {
261 let mut collected = Vec::new();
262 for node in nodes {
263 match node {
264 BuildNode::Directory { name, children } => {
265 let path = join_path(parent_path, name);
266 let slot = entries.len();
267 entries.push(FileTreeEntry {
268 depth,
269 kind: FileTreeEntryKind::Directory {
270 path: path.clone(),
271 name: name.clone(),
272 expanded: true,
273 staged: StageState::Unstaged,
274 file_indices: Vec::new(),
275 },
276 });
277 let descendants = flatten_into(children, depth + 1, &path, files, entries);
278 if let FileTreeEntryKind::Directory { staged, file_indices, .. } = &mut entries[slot].kind {
279 *staged = aggregate_stage_state(&descendants, files);
280 file_indices.clone_from(&descendants);
281 }
282 collected.extend(descendants);
283 }
284 BuildNode::File { file_index, name } => {
285 let file = &files[*file_index];
286 entries.push(FileTreeEntry {
287 depth,
288 kind: FileTreeEntryKind::File {
289 path: join_path(parent_path, name),
290 file_index: *file_index,
291 name: name.clone(),
292 status: file.status,
293 staged: file.staged,
294 additions: file.additions(),
295 deletions: file.deletions(),
296 },
297 });
298 collected.push(*file_index);
299 }
300 }
301 }
302 collected
303}
304
305fn aggregate_stage_state(file_indices: &[usize], files: &[FileDiff]) -> StageState {
306 let mut states = file_indices.iter().map(|&index| files[index].staged);
307 let Some(first) = states.next() else {
308 return StageState::Unstaged;
309 };
310 if states.all(|state| state == first) { first } else { StageState::PartiallyStaged }
311}
312
313fn entry_path(entry: &FileTreeEntry) -> &str {
314 match &entry.kind {
315 FileTreeEntryKind::Directory { path, .. } | FileTreeEntryKind::File { path, .. } => path,
316 }
317}
318
319fn join_path(parent: &str, name: &str) -> String {
320 if parent.is_empty() { name.to_string() } else { format!("{parent}/{name}") }
321}
322
323#[cfg(test)]
324mod tests {
325 use super::*;
326 use crate::git_diff::{FileDiff, FileStatus, Hunk, PatchLine, PatchLineKind, StageState};
327
328 fn file(path: &str, status: FileStatus, additions: usize, deletions: usize) -> FileDiff {
329 let mut lines = Vec::new();
330 for i in 0..additions {
331 lines.push(PatchLine {
332 kind: PatchLineKind::Added,
333 text: format!("added {i}"),
334 old_line_no: None,
335 new_line_no: Some(i + 1),
336 });
337 }
338 for i in 0..deletions {
339 lines.push(PatchLine {
340 kind: PatchLineKind::Removed,
341 text: format!("removed {i}"),
342 old_line_no: Some(i + 1),
343 new_line_no: None,
344 });
345 }
346 FileDiff {
347 old_path: None,
348 path: path.to_string(),
349 status,
350 staged: StageState::Unstaged,
351 hunks: if lines.is_empty() {
352 vec![]
353 } else {
354 vec![Hunk {
355 header: "@@ -1 +1 @@".to_string(),
356 old_start: 1,
357 old_count: deletions,
358 new_start: 1,
359 new_count: additions,
360 lines,
361 }]
362 },
363 binary: false,
364 }
365 }
366
367 fn modified(path: &str) -> FileDiff {
368 file(path, FileStatus::Modified, 1, 1)
369 }
370
371 fn added(path: &str) -> FileDiff {
372 file(path, FileStatus::Added, 2, 0)
373 }
374
375 #[test]
376 fn from_files_groups_by_directory() {
377 let files = vec![modified("src/a.rs"), modified("src/b.rs"), modified("lib/c.rs")];
378 let tree = FileTree::from_files(&files);
379 let entries = tree.visible_entries();
380 assert_eq!(entries.len(), 5);
381 assert!(matches!(&entries[0].kind, FileTreeEntryKind::Directory { name, .. } if name == "lib"));
382 assert!(matches!(&entries[1].kind, FileTreeEntryKind::File { name, .. } if name == "c.rs"));
383 assert!(matches!(&entries[2].kind, FileTreeEntryKind::Directory { name, .. } if name == "src"));
384 assert!(matches!(&entries[3].kind, FileTreeEntryKind::File { name, .. } if name == "a.rs"));
385 assert!(matches!(&entries[4].kind, FileTreeEntryKind::File { name, .. } if name == "b.rs"));
386 }
387
388 #[test]
389 fn visible_entries_respects_collapse() {
390 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
391 let mut tree = FileTree::from_files(&files);
392 assert_eq!(tree.visible_entries().len(), 3);
393
394 tree.collapse_or_parent();
395 assert_eq!(tree.visible_entries().len(), 1);
396 }
397
398 #[test]
399 fn navigate_clamps_at_bounds() {
400 let files = vec![modified("a.rs"), modified("b.rs")];
401 let mut tree = FileTree::from_files(&files);
402 assert_eq!(tree.selected_visible, 0);
403 tree.navigate(1);
404 assert_eq!(tree.selected_visible, 1);
405 tree.navigate(1);
406 assert_eq!(tree.selected_visible, 1, "navigating past the last entry should stay on it");
407 tree.navigate(-1);
408 assert_eq!(tree.selected_visible, 0);
409 tree.navigate(-1);
410 assert_eq!(tree.selected_visible, 0, "navigating before the first entry should stay on it");
411 }
412
413 #[test]
414 fn collapse_or_parent_collapses_dir() {
415 let files = vec![modified("src/a.rs")];
416 let mut tree = FileTree::from_files(&files);
417 tree.selected_visible = 0;
418 assert_eq!(tree.visible_entries().len(), 2);
419 tree.collapse_or_parent();
420 assert_eq!(tree.visible_entries().len(), 1);
421 }
422
423 #[test]
424 fn collapse_or_parent_moves_to_parent_from_file() {
425 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
426 let mut tree = FileTree::from_files(&files);
427 tree.selected_visible = 1;
428 tree.collapse_or_parent();
429 assert_eq!(tree.selected_visible, 0);
430 }
431
432 #[test]
433 fn expand_or_enter_returns_true_for_file() {
434 let files = vec![modified("a.rs")];
435 let mut tree = FileTree::from_files(&files);
436 assert!(tree.expand_or_enter());
437 }
438
439 #[test]
440 fn expand_or_enter_expands_collapsed_dir() {
441 let files = vec![modified("src/a.rs")];
442 let mut tree = FileTree::from_files(&files);
443 tree.collapse_or_parent();
444 assert_eq!(tree.visible_entries().len(), 1);
445 let result = tree.expand_or_enter();
446 assert!(!result);
447 assert_eq!(tree.visible_entries().len(), 2);
448 }
449
450 #[test]
451 fn path_compression_for_single_child_dirs() {
452 let files = vec![modified("src/deep/nested/file.rs")];
453 let tree = FileTree::from_files(&files);
454 let entries = tree.visible_entries();
455 assert_eq!(entries.len(), 2);
456 match &entries[0].kind {
457 FileTreeEntryKind::Directory { name, .. } => {
458 assert_eq!(name, "src/deep/nested");
459 }
460 FileTreeEntryKind::File { .. } => panic!("expected directory"),
461 }
462 }
463
464 #[test]
465 fn selected_file_index_returns_none_for_dir() {
466 let files = vec![modified("src/a.rs")];
467 let tree = FileTree::from_files(&files);
468 assert!(tree.selected_file_index().is_none());
469 }
470
471 #[test]
472 fn selected_file_index_returns_index_for_file() {
473 let files = vec![modified("a.rs")];
474 let tree = FileTree::from_files(&files);
475 assert_eq!(tree.selected_file_index(), Some(0));
476 }
477
478 #[test]
479 fn flat_files_no_grouping() {
480 let files = vec![modified("a.rs"), added("b.rs")];
481 let tree = FileTree::from_files(&files);
482 let entries = tree.visible_entries();
483 assert_eq!(entries.len(), 2);
484 assert!(matches!(&entries[0].kind, FileTreeEntryKind::File { name, .. } if name == "a.rs"));
485 assert!(matches!(&entries[1].kind, FileTreeEntryKind::File { name, .. } if name == "b.rs"));
486 }
487
488 #[test]
489 fn selected_visible_clamped_after_collapse() {
490 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
491 let mut tree = FileTree::from_files(&files);
492 tree.selected_visible = 2;
493 tree.collapse_or_parent();
494 assert!(tree.selected_visible < tree.visible_entries().len());
495 }
496
497 #[test]
498 fn rebuild_preserves_selected_directory() {
499 let files = vec![modified("lib/a.rs"), modified("src/b.rs")];
500 let mut tree = FileTree::from_files(&files);
501 tree.navigate(2);
502
503 tree.rebuild_from_files(&files);
504
505 assert!(
506 matches!(&tree.visible_entries()[tree.selected_visible()].kind, FileTreeEntryKind::Directory { name, .. } if name == "src")
507 );
508 }
509
510 #[test]
511 fn rebuild_preserves_collapsed_directories() {
512 let files = vec![modified("lib/a.rs"), modified("src/b.rs")];
513 let mut tree = FileTree::from_files(&files);
514 tree.collapse_or_parent();
515
516 tree.rebuild_from_files(&files);
517
518 assert!(
519 matches!(&tree.visible_entries()[0].kind, FileTreeEntryKind::Directory { name, expanded: false, .. } if name == "lib")
520 );
521 assert!(
522 !tree
523 .visible_entries()
524 .iter()
525 .any(|entry| matches!(&entry.kind, FileTreeEntryKind::File { name, .. } if name == "a.rs"))
526 );
527 }
528
529 #[test]
530 fn rebuild_preserves_collapsed_directory_hidden_under_collapsed_parent() {
531 let files = vec![modified("src/inner/a.rs"), modified("src/b.rs"), modified("src/inner/other/c.rs")];
532 let mut tree = FileTree::from_files(&files);
533 tree.navigate(1);
534 tree.collapse_or_parent();
535 tree.collapse_or_parent();
536 tree.collapse_or_parent();
537 assert_eq!(tree.visible_entries().len(), 1);
538
539 tree.rebuild_from_files(&files);
540 tree.expand_or_enter();
541
542 assert!(
543 matches!(&tree.visible_entries()[1].kind, FileTreeEntryKind::Directory { name, expanded: false, .. } if name == "inner"),
544 "nested collapsed directory should stay collapsed after rebuild"
545 );
546 }
547
548 #[test]
549 fn directory_stage_state_aggregates_descendant_files() {
550 let mut files = vec![modified("src/a.rs"), modified("src/b.rs")];
551 files[1].staged = StageState::Staged;
552 let tree = FileTree::from_files(&files);
553
554 assert!(matches!(
555 &tree.visible_entries()[0].kind,
556 FileTreeEntryKind::Directory { staged: StageState::PartiallyStaged, .. }
557 ));
558 }
559
560 #[test]
561 fn selected_file_indices_for_directory_includes_hidden_descendants() {
562 let files = vec![modified("src/nested/a.rs"), modified("src/b.rs"), modified("top.rs")];
563 let mut tree = FileTree::from_files(&files);
564 tree.collapse_or_parent();
565
566 let mut indices = tree.selected_file_indices();
567 indices.sort_unstable();
568 assert_eq!(indices, vec![0, 1]);
569 }
570
571 #[test]
572 fn expand_already_expanded_dir_moves_to_first_child() {
573 let files = vec![modified("src/a.rs"), modified("src/b.rs")];
574 let mut tree = FileTree::from_files(&files);
575 tree.selected_visible = 0;
576 let result = tree.expand_or_enter();
577 assert!(!result);
578 assert_eq!(tree.selected_visible, 1);
579 }
580}