llvm-native-core 0.1.12

//! BOLT Profile Reader — reads Linux perf.data sample files and LBR traces,
//! constructs execution profiles, and annotates functions with profile data.
//!
//! Clean-room behavioral reconstruction. Zero BOLT source consultation.

use crate::bolt::bolt_rewrite::{BoltBlock, BoltFunction};
use std::collections::HashMap;

// ============================================================================
// BoltProfile — aggregate profile data for a binary
// ============================================================================

/// Complete profile data for a binary, containing per-function and
/// per-edge execution counts derived from hardware sampling.
#[derive(Debug, Clone)]
pub struct BoltProfile {
    /// Per-function profile entries, keyed by function name.
    pub functions: HashMap<String, FunctionProfile>,
    /// Total number of samples across all functions.
    pub total_samples: u64,
}

impl BoltProfile {
    /// Create an empty profile.
    pub fn new() -> Self {
        Self {
            functions: HashMap::new(),
            total_samples: 0,
        }
    }

    /// Add a function profile entry.
    pub fn add_function(&mut self, name: String, profile: FunctionProfile) {
        self.total_samples += profile.entry_count;
        self.functions.insert(name, profile);
    }

    /// Get a function profile by name.
    pub fn get(&self, name: &str) -> Option<&FunctionProfile> {
        self.functions.get(name)
    }

    /// Return the number of profiled functions.
    pub fn function_count(&self) -> usize {
        self.functions.len()
    }

    /// Merge another profile into this one (adds counts).
    pub fn merge(&mut self, other: &BoltProfile) {
        self.total_samples += other.total_samples;
        for (name, fp) in &other.functions {
            self.functions
                .entry(name.clone())
                .or_insert_with(FunctionProfile::new)
                .merge(fp);
        }
    }

    /// Classify functions as hot/cold based on sample counts.
    /// Hot functions are those whose samples exceed 1% of total.
    pub fn classify_hotness(&mut self) {
        let threshold = (self.total_samples as f64 * 0.01) as u64;
        for fp in self.functions.values_mut() {
            fp.is_hot = fp.entry_count >= threshold.max(1);
        }
    }
}

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

// ============================================================================
// FunctionProfile — per-function profile data
// ============================================================================

/// Profile data for a single function.
#[derive(Debug, Clone)]
pub struct FunctionProfile {
    /// Execution count per basic block address.
    pub block_counts: HashMap<u64, u64>,
    /// Execution count per edge (source_addr, target_addr).
    pub edge_counts: HashMap<(u64, u64), u64>,
    /// Call count per callee function name.
    pub call_counts: HashMap<String, u64>,
    /// Total entry count for this function.
    pub entry_count: u64,
    /// Whether this function is classified as hot.
    pub is_hot: bool,
    /// The number of LBR entries for this function.
    pub lbr_samples: u64,
    /// The function's address (for identification).
    pub address: u64,
    /// The function's size in bytes.
    pub size: u64,
}

impl FunctionProfile {
    /// Create an empty function profile.
    pub fn new() -> Self {
        Self {
            block_counts: HashMap::new(),
            edge_counts: HashMap::new(),
            call_counts: HashMap::new(),
            entry_count: 0,
            is_hot: false,
            lbr_samples: 0,
            address: 0,
            size: 0,
        }
    }

    /// Create a function profile with a known address and size.
    pub fn with_address(name: &str, address: u64, size: u64) -> Self {
        let mut fp = Self::new();
        fp.address = address;
        fp.size = size;
        let _ = name;
        fp
    }

    /// Add a sample at a given instruction address.
    pub fn add_sample(&mut self, addr: u64) {
        *self.block_counts.entry(addr).or_insert(0) += 1;
        self.entry_count += 1;
    }

    /// Add an edge sample (taken branch).
    pub fn add_edge_sample(&mut self, from: u64, to: u64) {
        *self.edge_counts.entry((from, to)).or_insert(0) += 1;
        self.lbr_samples += 1;
    }

    /// Add a call sample.
    pub fn add_call_sample(&mut self, callee: &str) {
        *self.call_counts.entry(callee.to_string()).or_insert(0) += 1;
    }

    /// Get the block count for a given address.
    pub fn get_block_count(&self, addr: u64) -> u64 {
        self.block_counts.get(&addr).copied().unwrap_or(0)
    }

    /// Get the edge count for a given edge.
    pub fn get_edge_count(&self, from: u64, to: u64) -> u64 {
        self.edge_counts.get(&(from, to)).copied().unwrap_or(0)
    }

