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.is_head
194 .cmp(&lhs.is_head)
195 .then(rhs.unix_timestamp.cmp(&lhs.unix_timestamp))
196 });
197 branches.truncate(RECENT_BRANCHES_COUNT);
198 }
199 branches.sort_unstable_by(|lhs, rhs| {
200 rhs.is_head.cmp(&lhs.is_head).then(lhs.name.cmp(&rhs.name))
201 });
202 }
203 branches
204 .into_iter()
205 .enumerate()
206 .map(|(ix, command)| StringMatchCandidate::new(ix, &command.name))
207 .collect::<Vec<StringMatchCandidate>>()
208 });
209 let Some(candidates) = candidates.log_err() else {
210 return;
211 };
212 let matches: Vec<BranchEntry> = if query.is_empty() {
213 candidates
214 .into_iter()
215 .map(|candidate| BranchEntry::History(candidate.string))
216 .collect()
217 } else {
218 fuzzy::match_strings(
219 &candidates,
220 &query,
221 true,
222 10000,
223 &Default::default(),
224 cx.background_executor().clone(),
225 )
226 .await
227 .iter()
228 .cloned()
229 .map(BranchEntry::Branch)
230 .collect()
231 };
232 picker
233 .update(&mut cx, |picker, _| {
234 let delegate = &mut picker.delegate;
235 delegate.matches = matches;
236 if delegate.matches.is_empty() {
237 if !query.is_empty() {
238 delegate.matches.push(BranchEntry::NewBranch {
239 name: query.trim().replace(' ', "-"),
240 });
241 }
242
243 delegate.selected_index = 0;
244 } else {
245 delegate.selected_index =
246 core::cmp::min(delegate.selected_index, delegate.matches.len() - 1);
247 }
248 delegate.last_query = query;
249 })
250 .log_err();
251 })
252 }
253
254 fn confirm(&mut self, _: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
255 let Some(branch) = self.matches.get(self.selected_index()) else {
256 return;
257 };
258 cx.spawn_in(window, {
259 let branch = branch.clone();
260 |picker, mut cx| async move {
261 let branch_change_task = picker.update(&mut cx, |this, cx| {
262 let workspace = this
263 .delegate
264 .workspace
265 .upgrade()
266 .ok_or_else(|| anyhow!("workspace was dropped"))?;
267
268 let project = workspace.read(cx).project().read(cx);
269 let branch_to_checkout = match branch {
270 BranchEntry::Branch(branch) => branch.string,
271 BranchEntry::History(string) => string,
272 BranchEntry::NewBranch { name: branch_name } => branch_name,
273 };
274 let worktree = project
275 .visible_worktrees(cx)
276 .next()
277 .context("worktree disappeared")?;
278 let repository = ProjectPath::root_path(worktree.read(cx).id());
279
280 anyhow::Ok(project.update_or_create_branch(repository, branch_to_checkout, cx))
281 })??;
282
283 branch_change_task.await?;
284
285 picker.update(&mut cx, |_, cx| {
286 cx.emit(DismissEvent);
287
288 Ok::<(), anyhow::Error>(())
289 })
290 }
291 })
292 .detach_and_prompt_err("Failed to change branch", window, cx, |_, _, _| None);
293 }
294
295 fn dismissed(&mut self, _: &mut Window, cx: &mut Context<Picker<Self>>) {
296 cx.emit(DismissEvent);
297 }
298
299 fn render_match(
300 &self,
301 ix: usize,
302 selected: bool,
303 _window: &mut Window,
304 _cx: &mut Context<Picker<Self>>,
305 ) -> Option<Self::ListItem> {
306 let hit = &self.matches[ix];
307 let shortened_branch_name =
308 util::truncate_and_trailoff(&hit.name(), self.branch_name_trailoff_after);
309
310 Some(
311 ListItem::new(SharedString::from(format!("vcs-menu-{ix}")))
312 .inset(true)
313 .spacing(ListItemSpacing::Sparse)
314 .toggle_state(selected)
315 .when(matches!(hit, BranchEntry::History(_)), |el| {
316 el.end_slot(
317 Icon::new(IconName::HistoryRerun)
318 .color(Color::Muted)
319 .size(IconSize::Small),
320 )
321 })
322 .map(|el| match hit {
323 BranchEntry::Branch(branch) => {
324 let highlights: Vec<_> = branch
325 .positions
326 .iter()
327 .filter(|index| index < &&self.branch_name_trailoff_after)
328 .copied()
329 .collect();
330
331 el.child(HighlightedLabel::new(shortened_branch_name, highlights))
332 }
333 BranchEntry::History(_) => el.child(Label::new(shortened_branch_name)),
334 BranchEntry::NewBranch { name } => {
335 el.child(Label::new(format!("Create branch '{name}'")))
336 }
337 }),
338 )
339 }
340}