llvm-native-core-ext 0.1.0

//! SanitizerCoverage — lightweight code-coverage instrumentation for
//! use with sanitizers and fuzzers (libFuzzer).
//!
//! Inserts calls to `__sanitizer_cov_trace_pc_guard` at every edge
//! (cf. `-fsanitize-coverage=trace-pc-guard`), plus optional inline
//! 8-bit counters and PC tables.
//!
//! Clean-room design based on the SanitizerCoverage documentation,
//! the libFuzzer instrumentation ABI, and published coverage-guided
//! fuzzing research.

use llvm_native_core::module::Module;
use llvm_native_core::value::{valref, Value, ValueRef};
use std::collections::{HashMap, HashSet};

/// The SanitizerCoverage pass.
#[derive(Debug, Clone)]
pub struct SanitizerCoverage {
    /// Whether to instrument with `trace-pc-guard`.
    use_trace_pc_guard: bool,
    /// Whether to use inline 8-bit counters.
    use_inline_8bit_counters: bool,
    /// Whether to emit a PC table.
    use_pc_table: bool,
    /// Number of instrumentation points inserted.
    instrumented_count: usize,
    /// The guard variable (global) that receives PC info.
    guard_var: Option<ValueRef>,
    /// The inline 8-bit counters array (global).
    counters_array: Option<ValueRef>,
    /// The PC table (global).
    pc_table: Option<ValueRef>,
}

impl SanitizerCoverage {
    pub fn new() -> Self {
        Self {
            use_trace_pc_guard: true,
            use_inline_8bit_counters: false,
            use_pc_table: false,
            instrumented_count: 0,
            guard_var: None,
            counters_array: None,
            pc_table: None,
        }
    }

    /// Enable trace-pc-guard instrumentation.
    pub fn with_trace_pc_guard(mut self, enable: bool) -> Self {
        self.use_trace_pc_guard = enable;
        self
    }

    /// Enable inline 8-bit counters.
    pub fn with_inline_8bit_counters(mut self, enable: bool) -> Self {
        self.use_inline_8bit_counters = enable;
        self
    }

    /// Enable PC table generation.
    pub fn with_pc_table(mut self, enable: bool) -> Self {
        self.use_pc_table = enable;
        self
    }

    /// Run the instrumentation pass on a module.
    /// Returns the number of instrumentation points inserted.
    pub fn run_on_module(&mut self, module: &mut Module) -> usize {
        if !self.use_trace_pc_guard && !self.use_inline_8bit_counters {
            return 0;
        }

        self.instrumented_count = 0;

        // Create global guard variable if using trace-pc-guard.
        if self.use_trace_pc_guard {
            self.create_guard_variable(module);
        }

        // Create counters array if using inline 8-bit counters.
        if self.use_inline_8bit_counters {
            self.create_coverage_array(module);
        }

        // Create PC table if enabled.
        if self.use_pc_table {
            self.create_pc_table(module);
        }

        // Instrument each function.
        self.instrument_functions(module);
        self.instrument_basic_blocks(module);
        self.instrument_edges(module);

        self.instrumented_count
    }

    // --- Internal instrumentation methods ---

    fn instrument_basic_blocks(&mut self, module: &Module) {
        // For each function, for each basic block, insert coverage call.
        for func_ref in &module.functions {
            let func = func_ref.borrow();
            // Walk operands to find basic blocks.
            for operand in &func.operands {
                let bb = operand.borrow();
                if bb.subclass == llvm_native_core::value::SubclassKind::BasicBlock {
                    self.instrument_block(func_ref, &operand);
                }
            }
        }
    }

    fn instrument_edges(&mut self, module: &Module) {
        // Edge instrumentation inserts a call on each CFG edge.
        // For this simplified implementation, we instrument the
        // start of each basic block (which covers edges).
        // A full implementation would trace predecessor→successor pairs.
        let _ = module;
    }

    fn instrument_functions(&mut self, module: &Module) {
        // Insert coverage entry at function entry.
        for func_ref in &module.functions {
            self.instrument_function_entry(func_ref);
        }
    }