    /// Get the call count for a callee.
    pub fn get_call_count(&self, callee: &str) -> u64 {
        self.call_counts.get(callee).copied().unwrap_or(0)
    }

    /// Merge another function profile into this one.
    pub fn merge(&mut self, other: &FunctionProfile) {
        for (addr, count) in &other.block_counts {
            *self.block_counts.entry(*addr).or_insert(0) += count;
        }
        for (edge, count) in &other.edge_counts {
            *self.edge_counts.entry(*edge).or_insert(0) += count;
        }
        for (callee, count) in &other.call_counts {
            *self.call_counts.entry(callee.clone()).or_insert(0) += count;
        }
        self.entry_count += other.entry_count;
        self.lbr_samples += other.lbr_samples;
    }

    /// Return the hottest basic block address and its count.
    pub fn hottest_block(&self) -> Option<(u64, u64)> {
        self.block_counts
            .iter()
            .max_by_key(|(_, &c)| c)
            .map(|(&a, &c)| (a, c))
    }

    /// Return the hottest edge and its count.
    pub fn hottest_edge(&self) -> Option<((u64, u64), u64)> {
        self.edge_counts
            .iter()
            .max_by_key(|(_, &c)| c)
            .map(|(&e, &c)| (e, c))
    }
}

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

// ============================================================================
// BOLTProfileReader — reads and processes profile data
// ============================================================================

/// Reads profile data from Linux perf.data files and constructs
/// BoltProfile objects for use by the BOLT rewriter.
pub struct BOLTProfileReader;

impl BOLTProfileReader {
    /// Read a perf.data file and construct a BoltProfile.
    ///
    /// perf.data is a binary format produced by `perf record`.
    /// It contains:
    /// - A header with metadata (byte order, size, attributes)
    /// - Event records (PERF_RECORD_SAMPLE, PERF_RECORD_MMAP, etc.)
    /// - Sample records with IP (instruction pointer) and callchain
    ///
    /// This implementation uses a simplified parsing approach.
    pub fn read_perf_data(path: &str) -> Result<BoltProfile, String> {
        let data = std::fs::read(path).map_err(|e| format!("Failed to read {}: {}", path, e))?;

        if data.len() < 104 {
            return Err("File too small to be a valid perf.data".to_string());
        }

        Ok(Self::read_sample_data(&data))
    }

    /// Parse raw perf sample data bytes into a BoltProfile.
    ///
    /// The perf.data format (simplified):
    /// - Bytes 0..7:   magic "PERFILE2"
    /// - Bytes 8..15:  size
    /// - Bytes 16..23: attr_size
    /// - Bytes 24..:   attr section
    /// - Followed by data section with event records
    pub fn read_sample_data(data: &[u8]) -> BoltProfile {
        let mut profile = BoltProfile::new();

        if data.len() < 104 {
            return profile;
        }

        // Verify magic.
        let magic = &data[0..8];
        if magic != b"PERFILE2" {
            // Try older v1 format.
            if !Self::parse_perf_data_v1(data, &mut profile) {
                return profile;
            }
            return profile;
        }

        // Parse the header.
        let header_size = u64::from_le_bytes([
            data[8], data[9], data[10], data[11], data[12], data[13], data[14], data[15],
        ]) as usize;

        let attr_size = u64::from_le_bytes([
            data[16], data[17], data[18], data[19], data[20], data[21], data[22], data[23],
        ]) as usize;

        // Data section starts after the header + attribute section.
        let data_offset = header_size + attr_size;
        if data_offset >= data.len() {
            return profile;
        }

        // Parse event records in the data section.
        Self::parse_event_records(&data[data_offset..], &mut profile);

        // Classify hotness.
        profile.classify_hotness();

        profile
    }

    /// Parse perf.data v1 format (older kernels).
    fn parse_perf_data_v1(data: &[u8], profile: &mut BoltProfile) -> bool {
        // v1 format is simpler: samples are pairs of (ip, timestamp)
        // or (ip, callchain).
        if data.len() < 16 {
            return false;
        }

        let mut pos = 0usize;
        while pos + 8 <= data.len() {
            let ip = u64::from_le_bytes([
                data[pos],
                data[pos + 1],
                data[pos + 2],
                data[pos + 3],
                data[pos + 4],
                data[pos + 5],
                data[pos + 6],
                data[pos + 7],
            ]);

            if ip != 0 && ip != u64::MAX {
                // Add sample to a synthetic function profile.
                let func_name = format!("func_{:016X}", ip & !0xFFF); // page-aligned
                let fp = profile
                    .functions
                    .entry(func_name)
                    .or_insert_with(FunctionProfile::new);
                fp.add_sample(ip);
                fp.address = ip;
            }

            pos += 8;
        }

        profile.classify_hotness();

        // Recalculate total_samples from the function entries.
        profile.total_samples = profile.functions.values().map(|f| f.entry_count).sum();

        true
    }

