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 | RepositoryEvent::BranchChanged,
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 fn set_repo(&mut self, repo: Option<Entity<Repository>>, cx: &mut Context<Self>) {
114 self.repo = repo;
115 self.tree_diff = None;
116 self.base_commit = None;
117 self.head_commit = None;
118 cx.emit(BranchDiffEvent::FileListChanged);
119 *self.update_needed.borrow_mut() = ();
120 }
121
122 pub async fn handle_status_updates(
123 this: WeakEntity<Self>,
124 mut recv: postage::watch::Receiver<()>,
125 cx: &mut AsyncWindowContext,
126 ) {
127 Self::reload_tree_diff(this.clone(), cx).await.log_err();
128 while recv.next().await.is_some() {
129 let Ok(needs_update) = this.update(cx, |this, cx| {
130 let mut needs_update = false;
131
132 if this.repo.is_none() {
133 let active_repo = this
134 .project
135 .read(cx)
136 .git_store()
137 .read(cx)
138 .active_repository();
139 if active_repo.is_some() {
140 this.repo = active_repo;
141 needs_update = true;
142 }
143 } else if let Some(repo) = this.repo.as_ref() {
144 repo.update(cx, |repo, _| {
145 if let Some(branch) = &repo.branch
146 && let DiffBase::Merge { base_ref } = &this.diff_base
147 && let Some(commit) = branch.most_recent_commit.as_ref()
148 && &branch.ref_name == base_ref
149 && this.base_commit.as_ref() != Some(&commit.sha)
150 {
151 this.base_commit = Some(commit.sha.clone());
152 needs_update = true;
153 }
154
155 if repo.head_commit.as_ref().map(|c| &c.sha) != this.head_commit.as_ref() {
156 this.head_commit = repo.head_commit.as_ref().map(|c| c.sha.clone());
157 needs_update = true;
158 }
159 })
160 }
161 needs_update
162 }) else {
163 return;
164 };
165
166 if needs_update {
167 Self::reload_tree_diff(this.clone(), cx).await.log_err();
168 }
169 }
170 }
171
172 pub fn status_for_buffer_id(&self, buffer_id: BufferId, cx: &App) -> Option<FileStatus> {
173 let (repo, path) = self
174 .project
175 .read(cx)
176 .git_store()
177 .read(cx)
178 .repository_and_path_for_buffer_id(buffer_id, cx)?;
179 if self.repo() == Some(&repo) {
180 return self.merge_statuses(
181 repo.read(cx)
182 .status_for_path(&path)
183 .map(|status| status.status),
184 self.tree_diff
185 .as_ref()
186 .and_then(|diff| diff.entries.get(&path)),
187 );
188 }
189 None
190 }
191
192 pub fn merge_statuses(
193 &self,
194 diff_from_head: Option<FileStatus>,
195 diff_from_merge_base: Option<&TreeDiffStatus>,
196 ) -> Option<FileStatus> {
197 match (diff_from_head, diff_from_merge_base) {
198 (None, None) => None,
199 (Some(diff_from_head), None) => Some(diff_from_head),
200 (Some(diff_from_head @ FileStatus::Unmerged(_)), _) => Some(diff_from_head),
201
202 // file does not exist in HEAD
203 // but *does* exist in work-tree
204 // and *does* exist in merge-base
205 (
206 Some(FileStatus::Untracked)
207 | Some(FileStatus::Tracked(TrackedStatus {
208 index_status: StatusCode::Added,
209 worktree_status: _,
210 })),
211 Some(_),
212 ) => Some(FileStatus::Tracked(TrackedStatus {
213 index_status: StatusCode::Modified,
214 worktree_status: StatusCode::Modified,
215 })),
216
217 // file exists in HEAD
218 // but *does not* exist in work-tree
219 (Some(diff_from_head), Some(diff_from_merge_base)) if diff_from_head.is_deleted() => {
220 match diff_from_merge_base {
221 TreeDiffStatus::Added => None, // unchanged, didn't exist in merge base or worktree
222 _ => Some(diff_from_head),
223 }
224 }
225
226 // file exists in HEAD
227 // and *does* exist in work-tree
228 (Some(FileStatus::Tracked(_)), Some(tree_status)) => {
229 Some(FileStatus::Tracked(TrackedStatus {
230 index_status: match tree_status {
231 TreeDiffStatus::Added { .. } => StatusCode::Added,
232 _ => StatusCode::Modified,
233 },
234 worktree_status: match tree_status {
235 TreeDiffStatus::Added => StatusCode::Added,
236 _ => StatusCode::Modified,
237 },
238 }))
239 }
240
241 (_, Some(diff_from_merge_base)) => {
242 Some(diff_status_to_file_status(diff_from_merge_base))
243 }
244 }
245 }
246
247 pub async fn reload_tree_diff(
248 this: WeakEntity<Self>,
249 cx: &mut AsyncWindowContext,
250 ) -> Result<()> {
251 let task = this.update(cx, |this, cx| {
252 let DiffBase::Merge { base_ref } = this.diff_base.clone() else {
253 return None;
254 };
255 let Some(repo) = this.repo.as_ref() else {
256 this.tree_diff.take();
257 return None;
258 };
259 repo.update(cx, |repo, cx| {
260 Some(repo.diff_tree(
261 DiffTreeType::MergeBase {
262 base: base_ref,
263 head: "HEAD".into(),
264 },
265 cx,
266 ))
267 })
268 })?;
269 let Some(task) = task else { return Ok(()) };
270
271 let diff = task.await??;
272 this.update(cx, |this, cx| {
273 this.tree_diff = Some(diff);
274 cx.emit(BranchDiffEvent::FileListChanged);
275 cx.notify();
276 })
277 }
278
279 pub fn repo(&self) -> Option<&Entity<Repository>> {
280 self.repo.as_ref()
281 }
282
283 #[instrument(skip_all)]
284 pub fn load_buffers(&mut self, cx: &mut Context<Self>) -> Vec<DiffBuffer> {
285 let mut output = Vec::default();
286 let Some(repo) = self.repo.clone() else {
287 return output;
288 };
289
290 self.project.update(cx, |_project, cx| {
291 let mut seen = HashSet::default();
292
293 for item in repo.read(cx).cached_status() {
294 seen.insert(item.repo_path.clone());
295 let branch_diff = self
296 .tree_diff
297 .as_ref()
298 .and_then(|t| t.entries.get(&item.repo_path))
299 .cloned();
300 let Some(status) = self.merge_statuses(Some(item.status), branch_diff.as_ref())
301 else {
302 continue;
303 };
304 if !status.has_changes() {
305 continue;
306 }
307
308 let Some(project_path) =
309 repo.read(cx).repo_path_to_project_path(&item.repo_path, cx)
310 else {
311 continue;
312 };
313 let task = Self::load_buffer(branch_diff, project_path, repo.clone(), cx);
314
315 output.push(DiffBuffer {
316 repo_path: item.repo_path.clone(),
317 load: task,
318 file_status: item.status,
319 });
320 }
321 let Some(tree_diff) = self.tree_diff.as_ref() else {
322 return;
323 };
324
325 for (path, branch_diff) in tree_diff.entries.iter() {
326 if seen.contains(&path) {
327 continue;
328 }
329
330 let Some(project_path) = repo.read(cx).repo_path_to_project_path(&path, cx) else {
331 continue;
332 };
333 let task =
334 Self::load_buffer(Some(branch_diff.clone()), project_path, repo.clone(), cx);
335
336 let file_status = diff_status_to_file_status(branch_diff);
337
338 output.push(DiffBuffer {
339 repo_path: path.clone(),
340 load: task,
341 file_status,
342 });
343 }
344 });
345 output
346 }
347
348 #[instrument(skip_all)]
349 fn load_buffer(
350 branch_diff: Option<git::status::TreeDiffStatus>,
351 project_path: crate::ProjectPath,
352 repo: Entity<Repository>,
353 cx: &Context<'_, Project>,
354 ) -> Task<Result<(Entity<Buffer>, Entity<BufferDiff>)>> {
355 let task = cx.spawn(async move |project, cx| {
356 let buffer = project
357 .update(cx, |project, cx| project.open_buffer(project_path, cx))?
358 .await?;
359
360 let changes = if let Some(entry) = branch_diff {
361 let oid = match entry {
362 git::status::TreeDiffStatus::Added { .. } => None,
363 git::status::TreeDiffStatus::Modified { old, .. }
364 | git::status::TreeDiffStatus::Deleted { old } => Some(old),
365 };
366 project
367 .update(cx, |project, cx| {
368 project.git_store().update(cx, |git_store, cx| {
369 git_store.open_diff_since(oid, buffer.clone(), repo, cx)
370 })
371 })?
372 .await?
373 } else {
374 project
375 .update(cx, |project, cx| {
376 project.open_uncommitted_diff(buffer.clone(), cx)
377 })?
378 .await?
379 };
380 Ok((buffer, changes))
381 });
382 task
383 }
384}
385
386fn diff_status_to_file_status(branch_diff: &git::status::TreeDiffStatus) -> FileStatus {
387 let file_status = match branch_diff {
388 git::status::TreeDiffStatus::Added { .. } => FileStatus::Tracked(TrackedStatus {
389 index_status: StatusCode::Added,
390 worktree_status: StatusCode::Added,
391 }),
392 git::status::TreeDiffStatus::Modified { .. } => FileStatus::Tracked(TrackedStatus {
393 index_status: StatusCode::Modified,
394 worktree_status: StatusCode::Modified,
395 }),
396 git::status::TreeDiffStatus::Deleted { .. } => FileStatus::Tracked(TrackedStatus {
397 index_status: StatusCode::Deleted,
398 worktree_status: StatusCode::Deleted,
399 }),
400 };
401 file_status
402}
403
404#[derive(Debug)]
405pub struct DiffBuffer {
406 pub repo_path: RepoPath,
407 pub file_status: FileStatus,
408 pub load: Task<Result<(Entity<Buffer>, Entity<BufferDiff>)>>,
409}