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