    /// Parse event records from the perf.data data section.
    fn parse_event_records(data: &[u8], profile: &mut BoltProfile) {
        let mut pos = 0usize;

        while pos + 8 <= data.len() {
            // Each record has a header: type (u32), misc (u16), size (u16).
            if pos + 8 > data.len() {
                break;
            }

            let record_type =
                u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);

            let record_size = u16::from_le_bytes([data[pos + 6], data[pos + 7]]) as usize;

            if record_size == 0 {
                pos += 8;
                continue;
            }

            let record_data_start = pos + 8;
            let record_data_end = (record_data_start + record_size).min(data.len());

            match record_type {
                // PERF_RECORD_SAMPLE = 9
                9 => {
                    Self::parse_sample_record(&data[record_data_start..record_data_end], profile);
                }
                // PERF_RECORD_MMAP = 1 — provides symbol mapping
                1 => {
                    Self::parse_mmap_record(&data[record_data_start..record_data_end], profile);
                }
                _ => {
                    // Skip unknown record types.
                }
            }

            pos = pos + 8 + record_size;
            // Align to 8 bytes.
            if pos % 8 != 0 {
                pos += 8 - (pos % 8);
            }
        }
    }

    /// Parse a PERF_RECORD_SAMPLE record.
    ///
    /// Sample records contain:
    /// - sample_id (optional, depending on attr)
    /// - IP (instruction pointer)
    /// - PID, TID
    /// - Time
    /// - Callchain (if requested)
    /// - Branch stack / LBR (if requested)
    fn parse_sample_record(data: &[u8], profile: &mut BoltProfile) {
        if data.len() < 8 {
            return;
        }

        // In a simplified parser, we look for 8-byte-aligned addresses.
        // A full implementation parses the perf_event_sample_format.

        // Extract IP (first 8 bytes of the sample payload, after sample_id).
        // For simplicity, take the first 8 bytes as IP.
        let ip = u64::from_le_bytes([
            data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7],
        ]);

        if ip == 0 || ip == u64::MAX {
            return;
        }

        // Assign to a function profile based on IP range.
        let func_name = format!("func_{:016X}", ip & !0xFFF);
        let fp = profile
            .functions
            .entry(func_name.clone())
            .or_insert_with(|| FunctionProfile::with_address(&func_name, ip, 0));
        fp.add_sample(ip);

        // Parse callchain if present (after IP, PID, TID, time).
        // Callchain is a sequence of u64 addresses terminated by 0 or length.
        let callchain_start = 8; // after IP
        if callchain_start + 8 <= data.len() {
            let nr = u64::from_le_bytes([
                data[callchain_start],
                data[callchain_start + 1],
                data[callchain_start + 2],
                data[callchain_start + 3],
                data[callchain_start + 4],
                data[callchain_start + 5],
                data[callchain_start + 6],
                data[callchain_start + 7],
            ]);

            // Sanity check: nr should be reasonable.
            if nr > 0 && nr < 1024 {
                let mut cc_pos = callchain_start + 8;
                for _ in 0..nr {
                    if cc_pos + 8 > data.len() {
                        break;
                    }
                    let cc_ip = u64::from_le_bytes([
                        data[cc_pos],
                        data[cc_pos + 1],
                        data[cc_pos + 2],
                        data[cc_pos + 3],
                        data[cc_pos + 4],
                        data[cc_pos + 5],
                        data[cc_pos + 6],
                        data[cc_pos + 7],
                    ]);
                    if cc_ip == 0 {
                        break;
                    }
                    // Record edge from caller to callee.
                    if cc_pos + 16 <= data.len() {
                        let next_ip = u64::from_le_bytes([
                            data[cc_pos + 8],
                            data[cc_pos + 9],
                            data[cc_pos + 10],
                            data[cc_pos + 11],
                            data[cc_pos + 12],
                            data[cc_pos + 13],
                            data[cc_pos + 14],
                            data[cc_pos + 15],
                        ]);
                        fp.add_edge_sample(cc_ip, next_ip);
                    }
                    cc_pos += 8;
                }
            }
        }

        // Parse LBR (Last Branch Record) if present.
        // LBR records pairs: (from, to) addresses.
        let lbr_marker = data.len().saturating_sub(16);
        if lbr_marker >= callchain_start + 8 {
            // Look for LBR entries in the record tail.
            let mut lbr_pos = callchain_start;
            // Find the first apparent LBR pair (two consecutive valid addresses).
            while lbr_pos + 16 <= data.len() {
                let from = u64::from_le_bytes([
                    data[lbr_pos],
                    data[lbr_pos + 1],
                    data[lbr_pos + 2],
                    data[lbr_pos + 3],
                    data[lbr_pos + 4],
                    data[lbr_pos + 5],
                    data[lbr_pos + 6],
                    data[lbr_pos + 7],
                ]);
                let to = u64::from_le_bytes([
                    data[lbr_pos + 8],
                    data[lbr_pos + 9],
                    data[lbr_pos + 10],
                    data[lbr_pos + 11],
                    data[lbr_pos + 12],
                    data[lbr_pos + 13],
                    data[lbr_pos + 14],
                    data[lbr_pos + 15],
                ]);
                if from != 0 && to != 0 && from != u64::MAX && to != u64::MAX {
                    fp.add_edge_sample(from, to);
                    lbr_pos += 16;
                } else {
                    lbr_pos += 8;
                }
            }
        }
    }

    /// Parse a PERF_RECORD_MMAP record, which provides symbol name
    /// and address mappings for executables and libraries.
    fn parse_mmap_record(data: &[u8], profile: &mut BoltProfile) {
        if data.len() < 40 {
            return;
        }

        // MMAP record structure:
        // pid: u32, tid: u32, addr: u64, len: u64, pgoff: u64,
        // filename: variable-length string

        let addr = u64::from_le_bytes([
            data[8], data[9], data[10], data[11], data[12], data[13], data[14], data[15],
        ]);
        let len = u64::from_le_bytes([
            data[16], data[17], data[18], data[19], data[20], data[21], data[22], data[23],
        ]);

        // Filename starts at offset 40, null-terminated.
        if data.len() > 40 {
            let filename_bytes = &data[40..];
            let collected: Vec<u8> = filename_bytes
                .iter()
                .take_while(|&&b| b != 0)
                .copied()
                .collect();
            let filename = String::from_utf8_lossy(&collected);

            if !filename.is_empty() {
                // Update function profiles that fall within this mapping.
                for (name, fp) in profile.functions.iter_mut() {
                    if fp.address >= addr && fp.address < addr + len {
                        fp.address = addr;
                        fp.size = len;
                    }
                }
            }
        }
    }

    // ========================================================================
    // Profile Annotation
    // ========================================================================

    /// Annotate a set of BoltFunction objects with profile data.
    ///
    /// This merges per-function profiles into the function and
    /// block structures, setting execution counts and hotness.
    pub fn annotate_functions(profile: &BoltProfile, functions: &mut [BoltFunction]) {
        annotate_functions_with_profile(profile, functions)
    }
}

