ferric_ai/parser/

tree.rs

// Tree structures for flamegraph analysis

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use uuid::Uuid;

use super::nodes::{ParserNode, RgbColor};

/// A hierarchical tree structure containing TreeNodes with UUID-based relationships
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Tree {
    /// All nodes in the tree indexed by their UUID
    pub nodes: HashMap<Uuid, TreeNode>,
    /// Root node UUIDs (nodes with no parent)
    pub roots: Vec<Uuid>,
    /// Total samples from the flamegraph SVG for percentage calculations
    pub total_samples: u64,
}

/// A node in the call tree with UUID-based parent/child relationships
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TreeNode {
    /// Unique identifier for this node
    pub id: Uuid,
    /// Parent node UUID (None for root nodes)
    pub parent: Option<Uuid>,
    /// Child node UUIDs (direct children only)
    pub children: Vec<Uuid>,
    /// Descendant node UUIDs (all levels below)
    pub descendants: Vec<Uuid>,

    // === Core Function Identity ===
    /// Function name (cleaned and demangled)
    pub function_name: String,
    /// Raw function signature from flamegraph title
    pub raw_signature: String,

    // === Performance Metrics ===
    /// CPU percentage this function consumes (including children)
    pub cpu_percent: f64,
    /// Total sample count for this function (including children)
    pub sample_count: u64,
    /// Total sample count for this function (excluding children)
    pub self_sample_count: u64,

    // === Visual/Color Information ===
    /// RGB color from flamegraph (indicates heat level)
    pub color: Option<RgbColor>,

    // === Call Function Context ===
    pub source_library: Option<String>,
    pub is_alloc: bool,
}

impl Tree {
    /// Create a new Tree from a vector of ParserNodes
    pub fn from_parser_nodes(parser_nodes: Vec<ParserNode>, _total_samples: u64) -> Self {
        let mut tree = Tree {
            nodes: HashMap::new(),
            roots: Vec::new(),
            total_samples: 0, // Will be calculated from root nodes
        };

        // First pass: Create all TreeNodes with UUIDs
        let mut node_map: HashMap<Uuid,String> = HashMap::new();

        for parser_node in &parser_nodes {
            let id = Uuid::new_v4();
            let tree_node = TreeNode {
                id,
                parent: None, // Will be set in second pass
                children: Vec::new(), // Will be populated in second pass
                descendants: Vec::new(), // Will be populated in third pass
                function_name: parser_node.function_name.clone(),
                raw_signature: parser_node.raw_signature.clone(),
                cpu_percent: parser_node.cpu_percent,
                sample_count: parser_node.sample_count,
                self_sample_count: 0, // PLACEHOLDER
                color: parser_node.color.clone(),
                source_library: None, // PLACEHOLDER
                is_alloc: false, // PLACEHOLDER
            };

            node_map.insert(id, parser_node.unique_key.clone());
            tree.nodes.insert(id, tree_node);
        }

        // Second pass: Build parent-child relationships using positioning
        tree.build_relationships(parser_nodes, &node_map);

        // Third pass: Calculate self_sample_count for all nodes
        tree.calculate_self_sample_counts();

        // Fourth pass: Calculate total_samples from root nodes
        tree.calculate_total_samples();

        tree
    }

