1use std::collections::{HashMap, HashSet};
2use std::rc::Rc;
3
4use collections::HashSet as CollectionsHashSet;
5use std::path::PathBuf;
6use std::sync::Arc;
7
8use fuzzy::StringMatchCandidate;
9use git::repository::Branch as GitBranch;
10use gpui::{
11 AnyElement, App, Context, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable,
12 IntoElement, ParentElement, Render, SharedString, Styled, Task, Window, rems,
13};
14use picker::{Picker, PickerDelegate, PickerEditorPosition};
15use project::Project;
16use ui::{
17 Divider, DocumentationAside, HighlightedLabel, Icon, IconName, Label, LabelCommon, ListItem,
18 ListItemSpacing, prelude::*,
19};
20use util::ResultExt as _;
21
22use crate::{NewWorktreeBranchTarget, StartThreadIn};
23
24pub(crate) struct ThreadBranchPicker {
25 picker: Entity<Picker<ThreadBranchPickerDelegate>>,
26 focus_handle: FocusHandle,
27 _subscription: gpui::Subscription,
28}
29
30impl ThreadBranchPicker {
31 pub fn new(
32 project: Entity<Project>,
33 current_target: &StartThreadIn,
34 window: &mut Window,
35 cx: &mut Context<Self>,
36 ) -> Self {
37 let project_worktree_paths: HashSet<PathBuf> = project
38 .read(cx)
39 .visible_worktrees(cx)
40 .map(|worktree| worktree.read(cx).abs_path().to_path_buf())
41 .collect();
42
43 let has_multiple_repositories = project.read(cx).repositories(cx).len() > 1;
44 let current_branch_name = project
45 .read(cx)
46 .active_repository(cx)
47 .and_then(|repo| {
48 repo.read(cx)
49 .branch
50 .as_ref()
51 .map(|branch| branch.name().to_string())
52 })
53 .unwrap_or_else(|| "HEAD".to_string());
54
55 let repository = if has_multiple_repositories {
56 None
57 } else {
58 project.read(cx).active_repository(cx)
59 };
60 let branches_request = repository
61 .clone()
62 .map(|repo| repo.update(cx, |repo, _| repo.branches()));
63 let default_branch_request = repository
64 .clone()
65 .map(|repo| repo.update(cx, |repo, _| repo.default_branch(false)));
66 let worktrees_request = repository.map(|repo| repo.update(cx, |repo, _| repo.worktrees()));
67
68 let (worktree_name, branch_target) = match current_target {
69 StartThreadIn::NewWorktree {
70 worktree_name,
71 branch_target,
72 } => (worktree_name.clone(), branch_target.clone()),
73 _ => (None, NewWorktreeBranchTarget::default()),
74 };
75
76 let delegate = ThreadBranchPickerDelegate {
77 matches: vec![ThreadBranchEntry::CurrentBranch],
78 all_branches: None,
79 occupied_branches: None,
80 selected_index: 0,
81 worktree_name,
82 branch_target,
83 project_worktree_paths,
84 current_branch_name,
85 default_branch_name: None,
86 has_multiple_repositories,
87 };
88
89 let picker = cx.new(|cx| {
90 Picker::list(delegate, window, cx)
91 .list_measure_all()
92 .modal(false)
93 .max_height(Some(rems(20.).into()))
94 });
95
96 let focus_handle = picker.focus_handle(cx);
97
98 if let (Some(branches_request), Some(default_branch_request), Some(worktrees_request)) =
99 (branches_request, default_branch_request, worktrees_request)
100 {
101 let picker_handle = picker.downgrade();
102 cx.spawn_in(window, async move |_this, cx| {
103 let branches = branches_request.await??;
104 let default_branch = default_branch_request.await.ok().and_then(Result::ok).flatten();
105 let worktrees = worktrees_request.await??;
106
107 let remote_upstreams: CollectionsHashSet<_> = branches
108 .iter()
109 .filter_map(|branch| {
110 branch
111 .upstream
112 .as_ref()
113 .filter(|upstream| upstream.is_remote())
114 .map(|upstream| upstream.ref_name.clone())
115 })
116 .collect();
117
118 let mut occupied_branches = HashMap::new();
119 for worktree in worktrees {
120 let Some(branch_name) = worktree.branch_name().map(ToOwned::to_owned) else {
121 continue;
122 };
123
124 let reason = if picker_handle
125 .read_with(cx, |picker, _| {
126 picker
127 .delegate
128 .project_worktree_paths
129 .contains(&worktree.path)
130 })
131 .unwrap_or(false)
132 {
133 format!(
134 "This branch is already checked out in the current project worktree at {}.",
135 worktree.path.display()
136 )
137 } else {
138 format!(
139 "This branch is already checked out in a linked worktree at {}.",
140 worktree.path.display()
141 )
142 };
143
144 occupied_branches.insert(branch_name, reason);
145 }
146
147 let mut all_branches: Vec<_> = branches
148 .into_iter()
149 .filter(|branch| !remote_upstreams.contains(&branch.ref_name))
150 .collect();
151 all_branches.sort_by_key(|branch| {
152 (
153 branch.is_remote(),
154 !branch.is_head,
155 branch
156 .most_recent_commit
157 .as_ref()
158 .map(|commit| 0 - commit.commit_timestamp),
159 )
160 });
161
162 picker_handle.update_in(cx, |picker, window, cx| {
163 picker.delegate.all_branches = Some(all_branches);
164 picker.delegate.occupied_branches = Some(occupied_branches);
165 picker.delegate.default_branch_name = default_branch.map(|branch| branch.to_string());
166 picker.refresh(window, cx);
167 })?;
168
169 anyhow::Ok(())
170 })
171 .detach_and_log_err(cx);
172 }
173
174 let subscription = cx.subscribe(&picker, |_, _, _, cx| {
175 cx.emit(DismissEvent);
176 });
177
178 Self {
179 picker,
180 focus_handle,
181 _subscription: subscription,
182 }
183 }
184}
185
186impl Focusable for ThreadBranchPicker {
187 fn focus_handle(&self, _cx: &App) -> FocusHandle {
188 self.focus_handle.clone()
189 }
190}
191
192impl EventEmitter<DismissEvent> for ThreadBranchPicker {}
193
194impl Render for ThreadBranchPicker {
195 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
196 v_flex()
197 .w(rems(22.))
198 .elevation_3(cx)
199 .child(self.picker.clone())
200 .on_mouse_down_out(cx.listener(|_, _, _, cx| {
201 cx.emit(DismissEvent);
202 }))
203 }
204}
205
206#[derive(Clone)]
207enum ThreadBranchEntry {
208 CurrentBranch,
209 DefaultBranch,
210 Separator,
211 ExistingBranch {
212 branch: GitBranch,
213 positions: Vec<usize>,
214 },
215 CreateNamed {
216 name: String,
217 },
218}
219
220pub(crate) struct ThreadBranchPickerDelegate {
221 matches: Vec<ThreadBranchEntry>,
222 all_branches: Option<Vec<GitBranch>>,
223 occupied_branches: Option<HashMap<String, String>>,
224 selected_index: usize,
225 worktree_name: Option<String>,
226 branch_target: NewWorktreeBranchTarget,
227 project_worktree_paths: HashSet<PathBuf>,
228 current_branch_name: String,
229 default_branch_name: Option<String>,
230 has_multiple_repositories: bool,
231}
232
233impl ThreadBranchPickerDelegate {
234 fn new_worktree_action(&self, branch_target: NewWorktreeBranchTarget) -> StartThreadIn {
235 StartThreadIn::NewWorktree {
236 worktree_name: self.worktree_name.clone(),
237 branch_target,
238 }
239 }
240
241 fn selected_entry_name(&self) -> Option<&str> {
242 match &self.branch_target {
243 NewWorktreeBranchTarget::CurrentBranch => None,
244 NewWorktreeBranchTarget::ExistingBranch { name } => Some(name),
245 NewWorktreeBranchTarget::CreateBranch {
246 from_ref: Some(from_ref),
247 ..
248 } => Some(from_ref),
249 NewWorktreeBranchTarget::CreateBranch { name, .. } => Some(name),
250 }
251 }
252
253 fn prefer_create_entry(&self) -> bool {
254 matches!(
255 &self.branch_target,
256 NewWorktreeBranchTarget::CreateBranch { from_ref: None, .. }
257 )
258 }
259
260 fn fixed_matches(&self) -> Vec<ThreadBranchEntry> {
261 let mut matches = vec![ThreadBranchEntry::CurrentBranch];
262 if !self.has_multiple_repositories
263 && self
264 .default_branch_name
265 .as_ref()
266 .is_some_and(|default_branch_name| default_branch_name != &self.current_branch_name)
267 {
268 matches.push(ThreadBranchEntry::DefaultBranch);
269 }
270 matches
271 }
272
273 fn is_branch_occupied(&self, branch_name: &str) -> bool {
274 self.occupied_branches
275 .as_ref()
276 .is_some_and(|occupied| occupied.contains_key(branch_name))
277 }
278
279 fn branch_aside_text(&self, branch_name: &str, is_remote: bool) -> Option<SharedString> {
280 if self.is_branch_occupied(branch_name) {
281 Some(
282 format!(
283 "This branch is already checked out in another worktree. \
284 A new branch will be created from {branch_name}."
285 )
286 .into(),
287 )
288 } else if is_remote {
289 Some("A new local branch will be created from this remote branch.".into())
290 } else {
291 None
292 }
293 }
294
295 fn entry_aside_text(&self, entry: &ThreadBranchEntry) -> Option<SharedString> {
296 match entry {
297 ThreadBranchEntry::CurrentBranch => Some(SharedString::from(
298 "A new branch will be created from the current branch.",
299 )),
300 ThreadBranchEntry::DefaultBranch => {
301 let default_branch_name = self
302 .default_branch_name
303 .as_ref()
304 .filter(|name| *name != &self.current_branch_name)?;
305 self.branch_aside_text(default_branch_name, false)
306 }
307 ThreadBranchEntry::ExistingBranch { branch, .. } => {
308 self.branch_aside_text(branch.name(), branch.is_remote())
309 }
310 _ => None,
311 }
312 }
313
314 fn sync_selected_index(&mut self) {
315 let selected_entry_name = self.selected_entry_name().map(ToOwned::to_owned);
316 let prefer_create = self.prefer_create_entry();
317
318 if prefer_create {
319 if let Some(ref selected_entry_name) = selected_entry_name {
320 if let Some(index) = self.matches.iter().position(|entry| {
321 matches!(
322 entry,
323 ThreadBranchEntry::CreateNamed { name } if name == selected_entry_name
324 )
325 }) {
326 self.selected_index = index;
327 return;
328 }
329 }
330 } else if let Some(ref selected_entry_name) = selected_entry_name {
331 if selected_entry_name == &self.current_branch_name {
332 if let Some(index) = self
333 .matches
334 .iter()
335 .position(|entry| matches!(entry, ThreadBranchEntry::CurrentBranch))
336 {
337 self.selected_index = index;
338 return;
339 }
340 }
341
342 if self
343 .default_branch_name
344 .as_ref()
345 .is_some_and(|default_branch_name| default_branch_name == selected_entry_name)
346 {
347 if let Some(index) = self
348 .matches
349 .iter()
350 .position(|entry| matches!(entry, ThreadBranchEntry::DefaultBranch))
351 {
352 self.selected_index = index;
353 return;
354 }
355 }
356
357 if let Some(index) = self.matches.iter().position(|entry| {
358 matches!(
359 entry,
360 ThreadBranchEntry::ExistingBranch { branch, .. }
361 if branch.name() == selected_entry_name.as_str()
362 )
363 }) {
364 self.selected_index = index;
365 return;
366 }
367 }
368
369 if self.matches.len() > 1
370 && self
371 .matches
372 .iter()
373 .skip(1)
374 .all(|entry| matches!(entry, ThreadBranchEntry::CreateNamed { .. }))
375 {
376 self.selected_index = 1;
377 return;
378 }
379
380 self.selected_index = 0;
381 }
382}
383
384impl PickerDelegate for ThreadBranchPickerDelegate {
385 type ListItem = AnyElement;
386
387 fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
388 "Search branches…".into()
389 }
390
391 fn editor_position(&self) -> PickerEditorPosition {
392 PickerEditorPosition::Start
393 }
394
395 fn match_count(&self) -> usize {
396 self.matches.len()
397 }
398
399 fn selected_index(&self) -> usize {
400 self.selected_index
401 }
402
403 fn set_selected_index(
404 &mut self,
405 ix: usize,
406 _window: &mut Window,
407 _cx: &mut Context<Picker<Self>>,
408 ) {
409 self.selected_index = ix;
410 }
411
412 fn can_select(&self, ix: usize, _window: &mut Window, _cx: &mut Context<Picker<Self>>) -> bool {
413 !matches!(self.matches.get(ix), Some(ThreadBranchEntry::Separator))
414 }
415
416 fn update_matches(
417 &mut self,
418 query: String,
419 window: &mut Window,
420 cx: &mut Context<Picker<Self>>,
421 ) -> Task<()> {
422 if self.has_multiple_repositories {
423 let mut matches = self.fixed_matches();
424
425 if query.is_empty() {
426 if let Some(name) = self.selected_entry_name().map(ToOwned::to_owned) {
427 if self.prefer_create_entry() {
428 matches.push(ThreadBranchEntry::Separator);
429 matches.push(ThreadBranchEntry::CreateNamed { name });
430 }
431 }
432 } else {
433 matches.push(ThreadBranchEntry::Separator);
434 matches.push(ThreadBranchEntry::CreateNamed {
435 name: query.replace(' ', "-"),
436 });
437 }
438
439 self.matches = matches;
440 self.sync_selected_index();
441 return Task::ready(());
442 }
443
444 let Some(all_branches) = self.all_branches.clone() else {
445 self.matches = self.fixed_matches();
446 self.selected_index = 0;
447 return Task::ready(());
448 };
449
450 if query.is_empty() {
451 let mut matches = self.fixed_matches();
452 let filtered_branches: Vec<_> = all_branches
453 .into_iter()
454 .filter(|branch| {
455 branch.name() != self.current_branch_name
456 && self
457 .default_branch_name
458 .as_ref()
459 .is_none_or(|default_branch_name| branch.name() != default_branch_name)
460 })
461 .collect();
462
463 if !filtered_branches.is_empty() {
464 matches.push(ThreadBranchEntry::Separator);
465 }
466 for branch in filtered_branches {
467 matches.push(ThreadBranchEntry::ExistingBranch {
468 branch,
469 positions: Vec::new(),
470 });
471 }
472
473 if let Some(selected_entry_name) = self.selected_entry_name().map(ToOwned::to_owned) {
474 let has_existing = matches.iter().any(|entry| {
475 matches!(
476 entry,
477 ThreadBranchEntry::ExistingBranch { branch, .. }
478 if branch.name() == selected_entry_name
479 )
480 });
481 if self.prefer_create_entry() && !has_existing {
482 matches.push(ThreadBranchEntry::CreateNamed {
483 name: selected_entry_name,
484 });
485 }
486 }
487
488 self.matches = matches;
489 self.sync_selected_index();
490 return Task::ready(());
491 }
492
493 let candidates: Vec<_> = all_branches
494 .iter()
495 .enumerate()
496 .map(|(ix, branch)| StringMatchCandidate::new(ix, branch.name()))
497 .collect();
498 let executor = cx.background_executor().clone();
499 let query_clone = query.clone();
500 let normalized_query = query.replace(' ', "-");
501
502 let task = cx.background_executor().spawn(async move {
503 fuzzy::match_strings(
504 &candidates,
505 &query_clone,
506 true,
507 true,
508 10000,
509 &Default::default(),
510 executor,
511 )
512 .await
513 });
514
515 let all_branches_clone = all_branches;
516 cx.spawn_in(window, async move |picker, cx| {
517 let fuzzy_matches = task.await;
518
519 picker
520 .update_in(cx, |picker, _window, cx| {
521 let mut matches = picker.delegate.fixed_matches();
522 let mut has_dynamic_entries = false;
523
524 for candidate in &fuzzy_matches {
525 let branch = all_branches_clone[candidate.candidate_id].clone();
526 if branch.name() == picker.delegate.current_branch_name
527 || picker.delegate.default_branch_name.as_ref().is_some_and(
528 |default_branch_name| branch.name() == default_branch_name,
529 )
530 {
531 continue;
532 }
533 if !has_dynamic_entries {
534 matches.push(ThreadBranchEntry::Separator);
535 has_dynamic_entries = true;
536 }
537 matches.push(ThreadBranchEntry::ExistingBranch {
538 branch,
539 positions: candidate.positions.clone(),
540 });
541 }
542
543 if fuzzy_matches.is_empty() {
544 if !has_dynamic_entries {
545 matches.push(ThreadBranchEntry::Separator);
546 }
547 matches.push(ThreadBranchEntry::CreateNamed {
548 name: normalized_query.clone(),
549 });
550 }
551
552 picker.delegate.matches = matches;
553 if let Some(index) =
554 picker.delegate.matches.iter().position(|entry| {
555 matches!(entry, ThreadBranchEntry::ExistingBranch { .. })
556 })
557 {
558 picker.delegate.selected_index = index;
559 } else if !fuzzy_matches.is_empty() {
560 picker.delegate.selected_index = 0;
561 } else if let Some(index) =
562 picker.delegate.matches.iter().position(|entry| {
563 matches!(entry, ThreadBranchEntry::CreateNamed { .. })
564 })
565 {
566 picker.delegate.selected_index = index;
567 } else {
568 picker.delegate.sync_selected_index();
569 }
570 cx.notify();
571 })
572 .log_err();
573 })
574 }
575
576 fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
577 let Some(entry) = self.matches.get(self.selected_index) else {
578 return;
579 };
580
581 match entry {
582 ThreadBranchEntry::Separator => return,
583 ThreadBranchEntry::CurrentBranch => {
584 window.dispatch_action(
585 Box::new(self.new_worktree_action(NewWorktreeBranchTarget::CurrentBranch)),
586 cx,
587 );
588 }
589 ThreadBranchEntry::DefaultBranch => {
590 let Some(default_branch_name) = self.default_branch_name.clone() else {
591 return;
592 };
593 window.dispatch_action(
594 Box::new(
595 self.new_worktree_action(NewWorktreeBranchTarget::ExistingBranch {
596 name: default_branch_name,
597 }),
598 ),
599 cx,
600 );
601 }
602 ThreadBranchEntry::ExistingBranch { branch, .. } => {
603 let branch_target = if branch.is_remote() {
604 let branch_name = branch
605 .ref_name
606 .as_ref()
607 .strip_prefix("refs/remotes/")
608 .and_then(|stripped| stripped.split_once('/').map(|(_, name)| name))
609 .unwrap_or(branch.name())
610 .to_string();
611 NewWorktreeBranchTarget::CreateBranch {
612 name: branch_name,
613 from_ref: Some(branch.name().to_string()),
614 }
615 } else {
616 NewWorktreeBranchTarget::ExistingBranch {
617 name: branch.name().to_string(),
618 }
619 };
620 window.dispatch_action(Box::new(self.new_worktree_action(branch_target)), cx);
621 }
622 ThreadBranchEntry::CreateNamed { name } => {
623 window.dispatch_action(
624 Box::new(
625 self.new_worktree_action(NewWorktreeBranchTarget::CreateBranch {
626 name: name.clone(),
627 from_ref: None,
628 }),
629 ),
630 cx,
631 );
632 }
633 }
634
635 cx.emit(DismissEvent);
636 }
637
638 fn dismissed(&mut self, _window: &mut Window, _cx: &mut Context<Picker<Self>>) {}
639
640 fn render_match(
641 &self,
642 ix: usize,
643 selected: bool,
644 _window: &mut Window,
645 cx: &mut Context<Picker<Self>>,
646 ) -> Option<Self::ListItem> {
647 let entry = self.matches.get(ix)?;
648
649 match entry {
650 ThreadBranchEntry::Separator => Some(
651 div()
652 .py(DynamicSpacing::Base04.rems(cx))
653 .child(Divider::horizontal())
654 .into_any_element(),
655 ),
656 ThreadBranchEntry::CurrentBranch => {
657 let branch_name = if self.has_multiple_repositories {
658 SharedString::from("current branches")
659 } else {
660 SharedString::from(self.current_branch_name.clone())
661 };
662
663 Some(
664 ListItem::new("current-branch")
665 .inset(true)
666 .spacing(ListItemSpacing::Sparse)
667 .toggle_state(selected)
668 .child(Label::new(branch_name))
669 .into_any_element(),
670 )
671 }
672 ThreadBranchEntry::DefaultBranch => {
673 let default_branch_name = self
674 .default_branch_name
675 .as_ref()
676 .filter(|name| *name != &self.current_branch_name)?;
677 let is_occupied = self.is_branch_occupied(default_branch_name);
678
679 let item = ListItem::new("default-branch")
680 .inset(true)
681 .spacing(ListItemSpacing::Sparse)
682 .toggle_state(selected)
683 .child(Label::new(default_branch_name.clone()));
684
685 Some(
686 if is_occupied {
687 item.start_slot(Icon::new(IconName::GitBranchPlus).color(Color::Muted))
688 } else {
689 item
690 }
691 .into_any_element(),
692 )
693 }
694 ThreadBranchEntry::ExistingBranch {
695 branch, positions, ..
696 } => {
697 let branch_name = branch.name().to_string();
698 let needs_new_branch = self.is_branch_occupied(&branch_name) || branch.is_remote();
699
700 Some(
701 ListItem::new(SharedString::from(format!("branch-{ix}")))
702 .inset(true)
703 .spacing(ListItemSpacing::Sparse)
704 .toggle_state(selected)
705 .child(
706 h_flex()
707 .min_w_0()
708 .gap_1()
709 .child(
710 HighlightedLabel::new(branch_name, positions.clone())
711 .truncate(),
712 )
713 .when(needs_new_branch, |item| {
714 item.child(
715 Icon::new(IconName::GitBranchPlus)
716 .size(IconSize::Small)
717 .color(Color::Muted),
718 )
719 }),
720 )
721 .into_any_element(),
722 )
723 }
724 ThreadBranchEntry::CreateNamed { name } => Some(
725 ListItem::new("create-named-branch")
726 .inset(true)
727 .spacing(ListItemSpacing::Sparse)
728 .toggle_state(selected)
729 .child(Label::new(format!("Create Branch: \"{name}\"…")))
730 .into_any_element(),
731 ),
732 }
733 }
734
735 fn no_matches_text(&self, _window: &mut Window, _cx: &mut App) -> Option<SharedString> {
736 None
737 }
738
739 fn documentation_aside(
740 &self,
741 _window: &mut Window,
742 cx: &mut Context<Picker<Self>>,
743 ) -> Option<DocumentationAside> {
744 let entry = self.matches.get(self.selected_index)?;
745 let aside_text = self.entry_aside_text(entry)?;
746 let side = crate::ui::documentation_aside_side(cx);
747
748 Some(DocumentationAside::new(
749 side,
750 Rc::new(move |_| Label::new(aside_text.clone()).into_any_element()),
751 ))
752 }
753
754 fn documentation_aside_index(&self) -> Option<usize> {
755 let entry = self.matches.get(self.selected_index)?;
756 self.entry_aside_text(entry).map(|_| self.selected_index)
757 }
758}