ZIm/crates/git_ui/src/branch_picker.rs

use anyhow::{anyhow, Context as _};
use fuzzy::StringMatchCandidate;

use git::repository::Branch;
use gpui::{
    rems, App, Context, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable,
    InteractiveElement, IntoElement, Modifiers, ModifiersChangedEvent, ParentElement, Render,
    SharedString, Styled, Subscription, Task, Window,
};
use picker::{Picker, PickerDelegate};
use project::git::Repository;
use std::sync::Arc;
use time::OffsetDateTime;
use time_format::format_local_timestamp;
use ui::{
    prelude::*, HighlightedLabel, KeyBinding, ListItem, ListItemSpacing, PopoverMenuHandle, Tooltip,
};
use util::ResultExt;
use workspace::notifications::DetachAndPromptErr;
use workspace::{ModalView, Workspace};

pub fn init(cx: &mut App) {
    cx.observe_new(|workspace: &mut Workspace, _, _| {
        workspace.register_action(open);
        workspace.register_action(switch);
        workspace.register_action(checkout_branch);
    })
    .detach();
}

pub fn checkout_branch(
    workspace: &mut Workspace,
    _: &zed_actions::git::CheckoutBranch,
    window: &mut Window,
    cx: &mut Context<Workspace>,
) {
    open(workspace, &zed_actions::git::Branch, window, cx);
}

pub fn switch(
    workspace: &mut Workspace,
    _: &zed_actions::git::Switch,
    window: &mut Window,
    cx: &mut Context<Workspace>,
) {
    open(workspace, &zed_actions::git::Branch, window, cx);
}

pub fn open(
    workspace: &mut Workspace,
    _: &zed_actions::git::Branch,
    window: &mut Window,
    cx: &mut Context<Workspace>,
) {
    let repository = workspace.project().read(cx).active_repository(cx).clone();
    let style = BranchListStyle::Modal;
    workspace.toggle_modal(window, cx, |window, cx| {
        BranchList::new(repository, style, rems(34.), window, cx)
    })
}

pub fn popover(
    repository: Option<Entity<Repository>>,
    window: &mut Window,
    cx: &mut App,
) -> Entity<BranchList> {
    cx.new(|cx| {
        let list = BranchList::new(repository, BranchListStyle::Popover, rems(20.), window, cx);
        list.focus_handle(cx).focus(window);
        list
    })
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum BranchListStyle {
    Modal,
    Popover,
}

pub struct BranchList {
    width: Rems,
    pub popover_handle: PopoverMenuHandle<Self>,
    pub picker: Entity<Picker<BranchListDelegate>>,
    _subscription: Subscription,
}

impl BranchList {
    fn new(
        repository: Option<Entity<Repository>>,
        style: BranchListStyle,
        width: Rems,
        window: &mut Window,
        cx: &mut Context<Self>,
    ) -> Self {
        let popover_handle = PopoverMenuHandle::default();
        let all_branches_request = repository
            .clone()
            .map(|repository| repository.read(cx).branches());

        cx.spawn_in(window, |this, mut cx| async move {
            let mut all_branches = all_branches_request
                .context("No active repository")?
                .await??;

            all_branches.sort_by_key(|branch| {
                branch
                    .most_recent_commit
                    .as_ref()
                    .map(|commit| 0 - commit.commit_timestamp)
            });

            this.update_in(&mut cx, |this, window, cx| {
                this.picker.update(cx, |picker, cx| {
                    picker.delegate.all_branches = Some(all_branches);
                    picker.refresh(window, cx);
                })
            })?;

            anyhow::Ok(())
        })
        .detach_and_log_err(cx);

        let delegate = BranchListDelegate::new(repository.clone(), style);
        let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx));

        let _subscription = cx.subscribe(&picker, |_, _, _, cx| {
            cx.emit(DismissEvent);
        });

        Self {
            picker,
            width,
            popover_handle,
            _subscription,
        }
    }

    fn handle_modifiers_changed(
        &mut self,
        ev: &ModifiersChangedEvent,
        _: &mut Window,
        cx: &mut Context<Self>,
    ) {
        self.picker
            .update(cx, |picker, _| picker.delegate.modifiers = ev.modifiers)
    }
}
impl ModalView for BranchList {}
impl EventEmitter<DismissEvent> for BranchList {}

impl Focusable for BranchList {
    fn focus_handle(&self, cx: &App) -> FocusHandle {
        self.picker.focus_handle(cx)
    }
}

impl Render for BranchList {
    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
        v_flex()
            .w(self.width)
            .on_modifiers_changed(cx.listener(Self::handle_modifiers_changed))
            .child(self.picker.clone())
            .on_mouse_down_out({
                cx.listener(move |this, _, window, cx| {
                    this.picker.update(cx, |this, cx| {
                        this.cancel(&Default::default(), window, cx);
                    })
                })
            })
    }
}

