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_branch_name(&self, entry: &ThreadBranchEntry) -> Option<SharedString> {
296 match entry {
297 ThreadBranchEntry::CurrentBranch => {
298 Some(SharedString::from(self.current_branch_name.clone()))
299 }
300 ThreadBranchEntry::DefaultBranch => {
301 self.default_branch_name.clone().map(SharedString::from)
302 }
303 ThreadBranchEntry::ExistingBranch { branch, .. } => {
304 Some(SharedString::from(branch.name().to_string()))
305 }
306 _ => None,
307 }
308 }
309
310 fn entry_aside_text(&self, entry: &ThreadBranchEntry) -> Option<SharedString> {
311 match entry {
312 ThreadBranchEntry::CurrentBranch => Some(SharedString::from(
313 "A new branch will be created from the current branch.",
314 )),
315 ThreadBranchEntry::DefaultBranch => {
316 let default_branch_name = self
317 .default_branch_name
318 .as_ref()
319 .filter(|name| *name != &self.current_branch_name)?;
320 self.branch_aside_text(default_branch_name, false)
321 }
322 ThreadBranchEntry::ExistingBranch { branch, .. } => {
323 self.branch_aside_text(branch.name(), branch.is_remote())
324 }
325 _ => None,
326 }
327 }
328
329 fn sync_selected_index(&mut self) {
330 let selected_entry_name = self.selected_entry_name().map(ToOwned::to_owned);
331 let prefer_create = self.prefer_create_entry();
332
333 if prefer_create {
334 if let Some(ref selected_entry_name) = selected_entry_name {
335 if let Some(index) = self.matches.iter().position(|entry| {
336 matches!(
337 entry,
338 ThreadBranchEntry::CreateNamed { name } if name == selected_entry_name
339 )
340 }) {
341 self.selected_index = index;
342 return;
343 }
344 }
345 } else if let Some(ref selected_entry_name) = selected_entry_name {
346 if selected_entry_name == &self.current_branch_name {
347 if let Some(index) = self
348 .matches
349 .iter()
350 .position(|entry| matches!(entry, ThreadBranchEntry::CurrentBranch))
351 {
352 self.selected_index = index;
353 return;
354 }
355 }
356
357 if self
358 .default_branch_name
359 .as_ref()
360 .is_some_and(|default_branch_name| default_branch_name == selected_entry_name)
361 {
362 if let Some(index) = self
363 .matches
364 .iter()
365 .position(|entry| matches!(entry, ThreadBranchEntry::DefaultBranch))
366 {
367 self.selected_index = index;
368 return;
369 }
370 }
371
372 if let Some(index) = self.matches.iter().position(|entry| {
373 matches!(
374 entry,
375 ThreadBranchEntry::ExistingBranch { branch, .. }
376 if branch.name() == selected_entry_name.as_str()
377 )
378 }) {
379 self.selected_index = index;
380 return;
381 }
382 }
383
384 if self.matches.len() > 1
385 && self
386 .matches
387 .iter()
388 .skip(1)
389 .all(|entry| matches!(entry, ThreadBranchEntry::CreateNamed { .. }))
390 {
391 self.selected_index = 1;
392 return;
393 }
394
395 self.selected_index = 0;
396 }
397}
398
399impl PickerDelegate for ThreadBranchPickerDelegate {
400 type ListItem = AnyElement;
401
402 fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
403 "Search branches…".into()
404 }
405
406 fn editor_position(&self) -> PickerEditorPosition {
407 PickerEditorPosition::Start
408 }
409
410 fn match_count(&self) -> usize {
411 self.matches.len()
412 }
413
414 fn selected_index(&self) -> usize {
415 self.selected_index
416 }
417
418 fn set_selected_index(
419 &mut self,
420 ix: usize,
421 _window: &mut Window,
422 _cx: &mut Context<Picker<Self>>,
423 ) {
424 self.selected_index = ix;
425 }
426
427 fn can_select(&self, ix: usize, _window: &mut Window, _cx: &mut Context<Picker<Self>>) -> bool {
428 !matches!(self.matches.get(ix), Some(ThreadBranchEntry::Separator))
429 }
430
431 fn update_matches(
432 &mut self,
433 query: String,
434 window: &mut Window,
435 cx: &mut Context<Picker<Self>>,
436 ) -> Task<()> {
437 if self.has_multiple_repositories {
438 let mut matches = self.fixed_matches();
439
440 if query.is_empty() {
441 if let Some(name) = self.selected_entry_name().map(ToOwned::to_owned) {
442 if self.prefer_create_entry() {
443 matches.push(ThreadBranchEntry::Separator);
444 matches.push(ThreadBranchEntry::CreateNamed { name });
445 }
446 }
447 } else {
448 matches.push(ThreadBranchEntry::Separator);
449 matches.push(ThreadBranchEntry::CreateNamed {
450 name: query.replace(' ', "-"),
451 });
452 }
453
454 self.matches = matches;
455 self.sync_selected_index();
456 return Task::ready(());
457 }
458
459 let Some(all_branches) = self.all_branches.clone() else {
460 self.matches = self.fixed_matches();
461 self.selected_index = 0;
462 return Task::ready(());
463 };
464
465 if query.is_empty() {
466 let mut matches = self.fixed_matches();
467 let filtered_branches: Vec<_> = all_branches
468 .into_iter()
469 .filter(|branch| {
470 branch.name() != self.current_branch_name
471 && self
472 .default_branch_name
473 .as_ref()
474 .is_none_or(|default_branch_name| branch.name() != default_branch_name)
475 })
476 .collect();
477
478 if !filtered_branches.is_empty() {
479 matches.push(ThreadBranchEntry::Separator);
480 }
481 for branch in filtered_branches {
482 matches.push(ThreadBranchEntry::ExistingBranch {
483 branch,
484 positions: Vec::new(),
485 });
486 }
487
488 if let Some(selected_entry_name) = self.selected_entry_name().map(ToOwned::to_owned) {
489 let has_existing = matches.iter().any(|entry| {
490 matches!(
491 entry,
492 ThreadBranchEntry::ExistingBranch { branch, .. }
493 if branch.name() == selected_entry_name
494 )
495 });
496 if self.prefer_create_entry() && !has_existing {
497 matches.push(ThreadBranchEntry::CreateNamed {
498 name: selected_entry_name,
499 });
500 }
501 }
502
503 self.matches = matches;
504 self.sync_selected_index();
505 return Task::ready(());
506 }
507
508 let candidates: Vec<_> = all_branches
509 .iter()
510 .enumerate()
511 .map(|(ix, branch)| StringMatchCandidate::new(ix, branch.name()))
512 .collect();
513 let executor = cx.background_executor().clone();
514 let query_clone = query.clone();
515 let normalized_query = query.replace(' ', "-");
516
517 let task = cx.background_executor().spawn(async move {
518 fuzzy::match_strings(
519 &candidates,
520 &query_clone,
521 true,
522 true,
523 10000,
524 &Default::default(),
525 executor,
526 )
527 .await
528 });
529
530 let all_branches_clone = all_branches;
531 cx.spawn_in(window, async move |picker, cx| {
532 let fuzzy_matches = task.await;
533
534 picker
535 .update_in(cx, |picker, _window, cx| {
536 let mut matches = picker.delegate.fixed_matches();
537 let mut has_dynamic_entries = false;
538
539 for candidate in &fuzzy_matches {
540 let branch = all_branches_clone[candidate.candidate_id].clone();
541 if branch.name() == picker.delegate.current_branch_name
542 || picker.delegate.default_branch_name.as_ref().is_some_and(
543 |default_branch_name| branch.name() == default_branch_name,
544 )
545 {
546 continue;
547 }
548 if !has_dynamic_entries {
549 matches.push(ThreadBranchEntry::Separator);
550 has_dynamic_entries = true;
551 }
552 matches.push(ThreadBranchEntry::ExistingBranch {
553 branch,
554 positions: candidate.positions.clone(),
555 });
556 }
557
558 if fuzzy_matches.is_empty() {
559 if !has_dynamic_entries {
560 matches.push(ThreadBranchEntry::Separator);
561 }
562 matches.push(ThreadBranchEntry::CreateNamed {
563 name: normalized_query.clone(),
564 });
565 }
566
567 picker.delegate.matches = matches;
568 if let Some(index) =
569 picker.delegate.matches.iter().position(|entry| {
570 matches!(entry, ThreadBranchEntry::ExistingBranch { .. })
571 })
572 {
573 picker.delegate.selected_index = index;
574 } else if !fuzzy_matches.is_empty() {
575 picker.delegate.selected_index = 0;
576 } else if let Some(index) =
577 picker.delegate.matches.iter().position(|entry| {
578 matches!(entry, ThreadBranchEntry::CreateNamed { .. })
579 })
580 {
581 picker.delegate.selected_index = index;
582 } else {
583 picker.delegate.sync_selected_index();
584 }
585 cx.notify();
586 })
587 .log_err();
588 })
589 }
590
591 fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
592 let Some(entry) = self.matches.get(self.selected_index) else {
593 return;
594 };
595
596 match entry {
597 ThreadBranchEntry::Separator => return,
598 ThreadBranchEntry::CurrentBranch => {
599 window.dispatch_action(
600 Box::new(self.new_worktree_action(NewWorktreeBranchTarget::CurrentBranch)),
601 cx,
602 );
603 }
604 ThreadBranchEntry::DefaultBranch => {
605 let Some(default_branch_name) = self.default_branch_name.clone() else {
606 return;
607 };
608 window.dispatch_action(
609 Box::new(
610 self.new_worktree_action(NewWorktreeBranchTarget::ExistingBranch {
611 name: default_branch_name,
612 }),
613 ),
614 cx,
615 );
616 }
617 ThreadBranchEntry::ExistingBranch { branch, .. } => {
618 let branch_target = if branch.is_remote() {
619 let branch_name = branch
620 .ref_name
621 .as_ref()
622 .strip_prefix("refs/remotes/")
623 .and_then(|stripped| stripped.split_once('/').map(|(_, name)| name))
624 .unwrap_or(branch.name())
625 .to_string();
626 NewWorktreeBranchTarget::CreateBranch {
627 name: branch_name,
628 from_ref: Some(branch.name().to_string()),
629 }
630 } else {
631 NewWorktreeBranchTarget::ExistingBranch {
632 name: branch.name().to_string(),
633 }
634 };
635 window.dispatch_action(Box::new(self.new_worktree_action(branch_target)), cx);
636 }
637 ThreadBranchEntry::CreateNamed { name } => {
638 window.dispatch_action(
639 Box::new(
640 self.new_worktree_action(NewWorktreeBranchTarget::CreateBranch {
641 name: name.clone(),
642 from_ref: None,
643 }),
644 ),
645 cx,
646 );
647 }
648 }
649
650 cx.emit(DismissEvent);
651 }
652
653 fn dismissed(&mut self, _window: &mut Window, _cx: &mut Context<Picker<Self>>) {}
654
655 fn render_match(
656 &self,
657 ix: usize,
658 selected: bool,
659 _window: &mut Window,
660 cx: &mut Context<Picker<Self>>,
661 ) -> Option<Self::ListItem> {
662 let entry = self.matches.get(ix)?;
663
664 match entry {
665 ThreadBranchEntry::Separator => Some(
666 div()
667 .py(DynamicSpacing::Base04.rems(cx))
668 .child(Divider::horizontal())
669 .into_any_element(),
670 ),
671 ThreadBranchEntry::CurrentBranch => {
672 let branch_name = if self.has_multiple_repositories {
673 SharedString::from("current branches")
674 } else {
675 SharedString::from(self.current_branch_name.clone())
676 };
677
678 Some(
679 ListItem::new("current-branch")
680 .inset(true)
681 .spacing(ListItemSpacing::Sparse)
682 .toggle_state(selected)
683 .child(Label::new(branch_name))
684 .into_any_element(),
685 )
686 }
687 ThreadBranchEntry::DefaultBranch => {
688 let default_branch_name = self
689 .default_branch_name
690 .as_ref()
691 .filter(|name| *name != &self.current_branch_name)?;
692 let is_occupied = self.is_branch_occupied(default_branch_name);
693
694 let item = ListItem::new("default-branch")
695 .inset(true)
696 .spacing(ListItemSpacing::Sparse)
697 .toggle_state(selected)
698 .child(Label::new(default_branch_name.clone()));
699
700 Some(
701 if is_occupied {
702 item.start_slot(Icon::new(IconName::GitBranchPlus).color(Color::Muted))
703 } else {
704 item
705 }
706 .into_any_element(),
707 )
708 }
709 ThreadBranchEntry::ExistingBranch {
710 branch, positions, ..
711 } => {
712 let branch_name = branch.name().to_string();
713 let needs_new_branch = self.is_branch_occupied(&branch_name) || branch.is_remote();
714
715 Some(
716 ListItem::new(SharedString::from(format!("branch-{ix}")))
717 .inset(true)
718 .spacing(ListItemSpacing::Sparse)
719 .toggle_state(selected)
720 .child(
721 h_flex()
722 .min_w_0()
723 .gap_1()
724 .child(
725 HighlightedLabel::new(branch_name, positions.clone())
726 .truncate(),
727 )
728 .when(needs_new_branch, |item| {
729 item.child(
730 Icon::new(IconName::GitBranchPlus)
731 .size(IconSize::Small)
732 .color(Color::Muted),
733 )
734 }),
735 )
736 .into_any_element(),
737 )
738 }
739 ThreadBranchEntry::CreateNamed { name } => Some(
740 ListItem::new("create-named-branch")
741 .inset(true)
742 .spacing(ListItemSpacing::Sparse)
743 .toggle_state(selected)
744 .child(Label::new(format!("Create Branch: \"{name}\"…")))
745 .into_any_element(),
746 ),
747 }
748 }
749
750 fn no_matches_text(&self, _window: &mut Window, _cx: &mut App) -> Option<SharedString> {
751 None
752 }
753
754 fn documentation_aside(
755 &self,
756 _window: &mut Window,
757 cx: &mut Context<Picker<Self>>,
758 ) -> Option<DocumentationAside> {
759 let entry = self.matches.get(self.selected_index)?;
760 let branch_name = self.entry_branch_name(entry);
761 let aside_text = self.entry_aside_text(entry);
762
763 if branch_name.is_none() && aside_text.is_none() {
764 return None;
765 }
766
767 let side = crate::ui::documentation_aside_side(cx);
768
769 Some(DocumentationAside::new(
770 side,
771 Rc::new(move |cx| {
772 v_flex()
773 .gap_1()
774 .when_some(branch_name.clone(), |this, name| {
775 this.child(Label::new(name))
776 })
777 .when_some(aside_text.clone(), |this, text| {
778 this.child(
779 div()
780 .when(branch_name.is_some(), |this| {
781 this.pt_1()
782 .border_t_1()
783 .border_color(cx.theme().colors().border_variant)
784 })
785 .child(Label::new(text).color(Color::Muted)),
786 )
787 })
788 .into_any_element()
789 }),
790 ))
791 }
792
793 fn documentation_aside_index(&self) -> Option<usize> {
794 let entry = self.matches.get(self.selected_index)?;
795 if self.entry_branch_name(entry).is_some() || self.entry_aside_text(entry).is_some() {
796 Some(self.selected_index)
797 } else {
798 None
799 }
800 }
801}