use crate::style::Color;
use crate::utils::Selection;
use crate::widget::traits::{View, WidgetProps};
use crate::{impl_props_builders, impl_styled_view};
mod search;
mod types;
mod view;
pub use types::TreeNode;
type SelectCallback = Box<dyn Fn(&TreeNode)>;
pub struct Tree {
root: Vec<TreeNode>,
selection: Selection,
fg: Option<Color>,
bg: Option<Color>,
selected_fg: Option<Color>,
selected_bg: Option<Color>,
highlight_fg: Option<Color>,
indent: u16,
query: String,
searchable: bool,
matches: Vec<usize>,
current_match: usize,
props: WidgetProps,
multi_select: bool,
selected_indices: Vec<usize>,
on_select: Option<SelectCallback>,
}
impl Tree {
pub fn new() -> Self {
Self {
root: Vec::new(),
selection: Selection::new(0),
fg: None,
bg: None,
selected_fg: Some(Color::WHITE),
selected_bg: Some(Color::BLUE),
highlight_fg: Some(Color::YELLOW),
indent: 2,
query: String::new(),
searchable: false,
matches: Vec::new(),
current_match: 0,
props: WidgetProps::new(),
multi_select: false,
selected_indices: Vec::new(),
on_select: None,
}
}
pub fn searchable(mut self, enable: bool) -> Self {
self.searchable = enable;
self
}
pub fn highlight_fg(mut self, color: Color) -> Self {
self.highlight_fg = Some(color);
self
}
pub fn nodes(mut self, nodes: Vec<TreeNode>) -> Self {
self.root = nodes;
self.selection.set_len(self.count_visible());
self
}
pub fn node(mut self, node: TreeNode) -> Self {
self.root.push(node);
self.selection.set_len(self.count_visible());
self
}
pub fn selected(mut self, index: usize) -> Self {
self.selection.set(index);
self
}
pub fn fg(mut self, color: Color) -> Self {
self.fg = Some(color);
self
}
pub fn bg(mut self, color: Color) -> Self {
self.bg = Some(color);
self
}
pub fn selected_style(mut self, fg: Color, bg: Color) -> Self {
self.selected_fg = Some(fg);
self.selected_bg = Some(bg);
self
}
pub fn indent(mut self, indent: u16) -> Self {
self.indent = indent;
self
}
pub fn multi_select(mut self, enable: bool) -> Self {
self.multi_select = enable;
self
}
pub fn on_select(mut self, callback: impl Fn(&TreeNode) + 'static) -> Self {
self.on_select = Some(Box::new(callback));
self
}
pub fn selected_index(&self) -> usize {
self.selection.index
}
fn count_visible(&self) -> usize {
fn count_nodes(nodes: &[TreeNode]) -> usize {
let mut count = 0;
for node in nodes {
count += 1;
if node.expanded && !node.children.is_empty() {
count += count_nodes(&node.children);
}
}
count
}
count_nodes(&self.root)
}
fn get_node_at(&self, index: usize) -> Option<(&TreeNode, usize)> {
fn find_node<'a>(
nodes: &'a [TreeNode],
target: usize,
current: &mut usize,
depth: usize,
) -> Option<(&'a TreeNode, usize)> {
for node in nodes {
if *current == target {
return Some((node, depth));
}
*current += 1;
if node.expanded && !node.children.is_empty() {
if let Some(result) = find_node(&node.children, target, current, depth + 1) {
return Some(result);
}
}
}
None
}
let mut current = 0;
find_node(&self.root, index, &mut current, 0)
}
fn get_node_mut_at(&mut self, index: usize) -> Option<&mut TreeNode> {
fn find_node_mut<'a>(
nodes: &'a mut [TreeNode],
target: usize,
current: &mut usize,
) -> Option<&'a mut TreeNode> {
for node in nodes {
if *current == target {
return Some(node);
}
*current += 1;
if node.expanded && !node.children.is_empty() {
if let Some(result) = find_node_mut(&mut node.children, target, current) {
return Some(result);
}
}
}
None
}
let mut current = 0;
find_node_mut(&mut self.root, index, &mut current)
}
pub fn select_next(&mut self) {
self.selection.down();
}
pub fn select_prev(&mut self) {
self.selection.up();
}
pub fn select_first(&mut self) {
self.selection.first();
}
pub fn select_last(&mut self) {
self.selection.last();
}
pub fn toggle_expand(&mut self) {
if let Some(node) = self.get_node_mut_at(self.selection.index) {
if node.has_children() {
node.expanded = !node.expanded;
self.selection.set_len(self.count_visible());
}
}
}
pub fn expand(&mut self) {
if let Some(node) = self.get_node_mut_at(self.selection.index) {
if node.has_children() && !node.expanded {
node.expanded = true;
self.selection.set_len(self.count_visible());
}
}
}
pub fn collapse(&mut self) {
if let Some(node) = self.get_node_mut_at(self.selection.index) {
if node.expanded {
node.expanded = false;
self.selection.set_len(self.count_visible());
}
}
}
pub fn toggle_select(&mut self) {
let idx = self.selection.index;
if let Some(pos) = self.selected_indices.iter().position(|&i| i == idx) {
self.selected_indices.remove(pos);
} else {
self.selected_indices.push(idx);
}
if let Some(callback) = &self.on_select {
if let Some((node, _)) = self.get_node_at(idx) {
callback(node);
}
}
}
pub fn selected_nodes(&self) -> Vec<&TreeNode> {
self.selected_indices
.iter()
.filter_map(|&idx| self.get_node_at(idx).map(|(node, _)| node))
.collect()
}
pub fn selected_ids(&self) -> Vec<&str> {
self.selected_nodes()
.into_iter()
.filter_map(|node| node.id.as_deref())
.collect()
}
pub fn is_multi_selected(&self, index: usize) -> bool {
self.selected_indices.contains(&index)
}
fn find_parent_index(&self, index: usize) -> Option<usize> {
fn find_parent(
nodes: &[TreeNode],
target: usize,
current: &mut usize,
parent_index: Option<usize>,
) -> Option<usize> {
for node in nodes {
if *current == target {
return parent_index;
}
let my_index = *current;
*current += 1;
if node.expanded && !node.children.is_empty() {
if let Some(result) =
find_parent(&node.children, target, current, Some(my_index))
{
return Some(result);
}
}
}
None
}
let mut current = 0;
find_parent(&self.root, index, &mut current, None)
}
pub fn handle_key(&mut self, key: &crate::event::Key) -> bool {
use crate::event::Key;
if self.searchable && !self.query.is_empty() {
match key {
Key::Char('n') => return self.next_match(),
Key::Char('N') => return self.prev_match(),
Key::Escape => {
self.clear_query();
return true;
}
Key::Backspace => {
self.query.pop();
self.update_matches();
return true;
}
_ => {}
}
}
if self.searchable {
if let Key::Char(c) = key {
if c.is_alphanumeric() || *c == '_' || *c == '-' || *c == '.' || *c == '/' {
self.query.push(*c);
self.update_matches();
return true;
}
}
}
match key {
Key::Up | Key::Char('k') if !self.searchable => {
let old = self.selection.index;
self.select_prev();
old != self.selection.index
}
Key::Up if self.searchable => {
let old = self.selection.index;
self.select_prev();
old != self.selection.index
}
Key::Down | Key::Char('j') if !self.searchable => {
let old = self.selection.index;
self.select_next();
old != self.selection.index
}
Key::Down if self.searchable => {
let old = self.selection.index;
self.select_next();
old != self.selection.index
}
Key::Char(' ') if self.multi_select && !self.searchable => {
self.toggle_select();
true
}
Key::Enter | Key::Char(' ') if !self.searchable => {
let old_count = self.selection.len;
self.toggle_expand();
old_count != self.selection.len
}
Key::Enter if self.searchable => {
let old_count = self.selection.len;
self.toggle_expand();
old_count != self.selection.len
}
Key::Right | Key::Char('l') if !self.searchable => {
let old_count = self.selection.len;
self.expand();
old_count != self.selection.len
}
Key::Right if self.searchable => {
let old_count = self.selection.len;
self.expand();
old_count != self.selection.len
}
Key::Left | Key::Char('h') if !self.searchable => {
let is_expanded = self
.get_node_at(self.selection.index)
.map(|(n, _)| n.expanded && n.has_children())
.unwrap_or(false);
if is_expanded {
let old_count = self.selection.len;
self.collapse();
old_count != self.selection.len
} else if let Some(parent_idx) = self.find_parent_index(self.selection.index) {
let old = self.selection.index;
self.selection.set(parent_idx);
old != self.selection.index
} else {
false
}
}
Key::Left if self.searchable => {
let is_expanded = self
.get_node_at(self.selection.index)
.map(|(n, _)| n.expanded && n.has_children())
.unwrap_or(false);
if is_expanded {
let old_count = self.selection.len;
self.collapse();
old_count != self.selection.len
} else if let Some(parent_idx) = self.find_parent_index(self.selection.index) {
let old = self.selection.index;
self.selection.set(parent_idx);
old != self.selection.index
} else {
false
}
}
Key::Home => {
let old = self.selection.index;
self.select_first();
old != self.selection.index
}
Key::End => {
let old = self.selection.index;
self.select_last();
old != self.selection.index
}
_ => false,
}
}
pub fn len(&self) -> usize {
self.root.len()
}
pub fn is_empty(&self) -> bool {
self.root.is_empty()
}
pub fn visible_count(&self) -> usize {
self.selection.len
}
pub fn selected_label(&self) -> Option<&str> {
self.get_node_at(self.selection.index)
.map(|(n, _)| n.label.as_str())
}
}
impl Default for Tree {
fn default() -> Self {
Self::new()
}
}
impl View for Tree {
crate::impl_view_meta!("Tree");
fn render(&self, ctx: &mut crate::widget::traits::RenderContext) {
self.render_internal(ctx);
}
}
impl_styled_view!(Tree);
impl_props_builders!(Tree);
pub fn tree() -> Tree {
Tree::new()
}
pub fn tree_node(label: impl Into<String>) -> TreeNode {
TreeNode::new(label)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::layout::Rect;
use crate::render::Buffer;
use crate::widget::traits::RenderContext;
#[test]
fn test_tree_fg_bg_colors() {
let t = Tree::new().fg(Color::RED).bg(Color::BLUE);
assert_eq!(t.fg, Some(Color::RED));
assert_eq!(t.bg, Some(Color::BLUE));
}
#[test]
fn test_tree_indent() {
let t = Tree::new().indent(4);
assert_eq!(t.indent, 4);
}
#[test]
fn test_tree_render() {
let mut buffer = Buffer::new(40, 10);
let area = Rect::new(0, 0, 40, 10);
let mut ctx = RenderContext::new(&mut buffer, area);
let t = Tree::new().node(TreeNode::new("Files"));
t.render(&mut ctx);
assert_eq!(buffer.get(0, 0).unwrap().symbol, ' '); assert_eq!(buffer.get(1, 0).unwrap().symbol, 'F');
}
#[test]
fn test_tree_render_with_expanded_children() {
let mut buffer = Buffer::new(40, 10);
let area = Rect::new(0, 0, 40, 10);
let mut ctx = RenderContext::new(&mut buffer, area);
let t = Tree::new()
.node(
TreeNode::new("Parent")
.expanded(true)
.child(TreeNode::new("Child 1"))
.child(TreeNode::new("Child 2")),
)
.selected_style(Color::WHITE, Color::BLUE);
t.render(&mut ctx);
assert_eq!(buffer.get(0, 0).unwrap().symbol, '▼');
}
#[test]
fn test_tree_render_nested() {
let mut buffer = Buffer::new(40, 10);
let area = Rect::new(0, 0, 40, 10);
let mut ctx = RenderContext::new(&mut buffer, area);
let t = Tree::new().node(
TreeNode::new("Level 0").expanded(true).child(
TreeNode::new("Level 1")
.expanded(true)
.child(TreeNode::new("Level 2")),
),
);
t.render(&mut ctx);
}
#[test]
fn test_tree_render_empty() {
let mut buffer = Buffer::new(40, 10);
let area = Rect::new(0, 0, 40, 10);
let mut ctx = RenderContext::new(&mut buffer, area);
let t = Tree::new();
t.render(&mut ctx);
}
#[test]
fn test_tree_render_small_area() {
let mut buffer = Buffer::new(2, 1);
let area = Rect::new(0, 0, 2, 1);
let mut ctx = RenderContext::new(&mut buffer, area);
let t = Tree::new().node(TreeNode::new("Test"));
t.render(&mut ctx);
}
#[test]
fn test_tree_render_with_selection() {
let mut buffer = Buffer::new(40, 10);
let area = Rect::new(0, 0, 40, 10);
let mut ctx = RenderContext::new(&mut buffer, area);
let t = Tree::new()
.nodes(vec![TreeNode::new("First"), TreeNode::new("Second")])
.selected(1);
t.render(&mut ctx);
}
#[test]
fn test_tree_render_with_highlight() {
let mut buffer = Buffer::new(40, 10);
let area = Rect::new(0, 0, 40, 10);
let mut ctx = RenderContext::new(&mut buffer, area);
let mut t = Tree::new()
.nodes(vec![TreeNode::new("Hello World")])
.searchable(true)
.highlight_fg(Color::YELLOW);
t.set_query("hw");
t.render(&mut ctx);
}
}