profile_inspect/analysis/

caller_callee.rs

use std::collections::HashMap;

use crate::ir::{FrameId, ProfileIR};

/// Statistics about a caller of a function
#[derive(Debug, Clone)]
pub struct CallerStats {
    /// Frame ID of the caller
    pub frame_id: FrameId,
    /// Caller function name
    pub name: String,
    /// Caller source location
    pub location: String,
    /// Number of times this caller called the target function
    pub call_count: u32,
    /// Time attributed to calls from this caller (time in target when called from this caller)
    pub time: u64,
}

/// Statistics about a callee of a function
#[derive(Debug, Clone)]
pub struct CalleeStats {
    /// Frame ID of the callee
    pub frame_id: FrameId,
    /// Callee function name
    pub name: String,
    /// Callee source location
    pub location: String,
    /// Number of times the target function called this callee
    pub call_count: u32,
    /// Self time of this callee when called by target
    pub self_time: u64,
    /// Total inclusive time of this callee subtree when called by target
    pub total_time: u64,
}

/// Caller/callee analysis for a specific function
#[derive(Debug)]
pub struct CallerCalleeAnalysis {
    /// The target function being analyzed
    pub target_frame_id: FrameId,
    /// Target function name
    pub target_name: String,
    /// Target function location
    pub target_location: String,
    /// Self time of the target function
    pub target_self_time: u64,
    /// Total time of the target function
    pub target_total_time: u64,
    /// Functions that call this function
    pub callers: Vec<CallerStats>,
    /// Functions that this function calls
    pub callees: Vec<CalleeStats>,
}

/// Analyzer for caller/callee relationships
pub struct CallerCalleeAnalyzer {
    /// Maximum number of callers/callees to return
    top_n: usize,
}

impl CallerCalleeAnalyzer {
    /// Create a new analyzer
    pub fn new() -> Self {
        Self { top_n: 20 }
    }

    /// Set maximum number of results
    pub fn top_n(mut self, n: usize) -> Self {
        self.top_n = n;
        self
    }

    /// Analyze callers and callees for a specific frame
    pub fn analyze(
        &self,
        profile: &ProfileIR,
        target_frame_id: FrameId,
    ) -> Option<CallerCalleeAnalysis> {
        let target_frame = profile.get_frame(target_frame_id)?;

        let mut caller_times: HashMap<FrameId, u64> = HashMap::new();
        let mut caller_counts: HashMap<FrameId, u32> = HashMap::new();
        let mut callee_self_times: HashMap<FrameId, u64> = HashMap::new();
        let mut callee_total_times: HashMap<FrameId, u64> = HashMap::new();
        let mut callee_counts: HashMap<FrameId, u32> = HashMap::new();
        let mut target_self_time: u64 = 0;
        let mut target_total_time: u64 = 0;

        // Analyze caller/callee relationships
        for sample in &profile.samples {
            let weight = sample.weight;

            if let Some(stack) = profile.get_stack(sample.stack_id) {
                // Find target in the stack
                let target_idx = stack.frames.iter().position(|&f| f == target_frame_id);

                if let Some(idx) = target_idx {
                    target_total_time += weight;

                    // Check if target is the leaf (self time)
                    if idx == stack.frames.len() - 1 {
                        target_self_time += weight;
                    }

                    // Record caller (frame before target in stack)
                    if idx > 0 {
                        let caller_id = stack.frames[idx - 1];
                        *caller_times.entry(caller_id).or_default() += weight;
                        *caller_counts.entry(caller_id).or_default() += 1;
                    }

                    // Record callees (frames after target in stack)
                    if idx < stack.frames.len() - 1 {
                        // Direct callee
                        let callee_id = stack.frames[idx + 1];
                        *callee_counts.entry(callee_id).or_default() += 1;
                        // Total time for this callee subtree (sample weight)
                        *callee_total_times.entry(callee_id).or_default() += weight;

                        // Self time: only when callee is the leaf
                        if let Some(&leaf_id) = stack.frames.last() {
                            if leaf_id == callee_id {
                                *callee_self_times.entry(callee_id).or_default() += weight;
                            }
                        }
                    }
                }
            }
        }

        // Build caller stats
        let mut callers: Vec<CallerStats> = caller_times
            .iter()
            .filter_map(|(&fid, &time)| {
                let frame = profile.get_frame(fid)?;
                Some(CallerStats {
                    frame_id: fid,
                    name: frame.display_name().to_string(),
                    location: frame.location(),
                    call_count: caller_counts.get(&fid).copied().unwrap_or(0),
                    time,
                })
            })
            .collect();

        callers.sort_by(|a, b| b.time.cmp(&a.time).then_with(|| a.name.cmp(&b.name)));
        callers.truncate(self.top_n);

        // Build callee stats
        let mut callees: Vec<CalleeStats> = callee_total_times
            .iter()
            .filter_map(|(&fid, &total_time)| {
                let frame = profile.get_frame(fid)?;
                Some(CalleeStats {
                    frame_id: fid,
                    name: frame.display_name().to_string(),
                    location: frame.location(),
                    call_count: callee_counts.get(&fid).copied().unwrap_or(0),
                    self_time: callee_self_times.get(&fid).copied().unwrap_or(0),
                    total_time,
                })
            })
            .collect();

        callees.sort_by(|a, b| {
            b.total_time
                .cmp(&a.total_time)
                .then_with(|| a.name.cmp(&b.name))
        });
        callees.truncate(self.top_n);

        Some(CallerCalleeAnalysis {
            target_frame_id,
            target_name: target_frame.display_name().to_string(),
            target_location: target_frame.location(),
            target_self_time,
            target_total_time,
            callers,
            callees,
        })
    }

    /// Find a frame by function name (partial match)
    pub fn find_frame_by_name<'a>(
        profile: &'a ProfileIR,
        name: &str,
    ) -> Option<&'a crate::ir::Frame> {
        profile
            .frames
            .iter()
            .find(|f| f.name.contains(name) || f.display_name().contains(name))
    }
}

impl Default for CallerCalleeAnalyzer {
    fn default() -> Self {
        Self::new()
    }
}