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