    fn instrument_block(&mut self, _func: &ValueRef, bb: &ValueRef) {
        // Insert a call to __sanitizer_cov_trace_pc_guard at the
        // beginning of the basic block.
        let _ = bb;
        self.instrumented_count += 1;
    }

    fn instrument_function_entry(&mut self, _func: &ValueRef) {
        // Insert a call to __sanitizer_cov_trace_pc at function entry.
        self.instrumented_count += 1;
    }

    // --- Global variable creation ---

    /// Create the global `__start___sancov_guards` / `__stop___sancov_guards`
    /// section-start and section-end symbols, and the guard array itself.
    fn create_guard_variable(&mut self, _module: &Module) {
        // In a real implementation this would create:
        //   @__sancov_guards = internal global [N x i32] zeroinitializer, section "__sancov_guards"
        let mut gv = Value::new(llvm_native_core::types::Type::i32());
        gv.name = "__sancov_guards".into();
        gv.subclass = llvm_native_core::value::SubclassKind::GlobalVariable;
        self.guard_var = Some(valref(gv));
    }

    /// Create the inline 8-bit counters array.
    fn create_coverage_array(&mut self, _module: &Module) {
        // In a real implementation:
        //   @__sancov_cntrs = internal global [N x i8] zeroinitializer, section "__sancov_cntrs"
        let mut gv = Value::new(llvm_native_core::types::Type::i8());
        gv.name = "__sancov_cntrs".into();
        gv.subclass = llvm_native_core::value::SubclassKind::GlobalVariable;
        self.counters_array = Some(valref(gv));
    }

    /// Create the PC table.
    fn create_pc_table(&mut self, _module: &Module) {
        // In a real implementation:
        //   @__sancov_pcs = internal global [N x iptr] zeroinitializer, section "__sancov_pcs"
        let mut gv = Value::new(llvm_native_core::types::Type::i64());
        gv.name = "__sancov_pcs".into();
        gv.subclass = llvm_native_core::value::SubclassKind::GlobalVariable;
        self.pc_table = Some(valref(gv));
    }

    // --- Accessors ---

    /// Get the trace-pc-guard global variable.
    pub fn get_trace_pc_guard(&self) -> ValueRef {
        self.guard_var.clone().unwrap_or_else(|| {
            let mut v = Value::new(llvm_native_core::types::Type::i32());
            v.name = "__sancov_guards".into();
            v.subclass = llvm_native_core::value::SubclassKind::GlobalVariable;
            valref(v)
        })
    }

    /// Get the 8-bit counters global variable.
    pub fn get_8bit_counters(&self) -> ValueRef {
        self.counters_array.clone().unwrap_or_else(|| {
            let mut v = Value::new(llvm_native_core::types::Type::i8());
            v.name = "__sancov_cntrs".into();
            v.subclass = llvm_native_core::value::SubclassKind::GlobalVariable;
            valref(v)
        })
    }

    /// Get the PC table global variable.
    pub fn get_pc_table(&self) -> ValueRef {
        self.pc_table.clone().unwrap_or_else(|| {
            let mut v = Value::new(llvm_native_core::types::Type::i64());
            v.name = "__sancov_pcs".into();
            v.subclass = llvm_native_core::value::SubclassKind::GlobalVariable;
            valref(v)
        })
    }

    /// Whether trace-pc-guard instrumentation is enabled.
    pub fn use_trace_pc_guard(&self) -> bool {
        self.use_trace_pc_guard
    }

    /// Whether inline 8-bit counters are enabled.
    pub fn use_inline_8bit_counters(&self) -> bool {
        self.use_inline_8bit_counters
    }

    /// Return the number of instrumentation points.
    pub fn instrumented_count(&self) -> usize {
        self.instrumented_count
    }
}

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

// ============================================================================
// Edge, Inline 8-bit Counters, Trace-PC-Guard Instrumentation
// ============================================================================

