gpui-native-kit/crates/data/src/tree.rs

use std::rc::Rc;

use gpui::prelude::FluentBuilder as _;
use gpui::{
    AnyElement, App, ElementId, FocusHandle, InteractiveElement as _, IntoElement, KeyDownEvent,
    ParentElement as _, RenderOnce, Role, SharedString, StatefulInteractiveElement as _,
    Styled as _, Window, div,
};
use gpui_native_icons::{Icon, IconRole, icon};
use gpui_native_theme::ActiveTheme as _;
use gpui_native_utils::{nav_key, paint_selectable_row, selectable_label_color};

/// The change reporter, fired with the next selection by pointer and keys.
type ChangeHandler = Rc<dyn Fn(Option<usize>, &mut Window, &mut App)>;
/// The disclosure reporter, fired with the toggled row's index.
type DisclosureHandler = Rc<dyn Fn(usize, &mut Window, &mut App)>;
type ItemsFn = Box<dyn Fn(&mut Window, &mut App) -> Vec<TreeItem>>;

/// Tree item.
#[derive(Clone)]
pub struct TreeItem {
    label: SharedString,
    icon: Option<Icon>,
    depth: usize,
    expandable: bool,
    expanded: bool,
    disabled: bool,
}

impl TreeItem {
    /// Creates a leaf row with this label at the root depth.
    pub fn new<L>(label: L) -> Self
    where
        L: Into<SharedString>,
    {
        Self {
            label: label.into(),
            icon: None,
            depth: 0,
            expandable: false,
            expanded: false,
            disabled: false,
        }
    }

    /// Shows an icon leading the label.
    pub fn icon<I>(mut self, icon: I) -> Self
    where
        I: Into<Icon>,
    {
        self.icon = Some(icon.into());
        self
    }

    /// Nests the row one level per count.
    pub fn depth(mut self, depth: usize) -> Self {
        self.depth = depth;
        self
    }

    /// Marks the row as having children.
    pub fn expandable(mut self, expandable: bool) -> Self {
        self.expandable = expandable;
        self
    }

    /// Sets the row's disclosure state.
    pub fn expanded(mut self, expanded: bool) -> Self {
        self.expanded = expanded;
        self
    }

    /// Renders the row dimmed; clicks and keyboard skip it.
    pub fn disabled(mut self, disabled: bool) -> Self {
        self.disabled = disabled;
        self
    }
}

/// A tree.
#[derive(IntoElement)]
pub struct Tree {
    id: ElementId,
    selected: Option<usize>,
    a11y_label: Option<SharedString>,
    on_change: Option<ChangeHandler>,
    on_disclosure: Option<DisclosureHandler>,
    items: Option<ItemsFn>,
}

impl Tree {
    /// Creates a tree with nothing selected.
    pub fn new<I>(id: I) -> Self
    where
        I: Into<ElementId>,
    {
        Self {
            id: id.into(),
            selected: None,
            a11y_label: None,
            on_change: None,
            on_disclosure: None,
            items: None,
        }
    }

    /// Sets the current selection.
    pub fn selected(mut self, selected: Option<usize>) -> Self {
        self.selected = selected;
        self
    }

    /// Reports the next selection, from clicks and keys alike.
    pub fn on_change<F>(mut self, handler: F) -> Self
    where
        F: Fn(Option<usize>, &mut Window, &mut App) + 'static,
    {
        self.on_change = Some(Rc::new(handler));
        self
    }

    /// Reports the row whose disclosure the user toggled.
    pub fn on_disclosure<F>(mut self, handler: F) -> Self
    where
        F: Fn(usize, &mut Window, &mut App) + 'static,
    {
        self.on_disclosure = Some(Rc::new(handler));
        self
    }

    /// Sets the visible rows, built fresh each render.
    pub fn items<F>(mut self, items: F) -> Self
    where
        F: Fn(&mut Window, &mut App) -> Vec<TreeItem> + 'static,
    {
        self.items = Some(Box::new(items));
        self
    }

    /// Overrides the accessibility label, which defaults to the id.
    pub fn accessibility_label<L>(mut self, label: L) -> Self
    where
        L: Into<SharedString>,
    {
        self.a11y_label = Some(label.into());
        self
    }
}

