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