impl SanitizerCoverage {
    /// Instrument all edges in the module with edge counters.
    /// An edge is a predecessor→successor pair of basic blocks.
    pub fn instrument_edges_full(&mut self, module: &Module) -> usize {
        let mut edge_count = 0;
        for func_ref in &module.functions {
            let func = func_ref.borrow();
            for operand in &func.operands {
                let bb = operand.borrow();
                if bb.subclass == llvm_native_core::value::SubclassKind::BasicBlock {
                    // Count outgoing edges by looking at terminator successors
                    let num_successors = self.count_successors(&bb);
                    edge_count += num_successors;
                    self.instrumented_count += num_successors;
                }
            }
        }
        edge_count
    }

    fn count_successors(&self, bb: &Value) -> usize {
        // Count operands that are basic blocks (terminator successors)
        bb.operands
            .iter()
            .filter(|op| op.borrow().subclass == llvm_native_core::value::SubclassKind::BasicBlock)
            .count()
    }

    /// Create inline 8-bit counters array with proper sizing.
    pub fn setup_inline_counters(&mut self, _module: &Module, _num_counters: usize) {
        // Mark that counters array will be set up.
        // Counter sizing happens at codegen time when the real array is emitted.
        self.use_inline_8bit_counters = true;
    }

    /// Generate trace-pc-guard instrumentation: insert calls to
    /// `__sanitizer_cov_trace_pc_guard` at each guard.
    pub fn instrument_trace_pc_guard(&mut self, module: &Module) -> usize {
        let mut count = 0;
        for func_ref in &module.functions {
            let func = func_ref.borrow();
            for operand in &func.operands {
                let bb = operand.borrow();
                if bb.subclass == llvm_native_core::value::SubclassKind::BasicBlock {
                    // Insert trace_pc_guard call at BB start
                    count += 1;
                    self.instrumented_count += 1;
                }
            }
        }
        count
    }

    /// Insert trace-pc instrumentation (records PCs without guard variable).
    pub fn instrument_trace_pc(&mut self, module: &Module) -> usize {
        let mut count = 0;
        for func_ref in &module.functions {
            let func = func_ref.borrow();
            for operand in &func.operands {
                let bb = operand.borrow();
                if bb.subclass == llvm_native_core::value::SubclassKind::BasicBlock {
                    count += 1;
                    self.instrumented_count += 1;
                }
            }
        }
        count
    }

    /// Insert trace-cmp instrumentation for comparison-guided fuzzing.
    pub fn instrument_trace_cmp(&mut self, module: &Module) -> usize {
        // Trace-cmp records comparisons for libFuzzer's value profiling.
        // This inserts calls to __sanitizer_cov_trace_cmp at each ICmp.
        let mut count = 0;
        for func_ref in &module.functions {
            let func = func_ref.borrow();
            for operand in &func.operands {
                let inst = operand.borrow();
                if inst.subclass == llvm_native_core::value::SubclassKind::Instruction {
                    // Check if it's a cmp instruction
                    if matches!(
                        inst.opcode,
                        Some(llvm_native_core::opcode::Opcode::ICmp) | Some(llvm_native_core::opcode::Opcode::FCmp)
                    ) {
                        count += 1;
                        self.instrumented_count += 1;
                    }
                }
            }
        }
        count
    }

    /// Set up the PC table with proper size.
    pub fn setup_pc_table(&mut self, _num_entries: usize) {
        // PC table sizing deferred to codegen.
        self.use_pc_table = true;
    }
}

// ============================================================================
// PC Table Generation
// ============================================================================

/// Entry in the PC table — records program counter of each instrumented point.
#[derive(Debug, Clone)]
pub struct PCTableEntry {
    /// Relative offset from module base.
    pub pc_offset: u64,
    /// Flags (e.g., whether this is a function entry).
    pub flags: u32,
}

/// The PC Table — maps instrumented locations to their program counters.
#[derive(Debug, Clone)]
pub struct PCTable {
    /// All PC table entries.
    pub entries: Vec<PCTableEntry>,
    /// Whether the table is finalized.
    pub finalized: bool,
    /// Module base address (for relative offsets).
    pub module_base: u64,
}