    /// Build parent-child relationships based on flamegraph positioning
    fn build_relationships(&mut self, parser_nodes: Vec<ParserNode>, node_map: &HashMap<Uuid, String>) {
        // Create reverse lookup map from unique_key to UUID
        let key_to_uuid: HashMap<&String, Uuid> = node_map.iter().map(|(uuid, key)| (key, *uuid)).collect();

        // Group nodes by Y level for efficient processing
        let mut nodes_by_y: HashMap<u32, Vec<&ParserNode>> = HashMap::new();
        for node in &parser_nodes {
            let y_level = (node.y_position as u32 / 16) * 16; // Round to nearest 16px level
            nodes_by_y.entry(y_level).or_default().push(node);
        }

        let mut y_levels: Vec<u32> = nodes_by_y.keys().copied().collect();
        y_levels.sort();

        log::debug!("Found {} Y levels: {:?}", y_levels.len(), y_levels);

        // Track which nodes have been assigned parents
        let mut assigned_as_child = std::collections::HashSet::new();

        // Process each level, finding immediate parent relationships
        for (level_idx, &current_y) in y_levels.iter().enumerate() {
            let empty_vec = vec![];
            let current_nodes = nodes_by_y.get(&current_y).unwrap_or(&empty_vec);

            // Look for immediate parent level (next Y level down)
            if let Some(&parent_y) = y_levels.get(level_idx + 1) {
                let empty_parent_vec = vec![];
                let parent_nodes = nodes_by_y.get(&parent_y).unwrap_or(&empty_parent_vec);

                // For each node at current level, find its immediate parent
                for child_node in current_nodes {
                    let child_id = match key_to_uuid.get(&child_node.unique_key) {
                        Some(&id) => id,
                        None => continue,
                    };

                    // Skip if already has a parent
                    if assigned_as_child.contains(&child_id) {
                        continue;
                    }

                    // Find parent that contains this child's X range
                    for parent_node in parent_nodes {
                        if self.x_ranges_overlap(parent_node, child_node) {
                            let parent_id = match key_to_uuid.get(&parent_node.unique_key) {
                                Some(&id) => id,
                                None => continue,
                            };

                            // Set parent-child relationship
                            if let Some(child_tree_node) = self.nodes.get_mut(&child_id) {
                                child_tree_node.parent = Some(parent_id);
                                assigned_as_child.insert(child_id);
                            }

                            // Add child to parent's children list
                            if let Some(parent_tree_node) = self.nodes.get_mut(&parent_id) {
                                parent_tree_node.children.push(child_id);
                            }

                            break; // Found immediate parent, move to next child
                        }
                    }
                }
            }
        }

        // Identify root nodes (those without parents)
        self.roots.clear();
        for (&node_id, node) in &self.nodes {
            if !assigned_as_child.contains(&node_id) {
                self.roots.push(node_id);
                log::debug!("Root: {}", node.function_name);
            }
        }

        // Build descendants lists (all children recursively)
        self.build_descendants();

        log::debug!("Relationships complete: {} roots, {} nodes with children",
            self.roots.len(),
            self.nodes.values().filter(|n| !n.children.is_empty()).count());
    }

    /// Recursively build descendants list for all nodes
    fn build_descendants(&mut self) {
        // Create a temporary map to avoid borrow checker issues
        let mut descendants_map: HashMap<Uuid, Vec<Uuid>> = HashMap::new();

        for &node_id in self.nodes.keys() {
            let descendants = self.collect_all_descendants(node_id);
            descendants_map.insert(node_id, descendants);
        }

        // Update the nodes with their descendants
        for (node_id, descendants) in descendants_map {
            if let Some(node) = self.nodes.get_mut(&node_id) {
                node.descendants = descendants;
            }
        }
    }

    /// Recursively collect all descendants of a node
    fn collect_all_descendants(&self, node_id: Uuid) -> Vec<Uuid> {
        let mut descendants = Vec::new();

        if let Some(node) = self.nodes.get(&node_id) {
            for &child_id in &node.children {
                descendants.push(child_id);
                // Recursively add grandchildren, great-grandchildren, etc.
                descendants.extend(self.collect_all_descendants(child_id));
            }
        }

        descendants
    }

    /// Calculate self_sample_count for all nodes in the tree
    /// self_sample_count = inclusive sample_count - sum of children's inclusive sample_counts
    pub fn calculate_self_sample_counts(&mut self) {
        // Process all root nodes recursively
        for &root_id in self.roots.clone().iter() {
            self.calculate_self_sample_counts_recursive(root_id);
        }
    }

    /// Calculate total_samples from the sum of all root nodes' sample_count
    /// This reflects the actual total samples in the flamegraph more accurately
    /// than relying on the SVG total_samples attribute
    pub fn calculate_total_samples(&mut self) {
        self.total_samples = self.roots
            .iter()
            .filter_map(|&root_id| self.nodes.get(&root_id))
            .map(|node| node.sample_count)
            .sum();

        log::debug!("Calculated total_samples from {} root nodes: {}",
                   self.roots.len(), self.total_samples);
    }

    /// Recursively calculate self_sample_count starting from a node
    fn calculate_self_sample_counts_recursive(&mut self, node_id: Uuid) {
        // First, recursively calculate for all children (bottom-up approach)
        let child_ids: Vec<Uuid> = if let Some(node) = self.nodes.get(&node_id) {
            node.children.clone()
        } else {
            return;
        };

        for child_id in child_ids {
            self.calculate_self_sample_counts_recursive(child_id);
        }

        // Now calculate self_sample_count for this node
        if let Some(node) = self.nodes.get(&node_id) {
            let inclusive_count = node.sample_count;

            // Sum up all children's inclusive sample counts
            let children_inclusive_sum: u64 = node.children
                .iter()
                .filter_map(|&child_id| self.nodes.get(&child_id))
                .map(|child| child.sample_count)
                .sum();

            let self_count = inclusive_count.saturating_sub(children_inclusive_sum);

            // Update the node's self_sample_count
            if let Some(node_mut) = self.nodes.get_mut(&node_id) {
                node_mut.self_sample_count = self_count;
            }
        }
    }