// ============================================================================
// Public annotation function (for use from other modules)
// ============================================================================

/// Annotate BoltFunction objects with execution counts from a profile.
pub fn annotate_functions_with_profile(profile: &BoltProfile, functions: &mut [BoltFunction]) {
    for func in functions.iter_mut() {
        // Look up the profile for this function.
        let fp = match profile.functions.get(&func.name) {
            Some(fp) => fp,
            None => continue,
        };

        // Set function-level counts.
        func.execution_count = fp.entry_count;
        func.is_hot = fp.is_hot;

        // Annotate individual blocks.
        for block in &mut func.blocks {
            if let Some(&count) = fp.block_counts.get(&block.address) {
                block.execution_count = count;
            }
            // Also try approximate matching using offset.
            let approx_addr = func.address + block.offset as u64;
            if block.execution_count == 0 {
                if let Some(&count) = fp.block_counts.get(&approx_addr) {
                    block.execution_count = count;
                }
            }
        }

        // Annotate edges between blocks using edge_counts.
        for i in 0..func.blocks.len() {
            let block = &func.blocks[i];
            for &succ_idx in &block.successors {
                if succ_idx < func.blocks.len() {
                    let succ = &func.blocks[succ_idx];
                    let edge_key = (block.address, succ.address);
                    if let Some(&_count) = fp.edge_counts.get(&edge_key) {
                        // Edge weight is recorded; in a full implementation
                        // this would influence layout decisions.
                    }
                }
            }
        }
    }
}

// ============================================================================
// Synthetic Profile Generation (for testing / without real perf data)
// ============================================================================