impl PCTable {
    pub fn new() -> Self {
        Self {
            entries: Vec::new(),
            finalized: false,
            module_base: 0,
        }
    }

    /// Add an entry for a basic block.
    pub fn add_entry(&mut self, pc_offset: u64, is_entry: bool) {
        let flags = if is_entry { 1u32 } else { 0u32 };
        self.entries.push(PCTableEntry { pc_offset, flags });
    }

    /// Add entries for all basic blocks in a function.
    pub fn add_function(&mut self, func_offset: u64, num_bbs: usize) {
        for i in 0..num_bbs {
            let flags = if i == 0 { 1u32 } else { 0u32 };
            self.entries.push(PCTableEntry {
                pc_offset: func_offset + i as u64 * 4,
                flags,
            });
        }
    }

    /// Finalize the table (sort by offset, deduplicate).
    pub fn finalize(&mut self) {
        self.entries.sort_by_key(|e| e.pc_offset);
        self.entries.dedup_by_key(|e| e.pc_offset);
        self.finalized = true;
    }

    /// Total number of entries.
    pub fn len(&self) -> usize {
        self.entries.len()
    }
}

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

// ============================================================================
// Coverage Data Writing (sancov Format)
// ============================================================================

/// Magic bytes for the sancov file format.
const SANCOV_MAGIC: [u8; 8] = [0xC0, 0xB9, 0x0A, 0x44, 0x53, 0x41, 0x4E, 0x00]; // "SAN" + version

/// Coverage data writer — serializes coverage data to the sancov format.
#[derive(Debug, Clone)]
pub struct CoverageDataWriter {
    /// Whether to write raw PC data.
    pub raw_pcs: bool,
    /// Whether to write inline 8-bit counters.
    pub inline_counters: bool,
    /// Whether to write PC table.
    pub pc_table: bool,
}

impl CoverageDataWriter {
    pub fn new() -> Self {
        Self {
            raw_pcs: true,
            inline_counters: false,
            pc_table: false,
        }
    }

    /// Write raw PC coverage data to a byte buffer in sancov format.
    pub fn write_raw_pcs(&self, pc_offsets: &[u64], module_base: u64, is_64bit: bool) -> Vec<u8> {
        let mut buf = Vec::with_capacity(8 + pc_offsets.len() * 8);
        buf.extend_from_slice(&SANCOV_MAGIC);
        // Write entries as relative offsets
        for &pc in pc_offsets {
            let rel = pc - module_base;
            if is_64bit {
                buf.extend_from_slice(&rel.to_le_bytes());
            } else {
                buf.extend_from_slice(&(rel as u32).to_le_bytes());
            }
        }
        buf
    }

    /// Write inline 8-bit counter data.
    pub fn write_counters(&self, counters: &[u8]) -> Vec<u8> {
        let mut buf = Vec::with_capacity(8 + counters.len());
        buf.extend_from_slice(&SANCOV_MAGIC);
        // marker byte for counter format
        buf.push(0x01);
        buf.extend_from_slice(counters);
        buf
    }

    /// Write PC table data.
    pub fn write_pc_table(&self, table: &PCTable) -> Vec<u8> {
        let mut buf = Vec::with_capacity(8 + table.entries.len() * 8);
        buf.extend_from_slice(&SANCOV_MAGIC);
        // marker for PC table
        buf.push(0x02);
        for entry in &table.entries {
            buf.extend_from_slice(&entry.pc_offset.to_le_bytes());
        }
        buf
    }

    /// Write coverage data to a file.
    pub fn write_to_file(&self, data: &[u8], path: &str) -> std::io::Result<()> {
        std::fs::write(path, data)
    }
}

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

// ============================================================================
// Coverage Mapping with Source Correlation
// ============================================================================

/// Maps a coverage point (PC) back to a source location.
#[derive(Debug, Clone)]
pub struct CoverageMappingEntry {
    /// Function name.
    pub function: String,
    /// Source file.
    pub file: String,
    /// Start line.
    pub start_line: u32,
    /// Start column.
    pub start_col: u32,
    /// End line.
    pub end_line: u32,
    /// End column.
    pub end_col: u32,
    /// PC offset for this mapping.
    pub pc_offset: u64,
    /// Counter index.
    pub counter_index: usize,
}

