Skip to main content

envision/component/span_tree/
mod.rs

1//! A hierarchical span tree component for trace visualization.
2//!
3//! [`SpanTree`] displays hierarchical spans with horizontal timing bars
4//! aligned to a shared time axis, similar to the trace view in Jaeger or
5//! Zipkin. Each row shows a label on the left and a proportional duration
6//! bar on the right.
7//!
8//! State is stored in [`SpanTreeState`], updated via [`SpanTreeMessage`],
9//! and produces [`SpanTreeOutput`].
10//!
11//!
12//! # Example
13//!
14//! ```rust
15//! use envision::component::{
16//!     Component, SpanTree, SpanTreeState, SpanTreeMessage, SpanNode,
17//! };
18//! use ratatui::style::Color;
19//!
20//! let root = SpanNode::new("root", "frontend/request", 0.0, 1000.0)
21//!     .with_color(Color::Cyan)
22//!     .with_child(
23//!         SpanNode::new("db", "db/query", 100.0, 400.0)
24//!             .with_color(Color::Green),
25//!     );
26//! let mut state = SpanTreeState::new(vec![root]);
27//!
28//! // Navigate down
29//! SpanTree::update(&mut state, SpanTreeMessage::SelectDown);
30//! assert_eq!(state.selected_index(), Some(1));
31//! ```
32
33use std::collections::HashSet;
34
35use super::{Component, EventContext, RenderContext};
36use crate::input::{Event, Key};
37use crate::scroll::ScrollState;
38
39mod render;
40mod types;
41
42pub use types::{FlatSpan, SpanNode};
43
44/// Messages that can be sent to a SpanTree component.
45///
46/// # Example
47///
48/// ```rust
49/// use envision::component::{Component, SpanTree, SpanTreeState, SpanTreeMessage, SpanNode};
50///
51/// let root = SpanNode::new("r", "root", 0.0, 100.0);
52/// let mut state = SpanTreeState::new(vec![root]);
53/// SpanTree::update(&mut state, SpanTreeMessage::SelectDown);
54/// ```
55#[derive(Clone, Debug, PartialEq)]
56#[cfg_attr(
57    feature = "serialization",
58    derive(serde::Serialize, serde::Deserialize)
59)]
60pub enum SpanTreeMessage {
61    /// Replace all root spans.
62    SetRoots(Vec<SpanNode>),
63    /// Move selection up.
64    SelectUp,
65    /// Move selection down.
66    SelectDown,
67    /// Expand the selected node.
68    Expand,
69    /// Collapse the selected node.
70    Collapse,
71    /// Toggle expand/collapse of the selected node.
72    Toggle,
73    /// Expand all nodes.
74    ExpandAll,
75    /// Collapse all nodes.
76    CollapseAll,
77    /// Set the label column width.
78    SetLabelWidth(u16),
79}
80
81/// Output messages from a SpanTree component.
82///
83/// # Example
84///
85/// ```rust
86/// use envision::component::{
87///     Component, SpanTree, SpanTreeState, SpanTreeMessage, SpanTreeOutput, SpanNode,
88/// };
89///
90/// let root = SpanNode::new("r", "root", 0.0, 100.0)
91///     .with_child(SpanNode::new("c", "child", 10.0, 50.0));
92/// let mut state = SpanTreeState::new(vec![root]);
93///
94/// let output = SpanTree::update(&mut state, SpanTreeMessage::Collapse);
95/// assert_eq!(output, Some(SpanTreeOutput::Collapsed("r".into())));
96/// ```
97#[derive(Clone, Debug, PartialEq, Eq)]
98#[cfg_attr(
99    feature = "serialization",
100    derive(serde::Serialize, serde::Deserialize)
101)]
102pub enum SpanTreeOutput {
103    /// A span was selected. Contains the span id.
104    Selected(String),
105    /// A span was expanded. Contains the span id.
106    Expanded(String),
107    /// A span was collapsed. Contains the span id.
108    Collapsed(String),
109}
110
111/// State for a SpanTree component.
112///
113/// Contains the root spans, selection state, expanded set, and layout
114/// configuration.
115///
116/// # Example
117///
118/// ```rust
119/// use envision::component::{SpanTreeState, SpanNode};
120///
121/// let root = SpanNode::new("r", "root", 0.0, 1000.0)
122///     .with_child(SpanNode::new("c", "child", 100.0, 500.0));
123/// let state = SpanTreeState::new(vec![root]);
124///
125/// assert_eq!(state.roots().len(), 1);
126/// assert_eq!(state.global_start(), 0.0);
127/// assert_eq!(state.global_end(), 1000.0);
128/// assert_eq!(state.selected_index(), Some(0));
129/// ```
130#[derive(Clone, Debug, PartialEq)]
131#[cfg_attr(
132    feature = "serialization",
133    derive(serde::Serialize, serde::Deserialize)
134)]
135pub struct SpanTreeState {
136    /// Root spans.
137    roots: Vec<SpanNode>,
138    /// Selected index in flattened view.
139    selected_index: Option<usize>,
140    /// Set of expanded node IDs. All nodes are expanded by default.
141    expanded: HashSet<String>,
142    /// Scroll state for vertical scrolling.
143    scroll: ScrollState,
144    /// Earliest start time across all spans.
145    global_start: f64,
146    /// Latest end time across all spans.
147    global_end: f64,
148    /// Width allocated for labels.
149    label_width: u16,
150    /// Optional title.
151    title: Option<String>,
152}
153
154impl Default for SpanTreeState {
155    fn default() -> Self {
156        Self {
157            roots: Vec::new(),
158            selected_index: None,
159            expanded: HashSet::new(),
160            scroll: ScrollState::default(),
161            global_start: 0.0,
162            global_end: 0.0,
163            label_width: 30,
164            title: None,
165        }
166    }
167}
168
169impl SpanTreeState {
170    /// Creates a new span tree state with the given root spans.
171    ///
172    /// All nodes with children start expanded. The first node is selected
173    /// if the tree is non-empty.
174    ///
175    /// # Example
176    ///
177    /// ```rust
178    /// use envision::component::{SpanTreeState, SpanNode};
179    ///
180    /// let root = SpanNode::new("r", "root", 0.0, 100.0)
181    ///     .with_child(SpanNode::new("c", "child", 10.0, 50.0));
182    /// let state = SpanTreeState::new(vec![root]);
183    ///
184    /// assert_eq!(state.roots().len(), 1);
185    /// assert_eq!(state.selected_index(), Some(0));
186    /// assert_eq!(state.global_start(), 0.0);
187    /// assert_eq!(state.global_end(), 100.0);
188    /// ```
189    pub fn new(roots: Vec<SpanNode>) -> Self {
190        let mut expanded = HashSet::new();
191        for root in &roots {
192            Self::collect_expanded_ids(root, &mut expanded);
193        }
194
195        let (global_start, global_end) = Self::compute_global_range(&roots);
196        let selected_index = if roots.is_empty() { None } else { Some(0) };
197
198        let mut state = Self {
199            roots,
200            selected_index,
201            expanded,
202            scroll: ScrollState::default(),
203            global_start,
204            global_end,
205            label_width: 30,
206            title: None,
207        };
208        state.scroll.set_content_length(state.flatten().len());
209        state
210    }
211
212    /// Sets the title (builder pattern).
213    ///
214    /// # Example
215    ///
216    /// ```rust
217    /// use envision::component::{SpanTreeState, SpanNode};
218    ///
219    /// let state = SpanTreeState::new(vec![SpanNode::new("r", "root", 0.0, 10.0)])
220    ///     .with_title("Trace");
221    /// assert_eq!(state.title(), Some("Trace"));
222    /// ```
223    pub fn with_title(mut self, title: impl Into<String>) -> Self {
224        self.title = Some(title.into());
225        self
226    }
227
228    /// Sets the label column width (builder pattern).
229    ///
230    /// # Example
231    ///
232    /// ```rust
233    /// use envision::component::{SpanTreeState, SpanNode};
234    ///
235    /// let state = SpanTreeState::new(vec![SpanNode::new("r", "root", 0.0, 10.0)])
236    ///     .with_label_width(40);
237    /// assert_eq!(state.label_width(), 40);
238    /// ```
239    pub fn with_label_width(mut self, width: u16) -> Self {
240        self.label_width = width;
241        self
242    }
243
244    // ---- Accessors ----
245
246    /// Returns the root spans.
247    ///
248    /// # Example
249    ///
250    /// ```rust
251    /// use envision::component::{SpanTreeState, SpanNode};
252    ///
253    /// let state = SpanTreeState::new(vec![
254    ///     SpanNode::new("a", "svc-a", 0.0, 10.0),
255    ///     SpanNode::new("b", "svc-b", 5.0, 15.0),
256    /// ]);
257    /// assert_eq!(state.roots().len(), 2);
258    /// ```
259    pub fn roots(&self) -> &[SpanNode] {
260        &self.roots
261    }
262
263    /// Returns a mutable reference to the root spans.
264    ///
265    /// This is safe because span nodes are simple data containers.
266    /// The expanded set tracks nodes by string id, so mutating node
267    /// data does not corrupt expand/collapse state.
268    ///
269    /// # Example
270    ///
271    /// ```rust
272    /// use envision::component::{SpanTreeState, SpanNode};
273    /// use ratatui::style::Color;
274    ///
275    /// let root = SpanNode::new("r", "root", 0.0, 100.0);
276    /// let mut state = SpanTreeState::new(vec![root]);
277    /// state.roots_mut()[0].set_color(Color::Red);
278    /// assert_eq!(state.roots()[0].color(), Color::Red);
279    /// ```
280    /// **Note**: After modifying the collection, the scrollbar may be inaccurate
281    /// until the next render. Prefer dedicated methods (e.g., `push_event()`) when available.
282    pub fn roots_mut(&mut self) -> &mut Vec<SpanNode> {
283        &mut self.roots
284    }
285
286    /// Replaces all root spans and recomputes global times.
287    ///
288    /// # Example
289    ///
290    /// ```rust
291    /// use envision::component::{SpanTreeState, SpanNode};
292    ///
293    /// let mut state = SpanTreeState::new(vec![SpanNode::new("old", "old", 0.0, 10.0)]);
294    /// state.set_roots(vec![SpanNode::new("new", "new", 5.0, 50.0)]);
295    /// assert_eq!(state.roots().len(), 1);
296    /// assert_eq!(state.global_start(), 5.0);
297    /// assert_eq!(state.global_end(), 50.0);
298    /// ```
299    pub fn set_roots(&mut self, roots: Vec<SpanNode>) {
300        self.expanded.clear();
301        for root in &roots {
302            Self::collect_expanded_ids(root, &mut self.expanded);
303        }
304
305        let (global_start, global_end) = Self::compute_global_range(&roots);
306        self.global_start = global_start;
307        self.global_end = global_end;
308        self.roots = roots;
309        self.selected_index = if self.roots.is_empty() { None } else { Some(0) };
310        self.scroll.set_content_length(self.flatten().len());
311    }
312
313    /// Returns the currently selected flat index.
314    ///
315    /// # Example
316    ///
317    /// ```rust
318    /// use envision::component::{SpanTreeState, SpanNode};
319    ///
320    /// let state = SpanTreeState::new(vec![SpanNode::new("r", "root", 0.0, 10.0)]);
321    /// assert_eq!(state.selected_index(), Some(0));
322    ///
323    /// let empty = SpanTreeState::default();
324    /// assert_eq!(empty.selected_index(), None);
325    /// ```
326    pub fn selected_index(&self) -> Option<usize> {
327        self.selected_index
328    }
329
330    /// Returns the currently selected span in flattened view.
331    ///
332    /// # Example
333    ///
334    /// ```rust
335    /// use envision::component::{SpanTreeState, SpanNode};
336    ///
337    /// let root = SpanNode::new("r", "root", 0.0, 100.0);
338    /// let state = SpanTreeState::new(vec![root]);
339    /// let selected = state.selected_span().unwrap();
340    /// assert_eq!(selected.id(), "r");
341    /// assert_eq!(selected.label(), "root");
342    /// ```
343    pub fn selected_span(&self) -> Option<FlatSpan> {
344        let flat = self.flatten();
345        let idx = self.selected_index?;
346        flat.into_iter().nth(idx)
347    }
348
349    /// Returns the earliest start time across all spans.
350    pub fn global_start(&self) -> f64 {
351        self.global_start
352    }
353
354    /// Returns the latest end time across all spans.
355    pub fn global_end(&self) -> f64 {
356        self.global_end
357    }
358
359    /// Returns the label column width.
360    pub fn label_width(&self) -> u16 {
361        self.label_width
362    }
363
364    /// Returns the title, if set.
365    pub fn title(&self) -> Option<&str> {
366        self.title.as_deref()
367    }
368
369    /// Sets the title.
370    ///
371    /// # Example
372    ///
373    /// ```rust
374    /// use envision::component::SpanTreeState;
375    ///
376    /// let mut state = SpanTreeState::default();
377    /// state.set_title("Trace View");
378    /// assert_eq!(state.title(), Some("Trace View"));
379    /// ```
380    pub fn set_title(&mut self, title: impl Into<String>) {
381        self.title = Some(title.into());
382    }
383
384    /// Returns true if the tree is empty.
385    ///
386    /// # Example
387    ///
388    /// ```rust
389    /// use envision::component::SpanTreeState;
390    ///
391    /// assert!(SpanTreeState::default().is_empty());
392    /// ```
393    pub fn is_empty(&self) -> bool {
394        self.roots.is_empty()
395    }
396
397    /// Returns the set of expanded node IDs.
398    pub fn expanded_ids(&self) -> &HashSet<String> {
399        &self.expanded
400    }
401
402    // ---- Expand/Collapse ----
403
404    /// Expands a node by its ID.
405    ///
406    /// # Example
407    ///
408    /// ```rust
409    /// use envision::component::{SpanTreeState, SpanNode};
410    ///
411    /// let root = SpanNode::new("r", "root", 0.0, 100.0)
412    ///     .with_child(SpanNode::new("c", "child", 10.0, 50.0));
413    /// let mut state = SpanTreeState::new(vec![root]);
414    /// state.collapse("r");
415    /// assert!(!state.expanded_ids().contains("r"));
416    /// state.expand("r");
417    /// assert!(state.expanded_ids().contains("r"));
418    /// ```
419    pub fn expand(&mut self, id: &str) {
420        self.expanded.insert(id.to_string());
421        self.scroll.set_content_length(self.flatten().len());
422    }
423
424    /// Collapses a node by its ID.
425    ///
426    /// # Example
427    ///
428    /// ```rust
429    /// use envision::component::{SpanTreeState, SpanNode};
430    ///
431    /// let root = SpanNode::new("r", "root", 0.0, 100.0)
432    ///     .with_child(SpanNode::new("c", "child", 10.0, 50.0));
433    /// let mut state = SpanTreeState::new(vec![root]);
434    /// state.collapse("r");
435    /// assert!(!state.expanded_ids().contains("r"));
436    /// ```
437    pub fn collapse(&mut self, id: &str) {
438        self.expanded.remove(id);
439        // Clamp selection if it went out of bounds
440        let visible = self.flatten().len();
441        if let Some(idx) = self.selected_index {
442            if idx >= visible {
443                self.selected_index = Some(visible.saturating_sub(1));
444            }
445        }
446        self.scroll.set_content_length(visible);
447    }
448
449    /// Expands all nodes.
450    ///
451    /// # Example
452    ///
453    /// ```rust
454    /// use envision::component::{SpanTreeState, SpanNode};
455    ///
456    /// let root = SpanNode::new("r", "root", 0.0, 100.0)
457    ///     .with_child(SpanNode::new("c", "child", 10.0, 50.0));
458    /// let mut state = SpanTreeState::new(vec![root]);
459    /// state.collapse_all();
460    /// state.expand_all();
461    /// assert!(state.expanded_ids().contains("r"));
462    /// ```
463    pub fn expand_all(&mut self) {
464        self.expanded.clear();
465        for root in &self.roots {
466            Self::collect_expanded_ids(root, &mut self.expanded);
467        }
468        self.scroll.set_content_length(self.flatten().len());
469    }
470
471    /// Collapses all nodes.
472    ///
473    /// # Example
474    ///
475    /// ```rust
476    /// use envision::component::{SpanTreeState, SpanNode};
477    ///
478    /// let root = SpanNode::new("r", "root", 0.0, 100.0)
479    ///     .with_child(SpanNode::new("c", "child", 10.0, 50.0));
480    /// let mut state = SpanTreeState::new(vec![root]);
481    /// state.collapse_all();
482    /// assert!(state.expanded_ids().is_empty());
483    /// ```
484    pub fn collapse_all(&mut self) {
485        self.expanded.clear();
486        self.selected_index = if self.roots.is_empty() { None } else { Some(0) };
487        self.scroll.set_content_length(self.flatten().len());
488    }
489
490    // ---- Flatten ----
491
492    /// Flattens the visible hierarchy into a list of [`FlatSpan`] items.
493    ///
494    /// Only expanded nodes have their children included. The order is
495    /// depth-first, matching the visual tree order.
496    ///
497    /// # Example
498    ///
499    /// ```rust
500    /// use envision::component::{SpanTreeState, SpanNode};
501    ///
502    /// let root = SpanNode::new("r", "root", 0.0, 100.0)
503    ///     .with_child(SpanNode::new("c1", "child-1", 10.0, 50.0))
504    ///     .with_child(SpanNode::new("c2", "child-2", 50.0, 90.0));
505    /// let state = SpanTreeState::new(vec![root]);
506    /// let flat = state.flatten();
507    /// assert_eq!(flat.len(), 3);
508    /// assert_eq!(flat[0].id(), "r");
509    /// assert_eq!(flat[0].depth(), 0);
510    /// assert_eq!(flat[1].id(), "c1");
511    /// assert_eq!(flat[1].depth(), 1);
512    /// ```
513    pub fn flatten(&self) -> Vec<FlatSpan> {
514        let mut result = Vec::new();
515        for root in &self.roots {
516            self.flatten_node(root, 0, &mut result);
517        }
518        result
519    }
520
521    // ---- Instance methods ----
522
523    /// Updates the state with a message, returning any output.
524    pub fn update(&mut self, msg: SpanTreeMessage) -> Option<SpanTreeOutput> {
525        SpanTree::update(self, msg)
526    }
527
528    // ---- Private helpers ----
529
530    /// Recursively collects all node IDs with children into the expanded set.
531    fn collect_expanded_ids(node: &SpanNode, expanded: &mut HashSet<String>) {
532        if node.has_children() {
533            expanded.insert(node.id.clone());
534            for child in &node.children {
535                Self::collect_expanded_ids(child, expanded);
536            }
537        }
538    }
539
540    /// Computes the global start and end times across all spans.
541    fn compute_global_range(roots: &[SpanNode]) -> (f64, f64) {
542        let mut min_start = f64::INFINITY;
543        let mut max_end = f64::NEG_INFINITY;
544        for root in roots {
545            Self::compute_range_recursive(root, &mut min_start, &mut max_end);
546        }
547        if min_start.is_infinite() {
548            (0.0, 0.0)
549        } else {
550            (min_start, max_end)
551        }
552    }
553
554    /// Recursively computes min start and max end across all nodes.
555    fn compute_range_recursive(node: &SpanNode, min_start: &mut f64, max_end: &mut f64) {
556        if node.start < *min_start {
557            *min_start = node.start;
558        }
559        if node.end > *max_end {
560            *max_end = node.end;
561        }
562        for child in &node.children {
563            Self::compute_range_recursive(child, min_start, max_end);
564        }
565    }
566
567    /// Recursively flattens a node and its visible children.
568    fn flatten_node(&self, node: &SpanNode, depth: usize, result: &mut Vec<FlatSpan>) {
569        let is_expanded = self.expanded.contains(&node.id);
570        result.push(FlatSpan {
571            id: node.id.clone(),
572            label: node.label.clone(),
573            start: node.start,
574            end: node.end,
575            color: node.color,
576            status: node.status.clone(),
577            depth,
578            has_children: node.has_children(),
579            is_expanded,
580        });
581
582        if is_expanded {
583            for child in &node.children {
584                self.flatten_node(child, depth + 1, result);
585            }
586        }
587    }
588}
589
590/// A hierarchical span tree component for trace visualization.
591///
592/// Displays hierarchical spans with horizontal timing bars aligned to
593/// a shared time axis. Each row shows a label with tree
594/// expand/collapse indicators on the left, and a proportional
595/// duration bar on the right.
596///
597/// # Visual Format
598///
599/// ```text
600/// ┌─ Trace ─────────────────────────────────────────┐
601/// │ Label                     │ 0ms    500ms   1000ms│
602/// │───────────────────────────┼──────────────────────│
603/// │ ▾ frontend/request        │ ████████████████████ │
604/// │   ▾ api/handler           │   ██████████████     │
605/// │     db/query              │     ████             │
606/// │     cache/lookup          │          ███         │
607/// │   auth/validate           │ ███                  │
608/// └─────────────────────────────────────────────────┘
609/// ```
610///
611/// # Keyboard Navigation
612///
613/// - `Up/k`: Move selection up
614/// - `Down/j`: Move selection down
615/// - `Right/l`: Expand selected node
616/// - `Left/h`: Collapse selected node
617/// - `Space/Enter`: Toggle expand/collapse
618/// - `Shift+Right`: Increase label width
619/// - `Shift+Left`: Decrease label width
620///
621/// # Example
622///
623/// ```rust
624/// use envision::component::{
625///     Component, SpanTree, SpanTreeState, SpanTreeMessage, SpanNode,
626/// };
627/// use ratatui::style::Color;
628///
629/// let root = SpanNode::new("req", "frontend/request", 0.0, 1000.0)
630///     .with_color(Color::Cyan)
631///     .with_child(
632///         SpanNode::new("api", "api/handler", 50.0, 800.0)
633///             .with_color(Color::Yellow)
634///             .with_child(SpanNode::new("db", "db/query", 100.0, 400.0).with_color(Color::Green))
635///     );
636///
637/// let mut state = SpanTreeState::new(vec![root]);
638///
639/// // Navigate through spans
640/// SpanTree::update(&mut state, SpanTreeMessage::SelectDown);
641/// SpanTree::update(&mut state, SpanTreeMessage::Collapse);
642/// ```
643pub struct SpanTree;
644
645impl Component for SpanTree {
646    type State = SpanTreeState;
647    type Message = SpanTreeMessage;
648    type Output = SpanTreeOutput;
649
650    fn init() -> Self::State {
651        SpanTreeState::default()
652    }
653
654    fn update(state: &mut Self::State, msg: Self::Message) -> Option<Self::Output> {
655        match msg {
656            SpanTreeMessage::SetRoots(roots) => {
657                state.set_roots(roots);
658                None
659            }
660            SpanTreeMessage::SetLabelWidth(width) => {
661                state.label_width = width.clamp(10, 100);
662                None
663            }
664            SpanTreeMessage::ExpandAll => {
665                state.expand_all();
666                None
667            }
668            SpanTreeMessage::CollapseAll => {
669                state.collapse_all();
670                None
671            }
672            _ => {
673                let flat = state.flatten();
674                if flat.is_empty() {
675                    return None;
676                }
677
678                let selected = state.selected_index?;
679
680                match msg {
681                    SpanTreeMessage::SelectUp => {
682                        if selected > 0 {
683                            state.selected_index = Some(selected - 1);
684                            let span = &flat[selected - 1];
685                            return Some(SpanTreeOutput::Selected(span.id.clone()));
686                        }
687                        None
688                    }
689                    SpanTreeMessage::SelectDown => {
690                        if selected < flat.len() - 1 {
691                            state.selected_index = Some(selected + 1);
692                            let span = &flat[selected + 1];
693                            return Some(SpanTreeOutput::Selected(span.id.clone()));
694                        }
695                        None
696                    }
697                    SpanTreeMessage::Expand => {
698                        if let Some(span) = flat.get(selected) {
699                            if span.has_children && !span.is_expanded {
700                                let id = span.id.clone();
701                                state.expanded.insert(id.clone());
702                                state.scroll.set_content_length(state.flatten().len());
703                                return Some(SpanTreeOutput::Expanded(id));
704                            }
705                        }
706                        None
707                    }
708                    SpanTreeMessage::Collapse => {
709                        if let Some(span) = flat.get(selected) {
710                            if span.has_children && span.is_expanded {
711                                let id = span.id.clone();
712                                state.expanded.remove(&id);
713                                let new_flat = state.flatten();
714                                if selected >= new_flat.len() {
715                                    state.selected_index = Some(new_flat.len().saturating_sub(1));
716                                }
717                                state.scroll.set_content_length(new_flat.len());
718                                return Some(SpanTreeOutput::Collapsed(id));
719                            }
720                        }
721                        None
722                    }
723                    SpanTreeMessage::Toggle => {
724                        if let Some(span) = flat.get(selected) {
725                            if span.has_children {
726                                let id = span.id.clone();
727                                if span.is_expanded {
728                                    state.expanded.remove(&id);
729                                    let new_flat = state.flatten();
730                                    if selected >= new_flat.len() {
731                                        state.selected_index =
732                                            Some(new_flat.len().saturating_sub(1));
733                                    }
734                                    state.scroll.set_content_length(new_flat.len());
735                                    return Some(SpanTreeOutput::Collapsed(id));
736                                } else {
737                                    state.expanded.insert(id.clone());
738                                    state.scroll.set_content_length(state.flatten().len());
739                                    return Some(SpanTreeOutput::Expanded(id));
740                                }
741                            }
742                        }
743                        None
744                    }
745                    // Already handled above
746                    SpanTreeMessage::SetRoots(_)
747                    | SpanTreeMessage::SetLabelWidth(_)
748                    | SpanTreeMessage::ExpandAll
749                    | SpanTreeMessage::CollapseAll => unreachable!(),
750                }
751            }
752        }
753    }
754
755    fn handle_event(
756        state: &Self::State,
757        event: &Event,
758        ctx: &EventContext,
759    ) -> Option<Self::Message> {
760        if !ctx.focused || ctx.disabled {
761            return None;
762        }
763        if let Some(key) = event.as_key() {
764            let has_shift = key.modifiers.shift();
765            match key.code {
766                Key::Up | Key::Char('k') if !has_shift => Some(SpanTreeMessage::SelectUp),
767                Key::Down | Key::Char('j') if !has_shift => Some(SpanTreeMessage::SelectDown),
768                Key::Right | Key::Char('l') if has_shift => Some(SpanTreeMessage::SetLabelWidth(
769                    state.label_width.saturating_add(2),
770                )),
771                Key::Left | Key::Char('h') if has_shift => Some(SpanTreeMessage::SetLabelWidth(
772                    state.label_width.saturating_sub(2),
773                )),
774                Key::Right | Key::Char('l') => Some(SpanTreeMessage::Expand),
775                Key::Left | Key::Char('h') => Some(SpanTreeMessage::Collapse),
776                Key::Char(' ') | Key::Enter => Some(SpanTreeMessage::Toggle),
777                _ => None,
778            }
779        } else {
780            None
781        }
782    }
783
784    fn view(state: &Self::State, ctx: &mut RenderContext<'_, '_>) {
785        render::render_span_tree(
786            state,
787            ctx.frame,
788            ctx.area,
789            ctx.theme,
790            ctx.focused,
791            ctx.disabled,
792        );
793    }
794}
795
796#[cfg(test)]
797mod snapshot_tests;
798#[cfg(test)]
799mod tests;