hefesto_widgets/tree/

mod.rs

use std::collections::HashSet;

use ratatui::{
    buffer::Buffer,
    layout::Rect,
    style::Style,
    text::{Line, Span},
    widgets::ListItem,
};

use crate::{DEFAULT_HIGHLIGHT, HIGHLIGHT_SYMBOL, ScrollList, ScrollListState};

/// A node in a tree. Children are recursively nested.
/// `id` must be unique across the whole tree for expand/collapse tracking.
#[derive(Clone)]
pub struct TreeNode<'a> {
    pub text: Line<'a>,
    pub children: Vec<TreeNode<'a>>,
    pub id: usize,
}

/// Tree widget with expandable/collapsible branches.
///
/// Internally projects the tree into a flat `ScrollList` with indentation
/// and expand/collapse icons. The caller manages `TreeState::expanded`.
#[derive(Clone)]
pub struct Tree<'a> {
    nodes: Vec<TreeNode<'a>>,
    indent: u16,
    expand_icon: String,
    collapse_icon: String,
    leaf_icon: String,
    highlight_style: Style,
    highlight_symbol: String,
    borders: ratatui::widgets::Borders,
    border_style: Style,
}

impl<'a> Tree<'a> {
    pub fn new(nodes: Vec<TreeNode<'a>>) -> Self {
        Self {
            nodes,
            indent: 2,
            expand_icon: "▶ ".to_string(),
            collapse_icon: "▼ ".to_string(),
            leaf_icon: "  ".to_string(),
            highlight_style: DEFAULT_HIGHLIGHT,
            highlight_symbol: HIGHLIGHT_SYMBOL.to_string(),
            borders: ratatui::widgets::Borders::NONE,
            border_style: Style::new(),
        }
    }

    pub fn indent(mut self, n: u16) -> Self {
        self.indent = n;
        self
    }

    pub fn expand_icon(mut self, icon: &str) -> Self {
        self.expand_icon = icon.to_string();
        self
    }

    pub fn collapse_icon(mut self, icon: &str) -> Self {
        self.collapse_icon = icon.to_string();
        self
    }

    pub fn leaf_icon(mut self, icon: &str) -> Self {
        self.leaf_icon = icon.to_string();
        self
    }

    pub fn highlight_style(mut self, style: Style) -> Self {
        self.highlight_style = style;
        self
    }

    pub fn highlight_symbol(mut self, symbol: &str) -> Self {
        self.highlight_symbol = symbol.to_string();
        self
    }

    pub fn borders(mut self, borders: ratatui::widgets::Borders) -> Self {
        self.borders = borders;
        self
    }

    pub fn border_style(mut self, style: Style) -> Self {
        self.border_style = style;
        self
    }

    /// Count of visible items in the projected flat view.
    pub fn visible_count(&self, expanded: &HashSet<usize>) -> usize {
        let mut count = 0;
        Self::count_visible(&self.nodes, expanded, &mut count);
        count
    }

    fn count_visible(nodes: &[TreeNode], expanded: &HashSet<usize>, count: &mut usize) {
        for node in nodes {
            *count += 1;
            if !node.children.is_empty() && expanded.contains(&node.id) {
                Self::count_visible(&node.children, expanded, count);
            }
        }
    }

    /// Flatten the tree into visible `ListItem`s based on `expanded` set.
    fn project(&self, expanded: &HashSet<usize>) -> Vec<ListItem<'a>> {
        let mut items = Vec::new();
        for node in &self.nodes {
            Self::project_node(node, 0, self.indent, expanded, &self.expand_icon, &self.collapse_icon, &self.leaf_icon, &mut items);
        }
        items
    }

    fn project_node(
        node: &TreeNode<'a>,
        depth: u16,
        indent: u16,
        expanded: &HashSet<usize>,
        expand_icon: &str,
        collapse_icon: &str,
        leaf_icon: &str,
        items: &mut Vec<ListItem<'a>>,
    ) {
        let indent_str = " ".repeat((depth as usize) * indent as usize);
        let icon = if node.children.is_empty() {
            leaf_icon
        } else if expanded.contains(&node.id) {
            collapse_icon
        } else {
            expand_icon
        };

        let mut line = node.text.clone();
        let prefix = format!("{}{}", indent_str, icon);
        line.spans.insert(0, Span::raw(prefix));

        items.push(ListItem::new(line));

        if expanded.contains(&node.id) {
            for child in &node.children {
                Self::project_node(child, depth + 1, indent, expanded, expand_icon, collapse_icon, leaf_icon, items);
            }
        }
    }
}

/// State for the Tree widget.
#[derive(Clone, Default)]
pub struct TreeState {
    /// Set of node ids that are currently expanded.
    pub expanded: HashSet<usize>,
    pub scroll_state: ScrollListState,
}