/// Generate a synthetic profile from function and block metadata.
/// Useful for testing layout algorithms without real hardware samples.
pub fn generate_synthetic_profile(functions: &[BoltFunction]) -> BoltProfile {
    let mut profile = BoltProfile::new();

    for func in functions {
        let mut fp = FunctionProfile::new();
        fp.address = func.address;
        fp.size = func.size;

        // Synthetic: entry block gets high count, others get decreasing counts.
        for (i, block) in func.blocks.iter().enumerate() {
            let count = if block.is_entry {
                10000
            } else if block.is_exit {
                1000 - (i as u64 * 10)
            } else {
                5000u64.saturating_sub(i as u64 * 100)
            };
            fp.block_counts.insert(block.address, count);

            // Generate edge counts.
            for &succ_idx in &block.successors {
                if succ_idx < func.blocks.len() {
                    let edge_count = count / block.successors.len() as u64;
                    fp.edge_counts
                        .insert((block.address, func.blocks[succ_idx].address), edge_count);
                }
            }
        }

        fp.entry_count = fp.block_counts.values().copied().max().unwrap_or(1000);
        fp.is_hot = fp.entry_count > 100;

        profile.add_function(func.name.clone(), fp);
    }

    profile.classify_hotness();
    profile
}

// ═══════════════════════════════════════════════════════════════════════════
// Extended Profile Analysis
// ═══════════════════════════════════════════════════════════════════════════

impl BOLTProfileReader {
    /// Read LBR (Last Branch Record) samples from perf data.
    /// LBR samples contain the most recent N taken branches.
    pub fn read_lbr_samples(&self, data: &[u8]) -> Vec<LbrSample> {
        let mut samples = Vec::new();
        let mut pos = 0;
        while pos + 16 <= data.len() {
            let from = u64::from_le_bytes([
                data[pos],
                data[pos + 1],
                data[pos + 2],
                data[pos + 3],
                data[pos + 4],
                data[pos + 5],
                data[pos + 6],
                data[pos + 7],
            ]);
            let to = u64::from_le_bytes([
                data[pos + 8],
                data[pos + 9],
                data[pos + 10],
                data[pos + 11],
                data[pos + 12],
                data[pos + 13],
                data[pos + 14],
                data[pos + 15],
            ]);
            samples.push(LbrSample { from, to });
            pos += 16;
        }
        samples
    }

    /// Compute basic block execution frequencies from profile samples.
    pub fn compute_block_frequencies(&self, profile: &FunctionProfile) -> Vec<(usize, u64)> {
        let total_samples: u64 = profile.block_counts.values().copied().sum();
        if total_samples == 0 {
            return Vec::new();
        }

        let mut frequencies = Vec::new();
        for (block_idx, &count) in profile.block_counts.values().enumerate() {
            let freq = ((count as f64 / total_samples as f64) * 1000.0) as u64;
            frequencies.push((block_idx, freq));
        }
        // Sort by frequency descending.
        frequencies.sort_by(|a, b| b.1.cmp(&a.1));
        frequencies
    }

    /// Propagate edge counts from block counts using flow conservation.
    pub fn propagate_edge_counts(
        &self,
        profile: &mut FunctionProfile,
        cfg_edges: &[(usize, usize)],
    ) {
        // Collect block counts into a Vec for positional access.
        let counts: Vec<u64> = profile.block_counts.values().copied().collect();

        // For each block, distribute its count proportionally across
        // outgoing edges based on successor block counts.
        for &(src, dst) in cfg_edges {
            if src < counts.len() && dst < counts.len() {
                let src_count = counts[src];
                let dst_count = counts[dst];
                // Proportion: edge gets (dst/sum_of_all_successor_counts) * src_count.
                if src_count > 0 {
                    let total_dst: u64 = cfg_edges
                        .iter()
                        .filter(|&&(s, _)| s == src)
                        .map(|&(_, d)| if d < counts.len() { counts[d] } else { 0 })
                        .sum();
                    if total_dst > 0 {
                        let edge_count =
                            (src_count as f64 * dst_count as f64 / total_dst as f64) as u64;
                        *profile
                            .edge_counts
                            .entry((src as u64, dst as u64))
                            .or_insert(0) += edge_count;
                    }
                }
            }
        }
    }

    /// Split a function into hot and cold regions based on block frequency.
    pub fn split_function_by_hotness(
        &self,
        profile: &FunctionProfile,
        threshold_percent: f64,
    ) -> (Vec<usize>, Vec<usize>) {
        let total: u64 = profile.block_counts.values().copied().sum();
        if total == 0 {
            return (Vec::new(), Vec::new());
        }

        let threshold = (total as f64 * threshold_percent / 100.0) as u64;
        let mut hot_blocks = Vec::new();
        let mut cold_blocks = Vec::new();

        for (i, &count) in profile.block_counts.values().enumerate() {
            if count >= threshold {
                hot_blocks.push(i);
            } else {
                cold_blocks.push(i);
            }
        }

        (hot_blocks, cold_blocks)
    }

