lv_tui/widgets/

tree.rs

use crate::component::{Component, EventCx};
use crate::event::Event;
use crate::geom::{Pos, Rect, Size};
use crate::render::RenderCx;
use crate::style::Style;
use crate::text::Text;

/// A node in a [`Tree`] widget.
pub struct TreeNode {
    pub label: Text,
    pub children: Vec<TreeNode>,
    pub expanded: bool,
}

/// Flattened visible entry in a tree.
struct VisibleEntry {
    /// Index into the original tree's flat list (for selection tracking).
    node_path: Vec<usize>, // path of indices from root: [0, 2] = nodes[0].children[2]
    depth: usize,
    is_last: bool,
    has_children: bool,
    expanded: bool,
    label: Text,
}

/// A hierarchical tree view widget.
///
/// Supports expand/collapse via Enter/Space, keyboard navigation, and optional
/// guide lines.
pub struct Tree {
    nodes: Vec<TreeNode>,
    /// Index into the flattened visible list.
    selected: usize,
    show_guides: bool,
    scroll_offset: usize,
    rect: Rect,
    style: Style,
    select_style: Style,
}

impl Tree {
    /// Creates a new tree with the given root nodes.
    pub fn new(nodes: Vec<TreeNode>) -> Self {
        Self {
            nodes,
            selected: 0,
            show_guides: true,
            scroll_offset: 0,
            rect: Rect::default(),
            style: Style::default(),
            select_style: Style::default(),
        }
    }

    pub fn show_guides(mut self, show: bool) -> Self {
        self.show_guides = show;
        self
    }

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

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

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

    /// Flattens the tree into a list of visible entries based on expansion state.
    fn flatten(&self) -> Vec<VisibleEntry> {
        let mut entries = Vec::new();
        for (i, node) in self.nodes.iter().enumerate() {
            self.flatten_node(node, &vec![i], 0, i == self.nodes.len() - 1, &mut entries);
        }
        entries
    }

    fn flatten_node(
        &self,
        node: &TreeNode,
        path: &[usize],
        depth: usize,
        is_last: bool,
        entries: &mut Vec<VisibleEntry>,
    ) {
        let has_children = !node.children.is_empty();
        entries.push(VisibleEntry {
            node_path: path.to_vec(),
            depth,
            is_last,
            has_children,
            expanded: node.expanded,
            label: node.label.clone(),
        });

        if node.expanded {
            let child_count = node.children.len();
            for (i, child) in node.children.iter().enumerate() {
                let mut child_path = path.to_vec();
                child_path.push(i);
                self.flatten_node(
                    child,
                    &child_path,
                    depth + 1,
                    i == child_count - 1,
                    entries,
                );
            }
        }
    }

    /// Finds the node at the given path.
    fn node_at_path(&mut self, path: &[usize]) -> Option<&mut TreeNode> {
        if path.is_empty() {
            return None;
        }
        let mut node = self.nodes.get_mut(path[0])?;
        for &idx in &path[1..] {
            node = node.children.get_mut(idx)?;
        }
        Some(node)
    }
}

impl Component for Tree {
    fn render(&self, cx: &mut RenderCx) {
        let entries = self.flatten();
        if entries.is_empty() {
            return;
        }

        let visible_height = self.rect.height.max(1) as usize;
        let start = self.scroll_offset.min(entries.len().saturating_sub(1));
        let end = (start + visible_height).min(entries.len());

        for i in start..end {
            let entry = &entries[i];
            let is_selected = i == self.selected;
            let row_y = self.rect.y + (i - start) as u16;

            // Build prefix: indentation + guides + icon
            let mut prefix = String::new();

            if self.show_guides {
                // Build guide lines for ancestors
                if entry.depth > 0 {
                    // Determine which ancestor levels have continuation
                    let mut has_continuation = vec![false; entry.depth];
                    for a in 0..entry.depth {
                        // Check if there are more entries after this one at the same or deeper depth
                        // that belong to the same ancestor
                        for j in (i + 1)..entries.len() {
                            if entries[j].depth < a + 1 {
                                break;
                            }
                            if entries[j].depth == a + 1 {
                                has_continuation[a] = true;
                                break;
                            }
                        }
                    }

                    for a in 0..entry.depth {
                        if has_continuation[a] {
                            prefix.push_str("│ ");
                        } else {
                            prefix.push_str("  ");
                        }
                    }
                }

                // Branch character for the current entry
                if entry.depth > 0 {
                    if entry.is_last {
                        prefix.push_str("└─");
                    } else {
                        prefix.push_str("├─");
                    }
                }
            } else {
                for _ in 0..entry.depth {
                    prefix.push_str("  ");
                }
            }

            // Icon
            if entry.has_children {
                if entry.expanded {
                    prefix.push_str("▼ ");
                } else {
                    prefix.push_str("▶ ");
                }
            } else {
                prefix.push_str("  ");
            }

            // Render
            let style = if is_selected {
                &self.select_style
            } else {
                &self.style
            };

            cx.buffer.write_text(
                Pos {
                    x: self.rect.x,
                    y: row_y,
                },
                self.rect,
                &prefix,
                style,
            );

            let label_text = entry.label.first_text();
            let prefix_w = crate::widgets::textarea::str_width(&prefix);
            cx.buffer.write_text(
                Pos {
                    x: self.rect.x + prefix_w,
                    y: row_y,
                },
                self.rect,
                label_text,
                style,
            );
        }
    }

