1use collections::HashMap;
2use git::{repository::RepoPath, status::GitSummary};
3use std::{collections::BTreeMap, ops::Deref, path::Path};
4use sum_tree::Cursor;
5use text::Bias;
6use util::rel_path::RelPath;
7use worktree::{Entry, PathProgress, PathTarget, Traversal};
8
9use super::{RepositoryId, RepositorySnapshot, StatusEntry};
10
11/// Walks the worktree entries and their associated git statuses.
12pub struct GitTraversal<'a> {
13 traversal: Traversal<'a>,
14 current_entry_summary: Option<GitSummary>,
15 repo_root_to_snapshot: BTreeMap<&'a Path, &'a RepositorySnapshot>,
16 repo_location: Option<(
17 RepositoryId,
18 Cursor<'a, 'static, StatusEntry, PathProgress<'a>>,
19 )>,
20}
21
22impl<'a> GitTraversal<'a> {
23 pub fn new(
24 repo_snapshots: &'a HashMap<RepositoryId, RepositorySnapshot>,
25 traversal: Traversal<'a>,
26 ) -> GitTraversal<'a> {
27 let repo_root_to_snapshot = repo_snapshots
28 .values()
29 .map(|snapshot| (&*snapshot.work_directory_abs_path, snapshot))
30 .collect();
31 let mut this = GitTraversal {
32 traversal,
33 current_entry_summary: None,
34 repo_location: None,
35 repo_root_to_snapshot,
36 };
37 this.synchronize_statuses(true);
38 this
39 }
40
41 fn repo_root_for_path(&self, path: &Path) -> Option<(&'a RepositorySnapshot, RepoPath)> {
42 // We might need to perform a range search multiple times, as there may be a nested repository inbetween
43 // the target and our path. E.g:
44 // /our_root_repo/
45 // .git/
46 // other_repo/
47 // .git/
48 // our_query.txt
49 let query = path.ancestors();
50 for query in query {
51 let (_, snapshot) = self
52 .repo_root_to_snapshot
53 .range(Path::new("")..=query)
54 .last()?;
55
56 let stripped = snapshot
57 .abs_path_to_repo_path(path)
58 .map(|repo_path| (*snapshot, repo_path));
59 if stripped.is_some() {
60 return stripped;
61 }
62 }
63
64 None
65 }
66
67 fn synchronize_statuses(&mut self, reset: bool) {
68 self.current_entry_summary = None;
69
70 let Some(entry) = self.entry() else {
71 return;
72 };
73
74 let abs_path = self.traversal.snapshot().absolutize(&entry.path);
75
76 let Some((repo, repo_path)) = self.repo_root_for_path(&abs_path) else {
77 self.repo_location = None;
78 return;
79 };
80
81 // Update our state if we changed repositories.
82 if reset
83 || self
84 .repo_location
85 .as_ref()
86 .map(|(prev_repo_id, _)| *prev_repo_id)
87 != Some(repo.id)
88 {
89 self.repo_location = Some((repo.id, repo.statuses_by_path.cursor::<PathProgress>(())));
90 }
91
92 let Some((_, statuses)) = &mut self.repo_location else {
93 return;
94 };
95
96 if entry.is_dir() {
97 let mut statuses = statuses.clone();
98 statuses.seek_forward(&PathTarget::Path(&repo_path), Bias::Left);
99 let summary = statuses.summary(&PathTarget::Successor(&repo_path), Bias::Left);
100
101 self.current_entry_summary = Some(summary);
102 } else if entry.is_file() {
103 // For a file entry, park the cursor on the corresponding status
104 if statuses.seek_forward(&PathTarget::Path(&repo_path), Bias::Left) {
105 // TODO: Investigate statuses.item() being None here.
106 self.current_entry_summary = statuses.item().map(|item| item.status.into());
107 } else {
108 self.current_entry_summary = Some(GitSummary::UNCHANGED);
109 }
110 }
111 }
112
113 pub fn advance(&mut self) -> bool {
114 let found = self.traversal.advance_by(1);
115 self.synchronize_statuses(false);
116 found
117 }
118
119 pub fn advance_to_sibling(&mut self) -> bool {
120 let found = self.traversal.advance_to_sibling();
121 self.synchronize_statuses(false);
122 found
123 }
124
125 pub fn back_to_parent(&mut self) -> bool {
126 let found = self.traversal.back_to_parent();
127 self.synchronize_statuses(true);
128 found
129 }
130
131 pub fn start_offset(&self) -> usize {
132 self.traversal.start_offset()
133 }
134
135 pub fn end_offset(&self) -> usize {
136 self.traversal.end_offset()
137 }
138
139 pub fn entry(&self) -> Option<GitEntryRef<'a>> {
140 let entry = self.traversal.entry()?;
141 let git_summary = self.current_entry_summary.unwrap_or(GitSummary::UNCHANGED);
142 Some(GitEntryRef { entry, git_summary })
143 }
144}
145
146impl<'a> Iterator for GitTraversal<'a> {
147 type Item = GitEntryRef<'a>;
148
149 fn next(&mut self) -> Option<Self::Item> {
150 if let Some(item) = self.entry() {
151 self.advance();
152 Some(item)
153 } else {
154 None
155 }
156 }
157}
158
159pub struct ChildEntriesGitIter<'a> {
160 parent_path: &'a RelPath,
161 traversal: GitTraversal<'a>,
162}
163
164impl<'a> ChildEntriesGitIter<'a> {
165 pub fn new(
166 repo_snapshots: &'a HashMap<RepositoryId, RepositorySnapshot>,
167 worktree_snapshot: &'a worktree::Snapshot,
168 parent_path: &'a RelPath,
169 ) -> Self {
170 let mut traversal = GitTraversal::new(
171 repo_snapshots,
172 worktree_snapshot.traverse_from_path(true, true, true, parent_path),
173 );
174 traversal.advance();
175 ChildEntriesGitIter {
176 parent_path,
177 traversal,
178 }
179 }
180}
181
182impl<'a> Iterator for ChildEntriesGitIter<'a> {
183 type Item = GitEntryRef<'a>;
184
185 fn next(&mut self) -> Option<Self::Item> {
186 if let Some(item) = self.traversal.entry()
187 && item.path.starts_with(self.parent_path)
188 {
189 self.traversal.advance_to_sibling();
190 return Some(item);
191 }
192 None
193 }
194}
195
196#[derive(Debug, Clone, Copy)]
197pub struct GitEntryRef<'a> {
198 pub entry: &'a Entry,
199 pub git_summary: GitSummary,
200}
201
202impl GitEntryRef<'_> {
203 pub fn to_owned(self) -> GitEntry {
204 GitEntry {
205 entry: self.entry.clone(),
206 git_summary: self.git_summary,
207 }
208 }
209}
210
211impl Deref for GitEntryRef<'_> {
212 type Target = Entry;
213
214 fn deref(&self) -> &Self::Target {
215 self.entry
216 }
217}
218
219impl AsRef<Entry> for GitEntryRef<'_> {
220 fn as_ref(&self) -> &Entry {
221 self.entry
222 }
223}
224
225#[derive(Debug, Clone, PartialEq, Eq)]
226pub struct GitEntry {
227 pub entry: Entry,
228 pub git_summary: GitSummary,
229}
230
231impl GitEntry {
232 pub fn to_ref(&self) -> GitEntryRef<'_> {
233 GitEntryRef {
234 entry: &self.entry,
235 git_summary: self.git_summary,
236 }
237 }
238}
239
240impl Deref for GitEntry {
241 type Target = Entry;
242
243 fn deref(&self) -> &Self::Target {
244 &self.entry
245 }
246}
247
248impl AsRef<Entry> for GitEntry {
249 fn as_ref(&self) -> &Entry {
250 &self.entry
251 }
252}