git: Use repo snapshot to fetch picker branch list instead of spawning a job (#53661)

Anthony Eid created

This should fix a delay where the thread_branch_picker and git_picker
would take a while to load the branch list while other git jobs are
being processed before hand.

I also added a subscription in both pickers to repository so they
refresh their matches when the repo snapshot branch list changes.

Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Release Notes:

- N/A

Change summary

crates/agent_ui/src/thread_branch_picker.rs | 213 +++++++++++++---------
crates/git_ui/src/branch_picker.rs          | 178 +++++++++++-------
2 files changed, 232 insertions(+), 159 deletions(-)

Detailed changes

crates/agent_ui/src/thread_branch_picker.rs 🔗

@@ -1,18 +1,18 @@
-use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 
-use collections::HashSet as CollectionsHashSet;
+use collections::{HashMap, HashSet};
 use std::path::PathBuf;
 use std::sync::Arc;
 
 use fuzzy::StringMatchCandidate;
-use git::repository::Branch as GitBranch;
+use git::repository::{Branch as GitBranch, Worktree as GitWorktree};
 use gpui::{
     AnyElement, App, Context, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable,
-    IntoElement, ParentElement, Render, SharedString, Styled, Task, Window, rems,
+    IntoElement, ParentElement, Render, SharedString, Styled, Subscription, Task, Window, rems,
 };
 use picker::{Picker, PickerDelegate, PickerEditorPosition};
 use project::Project;
+use project::git_store::RepositoryEvent;
 use ui::{
     Divider, DocumentationAside, HighlightedLabel, Icon, IconName, Label, LabelCommon, ListItem,
     ListItemSpacing, prelude::*,
@@ -24,7 +24,7 @@ use crate::{NewWorktreeBranchTarget, StartThreadIn};
 pub(crate) struct ThreadBranchPicker {
     picker: Entity<Picker<ThreadBranchPickerDelegate>>,
     focus_handle: FocusHandle,
-    _subscription: gpui::Subscription,
+    _subscriptions: Vec<Subscription>,
 }
 
 impl ThreadBranchPicker {
@@ -57,13 +57,21 @@ impl ThreadBranchPicker {
         } else {
             project.read(cx).active_repository(cx)
         };
-        let branches_request = repository
-            .clone()
-            .map(|repo| repo.update(cx, |repo, _| repo.branches()));
+
+        let (all_branches, occupied_branches) = repository
+            .as_ref()
+            .map(|repo| {
+                let snapshot = repo.read(cx);
+                let branches = process_branches(&snapshot.branch_list);
+                let occupied =
+                    compute_occupied_branches(&snapshot.linked_worktrees, &project_worktree_paths);
+                (branches, occupied)
+            })
+            .unwrap_or_default();
+
         let default_branch_request = repository
             .clone()
             .map(|repo| repo.update(cx, |repo, _| repo.default_branch(false)));
-        let worktrees_request = repository.map(|repo| repo.update(cx, |repo, _| repo.worktrees()));
 
         let (worktree_name, branch_target) = match current_target {
             StartThreadIn::NewWorktree {
@@ -75,8 +83,8 @@ impl ThreadBranchPicker {
 
         let delegate = ThreadBranchPickerDelegate {
             matches: vec![ThreadBranchEntry::CurrentBranch],
-            all_branches: None,
-            occupied_branches: None,
+            all_branches,
+            occupied_branches,
             selected_index: 0,
             worktree_name,
             branch_target,
@@ -95,74 +103,50 @@ impl ThreadBranchPicker {
 
         let focus_handle = picker.focus_handle(cx);
 
-        if let (Some(branches_request), Some(default_branch_request), Some(worktrees_request)) =
-            (branches_request, default_branch_request, worktrees_request)
-        {
+        let mut subscriptions = Vec::new();
+
+        if let Some(repo) = &repository {
+            subscriptions.push(cx.subscribe_in(
+                repo,
+                window,
+                |this, repo, event: &RepositoryEvent, window, cx| match event {
+                    RepositoryEvent::BranchListChanged => {
+                        let all_branches = process_branches(&repo.read(cx).branch_list);
+                        this.picker.update(cx, |picker, cx| {
+                            picker.delegate.all_branches = all_branches;
+                            picker.refresh(window, cx);
+                        });
+                    }
+                    RepositoryEvent::GitWorktreeListChanged => {
+                        let project_worktree_paths =
+                            this.picker.read(cx).delegate.project_worktree_paths.clone();
+                        let occupied = compute_occupied_branches(
+                            &repo.read(cx).linked_worktrees,
+                            &project_worktree_paths,
+                        );
+                        this.picker.update(cx, |picker, cx| {
+                            picker.delegate.occupied_branches = occupied;
+                            picker.refresh(window, cx);
+                        });
+                    }
+                    _ => {}
+                },
+            ));
+        }
+
+        // Fetch default branch asynchronously since it requires a git operation
+        if let Some(default_branch_request) = default_branch_request {
             let picker_handle = picker.downgrade();
             cx.spawn_in(window, async move |_this, cx| {
-                let branches = branches_request.await??;
-                let default_branch = default_branch_request.await.ok().and_then(Result::ok).flatten();
-                let worktrees = worktrees_request.await??;
-
-                let remote_upstreams: CollectionsHashSet<_> = branches
-                    .iter()
-                    .filter_map(|branch| {
-                        branch
-                            .upstream
-                            .as_ref()
-                            .filter(|upstream| upstream.is_remote())
-                            .map(|upstream| upstream.ref_name.clone())
-                    })
-                    .collect();
-
-                let mut occupied_branches = HashMap::new();
-                for worktree in worktrees {
-                    let Some(branch_name) = worktree.branch_name().map(ToOwned::to_owned) else {
-                        continue;
-                    };
-
-                    let reason = if picker_handle
-                        .read_with(cx, |picker, _| {
-                            picker
-                                .delegate
-                                .project_worktree_paths
-                                .contains(&worktree.path)
-                        })
-                        .unwrap_or(false)
-                    {
-                        format!(
-                            "This branch is already checked out in the current project worktree at {}.",
-                            worktree.path.display()
-                        )
-                    } else {
-                        format!(
-                            "This branch is already checked out in a linked worktree at {}.",
-                            worktree.path.display()
-                        )
-                    };
-
-                    occupied_branches.insert(branch_name, reason);
-                }
-
-                let mut all_branches: Vec<_> = branches
-                    .into_iter()
-                    .filter(|branch| !remote_upstreams.contains(&branch.ref_name))
-                    .collect();
-                all_branches.sort_by_key(|branch| {
-                    (
-                        branch.is_remote(),
-                        !branch.is_head,
-                        branch
-                            .most_recent_commit
-                            .as_ref()
-                            .map(|commit| 0 - commit.commit_timestamp),
-                    )
-                });
+                let default_branch = default_branch_request
+                    .await
+                    .ok()
+                    .and_then(Result::ok)
+                    .flatten();
 
                 picker_handle.update_in(cx, |picker, window, cx| {
-                    picker.delegate.all_branches = Some(all_branches);
-                    picker.delegate.occupied_branches = Some(occupied_branches);
-                    picker.delegate.default_branch_name = default_branch.map(|branch| branch.to_string());
+                    picker.delegate.default_branch_name =
+                        default_branch.map(|branch| branch.to_string());
                     picker.refresh(window, cx);
                 })?;
 
@@ -171,14 +155,14 @@ impl ThreadBranchPicker {
             .detach_and_log_err(cx);
         }
 
-        let subscription = cx.subscribe(&picker, |_, _, _, cx| {
+        subscriptions.push(cx.subscribe(&picker, |_, _, _, cx| {
             cx.emit(DismissEvent);
-        });
+        }));
 
         Self {
             picker,
             focus_handle,
-            _subscription: subscription,
+            _subscriptions: subscriptions,
         }
     }
 }
@@ -219,8 +203,8 @@ enum ThreadBranchEntry {
 
 pub(crate) struct ThreadBranchPickerDelegate {
     matches: Vec<ThreadBranchEntry>,
-    all_branches: Option<Vec<GitBranch>>,
-    occupied_branches: Option<HashMap<String, String>>,
+    all_branches: Vec<GitBranch>,
+    occupied_branches: HashMap<String, String>,
     selected_index: usize,
     worktree_name: Option<String>,
     branch_target: NewWorktreeBranchTarget,
@@ -230,6 +214,65 @@ pub(crate) struct ThreadBranchPickerDelegate {
     has_multiple_repositories: bool,
 }
 
+fn process_branches(branches: &Arc<[GitBranch]>) -> Vec<GitBranch> {
+    let remote_upstreams: HashSet<_> = branches
+        .iter()
+        .filter_map(|branch| {
+            branch
+                .upstream
+                .as_ref()
+                .filter(|upstream| upstream.is_remote())
+                .map(|upstream| upstream.ref_name.clone())
+        })
+        .collect();
+
+    let mut result: Vec<GitBranch> = branches
+        .iter()
+        .filter(|branch| !remote_upstreams.contains(&branch.ref_name))
+        .cloned()
+        .collect();
+
+    result.sort_by_key(|branch| {
+        (
+            branch.is_remote(),
+            !branch.is_head,
+            branch
+                .most_recent_commit
+                .as_ref()
+                .map(|commit| 0 - commit.commit_timestamp),
+        )
+    });
+
+    result
+}
+
+fn compute_occupied_branches(
+    worktrees: &[GitWorktree],
+    project_worktree_paths: &HashSet<PathBuf>,
+) -> HashMap<String, String> {
+    let mut occupied_branches = HashMap::default();
+    for worktree in worktrees {
+        let Some(branch_name) = worktree.branch_name().map(ToOwned::to_owned) else {
+            continue;
+        };
+
+        let reason = if project_worktree_paths.contains(&worktree.path) {
+            format!(
+                "This branch is already checked out in the current project worktree at {}.",
+                worktree.path.display()
+            )
+        } else {
+            format!(
+                "This branch is already checked out in a linked worktree at {}.",
+                worktree.path.display()
+            )
+        };
+
+        occupied_branches.insert(branch_name, reason);
+    }
+    occupied_branches
+}
+
 impl ThreadBranchPickerDelegate {
     fn new_worktree_action(&self, branch_target: NewWorktreeBranchTarget) -> StartThreadIn {
         StartThreadIn::NewWorktree {
@@ -271,9 +314,7 @@ impl ThreadBranchPickerDelegate {
     }
 
     fn is_branch_occupied(&self, branch_name: &str) -> bool {
-        self.occupied_branches
-            .as_ref()
-            .is_some_and(|occupied| occupied.contains_key(branch_name))
+        self.occupied_branches.contains_key(branch_name)
     }
 
     fn branch_aside_text(&self, branch_name: &str, is_remote: bool) -> Option<SharedString> {
@@ -456,11 +497,7 @@ impl PickerDelegate for ThreadBranchPickerDelegate {
             return Task::ready(());
         }
 
-        let Some(all_branches) = self.all_branches.clone() else {
-            self.matches = self.fixed_matches();
-            self.selected_index = 0;
-            return Task::ready(());
-        };
+        let all_branches = self.all_branches.clone();
 
         if query.is_empty() {
             let mut matches = self.fixed_matches();

crates/git_ui/src/branch_picker.rs 🔗

@@ -11,7 +11,7 @@ use gpui::{
     SharedString, Styled, Subscription, Task, WeakEntity, Window, actions, rems,
 };
 use picker::{Picker, PickerDelegate, PickerEditorPosition};
-use project::git_store::Repository;
+use project::git_store::{Repository, RepositoryEvent};
 use project::project_settings::ProjectSettings;
 use settings::Settings;
 use std::sync::Arc;
@@ -113,7 +113,7 @@ pub struct BranchList {
     width: Rems,
     pub picker: Entity<Picker<BranchListDelegate>>,
     picker_focus_handle: FocusHandle,
-    _subscription: Option<Subscription>,
+    _subscriptions: Vec<Subscription>,
     embedded: bool,
 }
 
@@ -127,9 +127,10 @@ impl BranchList {
         cx: &mut Context<Self>,
     ) -> Self {
         let mut this = Self::new_inner(workspace, repository, style, width, false, window, cx);
-        this._subscription = Some(cx.subscribe(&this.picker, |_, _, _, cx| {
-            cx.emit(DismissEvent);
-        }));
+        this._subscriptions
+            .push(cx.subscribe(&this.picker, |_, _, _, cx| {
+                cx.emit(DismissEvent);
+            }));
         this
     }
 
@@ -142,18 +143,55 @@ impl BranchList {
         window: &mut Window,
         cx: &mut Context<Self>,
     ) -> Self {
-        let all_branches_request = repository
-            .clone()
-            .map(|repository| repository.update(cx, |repository, _| repository.branches()));
+        let all_branches = repository
+            .as_ref()
+            .map(|repo| process_branches(&repo.read(cx).branch_list))
+            .unwrap_or_default();
 
         let default_branch_request = repository.clone().map(|repository| {
             repository.update(cx, |repository, _| repository.default_branch(false))
         });
 
+        let mut delegate = BranchListDelegate::new(workspace, repository.clone(), style, cx);
+        delegate.all_branches = all_branches;
+
+        let picker = cx.new(|cx| {
+            Picker::uniform_list(delegate, window, cx)
+                .show_scrollbar(true)
+                .modal(!embedded)
+        });
+        let picker_focus_handle = picker.focus_handle(cx);
+
+        picker.update(cx, |picker, _| {
+            picker.delegate.focus_handle = picker_focus_handle.clone();
+        });
+
+        let mut subscriptions = Vec::new();
+
+        if let Some(repo) = &repository {
+            subscriptions.push(cx.subscribe_in(
+                repo,
+                window,
+                move |this, repo, event, window, cx| {
+                    if matches!(event, RepositoryEvent::BranchListChanged) {
+                        let branch_list = repo.read(cx).branch_list.clone();
+                        let all_branches = process_branches(&branch_list);
+                        this.picker.update(cx, |picker, cx| {
+                            picker.delegate.restore_selected_branch = picker
+                                .delegate
+                                .matches
+                                .get(picker.delegate.selected_index)
+                                .and_then(|entry| entry.as_branch().map(|b| b.ref_name.clone()));
+                            picker.delegate.all_branches = all_branches;
+                            picker.refresh(window, cx);
+                        });
+                    }
+                },
+            ));
+        }
+
+        // Fetch default branch asynchronously since it requires a git operation
         cx.spawn_in(window, async move |this, cx| {
-            let mut all_branches = all_branches_request
-                .context("No active repository")?
-                .await??;
             let default_branch = default_branch_request
                 .context("No active repository")?
                 .await
@@ -162,64 +200,21 @@ impl BranchList {
                 .flatten()
                 .flatten();
 
-            let all_branches = cx
-                .background_spawn(async move {
-                    let remote_upstreams: HashSet<_> = all_branches
-                        .iter()
-                        .filter_map(|branch| {
-                            branch
-                                .upstream
-                                .as_ref()
-                                .filter(|upstream| upstream.is_remote())
-                                .map(|upstream| upstream.ref_name.clone())
-                        })
-                        .collect();
-
-                    all_branches.retain(|branch| !remote_upstreams.contains(&branch.ref_name));
-
-                    all_branches.sort_by_key(|branch| {
-                        (
-                            !branch.is_head, // Current branch (is_head=true) comes first
-                            branch
-                                .most_recent_commit
-                                .as_ref()
-                                .map(|commit| 0 - commit.commit_timestamp),
-                        )
-                    });
-
-                    all_branches
-                })
-                .await;
-
-            let _ = this.update_in(cx, |this, window, cx| {
-                this.picker.update(cx, |picker, cx| {
+            let _ = this.update_in(cx, |this, _window, cx| {
+                this.picker.update(cx, |picker, _cx| {
                     picker.delegate.default_branch = default_branch;
-                    picker.delegate.all_branches = Some(all_branches);
-                    picker.refresh(window, cx);
-                })
+                });
             });
 
             anyhow::Ok(())
         })
         .detach_and_log_err(cx);
 
-        let delegate = BranchListDelegate::new(workspace, repository, style, cx);
-        let picker = cx.new(|cx| {
-            Picker::uniform_list(delegate, window, cx)
-                .show_scrollbar(true)
-                .modal(!embedded)
-        });
-        let picker_focus_handle = picker.focus_handle(cx);
-
-        picker.update(cx, |picker, _| {
-            picker.delegate.focus_handle = picker_focus_handle.clone();
-        });
-
         Self {
             picker,
             picker_focus_handle,
             width,
-            _subscription: None,
+            _subscriptions: subscriptions,
             embedded,
         }
     }
@@ -240,9 +235,10 @@ impl BranchList {
             window,
             cx,
         );
-        this._subscription = Some(cx.subscribe(&this.picker, |_, _, _, cx| {
-            cx.emit(DismissEvent);
-        }));
+        this._subscriptions
+            .push(cx.subscribe(&this.picker, |_, _, _, cx| {
+                cx.emit(DismissEvent);
+            }));
         this
     }
 
@@ -379,7 +375,7 @@ impl BranchFilter {
 pub struct BranchListDelegate {
     workspace: WeakEntity<Workspace>,
     matches: Vec<Entry>,
-    all_branches: Option<Vec<Branch>>,
+    all_branches: Vec<Branch>,
     default_branch: Option<SharedString>,
     repo: Option<Entity<Repository>>,
     style: BranchListStyle,
@@ -389,6 +385,7 @@ pub struct BranchListDelegate {
     branch_filter: BranchFilter,
     state: PickerState,
     focus_handle: FocusHandle,
+    restore_selected_branch: Option<SharedString>,
 }
 
 #[derive(Debug)]
@@ -403,6 +400,37 @@ enum PickerState {
     NewBranch,
 }
 
+fn process_branches(branches: &Arc<[Branch]>) -> Vec<Branch> {
+    let remote_upstreams: HashSet<_> = branches
+        .iter()
+        .filter_map(|branch| {
+            branch
+                .upstream
+                .as_ref()
+                .filter(|upstream| upstream.is_remote())
+                .map(|upstream| upstream.ref_name.clone())
+        })
+        .collect();
+
+    let mut result: Vec<Branch> = branches
+        .iter()
+        .filter(|branch| !remote_upstreams.contains(&branch.ref_name))
+        .cloned()
+        .collect();
+
+    result.sort_by_key(|branch| {
+        (
+            !branch.is_head,
+            branch
+                .most_recent_commit
+                .as_ref()
+                .map(|commit| 0 - commit.commit_timestamp),
+        )
+    });
+
+    result
+}
+
 impl BranchListDelegate {
     fn new(
         workspace: WeakEntity<Workspace>,
@@ -415,7 +443,7 @@ impl BranchListDelegate {
             matches: vec![],
             repo,
             style,
-            all_branches: None,
+            all_branches: Vec::new(),
             default_branch: None,
             selected_index: 0,
             last_query: Default::default(),
@@ -423,6 +451,7 @@ impl BranchListDelegate {
             branch_filter: BranchFilter::All,
             state: PickerState::List,
             focus_handle: cx.focus_handle(),
+            restore_selected_branch: None,
         }
     }
 
@@ -536,9 +565,10 @@ impl BranchListDelegate {
                 picker.delegate.matches.retain(|e| e != &entry);
 
                 if let Entry::Branch { branch, .. } = &entry {
-                    if let Some(all_branches) = &mut picker.delegate.all_branches {
-                        all_branches.retain(|e| e.ref_name != branch.ref_name);
-                    }
+                    picker
+                        .delegate
+                        .all_branches
+                        .retain(|e| e.ref_name != branch.ref_name);
                 }
 
                 if picker.delegate.matches.is_empty() {
@@ -666,9 +696,7 @@ impl PickerDelegate for BranchListDelegate {
         window: &mut Window,
         cx: &mut Context<Picker<Self>>,
     ) -> Task<()> {
-        let Some(all_branches) = self.all_branches.clone() else {
-            return Task::ready(());
-        };
+        let all_branches = self.all_branches.clone();
 
         let branch_filter = self.branch_filter;
         cx.spawn_in(window, async move |picker, cx| {
@@ -770,6 +798,14 @@ impl PickerDelegate for BranchListDelegate {
                     delegate.matches = matches;
                     if delegate.matches.is_empty() {
                         delegate.selected_index = 0;
+                    } else if let Some(ref_name) = delegate.restore_selected_branch.take() {
+                        delegate.selected_index = delegate
+                            .matches
+                            .iter()
+                            .position(|entry| {
+                                entry.as_branch().is_some_and(|b| b.ref_name == ref_name)
+                            })
+                            .unwrap_or(0);
                     } else {
                         delegate.selected_index =
                             core::cmp::min(delegate.selected_index, delegate.matches.len() - 1);
@@ -1408,7 +1444,7 @@ mod tests {
                         BranchListStyle::Modal,
                         cx,
                     );
-                    delegate.all_branches = Some(branches);
+                    delegate.all_branches = branches;
                     let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx));
                     let picker_focus_handle = picker.focus_handle(cx);
                     picker.update(cx, |picker, _| {
@@ -1423,7 +1459,7 @@ mod tests {
                         picker,
                         picker_focus_handle,
                         width: rems(34.),
-                        _subscription: Some(_subscription),
+                        _subscriptions: vec![_subscription],
                         embedded: false,
                     }
                 })