    fn measure(
        &self,
        _constraint: crate::layout::Constraint,
        _cx: &mut crate::component::MeasureCx,
    ) -> Size {
        let entries = self.flatten();
        let max_w = entries
            .iter()
            .map(|e| (e.depth * 2 + 2) as u16 + e.label.max_width())
            .max()
            .unwrap_or(0);
        Size {
            width: max_w,
            height: entries.len().max(1) as u16,
        }
    }

    fn event(&mut self, event: &Event, cx: &mut EventCx) {
        if matches!(event, Event::Focus | Event::Blur | Event::Tick) {
            return;
        }

        let entries = self.flatten();
        if entries.is_empty() {
            return;
        }

        if let Event::Key(key_event) = event {
            match &key_event.key {
                crate::event::Key::Up => {
                    if self.selected > 0 {
                        self.selected -= 1;
                        self.scroll_to_selected(entries.len());
                        cx.invalidate_paint();
                    }
                }
                crate::event::Key::Down => {
                    if self.selected + 1 < entries.len() {
                        self.selected += 1;
                        self.scroll_to_selected(entries.len());
                        cx.invalidate_paint();
                    }
                }
                crate::event::Key::Enter | crate::event::Key::Char(' ') => {
                    self.toggle_selected();
                    cx.invalidate_paint();
                }
                crate::event::Key::Right => {
                    self.expand_selected();
                    cx.invalidate_paint();
                }
                crate::event::Key::Left => {
                    self.collapse_or_parent(&entries);
                    cx.invalidate_paint();
                }
                crate::event::Key::Home => {
                    self.selected = 0;
                    self.scroll_offset = 0;
                    cx.invalidate_paint();
                }
                crate::event::Key::End => {
                    self.selected = entries.len().saturating_sub(1);
                    self.scroll_to_selected(entries.len());
                    cx.invalidate_paint();
                }
                _ => {}
            }
        }
    }

    fn layout(&mut self, rect: Rect, _cx: &mut crate::component::LayoutCx) {
        self.rect = rect;
    }

    fn focusable(&self) -> bool {
        false
    }

    fn style(&self) -> Style {
        self.style.clone()
    }
}

impl Tree {
    fn toggle_selected(&mut self) {
        let entries = self.flatten();
        if let Some(entry) = entries.get(self.selected) {
            if entry.has_children {
                if let Some(node) = self.node_at_path(&entry.node_path) {
                    node.expanded = !node.expanded;
                }
            }
        }
    }

    fn expand_selected(&mut self) {
        let entries = self.flatten();
        if let Some(entry) = entries.get(self.selected) {
            if entry.has_children && !entry.expanded {
                if let Some(node) = self.node_at_path(&entry.node_path) {
                    node.expanded = true;
                }
            }
        }
    }

    fn collapse_or_parent(&mut self, entries: &[VisibleEntry]) {
        if let Some(entry) = entries.get(self.selected) {
            if entry.has_children && entry.expanded {
                // Collapse current node
                if let Some(node) = self.node_at_path(&entry.node_path) {
                    node.expanded = false;
                }
            } else if entry.depth > 0 {
                // Jump to parent's position
                let parent_depth = entry.depth - 1;
                for (i, e) in entries.iter().enumerate() {
                    if e.node_path.len() == entry.node_path.len() - 1
                        && e.depth == parent_depth
                    {
                        self.selected = i;
                        break;
                    }
                }
            }
        }
    }

    fn scroll_to_selected(&mut self, total_visible: usize) {
        let visible_height = self.rect.height.max(1) as usize;
        if self.selected < self.scroll_offset {
            self.scroll_offset = self.selected;
        } else if self.selected >= self.scroll_offset + visible_height {
            self.scroll_offset = self.selected.saturating_sub(visible_height.saturating_sub(1));
        }
        self.scroll_offset = self.scroll_offset.min(
            total_visible.saturating_sub(visible_height),
        );
    }
}