use crate::components::{Box as RnkBox, Text};
use crate::core::{Color, Element, FlexDirection};
use std::collections::HashSet;
#[derive(Debug, Clone)]
pub struct TreeNode<T: Clone> {
pub id: String,
pub label: String,
pub data: Option<T>,
pub children: Vec<TreeNode<T>>,
}
impl<T: Clone> TreeNode<T> {
pub fn new(id: impl Into<String>, label: impl Into<String>) -> Self {
Self {
id: id.into(),
label: label.into(),
data: None,
children: Vec::new(),
}
}
pub fn leaf(id: impl Into<String>, label: impl Into<String>) -> Self {
Self::new(id, label)
}
pub fn with_data(id: impl Into<String>, label: impl Into<String>, data: T) -> Self {
Self {
id: id.into(),
label: label.into(),
data: Some(data),
children: Vec::new(),
}
}
pub fn child(mut self, child: TreeNode<T>) -> Self {
self.children.push(child);
self
}
pub fn children(mut self, children: impl IntoIterator<Item = TreeNode<T>>) -> Self {
self.children.extend(children);
self
}
pub fn is_leaf(&self) -> bool {
self.children.is_empty()
}
pub fn has_children(&self) -> bool {
!self.children.is_empty()
}
pub fn node_count(&self) -> usize {
1 + self.children.iter().map(|c| c.node_count()).sum::<usize>()
}
pub fn find(&self, id: &str) -> Option<&TreeNode<T>> {
if self.id == id {
return Some(self);
}
for child in &self.children {
if let Some(found) = child.find(id) {
return Some(found);
}
}
None
}
pub fn all_ids(&self) -> Vec<String> {
let mut ids = vec![self.id.clone()];
for child in &self.children {
ids.extend(child.all_ids());
}
ids
}
}
#[derive(Debug, Clone)]
pub struct TreeState {
expanded: HashSet<String>,
selected: Option<String>,
visible_nodes: Vec<String>,
cursor: usize,
}
impl TreeState {
pub fn new<T: Clone>(root: &TreeNode<T>) -> Self {
let mut state = Self {
expanded: HashSet::new(),
selected: None,
visible_nodes: Vec::new(),
cursor: 0,
};
state.rebuild_visible(root);
state
}
pub fn with_root_expanded<T: Clone>(root: &TreeNode<T>) -> Self {
let mut state = Self::new(root);
state.expanded.insert(root.id.clone());
state.rebuild_visible(root);
state
}
pub fn all_expanded<T: Clone>(root: &TreeNode<T>) -> Self {
let mut state = Self::new(root);
state.expand_all(root);
state.rebuild_visible(root);
state
}
pub fn is_expanded(&self, id: &str) -> bool {
self.expanded.contains(id)
}
pub fn expand(&mut self, id: &str) {
self.expanded.insert(id.to_string());
}
pub fn collapse(&mut self, id: &str) {
self.expanded.remove(id);
}
pub fn toggle(&mut self, id: &str) {
if self.expanded.contains(id) {
self.expanded.remove(id);
} else {
self.expanded.insert(id.to_string());
}
}
pub fn expand_all<T: Clone>(&mut self, root: &TreeNode<T>) {
for id in root.all_ids() {
self.expanded.insert(id);
}
}
pub fn collapse_all(&mut self) {
self.expanded.clear();
}
pub fn selected(&self) -> Option<&str> {
self.selected.as_deref()
}
pub fn set_selected(&mut self, id: Option<String>) {
self.selected = id;
}
pub fn cursor(&self) -> usize {
self.cursor
}
pub fn focused(&self) -> Option<&str> {
self.visible_nodes.get(self.cursor).map(|s| s.as_str())
}
pub fn cursor_up(&mut self) {
if self.cursor > 0 {
self.cursor -= 1;
}
}
pub fn cursor_down(&mut self) {
if self.cursor < self.visible_nodes.len().saturating_sub(1) {
self.cursor += 1;
}
}
pub fn cursor_first(&mut self) {
self.cursor = 0;
}
pub fn cursor_last(&mut self) {
self.cursor = self.visible_nodes.len().saturating_sub(1);
}
pub fn toggle_focused(&mut self) {
if let Some(id) = self.focused().map(|s| s.to_string()) {
self.toggle(&id);
}
}
pub fn select_focused(&mut self) {
self.selected = self.focused().map(|s| s.to_string());
}
pub fn rebuild_visible<T: Clone>(&mut self, root: &TreeNode<T>) {
self.visible_nodes.clear();
self.collect_visible(root);
if self.cursor >= self.visible_nodes.len() {
self.cursor = self.visible_nodes.len().saturating_sub(1);
}
}
fn collect_visible<T: Clone>(&mut self, node: &TreeNode<T>) {
self.visible_nodes.push(node.id.clone());
if self.is_expanded(&node.id) {
for child in &node.children {
self.collect_visible(child);
}
}
}
pub fn visible_count(&self) -> usize {
self.visible_nodes.len()
}
}
#[derive(Debug, Clone)]
pub struct TreeStyle {
pub indent: usize,
pub expanded_icon: String,
pub collapsed_icon: String,
pub leaf_icon: String,
pub connector: String,
pub last_connector: String,
pub vertical_line: String,
pub icon_color: Option<Color>,
pub selected_color: Option<Color>,
pub focused_color: Option<Color>,
pub focused_bg: Option<Color>,
pub show_lines: bool,
}
impl Default for TreeStyle {
fn default() -> Self {
Self {
indent: 2,
expanded_icon: "▼".to_string(),
collapsed_icon: "▶".to_string(),
leaf_icon: "•".to_string(),
connector: "├─".to_string(),
last_connector: "└─".to_string(),
vertical_line: "│ ".to_string(),
icon_color: Some(Color::Cyan),
selected_color: Some(Color::Green),
focused_color: Some(Color::Cyan),
focused_bg: None,
show_lines: true,
}
}
}
impl TreeStyle {
pub fn new() -> Self {
Self::default()
}
pub fn indent(mut self, indent: usize) -> Self {
self.indent = indent;
self
}
pub fn expanded_icon(mut self, icon: impl Into<String>) -> Self {
self.expanded_icon = icon.into();
self
}
pub fn collapsed_icon(mut self, icon: impl Into<String>) -> Self {
self.collapsed_icon = icon.into();
self
}
pub fn leaf_icon(mut self, icon: impl Into<String>) -> Self {
self.leaf_icon = icon.into();
self
}
pub fn icon_color(mut self, color: Color) -> Self {
self.icon_color = Some(color);
self
}
pub fn focused_color(mut self, color: Color) -> Self {
self.focused_color = Some(color);
self
}
pub fn show_lines(mut self, show: bool) -> Self {
self.show_lines = show;
self
}
pub fn folder_icons() -> Self {
Self {
expanded_icon: "📂".to_string(),
collapsed_icon: "📁".to_string(),
leaf_icon: "📄".to_string(),
..Default::default()
}
}
pub fn arrow_icons() -> Self {
Self {
expanded_icon: "▼".to_string(),
collapsed_icon: "▶".to_string(),
leaf_icon: " ".to_string(),
..Default::default()
}
}
pub fn plus_minus_icons() -> Self {
Self {
expanded_icon: "[-]".to_string(),
collapsed_icon: "[+]".to_string(),
leaf_icon: " - ".to_string(),
..Default::default()
}
}
pub fn minimal() -> Self {
Self {
show_lines: false,
expanded_icon: "▾".to_string(),
collapsed_icon: "▸".to_string(),
leaf_icon: " ".to_string(),
..Default::default()
}
}
}
#[derive(Debug, Clone)]
pub struct Tree<'a, T: Clone> {
root: &'a TreeNode<T>,
state: &'a TreeState,
style: TreeStyle,
focused: bool,
}
impl<'a, T: Clone> Tree<'a, T> {
pub fn new(root: &'a TreeNode<T>, state: &'a TreeState) -> Self {
Self {
root,
state,
style: TreeStyle::default(),
focused: true,
}
}
pub fn style(mut self, style: TreeStyle) -> Self {
self.style = style;
self
}
pub fn focused(mut self, focused: bool) -> Self {
self.focused = focused;
self
}
pub fn into_element(self) -> Element {
let mut container = RnkBox::new().flex_direction(FlexDirection::Column);
let elements = self.render_node(self.root, 0, vec![]);
for elem in elements {
container = container.child(elem);
}
container.into_element()
}
fn render_node(
&self,
node: &TreeNode<T>,
depth: usize,
parent_is_last: Vec<bool>,
) -> Vec<Element> {
let mut elements = Vec::new();
let is_focused = self.focused && self.state.focused() == Some(&node.id);
let is_selected = self.state.selected() == Some(&node.id);
let is_expanded = self.state.is_expanded(&node.id);
let mut prefix = String::new();
if self.style.show_lines && depth > 0 {
for &is_last in &parent_is_last[..parent_is_last.len().saturating_sub(1)] {
if is_last {
prefix.push_str(" ");
} else {
prefix.push_str(&self.style.vertical_line);
}
}
if let Some(&is_last) = parent_is_last.last() {
if is_last {
prefix.push_str(&self.style.last_connector);
} else {
prefix.push_str(&self.style.connector);
}
}
} else {
prefix = " ".repeat(depth * self.style.indent);
}
let icon = if node.is_leaf() {
&self.style.leaf_icon
} else if is_expanded {
&self.style.expanded_icon
} else {
&self.style.collapsed_icon
};
let line = format!("{}{} {}", prefix, icon, node.label);
let mut text = Text::new(&line);
if is_focused {
if let Some(color) = self.style.focused_color {
text = text.color(color);
}
if let Some(bg) = self.style.focused_bg {
text = text.background(bg);
}
text = text.bold();
} else if is_selected {
if let Some(color) = self.style.selected_color {
text = text.color(color);
}
}
elements.push(text.into_element());
if is_expanded {
let child_count = node.children.len();
for (i, child) in node.children.iter().enumerate() {
let is_last = i == child_count - 1;
let mut child_is_last = parent_is_last.clone();
child_is_last.push(is_last);
elements.extend(self.render_node(child, depth + 1, child_is_last));
}
}
elements
}
}
pub fn handle_tree_input<T: Clone>(
state: &mut TreeState,
root: &TreeNode<T>,
_input: &str,
key: &crate::hooks::Key,
) -> bool {
let mut handled = false;
if key.up_arrow {
state.cursor_up();
handled = true;
} else if key.down_arrow {
state.cursor_down();
handled = true;
} else if key.left_arrow {
if let Some(id) = state.focused().map(|s| s.to_string()) {
if state.is_expanded(&id) {
state.collapse(&id);
state.rebuild_visible(root);
}
}
handled = true;
} else if key.right_arrow {
if let Some(id) = state.focused().map(|s| s.to_string()) {
if !state.is_expanded(&id) {
state.expand(&id);
state.rebuild_visible(root);
}
}
handled = true;
} else if key.return_key || key.space {
if let Some(id) = state.focused().map(|s| s.to_string()) {
state.toggle(&id);
state.rebuild_visible(root);
}
handled = true;
} else if key.home {
state.cursor_first();
handled = true;
} else if key.end {
state.cursor_last();
handled = true;
}
handled
}
#[cfg(test)]
mod tests {
use super::*;
fn sample_tree() -> TreeNode<()> {
TreeNode::new("root", "Root")
.child(
TreeNode::new("a", "Node A")
.child(TreeNode::leaf("a1", "Leaf A1"))
.child(TreeNode::leaf("a2", "Leaf A2")),
)
.child(TreeNode::leaf("b", "Node B"))
.child(TreeNode::new("c", "Node C").child(TreeNode::leaf("c1", "Leaf C1")))
}
#[test]
fn test_tree_node_creation() {
let node = TreeNode::<()>::new("test", "Test Node");
assert_eq!(node.id, "test");
assert_eq!(node.label, "Test Node");
assert!(node.is_leaf());
}
#[test]
fn test_tree_node_with_children() {
let node = TreeNode::<()>::new("parent", "Parent")
.child(TreeNode::leaf("child1", "Child 1"))
.child(TreeNode::leaf("child2", "Child 2"));
assert!(!node.is_leaf());
assert!(node.has_children());
assert_eq!(node.children.len(), 2);
}
#[test]
fn test_tree_node_count() {
let tree = sample_tree();
assert_eq!(tree.node_count(), 7); }
#[test]
fn test_tree_find() {
let tree = sample_tree();
assert!(tree.find("a1").is_some());
assert!(tree.find("nonexistent").is_none());
}
#[test]
fn test_tree_state_new() {
let tree = sample_tree();
let state = TreeState::new(&tree);
assert_eq!(state.visible_count(), 1); assert_eq!(state.cursor(), 0);
assert!(!state.is_expanded("root"));
}
#[test]
fn test_tree_state_expand() {
let tree = sample_tree();
let mut state = TreeState::new(&tree);
state.expand("root");
state.rebuild_visible(&tree);
assert!(state.is_expanded("root"));
assert_eq!(state.visible_count(), 4); }
#[test]
fn test_tree_state_expand_all() {
let tree = sample_tree();
let mut state = TreeState::new(&tree);
state.expand_all(&tree);
state.rebuild_visible(&tree);
assert_eq!(state.visible_count(), 7); }
#[test]
fn test_tree_state_navigation() {
let tree = sample_tree();
let mut state = TreeState::all_expanded(&tree);
assert_eq!(state.cursor(), 0);
assert_eq!(state.focused(), Some("root"));
state.cursor_down();
assert_eq!(state.focused(), Some("a"));
state.cursor_down();
assert_eq!(state.focused(), Some("a1"));
state.cursor_up();
assert_eq!(state.focused(), Some("a"));
state.cursor_last();
assert_eq!(state.focused(), Some("c1"));
state.cursor_first();
assert_eq!(state.focused(), Some("root"));
}
#[test]
fn test_tree_state_toggle() {
let tree = sample_tree();
let mut state = TreeState::new(&tree);
state.expand("root");
state.rebuild_visible(&tree);
assert_eq!(state.visible_count(), 4);
state.toggle("root");
state.rebuild_visible(&tree);
assert_eq!(state.visible_count(), 1);
}
#[test]
fn test_tree_style_presets() {
let _default = TreeStyle::default();
let _folder = TreeStyle::folder_icons();
let _arrow = TreeStyle::arrow_icons();
let _plus_minus = TreeStyle::plus_minus_icons();
let _minimal = TreeStyle::minimal();
}
}