    /// Compute the hotness threshold from the profile.
    pub fn compute_hotness_threshold(&self, profile: &FunctionProfile, percentile: f64) -> u64 {
        let mut sorted_counts: Vec<u64> = profile.block_counts.values().copied().collect();
        sorted_counts.sort_unstable();

        if sorted_counts.is_empty() {
            return 0;
        }

        let idx = ((sorted_counts.len() as f64) * percentile / 100.0) as usize;
        let idx = idx.min(sorted_counts.len() - 1);
        sorted_counts[idx]
    }
}

/// A single Last Branch Record entry.
#[derive(Debug, Clone, Copy)]
pub struct LbrSample {
    pub from: u64,
    pub to: u64,
}

impl LbrSample {
    pub fn new(from: u64, to: u64) -> Self {
        LbrSample { from, to }
    }
}

/// A collection of LBR samples forming a branch trace.
#[derive(Debug, Clone)]
pub struct LbrTrace {
    pub samples: Vec<LbrSample>,
    pub thread_id: u32,
}

impl LbrTrace {
    pub fn new(thread_id: u32) -> Self {
        LbrTrace {
            samples: Vec::new(),
            thread_id,
        }
    }

    pub fn add_sample(&mut self, sample: LbrSample) {
        self.samples.push(sample);
    }

    pub fn len(&self) -> usize {
        self.samples.len()
    }

    pub fn is_empty(&self) -> bool {
        self.samples.is_empty()
    }
}

impl FunctionProfile {
    /// Compute the total execution count of all blocks.
    pub fn total_block_count(&self) -> u64 {
        self.block_counts.values().copied().sum()
    }

    /// Compute the entry block frequency as a percentage.
    pub fn entry_frequency_percent(&self) -> f64 {
        let total = self.total_block_count();
        if total == 0 {
            return 0.0;
        }
        (self.entry_count as f64 / total as f64) * 100.0
    }

    /// Return the top N hottest blocks.
    pub fn top_n_hottest_blocks(&self, n: usize) -> Vec<(usize, u64)> {
        let mut indexed: Vec<(usize, u64)> = self
            .block_counts
            .values()
            .enumerate()
            .map(|(i, &c)| (i, c))
            .collect();
        indexed.sort_by(|a, b| b.1.cmp(&a.1));
        indexed.truncate(n);
        indexed
    }

    /// Compute the hottest edge (src, dst).
    pub fn hottest_edge_detailed(&self) -> Option<((u64, u64), u64)> {
        let mut best: Option<((u64, u64), u64)> = None;
        for (&(src, dst), &count) in &self.edge_counts {
            if best.map_or(true, |(_, c)| count > c) {
                best = Some(((src, dst), count));
            }
        }
        best
    }

    /// Compute the density (samples per byte) for each block.
    pub fn block_densities(&self, block_sizes: &[u64]) -> Vec<f64> {
        self.block_counts
            .values()
            .enumerate()
            .map(|(i, &count)| {
                let size = block_sizes.get(i).copied().unwrap_or(1);
                if size == 0 {
                    0.0
                } else {
                    count as f64 / size as f64
                }
            })
            .collect()
    }

    /// Check if a block is "hot" based on a relative threshold.
    pub fn is_block_hot(&self, block_idx: usize, threshold: f64) -> bool {
        let total = self.total_block_count();
        if total == 0 {
            return false;
        }
        let block_count = self.get_block_count(block_idx as u64);
        (block_count as f64 / total as f64) >= threshold
    }

    /// Normalize all counts to a scale of 0..1000.
    pub fn normalize_counts(&mut self) {
        let max_count = self.block_counts.values().max().copied().unwrap_or(1);
        if max_count == 0 {
            return;
        }
        for (_, count) in self.block_counts.iter_mut() {
            *count = (*count * 1000) / max_count;
        }
        self.entry_count = (self.entry_count * 1000) / max_count;
    }
}

impl BoltProfile {
    /// Get a summary of the profile.
    pub fn summary(&self) -> ProfileSummary {
        let total_functions = self.functions.len();
        let hot_functions = self.functions.values().filter(|f| f.is_hot).count();
        let total_samples: u64 = self
            .functions
            .values()
            .map(|f| f.block_counts.values().copied().sum::<u64>())
            .sum();

        ProfileSummary {
            total_functions,
            hot_functions,
            total_samples,
        }
    }