    /// Check if two nodes have overlapping X coordinate ranges
    fn x_ranges_overlap(&self, parent: &ParserNode, child: &ParserNode) -> bool {
        let parent_start = parent.x_position;
        let parent_end = parent.x_position + parent.width;
        let child_start = child.x_position;
        let child_end = child.x_position + child.width;

        // Check for overlap: child must be within parent's X range
        child_start >= parent_start && child_end <= parent_end
    }


    /// Get a node by UUID
    pub fn get_node(&self, id: Uuid) -> Option<&TreeNode> {
        self.nodes.get(&id)
    }

    /// Get a mutable node by UUID
    pub fn get_node_mut(&mut self, id: Uuid) -> Option<&mut TreeNode> {
        self.nodes.get_mut(&id)
    }

    /// Get all root nodes
    pub fn get_roots(&self) -> Vec<&TreeNode> {
        self.roots
            .iter()
            .filter_map(|&id| self.nodes.get(&id))
            .collect()
    }

    /// Get children of a node
    pub fn get_children(&self, node_id: Uuid) -> Vec<&TreeNode> {
        if let Some(node) = self.nodes.get(&node_id) {
            node.children
                .iter()
                .filter_map(|&child_id| self.nodes.get(&child_id))
                .collect()
        } else {
            Vec::new()
        }
    }

    /// Get parent of a node
    pub fn get_parent(&self, node_id: Uuid) -> Option<&TreeNode> {
        if let Some(node) = self.nodes.get(&node_id) {
            node.parent.and_then(|parent_id| self.nodes.get(&parent_id))
        } else {
            None
        }
    }

    /// Get total number of nodes in the tree
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }
}

impl TreeNode {
    /// Check if this node is a root (has no parent)
    pub fn is_root(&self) -> bool {
        self.parent.is_none()
    }

    /// Check if this node is a leaf (has no children)
    pub fn is_leaf(&self) -> bool {
        self.children.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::nodes::FlamegraphParser;

    // #[test]
    // fn test_tree_from_parser_nodes() {
    //     // Create some test ParserNodes
    //     let parser_nodes = vec![
    //         ParserNode {
    //             todo!()
    //         },
    //         ParserNode {
    //             todo!()
    //         },
    //     ];

    //     let tree = Tree::from_parser_nodes(parser_nodes);

    //     // Verify we have 2 nodes
    //     assert_eq!(tree.node_count(), 2);

    //     // Find the main and child nodes
    //     let main_node = tree
    //         .nodes
    //         .values()
    //         .find(|node| node.function_name == "main")
    //         .expect("Should find main node");

    //     let child_node = tree
    //         .nodes
    //         .values()
    //         .find(|node| node.function_name == "child")
    //         .expect("Should find child node");

    //     // Verify main is root and child has main as parent
    //     assert!(main_node.is_root());
    //     assert_eq!(child_node.parent, Some(main_node.id));
    //     assert!(main_node.children.contains(&child_node.id));

    //     println!("Tree created successfully with {} nodes", tree.node_count());
    //     println!("Main node ID: {}", main_node.id);
    //     println!("Child node ID: {}", child_node.id);
    //     println!("Child's parent: {:?}", child_node.parent);
    // }

    #[test]
    fn test_tree_with_real_flamegraph() {
        crate::init_test_logging();
        let svg_content = include_str!("test_flamegraph.svg");
        let mut parser = FlamegraphParser::default();
        let parser_nodes = parser.parse_svg(svg_content).expect("Should parse SVG");

        let tree = Tree::from_parser_nodes(parser_nodes, 0);

        log::debug!("Created tree with {} nodes from real flamegraph", tree.node_count());
        log::debug!("Number of root nodes: {}", tree.roots.len());

        log::debug!("Root node IDs: {:?}", tree.roots);

        // Test the Display implementation
        log::debug!("{}", tree);

        // Verify we have nodes
        assert!(tree.node_count() > 0);
        // assert!(!tree.roots.is_empty());

        // Check that some nodes have proper relationships
        let nodes_with_children = tree
            .nodes
            .values()
            .filter(|node| !node.children.is_empty())
            .count();

        log::debug!("Nodes with children: {}", nodes_with_children);

        // Should have some parent-child relationships
        assert!(nodes_with_children > 0);
    }
}