1use anyhow::Result;
2use buffer_diff::BufferDiff;
3use collections::HashSet;
4use futures::StreamExt;
5use git::{
6 repository::RepoPath,
7 status::{DiffTreeType, FileStatus, StatusCode, TrackedStatus, TreeDiff, TreeDiffStatus},
8};
9use gpui::{
10 App, AsyncWindowContext, Context, Entity, EventEmitter, SharedString, Subscription, Task,
11 WeakEntity, Window,
12};
13
14use language::Buffer;
15use text::BufferId;
16use util::ResultExt;
17use ztracing::instrument;
18
19use crate::{
20 Project,
21 git_store::{GitStoreEvent, Repository, RepositoryEvent},
22};
23
24#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
25pub enum DiffBase {
26 Head,
27 Merge { base_ref: SharedString },
28}
29
30impl DiffBase {
31 pub fn is_merge_base(&self) -> bool {
32 matches!(self, DiffBase::Merge { .. })
33 }
34}
35
36pub struct BranchDiff {
37 diff_base: DiffBase,
38 repo: Option<Entity<Repository>>,
39 project: Entity<Project>,
40 base_commit: Option<SharedString>,
41 head_commit: Option<SharedString>,
42 tree_diff: Option<TreeDiff>,
43 _subscription: Subscription,
44 update_needed: postage::watch::Sender<()>,
45 _task: Task<()>,
46}
47
48pub enum BranchDiffEvent {
49 FileListChanged,
50}
51
52impl EventEmitter<BranchDiffEvent> for BranchDiff {}
53
54impl BranchDiff {
55 pub fn new(
56 source: DiffBase,
57 project: Entity<Project>,
58 window: &mut Window,
59 cx: &mut Context<Self>,
60 ) -> Self {
61 let git_store = project.read(cx).git_store().clone();
62 let repo = git_store.read(cx).active_repository();
63 let git_store_subscription = cx.subscribe_in(
64 &git_store,
65 window,
66 move |this, _git_store, event, _window, cx| {
67 let should_update = match event {
68 GitStoreEvent::ActiveRepositoryChanged(new_repo_id) => {
69 this.repo.is_none() && new_repo_id.is_some()
70 }
71 GitStoreEvent::RepositoryUpdated(
72 event_repo_id,
73 RepositoryEvent::StatusesChanged,
74 _,
75 ) => this
76 .repo
77 .as_ref()
78 .is_some_and(|r| r.read(cx).snapshot().id == *event_repo_id),
79 GitStoreEvent::ConflictsUpdated => this.repo.is_some(),
80 _ => false,
81 };
82
83 if should_update {
84 cx.emit(BranchDiffEvent::FileListChanged);
85 *this.update_needed.borrow_mut() = ();
86 }
87 },
88 );
89
90 let (send, recv) = postage::watch::channel::<()>();
91 let worker = window.spawn(cx, {
92 let this = cx.weak_entity();
93 async |cx| Self::handle_status_updates(this, recv, cx).await
94 });
95
96 Self {
97 diff_base: source,
98 repo,
99 project,
100 tree_diff: None,
101 base_commit: None,
102 head_commit: None,
103 _subscription: git_store_subscription,
104 _task: worker,
105 update_needed: send,
106 }
107 }
108
109 pub fn diff_base(&self) -> &DiffBase {
110 &self.diff_base
111 }
112
113 pub async fn handle_status_updates(
114 this: WeakEntity<Self>,
115 mut recv: postage::watch::Receiver<()>,
116 cx: &mut AsyncWindowContext,
117 ) {
118 Self::reload_tree_diff(this.clone(), cx).await.log_err();
119 while recv.next().await.is_some() {
120 let Ok(needs_update) = this.update(cx, |this, cx| {
121 let mut needs_update = false;
122
123 if this.repo.is_none() {
124 let active_repo = this
125 .project
126 .read(cx)
127 .git_store()
128 .read(cx)
129 .active_repository();
130 if active_repo.is_some() {
131 this.repo = active_repo;
132 needs_update = true;
133 }
134 } else if let Some(repo) = this.repo.as_ref() {
135 repo.update(cx, |repo, _| {
136 if let Some(branch) = &repo.branch
137 && let DiffBase::Merge { base_ref } = &this.diff_base
138 && let Some(commit) = branch.most_recent_commit.as_ref()
139 && &branch.ref_name == base_ref
140 && this.base_commit.as_ref() != Some(&commit.sha)
141 {
142 this.base_commit = Some(commit.sha.clone());
143 needs_update = true;
144 }
145
146 if repo.head_commit.as_ref().map(|c| &c.sha) != this.head_commit.as_ref() {
147 this.head_commit = repo.head_commit.as_ref().map(|c| c.sha.clone());
148 needs_update = true;
149 }
150 })
151 }
152 needs_update
153 }) else {
154 return;
155 };
156
157 if needs_update {
158 Self::reload_tree_diff(this.clone(), cx).await.log_err();
159 }
160 }
161 }
162
163 pub fn status_for_buffer_id(&self, buffer_id: BufferId, cx: &App) -> Option<FileStatus> {
164 let (repo, path) = self
165 .project
166 .read(cx)
167 .git_store()
168 .read(cx)
169 .repository_and_path_for_buffer_id(buffer_id, cx)?;
170 if self.repo() == Some(&repo) {
171 return self.merge_statuses(
172 repo.read(cx)
173 .status_for_path(&path)
174 .map(|status| status.status),
175 self.tree_diff
176 .as_ref()
177 .and_then(|diff| diff.entries.get(&path)),
178 );
179 }
180 None
181 }
182
183 pub fn merge_statuses(
184 &self,
185 diff_from_head: Option<FileStatus>,
186 diff_from_merge_base: Option<&TreeDiffStatus>,
187 ) -> Option<FileStatus> {
188 match (diff_from_head, diff_from_merge_base) {
189 (None, None) => None,
190 (Some(diff_from_head), None) => Some(diff_from_head),
191 (Some(diff_from_head @ FileStatus::Unmerged(_)), _) => Some(diff_from_head),
192
193 // file does not exist in HEAD
194 // but *does* exist in work-tree
195 // and *does* exist in merge-base
196 (
197 Some(FileStatus::Untracked)
198 | Some(FileStatus::Tracked(TrackedStatus {
199 index_status: StatusCode::Added,
200 worktree_status: _,
201 })),
202 Some(_),
203 ) => Some(FileStatus::Tracked(TrackedStatus {
204 index_status: StatusCode::Modified,
205 worktree_status: StatusCode::Modified,
206 })),
207
208 // file exists in HEAD
209 // but *does not* exist in work-tree
210 (Some(diff_from_head), Some(diff_from_merge_base)) if diff_from_head.is_deleted() => {
211 match diff_from_merge_base {
212 TreeDiffStatus::Added => None, // unchanged, didn't exist in merge base or worktree
213 _ => Some(diff_from_head),
214 }
215 }
216
217 // file exists in HEAD
218 // and *does* exist in work-tree
219 (Some(FileStatus::Tracked(_)), Some(tree_status)) => {
220 Some(FileStatus::Tracked(TrackedStatus {
221 index_status: match tree_status {
222 TreeDiffStatus::Added { .. } => StatusCode::Added,
223 _ => StatusCode::Modified,
224 },
225 worktree_status: match tree_status {
226 TreeDiffStatus::Added => StatusCode::Added,
227 _ => StatusCode::Modified,
228 },
229 }))
230 }
231
232 (_, Some(diff_from_merge_base)) => {
233 Some(diff_status_to_file_status(diff_from_merge_base))
234 }
235 }
236 }
237
238 pub async fn reload_tree_diff(
239 this: WeakEntity<Self>,
240 cx: &mut AsyncWindowContext,
241 ) -> Result<()> {
242 let task = this.update(cx, |this, cx| {
243 let DiffBase::Merge { base_ref } = this.diff_base.clone() else {
244 return None;
245 };
246 let Some(repo) = this.repo.as_ref() else {
247 this.tree_diff.take();
248 return None;
249 };
250 repo.update(cx, |repo, cx| {
251 Some(repo.diff_tree(
252 DiffTreeType::MergeBase {
253 base: base_ref,
254 head: "HEAD".into(),
255 },
256 cx,
257 ))
258 })
259 })?;
260 let Some(task) = task else { return Ok(()) };
261
262 let diff = task.await??;
263 this.update(cx, |this, cx| {
264 this.tree_diff = Some(diff);
265 cx.emit(BranchDiffEvent::FileListChanged);
266 cx.notify();
267 })
268 }
269
270 pub fn repo(&self) -> Option<&Entity<Repository>> {
271 self.repo.as_ref()
272 }
273
274 #[instrument(skip_all)]
275 pub fn load_buffers(&mut self, cx: &mut Context<Self>) -> Vec<DiffBuffer> {
276 let mut output = Vec::default();
277 let Some(repo) = self.repo.clone() else {
278 return output;
279 };
280
281 self.project.update(cx, |_project, cx| {
282 let mut seen = HashSet::default();
283
284 for item in repo.read(cx).cached_status() {
285 seen.insert(item.repo_path.clone());
286 let branch_diff = self
287 .tree_diff
288 .as_ref()
289 .and_then(|t| t.entries.get(&item.repo_path))
290 .cloned();
291 let Some(status) = self.merge_statuses(Some(item.status), branch_diff.as_ref())
292 else {
293 continue;
294 };
295 if !status.has_changes() {
296 continue;
297 }
298
299 let Some(project_path) =
300 repo.read(cx).repo_path_to_project_path(&item.repo_path, cx)
301 else {
302 continue;
303 };
304 let task = Self::load_buffer(branch_diff, project_path, repo.clone(), cx);
305
306 output.push(DiffBuffer {
307 repo_path: item.repo_path.clone(),
308 load: task,
309 file_status: item.status,
310 });
311 }
312 let Some(tree_diff) = self.tree_diff.as_ref() else {
313 return;
314 };
315
316 for (path, branch_diff) in tree_diff.entries.iter() {
317 if seen.contains(&path) {
318 continue;
319 }
320
321 let Some(project_path) = repo.read(cx).repo_path_to_project_path(&path, cx) else {
322 continue;
323 };
324 let task =
325 Self::load_buffer(Some(branch_diff.clone()), project_path, repo.clone(), cx);
326
327 let file_status = diff_status_to_file_status(branch_diff);
328
329 output.push(DiffBuffer {
330 repo_path: path.clone(),
331 load: task,
332 file_status,
333 });
334 }
335 });
336 output
337 }
338
339 #[instrument(skip_all)]
340 fn load_buffer(
341 branch_diff: Option<git::status::TreeDiffStatus>,
342 project_path: crate::ProjectPath,
343 repo: Entity<Repository>,
344 cx: &Context<'_, Project>,
345 ) -> Task<Result<(Entity<Buffer>, Entity<BufferDiff>)>> {
346 let task = cx.spawn(async move |project, cx| {
347 let buffer = project
348 .update(cx, |project, cx| project.open_buffer(project_path, cx))?
349 .await?;
350
351 let changes = if let Some(entry) = branch_diff {
352 let oid = match entry {
353 git::status::TreeDiffStatus::Added { .. } => None,
354 git::status::TreeDiffStatus::Modified { old, .. }
355 | git::status::TreeDiffStatus::Deleted { old } => Some(old),
356 };
357 project
358 .update(cx, |project, cx| {
359 project.git_store().update(cx, |git_store, cx| {
360 git_store.open_diff_since(oid, buffer.clone(), repo, cx)
361 })
362 })?
363 .await?
364 } else {
365 project
366 .update(cx, |project, cx| {
367 project.open_uncommitted_diff(buffer.clone(), cx)
368 })?
369 .await?
370 };
371 Ok((buffer, changes))
372 });
373 task
374 }
375}
376
377fn diff_status_to_file_status(branch_diff: &git::status::TreeDiffStatus) -> FileStatus {
378 let file_status = match branch_diff {
379 git::status::TreeDiffStatus::Added { .. } => FileStatus::Tracked(TrackedStatus {
380 index_status: StatusCode::Added,
381 worktree_status: StatusCode::Added,
382 }),
383 git::status::TreeDiffStatus::Modified { .. } => FileStatus::Tracked(TrackedStatus {
384 index_status: StatusCode::Modified,
385 worktree_status: StatusCode::Modified,
386 }),
387 git::status::TreeDiffStatus::Deleted { .. } => FileStatus::Tracked(TrackedStatus {
388 index_status: StatusCode::Deleted,
389 worktree_status: StatusCode::Deleted,
390 }),
391 };
392 file_status
393}
394
395#[derive(Debug)]
396pub struct DiffBuffer {
397 pub repo_path: RepoPath,
398 pub file_status: FileStatus,
399 pub load: Task<Result<(Entity<Buffer>, Entity<BufferDiff>)>>,
400}