#[derive(Debug, Clone)]
struct BranchEntry {
    branch: Branch,
    positions: Vec<usize>,
}

pub struct BranchListDelegate {
    matches: Vec<BranchEntry>,
    all_branches: Option<Vec<Branch>>,
    repo: Option<Entity<Repository>>,
    style: BranchListStyle,
    selected_index: usize,
    last_query: String,
    modifiers: Modifiers,
}

impl BranchListDelegate {
    fn new(repo: Option<Entity<Repository>>, style: BranchListStyle) -> Self {
        Self {
            matches: vec![],
            repo,
            style,
            all_branches: None,
            selected_index: 0,
            last_query: Default::default(),
            modifiers: Default::default(),
        }
    }

    fn has_exact_match(&self, target: &str) -> bool {
        self.matches.iter().any(|entry| entry.branch.name == target)
    }

    fn create_branch(
        &self,
        new_branch_name: SharedString,
        window: &mut Window,
        cx: &mut Context<Picker<Self>>,
    ) {
        let Some(repo) = self.repo.clone() else {
            return;
        };
        cx.spawn(|_, cx| async move {
            cx.update(|cx| repo.read(cx).create_branch(&new_branch_name))?
                .await??;
            cx.update(|cx| repo.read(cx).change_branch(&new_branch_name))?
                .await??;
            Ok(())
        })
        .detach_and_prompt_err("Failed to create branch", window, cx, |e, _, _| {
            Some(e.to_string())
        });
        cx.emit(DismissEvent);
    }
}

impl PickerDelegate for BranchListDelegate {
    type ListItem = ListItem;

    fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
        "Select branch...".into()
    }

    fn match_count(&self) -> usize {
        self.matches.len()
    }

    fn selected_index(&self) -> usize {
        self.selected_index
    }

    fn set_selected_index(
        &mut self,
        ix: usize,
        _window: &mut Window,
        _: &mut Context<Picker<Self>>,
    ) {
        self.selected_index = ix;
    }

    fn update_matches(
        &mut self,
        query: String,
        window: &mut Window,
        cx: &mut Context<Picker<Self>>,
    ) -> Task<()> {
        let Some(all_branches) = self.all_branches.clone() else {
            return Task::ready(());
        };

        const RECENT_BRANCHES_COUNT: usize = 10;
        cx.spawn_in(window, move |picker, mut cx| async move {
            let matches = if query.is_empty() {
                all_branches
                    .into_iter()
                    .take(RECENT_BRANCHES_COUNT)
                    .map(|branch| BranchEntry {
                        branch,
                        positions: Vec::new(),
                    })
                    .collect()
            } else {
                let candidates = all_branches
                    .iter()
                    .enumerate()
                    .map(|(ix, command)| StringMatchCandidate::new(ix, &command.name.clone()))
                    .collect::<Vec<StringMatchCandidate>>();
                fuzzy::match_strings(
                    &candidates,
                    &query,
                    true,
                    10000,
                    &Default::default(),
                    cx.background_executor().clone(),
                )
                .await
                .iter()
                .cloned()
                .map(|candidate| BranchEntry {
                    branch: all_branches[candidate.candidate_id].clone(),
                    positions: candidate.positions,
                })
                .collect()
            };
            picker
                .update(&mut cx, |picker, _| {
                    let delegate = &mut picker.delegate;
                    delegate.matches = matches;
                    if delegate.matches.is_empty() {
                        delegate.selected_index = 0;
                    } else {
                        delegate.selected_index =
                            core::cmp::min(delegate.selected_index, delegate.matches.len() - 1);
                    }
                    delegate.last_query = query;
                })
                .log_err();
        })
    }

    fn confirm(&mut self, secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
        let new_branch_name = SharedString::from(self.last_query.trim().replace(" ", "-"));
        if !new_branch_name.is_empty()
            && !self.has_exact_match(&new_branch_name)
            && ((self.selected_index == 0 && self.matches.len() == 0) || secondary)
        {
            self.create_branch(new_branch_name, window, cx);
            return;
        }

        let Some(entry) = self.matches.get(self.selected_index()) else {
            return;
        };

        let current_branch = self.repo.as_ref().map(|repo| {
            repo.update(cx, |repo, _| {
                repo.current_branch().map(|branch| branch.name.clone())
            })
        });

        if current_branch
            .flatten()
            .is_some_and(|current_branch| current_branch == entry.branch.name)
        {
            cx.emit(DismissEvent);
            return;
        }

        cx.spawn_in(window, {
            let branch = entry.branch.clone();
            |picker, mut cx| async move {
                let branch_change_task = picker.update(&mut cx, |this, cx| {
                    let repo = this
                        .delegate
                        .repo
                        .as_ref()
                        .ok_or_else(|| anyhow!("No active repository"))?
                        .clone();

                    let cx = cx.to_async();

                    anyhow::Ok(async move {
                        cx.update(|cx| repo.read(cx).change_branch(&branch.name))?
                            .await?
                    })
                })??;

                branch_change_task.await?;

                picker.update(&mut cx, |_, cx| {
                    cx.emit(DismissEvent);

                    anyhow::Ok(())
                })
            }
        })
        .detach_and_prompt_err("Failed to change branch", window, cx, |_, _, _| None);
    }

    fn dismissed(&mut self, _: &mut Window, cx: &mut Context<Picker<Self>>) {
        cx.emit(DismissEvent);
    }

    fn render_header(&self, _: &mut Window, _cx: &mut Context<Picker<Self>>) -> Option<AnyElement> {
        None
    }

    fn render_match(
        &self,
        ix: usize,
        selected: bool,
        _window: &mut Window,
        _cx: &mut Context<Picker<Self>>,
    ) -> Option<Self::ListItem> {
        let entry = &self.matches[ix];

        let (commit_time, commit_message) = entry
            .branch
            .most_recent_commit
            .as_ref()
            .map(|commit| {
                let commit_message = commit.subject.clone();
                let commit_time = OffsetDateTime::from_unix_timestamp(commit.commit_timestamp)
                    .unwrap_or_else(|_| OffsetDateTime::now_utc());
                let formatted_time = format_local_timestamp(
                    commit_time,
                    OffsetDateTime::now_utc(),
                    time_format::TimestampFormat::Relative,
                );
                (formatted_time, commit_message)
            })
            .unwrap_or_else(|| {
                (
                    "Unknown Date".to_string(),
                    SharedString::from("No commit message available"),
                )
            });

        Some(
            ListItem::new(SharedString::from(format!("vcs-menu-{ix}")))
                .inset(true)
                .spacing(match self.style {
                    BranchListStyle::Modal => ListItemSpacing::default(),
                    BranchListStyle::Popover => ListItemSpacing::ExtraDense,
                })
                .spacing(ListItemSpacing::Sparse)
                .toggle_state(selected)
                .child(
                    v_flex()
                        .w_full()
                        .child(
                            h_flex()
                                .w_full()
                                .flex_shrink()
                                .overflow_x_hidden()
                                .gap_2()
                                .justify_between()
                                .child(
                                    div().flex_shrink().overflow_x_hidden().child(
                                        HighlightedLabel::new(
                                            entry.branch.name.clone(),
                                            entry.positions.clone(),
                                        )
                                        .truncate(),
                                    ),
                                )
                                .child(
                                    Label::new(commit_time)
                                        .size(LabelSize::Small)
                                        .color(Color::Muted)
                                        .into_element(),
                                ),
                        )
                        .when(self.style == BranchListStyle::Modal, |el| {
                            el.child(
                                div().max_w_96().child(
                                    Label::new(
                                        commit_message
                                            .split('\n')
                                            .next()
                                            .unwrap_or_default()
                                            .to_string(),
                                    )
                                    .size(LabelSize::Small)
                                    .truncate()
                                    .color(Color::Muted),
                                ),
                            )
                        }),
                ),
        )
    }

    fn render_footer(
        &self,
        window: &mut Window,
        cx: &mut Context<Picker<Self>>,
    ) -> Option<AnyElement> {
        let new_branch_name = SharedString::from(self.last_query.trim().replace(" ", "-"));
        let handle = cx.weak_entity();
        Some(
            h_flex()
                .w_full()
                .p_2()
                .gap_2()
                .border_t_1()
                .border_color(cx.theme().colors().border_variant)
                .when(
                    !new_branch_name.is_empty() && !self.has_exact_match(&new_branch_name),
                    |el| {
                        el.child(
                            Button::new(
                                "create-branch",
                                format!("Create branch '{new_branch_name}'",),
                            )
                            .key_binding(KeyBinding::for_action(
                                &menu::SecondaryConfirm,
                                window,
                                cx,
                            ))
                            .toggle_state(
                                self.modifiers.secondary()
                                    || (self.selected_index == 0 && self.matches.len() == 0),
                            )
                            .tooltip(Tooltip::for_action_title(
                                "Create branch",
                                &menu::SecondaryConfirm,
                            ))
                            .on_click(move |_, window, cx| {
                                let new_branch_name = new_branch_name.clone();
                                if let Some(picker) = handle.upgrade() {
                                    picker.update(cx, |picker, cx| {
                                        picker.delegate.create_branch(new_branch_name, window, cx)
                                    });
                                }
                            }),
                        )
                    },
                )
                .into_any(),
        )
    }

    fn no_matches_text(&self, _window: &mut Window, _cx: &mut App) -> Option<SharedString> {
        None
    }
}