grafos_profile/

span_tree.rs

//! Hierarchical span tree reconstruction from flat span lists.
//!
//! Builds a tree from parent-child relationships using [`ResourceSpan::parent_span_id`].
//! Orphan spans (missing parents) become additional roots.

use alloc::vec::Vec;

use grafos_observe::span::ResourceSpan;

/// A node in the hierarchical span tree.
#[derive(Debug, Clone)]
pub struct SpanTreeNode {
    /// The span data at this node.
    pub span: ResourceSpan,
    /// Indices of child nodes in the parent [`SpanTree::nodes`] vector.
    pub children: Vec<usize>,
}

/// Hierarchical reconstruction of spans from parent-child relationships.
#[derive(Debug, Clone)]
pub struct SpanTree {
    /// All nodes in the tree (flat storage, children reference by index).
    pub nodes: Vec<SpanTreeNode>,
    /// Indices of root nodes (spans with no parent or orphaned parents).
    pub roots: Vec<usize>,
}

impl SpanTree {
    /// Build a span tree from a flat list of spans.
    ///
    /// Uses `parent_span_id` to establish parent-child relationships.
    /// Spans whose parent is not found in the list become additional roots.
    pub fn build(spans: &[ResourceSpan]) -> Self {
        let mut nodes: Vec<SpanTreeNode> = spans
            .iter()
            .map(|s| SpanTreeNode {
                span: s.clone(),
                children: Vec::new(),
            })
            .collect();

        // Build a map from span_id -> index
        let mut roots = Vec::new();

        // For each span, find its parent and add it as a child
        for i in 0..nodes.len() {
            let parent_id = nodes[i].span.parent_span_id;
            if let Some(pid) = parent_id {
                if pid.is_valid() {
                    // Find parent by span_id
                    let parent_idx = nodes
                        .iter()
                        .position(|n| n.span.trace_context.span_id == pid);
                    if parent_idx.is_some() {
                        // Has a parent in the list — will be wired up in second pass
                        continue;
                    }
                }
            }
            // No valid parent or parent not found -> root
            roots.push(i);
        }

        // Second pass: actually wire up children
        // Collect (parent_idx, child_idx) pairs first
        let mut child_links: Vec<(usize, usize)> = Vec::new();
        for i in 0..nodes.len() {
            if roots.contains(&i) {
                continue;
            }
            if let Some(pid) = nodes[i].span.parent_span_id {
                if pid.is_valid() {
                    if let Some(pi) = nodes
                        .iter()
                        .position(|n| n.span.trace_context.span_id == pid)
                    {
                        child_links.push((pi, i));
                    } else {
                        roots.push(i);
                    }
                }
            }
        }

        for (parent_idx, child_idx) in child_links {
            nodes[parent_idx].children.push(child_idx);
        }

        SpanTree { nodes, roots }
    }

    /// Total lease cost across all root subtrees.
    pub fn total_cost(&self) -> u64 {
        self.nodes.iter().map(|n| n.span.lease_cost_byte_secs).sum()
    }

    /// Depth-first iteration yielding (depth, node_index).
    pub fn dfs(&self) -> Vec<(usize, usize)> {
        let mut result = Vec::new();
        for &root_idx in &self.roots {
            self.dfs_visit(root_idx, 0, &mut result);
        }
        result
    }

    fn dfs_visit(&self, idx: usize, depth: usize, result: &mut Vec<(usize, usize)>) {
        result.push((depth, idx));
        for &child_idx in &self.nodes[idx].children {
            self.dfs_visit(child_idx, depth + 1, result);
        }
    }

    /// Compute the aggregated cost of a subtree rooted at the given node index.
    pub fn subtree_cost(&self, idx: usize) -> u64 {
        let mut cost = self.nodes[idx].span.lease_cost_byte_secs;
        for &child_idx in &self.nodes[idx].children {
            cost += self.subtree_cost(child_idx);
        }
        cost
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use grafos_observe::trace::{SpanId, TraceContext};

    fn test_ctx() -> TraceContext {
        let mut bytes = [0u8; 24];
        for (i, b) in bytes.iter_mut().enumerate() {
            *b = (i as u8).wrapping_add(0x42);
        }
        TraceContext::new_root(&bytes)
    }

    #[test]
    fn single_root() {
        let ctx = test_ctx();
        let span = ResourceSpan::new("root", ctx);
        let tree = SpanTree::build(&[span]);

        assert_eq!(tree.nodes.len(), 1);
        assert_eq!(tree.roots.len(), 1);
        assert_eq!(tree.roots[0], 0);
        assert!(tree.nodes[0].children.is_empty());
    }

    #[test]
    fn parent_child_hierarchy() {
        let root_ctx = test_ctx();
        let child_ctx = root_ctx.child(&[0xAA; 8]);
        let grandchild_ctx = child_ctx.child(&[0xBB; 8]);

        let mut root_span = ResourceSpan::new("root", root_ctx);
        root_span.lease_cost_byte_secs = 100;

        let mut child_span = ResourceSpan::new("child", child_ctx);
        child_span.parent_span_id = Some(root_ctx.span_id);
        child_span.lease_cost_byte_secs = 50;

        let mut grandchild_span = ResourceSpan::new("grandchild", grandchild_ctx);
        grandchild_span.parent_span_id = Some(child_ctx.span_id);
        grandchild_span.lease_cost_byte_secs = 25;

        let tree = SpanTree::build(&[root_span, child_span, grandchild_span]);

        assert_eq!(tree.roots.len(), 1);
        assert_eq!(tree.roots[0], 0);
        assert_eq!(tree.nodes[0].children, vec![1]);
        assert_eq!(tree.nodes[1].children, vec![2]);
        assert!(tree.nodes[2].children.is_empty());

        assert_eq!(tree.total_cost(), 175);
        assert_eq!(tree.subtree_cost(0), 175);
        assert_eq!(tree.subtree_cost(1), 75);
        assert_eq!(tree.subtree_cost(2), 25);
    }

    #[test]
    fn orphan_becomes_root() {
        let ctx1 = test_ctx();
        let ctx2 = ctx1.child(&[0xCC; 8]);

        let root_span = ResourceSpan::new("root", ctx1);

        // Orphan: has a parent_span_id that doesn't match any span in the list
        let mut orphan = ResourceSpan::new("orphan", ctx2);
        orphan.parent_span_id = Some(SpanId(0xDEAD));

        let tree = SpanTree::build(&[root_span, orphan]);

        assert_eq!(tree.roots.len(), 2);
        assert!(tree.roots.contains(&0));
        assert!(tree.roots.contains(&1));
    }

    #[test]
    fn dfs_order() {
        let root_ctx = test_ctx();
        let child_a_ctx = root_ctx.child(&[0xAA; 8]);
        let child_b_ctx = root_ctx.child(&[0xBB; 8]);

        let root = ResourceSpan::new("root", root_ctx);

        let mut child_a = ResourceSpan::new("child_a", child_a_ctx);
        child_a.parent_span_id = Some(root_ctx.span_id);

        let mut child_b = ResourceSpan::new("child_b", child_b_ctx);
        child_b.parent_span_id = Some(root_ctx.span_id);

        let tree = SpanTree::build(&[root, child_a, child_b]);
        let dfs = tree.dfs();

        assert_eq!(dfs.len(), 3);
        assert_eq!(dfs[0], (0, 0)); // root at depth 0
        assert_eq!(dfs[1], (1, 1)); // child_a at depth 1
        assert_eq!(dfs[2], (1, 2)); // child_b at depth 1
    }
}