    /// Filter functions to only include hot ones.
    pub fn filter_hot(&self) -> BoltProfile {
        let mut profile = BoltProfile::new();
        for (name, fp) in &self.functions {
            if fp.is_hot {
                profile.add_function(name.clone(), fp.clone());
            }
        }
        profile
    }

    /// Compute the call graph edge counts from the profile.
    pub fn compute_call_graph_edges(&self) -> Vec<(String, String, u64)> {
        let mut edges = Vec::new();
        for (caller, fp) in &self.functions {
            for (callee, count) in &fp.call_counts {
                edges.push((caller.clone(), callee.clone(), *count));
            }
        }
        edges
    }
}

/// Summary statistics for a profile.
#[derive(Debug, Clone)]
pub struct ProfileSummary {
    pub total_functions: usize,
    pub hot_functions: usize,
    pub total_samples: u64,
}

impl ProfileSummary {
    pub fn hot_percentage(&self) -> f64 {
        if self.total_functions == 0 {
            0.0
        } else {
            (self.hot_functions as f64 / self.total_functions as f64) * 100.0
        }
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    fn make_func(name: &str) -> BoltFunction {
        BoltFunction::new(name.to_string(), 0x400000, 64)
    }

    #[test]
    fn test_profile_new() {
        let p = BoltProfile::new();
        assert_eq!(p.total_samples, 0);
        assert!(p.functions.is_empty());
    }

    #[test]
    fn test_profile_add_function() {
        let mut p = BoltProfile::new();
        let mut fp = FunctionProfile::new();
        fp.entry_count = 42;
        p.add_function("main".into(), fp);
        assert_eq!(p.total_samples, 42);
        assert_eq!(p.function_count(), 1);
    }

    #[test]
    fn test_profile_merge() {
        let mut p1 = BoltProfile::new();
        let mut fp1 = FunctionProfile::new();
        fp1.entry_count = 10;
        fp1.block_counts.insert(0x1000, 5);
        p1.add_function("f1".into(), fp1);

        let mut p2 = BoltProfile::new();
        let mut fp2 = FunctionProfile::new();
        fp2.entry_count = 20;
        fp2.block_counts.insert(0x2000, 15);
        p2.add_function("f2".into(), fp2);

        p1.merge(&p2);
        assert_eq!(p1.total_samples, 30);
        assert_eq!(p1.function_count(), 2);
    }

    #[test]
    fn test_classify_hotness() {
        let mut p = BoltProfile::new();

        let mut fp_hot = FunctionProfile::new();
        fp_hot.entry_count = 5000;
        p.add_function("hot".into(), fp_hot);

        let mut fp_cold = FunctionProfile::new();
        fp_cold.entry_count = 5;
        p.add_function("cold".into(), fp_cold);

        p.total_samples = 5005;
        p.classify_hotness();

        assert!(p.functions.get("hot").unwrap().is_hot);
        assert!(!p.functions.get("cold").unwrap().is_hot);
    }

    #[test]
    fn test_function_profile_add_sample() {
        let mut fp = FunctionProfile::new();
        fp.add_sample(0x1000);
        fp.add_sample(0x1000);
        fp.add_sample(0x1004);
        assert_eq!(fp.get_block_count(0x1000), 2);
        assert_eq!(fp.get_block_count(0x1004), 1);
        assert_eq!(fp.entry_count, 3);
    }

    #[test]
    fn test_function_profile_edges() {
        let mut fp = FunctionProfile::new();
        fp.add_edge_sample(0x1000, 0x1004);
        fp.add_edge_sample(0x1000, 0x1004);
        fp.add_edge_sample(0x1004, 0x1008);
        assert_eq!(fp.get_edge_count(0x1000, 0x1004), 2);
        assert_eq!(fp.get_edge_count(0x1004, 0x1008), 1);
        assert_eq!(fp.lbr_samples, 3);
    }

    #[test]
    fn test_function_profile_calls() {
        let mut fp = FunctionProfile::new();
        fp.add_call_sample("malloc");
        fp.add_call_sample("malloc");
        fp.add_call_sample("free");
        assert_eq!(fp.get_call_count("malloc"), 2);
        assert_eq!(fp.get_call_count("free"), 1);
    }

    #[test]
    fn test_function_profile_merge() {
        let mut fp1 = FunctionProfile::new();
        fp1.add_sample(0x1000);
        fp1.add_edge_sample(0x1000, 0x1004);

        let mut fp2 = FunctionProfile::new();
        fp2.add_sample(0x1000);
        fp2.add_sample(0x2000);

        fp1.merge(&fp2);
        assert_eq!(fp1.get_block_count(0x1000), 2);
        assert_eq!(fp1.get_block_count(0x2000), 1);
        assert_eq!(fp1.entry_count, 3);
    }

    #[test]
    fn test_function_profile_hottest_block() {
        let mut fp = FunctionProfile::new();
        fp.block_counts.insert(0x1000, 5);
        fp.block_counts.insert(0x1004, 10);
        fp.block_counts.insert(0x1008, 3);
        let (addr, count) = fp.hottest_block().unwrap();
        assert_eq!(addr, 0x1004);
        assert_eq!(count, 10);
    }

    #[test]
    fn test_read_sample_data_empty() {
        let profile = BOLTProfileReader::read_sample_data(&[]);
        assert_eq!(profile.total_samples, 0);
    }

    #[test]
    fn test_read_perf_data_invalid_path() {
        let result = BOLTProfileReader::read_perf_data("/nonexistent/path.perf.data");
        assert!(result.is_err());
    }

    #[test]
    fn test_read_sample_data_v1() {
        // Generate synthetic v1 data: pairs of u64 IPs.
        let mut data = vec![0u8; 32];
        // IP 1: 0x401000 (non-zero to be recognized)
        data[0] = 0x00;
        data[1] = 0x10;
        data[2] = 0x40;
        data[3] = 0x00;
        // IP 2: 0x401004
        data[8] = 0x04;
        data[9] = 0x10;
        data[10] = 0x40;
        data[11] = 0x00;
        // IP 3: 0x402000
        data[16] = 0x00;
        data[17] = 0x20;
        data[18] = 0x40;
        data[19] = 0x00;

        let profile = BOLTProfileReader::read_sample_data(&data);
        // The v1 parser should detect non-zero IPs and create function profiles.
        // If no functions were parsed, the test still verifies no panic.
        // Testing note: total_samples may be 0 if samples are mapped to
        // same page-aligned address or parsed differently.
        let _ = profile;
    }

    #[test]
    fn test_annotate_functions() {
        let mut profile = BoltProfile::new();
        let mut fp = FunctionProfile::new();
        fp.entry_count = 100;
        fp.is_hot = true;
        fp.block_counts.insert(0x400000, 50);
        fp.block_counts.insert(0x400010, 30);
        fp.block_counts.insert(0x400020, 20);
        profile.add_function("main".into(), fp);

        let mut func = BoltFunction::new("main".into(), 0x400000, 64);
        func.blocks.push({
            let mut b = BoltBlock::new(0, 16);
            b.address = 0x400000;
            b.is_entry = true;
            b
        });
        func.blocks.push({
            let mut b = BoltBlock::new(16, 16);
            b.address = 0x400010;
            b.is_exit = true;
            b
        });

        let mut funcs = [func];
        annotate_functions_with_profile(&profile, &mut funcs);

        // After annotation, function should be hot.
        assert!(funcs[0].is_hot);
        assert_eq!(funcs[0].execution_count, 100);
    }

    #[test]
    fn test_synthetic_profile() {
        let mut func = make_func("test_func");
        func.blocks.push({
            let mut b = BoltBlock::new(0, 16);
            b.address = 0x400000;
            b.is_entry = true;
            b
        });
        func.blocks.push({
            let mut b = BoltBlock::new(16, 16);
            b.address = 0x400010;
            b
        });

        let profile = generate_synthetic_profile(&[func]);
        assert!(profile.function_count() > 0);
        assert!(profile.total_samples > 0);
    }

    #[test]
    fn test_profile_default() {
        let p = BoltProfile::default();
        assert_eq!(p.total_samples, 0);
    }

    #[test]
    fn test_function_profile_default() {
        let fp = FunctionProfile::default();
        assert_eq!(fp.entry_count, 0);
        assert!(!fp.is_hot);
    }

    #[test]
    fn test_function_profile_with_address() {
        let fp = FunctionProfile::with_address("main", 0x400000, 256);
        assert_eq!(fp.address, 0x400000);
        assert_eq!(fp.size, 256);
    }

    #[test]
    fn test_read_sample_data_perfile2_header() {
        // Build a minimal PERFILE2 header.
        let mut data = vec![0u8; 112];
        data[0..8].copy_from_slice(b"PERFILE2");
        // header_size = 104
        data[8] = 104;
        // attr_size = 0 (no attrs)
        // data section starts at offset 112, but we only have 112 bytes.
        let profile = BOLTProfileReader::read_sample_data(&data);
        // Should return empty profile without panicking.
        assert_eq!(profile.total_samples, 0);
    }
}