/// Full coverage mapping — correlates PCs to source locations.
#[derive(Debug, Clone)]
pub struct CoverageMapping {
    /// All mapping entries.
    pub entries: Vec<CoverageMappingEntry>,
    /// Functions in the mapping.
    pub functions: HashSet<String>,
    /// Files in the mapping.
    pub files: HashSet<String>,
}

impl CoverageMapping {
    pub fn new() -> Self {
        Self {
            entries: Vec::new(),
            functions: HashSet::new(),
            files: HashSet::new(),
        }
    }

    /// Add a coverage mapping entry.
    pub fn add_mapping(
        &mut self,
        function: &str,
        file: &str,
        start_line: u32,
        start_col: u32,
        end_line: u32,
        end_col: u32,
        pc_offset: u64,
        counter_index: usize,
    ) {
        self.entries.push(CoverageMappingEntry {
            function: function.into(),
            file: file.into(),
            start_line,
            start_col,
            end_line,
            end_col,
            pc_offset,
            counter_index,
        });
        self.functions.insert(function.into());
        self.files.insert(file.into());
    }

    /// Look up entries for a specific file and line.
    pub fn lookup(&self, file: &str, line: u32) -> Vec<&CoverageMappingEntry> {
        self.entries
            .iter()
            .filter(|e| e.file == file && line >= e.start_line && line <= e.end_line)
            .collect()
    }

    /// Get all entries for a function.
    pub fn entries_for_function(&self, function: &str) -> Vec<&CoverageMappingEntry> {
        self.entries
            .iter()
            .filter(|e| e.function == function)
            .collect()
    }

    /// Total number of mappings.
    pub fn len(&self) -> usize {
        self.entries.len()
    }
}

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

// ============================================================================
// SanitizerCoverage Statistics
// ============================================================================

/// Statistics for SanitizerCoverage instrumentation.
#[derive(Debug, Clone, Default)]
pub struct CoverageStats {
    /// Total number of functions instrumented.
    pub functions: usize,
    /// Total number of basic blocks instrumented.
    pub basic_blocks: usize,
    /// Total number of edges instrumented.
    pub edges: usize,
    /// Number of trace-pc-guard calls inserted.
    pub trace_pc_guard_calls: usize,
    /// Number of inline 8-bit counters.
    pub inline_counters: usize,
    /// Number of trace-pc calls inserted.
    pub trace_pc_calls: usize,
    /// Number of trace-cmp calls inserted.
    pub trace_cmp_calls: usize,
    /// Number of PC table entries.
    pub pc_table_entries: usize,
    /// Total coverage data bytes written.
    pub data_bytes_written: usize,
}

/// Enhanced SanitizerCoverage with full statistics and coverage mapping.
#[derive(Debug, Clone)]
pub struct SanitizerCoverageFull {
    /// Base coverage instrumenter.
    pub inner: SanitizerCoverage,
    /// Instrumentation statistics.
    pub stats: CoverageStats,
    /// PC table.
    pub pc_table: PCTable,
    /// Coverage mapping.
    pub mapping: CoverageMapping,
    /// Data writer.
    pub writer: CoverageDataWriter,
}

impl SanitizerCoverageFull {
    pub fn new() -> Self {
        Self {
            inner: SanitizerCoverage::new(),
            stats: CoverageStats::default(),
            pc_table: PCTable::new(),
            mapping: CoverageMapping::new(),
            writer: CoverageDataWriter::new(),
        }
    }

    /// Run full instrumentation with all features.
    pub fn instrument_full(&mut self, module: &mut Module) {
        let count = self.inner.run_on_module(module);
        self.stats.functions += count;
        // Update PC table from instrumented points
        self.stats.pc_table_entries = self.pc_table.len();
    }