impl TreeState {
    pub fn toggle(&mut self, id: usize) {
        if self.expanded.contains(&id) {
            self.expanded.remove(&id);
        } else {
            self.expanded.insert(id);
        }
    }

    /// Expand all nodes that have children.
    pub fn expand_all(&mut self, nodes: &[TreeNode]) {
        for node in nodes {
            if !node.children.is_empty() {
                self.expanded.insert(node.id);
                self.expand_all(&node.children);
            }
        }
    }

    /// Expand this node and all its descendants recursively.
    pub fn expand_subtree(&mut self, nodes: &[TreeNode], id: usize) {
        if let Some(node) = Self::find_node(nodes, id) {
            self.expanded.insert(id);
            for child in &node.children {
                self.expand_subtree(nodes, child.id);
            }
        }
    }

    /// Collapse from cursor position.
    ///
    /// If the cursor is on an expanded node with children → collapse that node.
    /// Otherwise find the nearest expanded ancestor and collapse it.
    pub fn collapse_at(&mut self, nodes: &[TreeNode], cursor_id: usize) {
        let is_expanded = self.expanded.contains(&cursor_id);
        let has_children = Self::find_node(nodes, cursor_id)
            .map(|n| !n.children.is_empty())
            .unwrap_or(false);

        if is_expanded && has_children {
            self.expanded.remove(&cursor_id);
            return;
        }

        if let Some(ancestor) = self.nearest_expanded_ancestor(nodes, cursor_id) {
            self.expanded.remove(&ancestor);
        }
    }

    /// Expand from cursor position.
    ///
    /// If cursor is on a collapsed node → expand just that node.
    /// If cursor is on an already-expanded node → expand the entire subtree recursively.
    pub fn expand_at(&mut self, nodes: &[TreeNode], cursor_id: usize) {
        if self.expanded.contains(&cursor_id) {
            self.expand_subtree(nodes, cursor_id);
        } else {
            self.expanded.insert(cursor_id);
        }
    }

    /// Find the nearest ancestor of `cursor_id` that is currently expanded,
    /// walking from the root downward. Returns `None` if no ancestor is expanded.
    pub fn nearest_expanded_ancestor(&self, nodes: &[TreeNode], cursor_id: usize) -> Option<usize> {
        let path = Self::find_path(nodes, cursor_id)?;
        let mut result = None;
        for id in &path {
            if *id == cursor_id {
                break;
            }
            if self.expanded.contains(id) {
                result = Some(*id);
            }
        }
        result
    }

    // ── private helpers ──

    fn find_node<'b>(nodes: &'b [TreeNode<'b>], target_id: usize) -> Option<&'b TreeNode<'b>> {
        for node in nodes {
            if node.id == target_id {
                return Some(node);
            }
            if let Some(found) = Self::find_node(&node.children, target_id) {
                return Some(found);
            }
        }
        None
    }

    fn find_path<'b>(nodes: &'b [TreeNode<'b>], target_id: usize) -> Option<Vec<usize>> {
        for node in nodes {
            if node.id == target_id {
                return Some(vec![node.id]);
            }
            if let Some(mut path) = Self::find_path(&node.children, target_id) {
                path.insert(0, node.id);
                return Some(path);
            }
        }
        None
    }

    /// Return the visible item index of the first node with `id`,
    /// or `None` if it's hidden behind a collapsed ancestor.
    pub fn visible_index_of(&self, nodes: &[TreeNode], target_id: usize) -> Option<usize> {
        let mut idx = 0usize;
        Self::walk(nodes, target_id, &self.expanded, true, &mut idx)
    }
    pub fn walk<'b>(nodes: &'b [TreeNode<'b>], target_id: usize, expanded: &HashSet<usize>, visible: bool, idx: &mut usize) -> Option<usize> {
        for node in nodes {
            if visible {
                if node.id == target_id {
                    return Some(*idx);
                }
                *idx += 1;
            }

            if !node.children.is_empty() && expanded.contains(&node.id) {
                if let Some(found) = Self::walk(&node.children, target_id, expanded, visible, idx) {
                    return Some(found);
                }
            }
        }
        None
    }
}

impl<'a> ratatui::widgets::StatefulWidget for Tree<'a> {
    type State = TreeState;

    fn render(self, area: Rect, buf: &mut Buffer, state: &mut Self::State) {
        let items = self.project(&state.expanded);
        ScrollList::new_list(items)
            .highlight_style(self.highlight_style)
            .highlight_symbol(&self.highlight_symbol)
            .borders(self.borders)
            .border_style(self.border_style)
            .render(area, buf, &mut state.scroll_state);
    }
}
#[cfg(test)]
mod tests;