1use anyhow::{anyhow, Context as _, Result};
2use fuzzy::{StringMatch, StringMatchCandidate};
3
4use git::repository::Branch;
5use gpui::{
6 rems, App, AsyncApp, Context, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable,
7 InteractiveElement, IntoElement, ParentElement, Render, SharedString, Styled, Subscription,
8 Task, WeakEntity, Window,
9};
10use picker::{Picker, PickerDelegate};
11use project::ProjectPath;
12use std::sync::Arc;
13use ui::{prelude::*, HighlightedLabel, ListItem, ListItemSpacing};
14use util::ResultExt;
15use workspace::notifications::DetachAndPromptErr;
16use workspace::{ModalView, Workspace};
17
18pub fn init(cx: &mut App) {
19 cx.observe_new(|workspace: &mut Workspace, _, _| {
20 workspace.register_action(open);
21 })
22 .detach();
23}
24
25pub fn open(
26 _: &mut Workspace,
27 _: &zed_actions::git::Branch,
28 window: &mut Window,
29 cx: &mut Context<Workspace>,
30) {
31 let this = cx.entity().clone();
32 cx.spawn_in(window, |_, mut cx| async move {
33 // Modal branch picker has a longer trailoff than a popover one.
34 let delegate = BranchListDelegate::new(this.clone(), 70, &cx).await?;
35
36 this.update_in(&mut cx, |workspace, window, cx| {
37 workspace.toggle_modal(window, cx, |window, cx| {
38 BranchList::new(delegate, 34., window, cx)
39 })
40 })?;
41
42 Ok(())
43 })
44 .detach_and_prompt_err("Failed to read branches", window, cx, |_, _, _| None)
45}
46
47pub struct BranchList {
48 pub picker: Entity<Picker<BranchListDelegate>>,
49 rem_width: f32,
50 _subscription: Subscription,
51}
52
53impl BranchList {
54 pub fn new(
55 delegate: BranchListDelegate,
56 rem_width: f32,
57 window: &mut Window,
58 cx: &mut Context<Self>,
59 ) -> Self {
60 let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx));
61 let _subscription = cx.subscribe(&picker, |_, _, _, cx| cx.emit(DismissEvent));
62 Self {
63 picker,
64 rem_width,
65 _subscription,
66 }
67 }
68}
69impl ModalView for BranchList {}
70impl EventEmitter<DismissEvent> for BranchList {}
71
72impl Focusable for BranchList {
73 fn focus_handle(&self, cx: &App) -> FocusHandle {
74 self.picker.focus_handle(cx)
75 }
76}
77
78impl Render for BranchList {
79 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
80 v_flex()
81 .w(rems(self.rem_width))
82 .child(self.picker.clone())
83 .on_mouse_down_out(cx.listener(|this, _, window, cx| {
84 this.picker.update(cx, |this, cx| {
85 this.cancel(&Default::default(), window, cx);
86 })
87 }))
88 }
89}
90
91#[derive(Debug, Clone)]
92enum BranchEntry {
93 Branch(StringMatch),
94 History(String),
95 NewBranch { name: String },
96}
97
98impl BranchEntry {
99 fn name(&self) -> &str {
100 match self {
101 Self::Branch(branch) => &branch.string,
102 Self::History(branch) => &branch,
103 Self::NewBranch { name } => &name,
104 }
105 }
106}
107
108pub struct BranchListDelegate {
109 matches: Vec<BranchEntry>,
110 all_branches: Vec<Branch>,
111 workspace: WeakEntity<Workspace>,
112 selected_index: usize,
113 last_query: String,
114 /// Max length of branch name before we truncate it and add a trailing `...`.
115 branch_name_trailoff_after: usize,
116}
117
118impl BranchListDelegate {
119 pub async fn new(
120 workspace: Entity<Workspace>,
121 branch_name_trailoff_after: usize,
122 cx: &AsyncApp,
123 ) -> Result<Self> {
124 let all_branches_request = cx.update(|cx| {
125 let project = workspace.read(cx).project().read(cx);
126 let first_worktree = project
127 .visible_worktrees(cx)
128 .next()
129 .context("No worktrees found")?;
130 let project_path = ProjectPath::root_path(first_worktree.read(cx).id());
131 anyhow::Ok(project.branches(project_path, cx))
132 })??;
133
134 let all_branches = all_branches_request.await?;
135
136 Ok(Self {
137 matches: vec![],
138 workspace: workspace.downgrade(),
139 all_branches,
140 selected_index: 0,
141 last_query: Default::default(),
142 branch_name_trailoff_after,
143 })
144 }
145
146 pub fn branch_count(&self) -> usize {
147 self.matches
148 .iter()
149 .filter(|item| matches!(item, BranchEntry::Branch(_)))
150 .count()
151 }
152}
153
154impl PickerDelegate for BranchListDelegate {
155 type ListItem = ListItem;
156
157 fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
158 "Select branch...".into()
159 }
160
161 fn match_count(&self) -> usize {
162 self.matches.len()
163 }
164
165 fn selected_index(&self) -> usize {
166 self.selected_index
167 }
168
169 fn set_selected_index(
170 &mut self,
171 ix: usize,
172 _window: &mut Window,
173 _: &mut Context<Picker<Self>>,
174 ) {
175 self.selected_index = ix;
176 }
177
178 fn update_matches(
179 &mut self,
180 query: String,
181 window: &mut Window,
182 cx: &mut Context<Picker<Self>>,
183 ) -> Task<()> {
184 cx.spawn_in(window, move |picker, mut cx| async move {
185 let candidates = picker.update(&mut cx, |picker, _| {
186 const RECENT_BRANCHES_COUNT: usize = 10;
187 let mut branches = picker.delegate.all_branches.clone();
188 if query.is_empty() {
189 if branches.len() > RECENT_BRANCHES_COUNT {
190 // Truncate list of recent branches
191 // Do a partial sort to show recent-ish branches first.
192 branches.select_nth_unstable_by(RECENT_BRANCHES_COUNT - 1, |lhs, rhs| {
193 rhs.priority_key().cmp(&lhs.priority_key())
194 });
195 branches.truncate(RECENT_BRANCHES_COUNT);
196 }
197 branches.sort_unstable_by(|lhs, rhs| {
198 rhs.is_head.cmp(&lhs.is_head).then(lhs.name.cmp(&rhs.name))
199 });
200 }
201 branches
202 .into_iter()
203 .enumerate()
204 .map(|(ix, command)| StringMatchCandidate::new(ix, &command.name))
205 .collect::<Vec<StringMatchCandidate>>()
206 });
207 let Some(candidates) = candidates.log_err() else {
208 return;
209 };
210 let matches: Vec<BranchEntry> = if query.is_empty() {
211 candidates
212 .into_iter()
213 .map(|candidate| BranchEntry::History(candidate.string))
214 .collect()
215 } else {
216 fuzzy::match_strings(
217 &candidates,
218 &query,
219 true,
220 10000,
221 &Default::default(),
222 cx.background_executor().clone(),
223 )
224 .await
225 .iter()
226 .cloned()
227 .map(BranchEntry::Branch)
228 .collect()
229 };
230 picker
231 .update(&mut cx, |picker, _| {
232 let delegate = &mut picker.delegate;
233 delegate.matches = matches;
234 if delegate.matches.is_empty() {
235 if !query.is_empty() {
236 delegate.matches.push(BranchEntry::NewBranch {
237 name: query.trim().replace(' ', "-"),
238 });
239 }
240
241 delegate.selected_index = 0;
242 } else {
243 delegate.selected_index =
244 core::cmp::min(delegate.selected_index, delegate.matches.len() - 1);
245 }
246 delegate.last_query = query;
247 })
248 .log_err();
249 })
250 }
251
252 fn confirm(&mut self, _: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
253 let Some(branch) = self.matches.get(self.selected_index()) else {
254 return;
255 };
256
257 let current_branch = self
258 .workspace
259 .update(cx, |workspace, cx| {
260 workspace
261 .project()
262 .read(cx)
263 .active_repository(cx)
264 .and_then(|repo| repo.read(cx).current_branch())
265 .map(|branch| branch.name.to_string())
266 })
267 .ok()
268 .flatten();
269
270 if current_branch == Some(branch.name().to_string()) {
271 cx.emit(DismissEvent);
272 return;
273 }
274
275 cx.spawn_in(window, {
276 let branch = branch.clone();
277 |picker, mut cx| async move {
278 let branch_change_task = picker.update(&mut cx, |this, cx| {
279 let workspace = this
280 .delegate
281 .workspace
282 .upgrade()
283 .ok_or_else(|| anyhow!("workspace was dropped"))?;
284
285 let project = workspace.read(cx).project().read(cx);
286 let branch_to_checkout = match branch {
287 BranchEntry::Branch(branch) => branch.string,
288 BranchEntry::History(string) => string,
289 BranchEntry::NewBranch { name: branch_name } => branch_name,
290 };
291 let worktree = project
292 .visible_worktrees(cx)
293 .next()
294 .context("worktree disappeared")?;
295 let repository = ProjectPath::root_path(worktree.read(cx).id());
296
297 anyhow::Ok(project.update_or_create_branch(repository, branch_to_checkout, cx))
298 })??;
299
300 branch_change_task.await?;
301
302 picker.update(&mut cx, |_, cx| {
303 cx.emit(DismissEvent);
304
305 Ok::<(), anyhow::Error>(())
306 })
307 }
308 })
309 .detach_and_prompt_err("Failed to change branch", window, cx, |_, _, _| None);
310 }
311
312 fn dismissed(&mut self, _: &mut Window, cx: &mut Context<Picker<Self>>) {
313 cx.emit(DismissEvent);
314 }
315
316 fn render_match(
317 &self,
318 ix: usize,
319 selected: bool,
320 _window: &mut Window,
321 _cx: &mut Context<Picker<Self>>,
322 ) -> Option<Self::ListItem> {
323 let hit = &self.matches[ix];
324 let shortened_branch_name =
325 util::truncate_and_trailoff(&hit.name(), self.branch_name_trailoff_after);
326
327 Some(
328 ListItem::new(SharedString::from(format!("vcs-menu-{ix}")))
329 .inset(true)
330 .spacing(ListItemSpacing::Sparse)
331 .toggle_state(selected)
332 .when(matches!(hit, BranchEntry::History(_)), |el| {
333 el.end_slot(
334 Icon::new(IconName::HistoryRerun)
335 .color(Color::Muted)
336 .size(IconSize::Small),
337 )
338 })
339 .map(|el| match hit {
340 BranchEntry::Branch(branch) => {
341 let highlights: Vec<_> = branch
342 .positions
343 .iter()
344 .filter(|index| index < &&self.branch_name_trailoff_after)
345 .copied()
346 .collect();
347
348 el.child(HighlightedLabel::new(shortened_branch_name, highlights))
349 }
350 BranchEntry::History(_) => el.child(Label::new(shortened_branch_name)),
351 BranchEntry::NewBranch { name } => {
352 el.child(Label::new(format!("Create branch '{name}'")))
353 }
354 }),
355 )
356 }
357}