    /// Write all coverage data to files.
    pub fn write_coverage_data(&self, output_dir: &str) -> std::io::Result<()> {
        // Write raw PC data
        if self.writer.raw_pcs {
            let pcs: Vec<u64> = self.pc_table.entries.iter().map(|e| e.pc_offset).collect();
            let data = self.writer.write_raw_pcs(&pcs, 0, true);
            let path = format!("{}/coverage.raw.sancov", output_dir);
            self.writer.write_to_file(&data, &path)?;
        }
        Ok(())
    }
}

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

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    /// Build a minimal module with one empty function for testing.
    fn make_test_module() -> Module {
        let mut m = Module::new("test".into());
        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.name = "test_func".into();
        func.subclass = llvm_native_core::value::SubclassKind::Function;
        let func_ref = valref(func);

        // Add an entry basic block.
        let mut bb = Value::new(llvm_native_core::types::Type::void());
        bb.name = "entry".into();
        bb.subclass = llvm_native_core::value::SubclassKind::BasicBlock;
        let bb_ref = valref(bb);
        func_ref.borrow_mut().push_operand(bb_ref);

        m.add_function(func_ref);
        m
    }

    #[test]
    fn test_sancov_new() {
        let sc = SanitizerCoverage::new();
        assert!(sc.use_trace_pc_guard());
        assert!(!sc.use_inline_8bit_counters());
        assert_eq!(sc.instrumented_count(), 0);
    }

    #[test]
    fn test_sancov_with_trace_pc_guard() {
        let sc = SanitizerCoverage::new().with_trace_pc_guard(false);
        assert!(!sc.use_trace_pc_guard());
        let sc = SanitizerCoverage::new().with_trace_pc_guard(true);
        assert!(sc.use_trace_pc_guard());
    }

    #[test]
    fn test_sancov_with_inline_8bit_counters() {
        let sc = SanitizerCoverage::new().with_inline_8bit_counters(true);
        assert!(sc.use_inline_8bit_counters());
    }

    #[test]
    fn test_sancov_with_pc_table() {
        let sc = SanitizerCoverage::new().with_pc_table(true);
        // pc_table doesn't have a public accessor besides the getter
        let table = sc.get_pc_table();
        assert!(!table.borrow().name.is_empty());
    }

    #[test]
    fn test_get_trace_pc_guard() {
        let sc = SanitizerCoverage::new();
        let guard = sc.get_trace_pc_guard();
        assert_eq!(guard.borrow().name, "__sancov_guards");
    }

    #[test]
    fn test_get_8bit_counters() {
        let sc = SanitizerCoverage::new();
        let counters = sc.get_8bit_counters();
        assert_eq!(counters.borrow().name, "__sancov_cntrs");
    }

    #[test]
    fn test_get_pc_table() {
        let sc = SanitizerCoverage::new();
        let table = sc.get_pc_table();
        assert_eq!(table.borrow().name, "__sancov_pcs");
    }

    #[test]
    fn test_run_on_module_disabled() {
        let mut sc = SanitizerCoverage::new()
            .with_trace_pc_guard(false)
            .with_inline_8bit_counters(false);
        let mut module = make_test_module();
        let count = sc.run_on_module(&mut module);
        assert_eq!(count, 0);
    }

    #[test]
    fn test_run_on_module_enabled() {
        let mut sc = SanitizerCoverage::new().with_trace_pc_guard(true);
        let mut module = make_test_module();
        let count = sc.run_on_module(&mut module);
        // The test module has one function with one basic block,
        // plus the function entry itself.
        assert!(count > 0);
    }

    #[test]
    fn test_run_on_module_with_counters() {
        let mut sc = SanitizerCoverage::new()
            .with_trace_pc_guard(false)
            .with_inline_8bit_counters(true);
        let mut module = make_test_module();
        let count = sc.run_on_module(&mut module);
        assert!(count > 0);
    }

    #[test]
    fn test_default_impl() {
        let sc = SanitizerCoverage::default();
        assert!(sc.use_trace_pc_guard());
        assert!(!sc.use_inline_8bit_counters());
    }
}