impl RenderOnce for Tree {
    fn render(self, window: &mut Window, cx: &mut App) -> impl IntoElement {
        let items = match self.items {
            Some(build) => build(window, cx),
            None => Vec::new(),
        };

        let selectable: Vec<usize> = items
            .iter()
            .enumerate()
            .filter(|(_, item)| !item.disabled)
            .map(|(index, _)| index)
            .collect();

        let focus = gpui_native_utils::focus_handle(&self.id, window, cx);
        let theme = cx.theme();

        let label = self
            .a11y_label
            .clone()
            .unwrap_or_else(|| SharedString::from(self.id.to_string()));

        let sibling_sets: Vec<(usize, usize)> = {
            let mut sets = vec![(0usize, 0usize); items.len()];
            let mut start = 0usize;

            while start < items.len() {
                let depth = items[start].depth;

                let mut end = start + 1;
                while end < items.len() && items[end].depth > depth {
                    end += 1;
                }

                let mut position = 0usize;
                for index in start..end {
                    if items[index].depth == depth {
                        position += 1;
                        sets[index] = (position, 0);
                    }
                }

                for index in start..end {
                    if items[index].depth == depth {
                        sets[index].1 = position;
                    }
                }

                start = end;
            }

            sets
        };

        let rows: Vec<AnyElement> = items
            .iter()
            .enumerate()
            .map(|(index, _item)| {
                let (position_in_set, size_of_set) = sibling_sets[index];
                render_row(
                    cx,
                    &items[index],
                    index,
                    position_in_set,
                    size_of_set,
                    self.selected == Some(index),
                    &self.on_change,
                    &self.on_disclosure,
                    &focus,
                )
            })
            .collect();

        div()
            .id(self.id)
            .flex()
            .flex_col()
            .gap_0p5()
            .text_size(theme.sizes.medium.text_size)
            .track_focus(&focus)
            .focusable()
            .tab_stop(true)
            .role(Role::Tree)
            .aria_label(label)
            .on_key_down({
                let selected = self.selected;
                let on_change = self.on_change.clone();
                let on_disclosure = self.on_disclosure.clone();
                move |event: &KeyDownEvent, window, cx| {
                    if nav_key(&selectable, selected, &event.keystroke.key, |next| {
                        if let Some(on_change) = &on_change {
                            on_change(next, window, cx);
                        }
                    }) {
                        cx.stop_propagation();
                        return;
                    }
                    match event.keystroke.key.as_str() {
                        "left" | "right" => {
                            let expanding = event.keystroke.key == "right";
                            if let Some(index) = selected
                                && let Some(item) = items.get(index)
                                && item.expandable
                                && !item.disabled
                                && item.expanded != expanding
                                && let Some(on_disclosure) = &on_disclosure
                            {
                                on_disclosure(index, window, cx);
                            }
                        }
                        _ => return,
                    }
                    cx.stop_propagation();
                }
            })
            .children(rows)
    }
}

#[allow(clippy::too_many_arguments)]
fn render_row(
    cx: &App,
    item: &TreeItem,
    index: usize,
    position_in_set: usize,
    size_of_set: usize,
    selected: bool,
    on_change: &Option<ChangeHandler>,
    on_disclosure: &Option<DisclosureHandler>,
    focus: &FocusHandle,
) -> AnyElement {
    let theme = cx.theme();
    let label_color = selectable_label_color(cx, selected, item.disabled);
    let disabled = item.disabled;
    // One indent step per depth level, on the theme's horizontal padding.
    let indent = theme.sizes.medium.padding_x * (item.depth as f32 + 1.0);

    div()
        .id(("tree-item", index))
        .flex()
        .items_center()
        .gap_2()
        .pl(indent)
        .pr_2()
        .py_1()
        .rounded(theme.radius.sm)
        .role(Role::TreeItem)
        .aria_selected(selected)
        .aria_level(item.depth + 1)
        .aria_position_in_set(position_in_set)
        .aria_size_of_set(size_of_set)
        .when(item.expandable, |row| row.aria_expanded(item.expanded))
        .map(|row| paint_selectable_row(cx, row, selected, disabled))
        .when(!disabled, |row| {
            row.cursor_pointer().on_click({
                let on_change = on_change.clone();
                let focus = focus.clone();
                move |_, window, cx| {
                    window.focus(&focus, cx);
                    if let Some(on_change) = &on_change {
                        on_change(Some(index), window, cx);
                    }
                }
            })
        })
        .child(render_disclosure(
            cx,
            item,
            index,
            disabled,
            selected,
            on_disclosure,
        ))
        .when_some(item.icon.clone(), |row, icon| {
            row.child(icon.size(theme.sizes.medium.icon_size).color(label_color))
        })
        .child(
            div()
                .flex_1()
                .text_color(label_color)
                .child(item.label.clone()),
        )
        .into_any_element()
}

/// The disclosure slot.
fn render_disclosure(
    cx: &App,
    item: &TreeItem,
    index: usize,
    disabled: bool,
    selected: bool,
    on_disclosure: &Option<DisclosureHandler>,
) -> AnyElement {
    let theme = cx.theme();
    let glyph_color = if disabled {
        theme.colors.muted_foreground
    } else if selected {
        theme.colors.primary_foreground
    } else {
        theme.colors.muted_foreground
    };

    let cell = div()
        .flex_none()
        .size(theme.sizes.medium.icon_size)
        .flex()
        .items_center()
        .justify_center()
        .child(
            icon(if item.expanded {
                IconRole::ChevronDown
            } else {
                IconRole::ChevronRight
            })
            .size(theme.sizes.medium.dense_icon_size)
            .color(glyph_color),
        );

    if item.expandable && !disabled {
        cell.id(("tree-disclosure", index))
            .cursor_pointer()
            .on_click({
                let on_disclosure = on_disclosure.clone();
                move |_, window, cx| {
                    if let Some(on_disclosure) = &on_disclosure {
                        on_disclosure(index, window, cx);
                    }
                    // The toggle is not a selection: keep it off the row.
                    cx.stop_propagation();
                }
            })
            .into_any_element()
    } else {
        cell.into_any_element()
    }
}