llvm-native-core-ext 0.1.0

//! DataFlow Sanitizer (DFSan) — dynamic taint/data-flow tracking.
//! Clean-room behavioral reconstruction.
//!
//! @llvm_behavior: DFSan instruments LLVM IR to track data flow at runtime.
//! Each byte of program memory has an associated "shadow" label (a
//! bitmask) that records which input sources influenced its value.
//! Labels propagate through arithmetic, memory, and control flow
//! operations by unioning the labels of all operands.
//!
//! Design:
//! - Shadow memory: a dedicated region mapping application memory to labels
//! - ABI list: functions whose labels follow defined propagation rules
//! - Instrumentation: loads/stores/calls are instrumented to track labels
//! - Union labels: combine labels from multiple sources via bitwise OR
//!
//! This is a dynamic analysis tool; this pass inserts the instrumentation
//! calls into the IR that the dfsan runtime library handles at execution.

use llvm_native_core::module::Module;
use llvm_native_core::opcode::Opcode;
use llvm_native_core::types::Type;
use llvm_native_core::value::{valref, SubclassKind, Value, ValueRef};
use std::collections::{HashMap, HashSet};

// ============================================================================
// Label and Shadow Constants
// ============================================================================

/// Default shadow memory offset (application ↔ shadow mapping).
const DEFAULT_SHADOW_OFFSET: u64 = 0x200000000000;

/// Shadow memory scale: shadow address = (app_addr - offset) * scale.
const DEFAULT_SHADOW_SCALE: u64 = 1;

/// Size of a label in bytes (u32 = 4 bytes).
const LABEL_SIZE: u32 = 4;

/// A label is a u32 bitmask: each bit represents a distinct taint source.
pub type Label = u32;

/// Special label value: untainted / clean.
const LABEL_CLEAN: Label = 0;

/// Special label value: all sources tainted.
const LABEL_ALL: Label = !0u32;

// ============================================================================
// ABI List Entry
// ============================================================================

/// Describes how labels propagate through a specific function.
#[derive(Debug, Clone)]
pub struct ABIListEntry {
    /// The function name (mangled).
    pub name: String,
    /// How arguments propagate: arg_index -> label source.
    /// If empty, all args contribute equally (union of all arg labels).
    pub arg_rules: Vec<ArgLabelRule>,
    /// Whether the return value gets the union of arg labels.
    pub ret_union: bool,
    /// Whether to discard labels on entry (discard = reset to clean).
    pub discard: bool,
    /// Whether this is a custom function handled by the runtime.
    pub custom: bool,
}

#[derive(Debug, Clone)]
pub enum ArgLabelRule {
    /// This argument's label contributes to the result.
    Propagate,
    /// This argument's label is ignored.
    Ignore,
    /// This argument's label goes to a specific shadow slot.
    ShadowOutput(usize),
}

// ============================================================================
// Shadow Variable Info
// ============================================================================

/// Information about a shadow variable (label storage for an IR value).
#[derive(Debug, Clone)]
struct ShadowInfo {
    /// The original value being shadowed.
    original: ValueRef,
    /// The shadow storage (an alloca or global holding the label).
    shadow: ValueRef,
    /// Size in bytes that this shadow covers.
    size: u32,
}

// ============================================================================
// DataFlow Sanitizer Pass
// ============================================================================

/// DataFlow Sanitizer pass — instruments IR for data-flow tracking.
#[derive(Debug, Clone)]
pub struct DataFlowSanitizer {
    /// Shadow memory offset for the platform.
    pub shadow_offset: u64,
    /// Shadow memory scale factor.
    pub shadow_scale: u64,
    /// ABI list: known functions with custom label propagation.
    pub abi_list: HashMap<String, ABIListEntry>,
    /// Number of instrumentation points inserted.
    pub instrumented_count: usize,
    /// Label counter for generating unique label bits.
    next_label_bit: u32,
    /// Shadow variables created for IR values.
    shadow_vars: Vec<ShadowInfo>,
}

impl DataFlowSanitizer {
    /// Create a new DataFlow Sanitizer with default settings.
    pub fn new() -> Self {
        Self {
            shadow_offset: DEFAULT_SHADOW_OFFSET,
            shadow_scale: DEFAULT_SHADOW_SCALE,
            abi_list: HashMap::new(),
            instrumented_count: 0,
            next_label_bit: 0,
            shadow_vars: Vec::new(),
        }
    }

    /// Instrument a module for data-flow tracking.
    ///
    /// Walks all functions and inserts shadow propagation code for:
    /// - Memory operations (load, store, alloca)
    /// - Arithmetic operations (add, sub, mul, bitwise ops)
    /// - Calls to known functions (ABI list)
    /// - Function entry/exit to manage label state
    ///
    /// Returns the number of instrumentation points inserted.
    pub fn instrument(&mut self, module: &mut Module) -> usize {
        self.instrumented_count = 0;
        self.shadow_vars.clear();

        // Phase 1: Add shadow global variables for all globals
        self.add_shadow_variables(module);

        // Phase 2: Instrument each function
        for func in module.functions.clone() {
            self.instrument_function(&func);
        }

        self.instrumented_count
    }

    // ========================================================================
    // Shadow variable creation for globals
    // ========================================================================

    /// Add shadow variables for each global in the module.
    fn add_shadow_variables(&mut self, module: &Module) {
        for global in &module.globals {
            let g = global.borrow();
            let shadow_name = format!("__dfsan_shadow_{}", g.name);
            let size = g.ty.size_bytes();

            // Create a shadow global (array of labels)
            let elem_ty = Type::i32();
            let shadow_ty = Type::new(llvm_native_core::types::TypeKind::Array {
                len: size as u64 / LABEL_SIZE as u64 + 1,
                element_type_id: elem_ty.id,
            });

            let mut shadow = Value::new(shadow_ty)
                .with_subclass(SubclassKind::GlobalVariable)
                .named(&shadow_name);

            let shadow_ref = valref(shadow);
            self.shadow_vars.push(ShadowInfo {
                original: global.clone(),
                shadow: shadow_ref,
                size: size as u32,
            });
        }
    }

    // ========================================================================
    // Function instrumentation
    // ========================================================================

    /// Instrument a single function for data-flow tracking.
    fn instrument_function(&mut self, func: &ValueRef) {
        let f = func.borrow();

        // Collect blocks to instrument (clone to avoid borrow issues)
        let blocks: Vec<ValueRef> = f
            .operands
            .iter()
            .filter(|op| op.borrow().is_basic_block())
            .cloned()
            .collect();

        // Insert function entry instrumentation: set up shadow for args
        if !blocks.is_empty() {
            self.instrument_function_entry(func, &blocks[0]);
        }

        // Walk each block and instrument instructions
        for bb in &blocks {
            let block = bb.borrow();
            let insts: Vec<ValueRef> = block
                .operands
                .iter()
                .filter(|op| op.borrow().is_instruction())
                .cloned()
                .collect();

            for inst in &insts {
                let opcode = inst.borrow().get_opcode();
                match opcode {
                    Some(Opcode::Load) => self.instrument_load(inst),
                    Some(Opcode::Store) => self.instrument_store(inst),
                    Some(Opcode::Call) => self.instrument_call(inst),
                    Some(_) => self.propagate_shadow(inst),
                    None => {}
                }
            }
        }
    }

    /// Instrument function entry: create shadow for arguments.
    fn instrument_function_entry(&mut self, func: &ValueRef, _entry_bb: &ValueRef) {
        let f = func.borrow();

        // Arguments (operands that are not basic blocks, stored before BBs)
        for arg_val in &f.operands {
            let arg = arg_val.borrow();
            if arg.subclass == SubclassKind::Argument {
                // Create a shadow alloca for this argument
                let shadow = self.create_shadow_alloca(arg_val, arg.ty.size_bytes() as u32);
                self.shadow_vars.push(ShadowInfo {
                    original: arg_val.clone(),
                    shadow,
                    size: arg.ty.size_bytes() as u32,
                });
                self.instrumented_count += 1;
            }
        }
    }

    /// Instrument a load instruction to propagate shadow labels.
    fn instrument_load(&mut self, inst: &ValueRef) {
        let i = inst.borrow();
        if i.operands.is_empty() {
            return;
        }

        let ptr = &i.operands[0];
        let load_size = i.ty.size_bytes() as u32;

        // Create a shadow load: load the label from shadow memory
        let shadow = self.create_shadow_alloca(inst, load_size);
        self.shadow_vars.push(ShadowInfo {
            original: inst.clone(),
            shadow,
            size: load_size,
        });
        self.instrumented_count += 1;
    }

    /// Instrument a store instruction to propagate shadow labels.
    fn instrument_store(&mut self, inst: &ValueRef) {
        let i = inst.borrow();
        if i.operands.len() < 2 {
            return;
        }

        let value = &i.operands[0];
        let ptr = &i.operands[1];

        // Find the shadow for the stored value (if any)
        let stored_label = self.get_shadow_label(value);

        // Insert a shadow store: write the label to shadow memory
        // In a real implementation this would generate a store to the
        // shadow address. Here we record the instrumentation point.
        if stored_label != LABEL_CLEAN {
            let _ = ptr; // Would compute shadow address
            self.instrumented_count += 1;
        }
    }

    /// Instrument a call instruction.
    fn instrument_call(&mut self, inst: &ValueRef) {
        let i = inst.borrow();
        if i.operands.is_empty() {
            return;
        }

        let callee = &i.operands[0];
        let callee_name = callee.borrow().name.clone();

        // Check ABI list for custom propagation
        if let Some(entry) = self.abi_list.get(&callee_name) {
            if entry.custom {
                // Custom function, handled by runtime
                self.instrumented_count += 1;
                return;
            }
        }

        // Default: union all argument labels into the return
        let mut union_label: Label = LABEL_CLEAN;
        for (idx, arg) in i.operands.iter().enumerate().skip(1) {
            let arg_label = self.get_shadow_label(arg);
            union_label = self.create_union_shadow_label(union_label, arg_label);
        }

        if union_label != LABEL_CLEAN {
            let return_size = i.ty.size_bytes() as u32;
            let _shadow = self.create_shadow_alloca(inst, return_size);
            self.instrumented_count += 1;
        }
    }

    /// Propagate shadow labels through a general instruction.
    ///
    /// For arithmetic/logic instructions, the shadow of the result is
    /// the union of the shadows of all operands (bitwise OR).
    fn propagate_shadow(&mut self, inst: &ValueRef) {
        let i = inst.borrow();
        let inst_size = i.ty.size_bytes() as u32;

        if i.operands.is_empty() || i.ty.is_void() {
            return;
        }

        let mut union_label: Label = LABEL_CLEAN;
        for arg in &i.operands {
            let lbl = self.get_shadow_label(arg);
            union_label = self.create_union_shadow_label(union_label, lbl);
        }

        if union_label != LABEL_CLEAN {
            let shadow = self.create_shadow_alloca(inst, inst_size);
            self.shadow_vars.push(ShadowInfo {
                original: inst.clone(),
                shadow,
                size: inst_size,
            });
            self.instrumented_count += 1;
        }
    }

    // ========================================================================
    // Label utilities
    // ========================================================================

    /// Create a union of two shadow labels (bitwise OR).
    pub fn create_union_shadow_label(&self, a: u32, b: u32) -> u32 {
        a | b
    }

    /// Get the shadow label for a value, looking it up in the shadow map.
    fn get_shadow_label(&self, val: &ValueRef) -> Label {
        for si in &self.shadow_vars {
            if Rc::ptr_eq(&si.original, val) {
                return LABEL_ALL; // In a real impl, would read from shadow memory
            }
        }
        LABEL_CLEAN
    }

    /// Create a shadow alloca for a value's label.
    fn create_shadow_alloca(&self, _val: &ValueRef, _size: u32) -> ValueRef {
        // Create a label-sized alloca for holding the shadow label
        let label_ty = Type::i32();
        let mut alloca = Value::new(label_ty)
            .with_subclass(SubclassKind::AllocaInst)
            .named("dfsan_label");
        alloca.opcode = Some(Opcode::Alloca);
        valref(alloca)
    }

    // ========================================================================
    // ABI List
    // ========================================================================

    /// Build the default ABI list with common C library functions.
    pub fn create_abi_list(&self) -> Vec<String> {
        let mut list = Vec::new();

        // Memory functions — labels propagate through memory
        list.push("memcpy".to_string());
        list.push("memmove".to_string());
        list.push("memset".to_string());
        list.push("memcmp".to_string());

        // String functions
        list.push("strlen".to_string());
        list.push("strcpy".to_string());
        list.push("strncpy".to_string());
        list.push("strcmp".to_string());
        list.push("strncmp".to_string());
        list.push("strchr".to_string());
        list.push("strrchr".to_string());
        list.push("strstr".to_string());

        // I/O functions — labels propagate through read data
        list.push("read".to_string());
        list.push("write".to_string());
        list.push("fread".to_string());
        list.push("fwrite".to_string());
        list.push("pread".to_string());
        list.push("pwrite".to_string());
        list.push("recv".to_string());
        list.push("send".to_string());
        list.push("recvfrom".to_string());
        list.push("sendto".to_string());

        // Allocation functions — memory is clean on allocation
        list.push("malloc".to_string());
        list.push("calloc".to_string());
        list.push("realloc".to_string());
        list.push("free".to_string());
        list.push("mmap".to_string());
        list.push("munmap".to_string());

        // POSIX
        list.push("getenv".to_string());
        list.push("atoi".to_string());
        list.push("atol".to_string());
        list.push("atoll".to_string());
        list.push("strtol".to_string());
        list.push("strtoll".to_string());
        list.push("strtoul".to_string());
        list.push("strtoull".to_string());
        list.push("strtod".to_string());
        list.push("strtof".to_string());

        list
    }

    /// Register a custom function in the ABI list.
    pub fn register_abi_function(&mut self, name: &str, entry: ABIListEntry) {
        self.abi_list.insert(name.to_string(), entry);
    }

    /// Check if a function is in the ABI list.
    pub fn is_in_abi_list(&self, name: &str) -> bool {
        self.abi_list.contains_key(name)
    }

    /// Allocate a new unique label bit for a source.
    pub fn allocate_label(&mut self) -> Label {
        if self.next_label_bit >= 32 {
            // All bits exhausted
            return LABEL_ALL;
        }
        let bit = 1u32 << self.next_label_bit;
        self.next_label_bit += 1;
        bit
    }
}

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

// ============================================================================
// Helper: Type size_bytes
// ============================================================================

/// Extension trait for Type to get approximate byte size.
trait TypeSize {
    fn size_bytes(&self) -> usize;
    fn is_void(&self) -> bool;
}

impl TypeSize for Type {
    fn size_bytes(&self) -> usize {
        match &self.kind {
            llvm_native_core::types::TypeKind::Void => 0,
            llvm_native_core::types::TypeKind::Half => 2,
            llvm_native_core::types::TypeKind::BFloat => 2,
            llvm_native_core::types::TypeKind::Float => 4,
            llvm_native_core::types::TypeKind::Double => 8,
            llvm_native_core::types::TypeKind::FP128 => 16,
            llvm_native_core::types::TypeKind::X86FP80 => 10,
            llvm_native_core::types::TypeKind::PPCFP128 => 16,
            llvm_native_core::types::TypeKind::Label => 0,
            llvm_native_core::types::TypeKind::Metadata => 0,
            llvm_native_core::types::TypeKind::X86AMX => 0,
            llvm_native_core::types::TypeKind::X86MMX => 8,
            llvm_native_core::types::TypeKind::Token => 0,
            llvm_native_core::types::TypeKind::Integer { bits } => ((bits + 7) / 8) as usize,
            llvm_native_core::types::TypeKind::Pointer { .. } => 8,
            llvm_native_core::types::TypeKind::Array { len, .. } => {
                // Approximate; full resolution needs type store
                (*len as usize) * 8 // assume element is pointer-sized
            }
            llvm_native_core::types::TypeKind::Struct { .. } => 8, // conservative default
            llvm_native_core::types::TypeKind::FixedVector { len, .. } => (*len as usize) * 8,
            llvm_native_core::types::TypeKind::ScalableVector { min_elems, .. } => (*min_elems as usize) * 8,
            llvm_native_core::types::TypeKind::Function { .. } => 0,
        }
    }

    fn is_void(&self) -> bool {
        matches!(self.kind, llvm_native_core::types::TypeKind::Void)
    }
}

// ============================================================================
// Label Propagation Rules for All Operations
// ============================================================================

impl DataFlowSanitizer {
    /// Propagate shadow labels through a binary operation.
    /// Returns the resulting shadow label (union of operand labels).
    pub fn propagate_binary_op(&self, lhs_label: Label, rhs_label: Label) -> Label {
        lhs_label | rhs_label
    }

    /// Propagate shadow labels through a cast operation.
    /// Casts preserve the label of the source value.
    pub fn propagate_cast(&self, src_label: Label) -> Label {
        src_label
    }

    /// Propagate shadow labels through a PHI node.
    /// Result is the union of all incoming labels.
    pub fn propagate_phi(&self, incoming_labels: &[Label]) -> Label {
        incoming_labels.iter().fold(LABEL_CLEAN, |acc, &l| acc | l)
    }

    /// Propagate shadow labels through a select instruction.
    /// Condition is ignored; result is union of true/false value labels.
    pub fn propagate_select(&self, true_label: Label, false_label: Label) -> Label {
        true_label | false_label
    }

    /// Propagate shadow labels through an extractvalue.
    pub fn propagate_extractvalue(&self, aggregate_label: Label) -> Label {
        aggregate_label
    }

    /// Propagate shadow labels through an insertvalue.
    pub fn propagate_insertvalue(&self, agg_label: Label, inserted_label: Label) -> Label {
        agg_label | inserted_label
    }

    /// Propagate shadow labels for a GEP (getelementptr).
    /// The index labels are ignored; result label is the pointer label.
    pub fn propagate_gep(&self, ptr_label: Label, _index_labels: &[Label]) -> Label {
        ptr_label
    }

    /// Propagate shadow labels for a comparison.
    /// Result is union of compared operand labels.
    pub fn propagate_cmp(&self, lhs_label: Label, rhs_label: Label) -> Label {
        lhs_label | rhs_label
    }
}

// ============================================================================
// Union Table for Combining Labels Efficiently
// ============================================================================

/// A union table for fast label combination.
/// Maps (label_a, label_b) → combined_label using a flat array.
/// This allows O(1) label unioning at runtime.
#[derive(Debug, Clone)]
pub struct UnionTable {
    /// Flat table: table[a * size + b] = a | b
    table: Vec<Label>,
    /// Number of distinct labels (size of one dimension).
    size: usize,
}

impl UnionTable {
    /// Create a new union table for up to `num_labels` distinct labels.
    pub fn new(num_labels: usize) -> Self {
        let size = num_labels;
        let mut table = vec![LABEL_CLEAN; size * size];
        for a in 0..size {
            for b in 0..size {
                table[a * size + b] = (a as Label) | (b as Label);
            }
        }
        Self { table, size }
    }

    /// Get the union of two labels.
    pub fn union(&self, a: Label, b: Label) -> Label {
        let ai = a as usize;
        let bi = b as usize;
        if ai < self.size && bi < self.size {
            self.table[ai * self.size + bi]
        } else {
            a | b
        }
    }

    /// Get the union of multiple labels.
    pub fn union_all(&self, labels: &[Label]) -> Label {
        labels
            .iter()
            .fold(LABEL_CLEAN, |acc, &l| self.union(acc, l))
    }

    /// Number of entries in the table.
    pub fn len(&self) -> usize {
        self.table.len()
    }
}

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

// ============================================================================
// ABI List Parsing
// ============================================================================

/// Parser for the DFSan ABI list file format.
/// Format: each line is "[discard] fun:name [uninstrumented|custom|custom_ret|
///          arg1=propagate|ignore|shadow ...]"
#[derive(Debug, Clone)]
pub struct ABIListParser {
    /// Parsed entries.
    pub entries: Vec<ABIListEntry>,
    /// Current line number during parsing.
    line_number: usize,
}

impl ABIListParser {
    pub fn new() -> Self {
        Self {
            entries: Vec::new(),
            line_number: 0,
        }
    }

    /// Parse an ABI list from a string (entire file contents).
    pub fn parse(&mut self, contents: &str) -> usize {
        for line in contents.lines() {
            self.line_number += 1;
            let trimmed = line.trim();
            // Skip empty lines and comments
            if trimmed.is_empty() || trimmed.starts_with('#') {
                continue;
            }
            if let Some(entry) = self.parse_line(trimmed) {
                self.entries.push(entry);
            }
        }
        self.entries.len()
    }

    /// Parse a single ABI list line.
    fn parse_line(&self, line: &str) -> Option<ABIListEntry> {
        let mut parts: Vec<&str> = line.split_whitespace().collect();
        if parts.is_empty() {
            return None;
        }

        let mut discard = false;
        let mut name = String::new();
        let mut arg_rules = Vec::new();
        let mut ret_union = false;
        let mut custom = false;

        // Handle [discard] prefix
        if parts[0].eq_ignore_ascii_case("[discard]") {
            discard = true;
            parts.remove(0);
        }

        for part in &parts {
            if part.starts_with("fun:") {
                name = part[4..].to_string();
            } else if part.eq_ignore_ascii_case("uninstrumented") {
                // All args and return are un-instrumented
                ret_union = false;
                arg_rules.clear();
            } else if part.eq_ignore_ascii_case("custom") {
                custom = true;
            } else if part.eq_ignore_ascii_case("custom_ret") {
                ret_union = true;
            } else if part.starts_with("arg") {
                // Parse argN=rule
                if let Some(eq_pos) = part.find('=') {
                    let rule_str = &part[eq_pos + 1..];
                    let rule = match rule_str {
                        "propagate" => ArgLabelRule::Propagate,
                        "ignore" => ArgLabelRule::Ignore,
                        "shadow" => ArgLabelRule::ShadowOutput(0),
                        _ => ArgLabelRule::Propagate,
                    };
                    arg_rules.push(rule);
                }
            }
        }

        if name.is_empty() {
            return None;
        }

        Some(ABIListEntry {
            name,
            arg_rules,
            ret_union,
            discard,
            custom,
        })
    }

    /// Look up an entry by function name.
    pub fn lookup(&self, func_name: &str) -> Option<&ABIListEntry> {
        self.entries.iter().find(|e| e.name == func_name)
    }
}

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

// ============================================================================
// Origin Tracking
// ============================================================================

/// Tracks the origin of tainted data — records which taint source
/// (e.g., file read, network input) first introduced a label.
#[derive(Debug, Clone)]
pub struct OriginTracker {
    /// Whether origin tracking is enabled.
    pub enabled: bool,
    /// Origin ID for each label. Maps label → origin_description.
    label_origins: HashMap<Label, String>,
    /// Next available origin ID.
    next_origin_id: u32,
}

impl OriginTracker {
    pub fn new() -> Self {
        Self {
            enabled: false,
            label_origins: HashMap::new(),
            next_origin_id: 0,
        }
    }

    /// Register a new taint source with a human-readable description.
    /// Returns the label assigned to this source.
    pub fn register_source(&mut self, description: &str) -> Label {
        let label = 1u32 << self.next_origin_id;
        self.next_origin_id += 1;
        self.label_origins.insert(label, description.into());
        label
    }

    /// Get the origin description for a given label.
    pub fn origin_of(&self, label: Label) -> Option<&str> {
        self.label_origins.get(&label).map(|s| s.as_str())
    }

    /// Get all origins that contribute to a combined label.
    pub fn origins_of(&self, label: Label) -> Vec<String> {
        let mut origins = Vec::new();
        let mut remaining = label;
        let mut bit = 0u32;
        while remaining != LABEL_CLEAN && bit < 32 {
            if remaining & 1 != 0 {
                let single_label = 1u32 << bit;
                if let Some(desc) = self.label_origins.get(&single_label) {
                    origins.push(desc.clone());
                }
            }
            remaining >>= 1;
            bit += 1;
        }
        origins
    }

    /// Number of registered sources.
    pub fn source_count(&self) -> usize {
        self.label_origins.len()
    }
}

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

// ============================================================================
// Custom Function Wrappers
// ============================================================================

/// A custom function wrapper — specifies how DFSan should handle
/// a particular function's shadow propagation.
#[derive(Debug, Clone)]
pub struct CustomFunctionWrapper {
    /// Function name.
    pub function_name: String,
    /// Number of arguments.
    pub arg_count: usize,
    /// Label propagation rule for each argument.
    pub arg_rules: Vec<ArgLabelRule>,
    /// Whether the return label is the union of all arg labels.
    pub ret_label_union: bool,
    /// Whether the function is fully custom (runtime handles it).
    pub is_custom: bool,
    /// Custom wrapper identifier.
    pub wrapper_id: u32,
}

/// Registry of custom function wrappers.
#[derive(Debug, Clone, Default)]
pub struct CustomWrapperRegistry {
    /// All registered wrappers.
    pub wrappers: Vec<CustomFunctionWrapper>,
    /// Mapping from function name to wrapper index.
    name_index: HashMap<String, usize>,
}

impl CustomWrapperRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a custom function wrapper.
    pub fn register(&mut self, wrapper: CustomFunctionWrapper) {
        self.name_index
            .insert(wrapper.function_name.clone(), self.wrappers.len());
        self.wrappers.push(wrapper);
    }

    /// Look up a wrapper by function name.
    pub fn lookup(&self, name: &str) -> Option<&CustomFunctionWrapper> {
        self.name_index
            .get(name)
            .and_then(|&idx| self.wrappers.get(idx))
    }

    /// Check if a function has a custom wrapper.
    pub fn is_custom(&self, name: &str) -> bool {
        self.name_index.contains_key(name)
    }

    /// Number of registered wrappers.
    pub fn count(&self) -> usize {
        self.wrappers.len()
    }

    /// Build wrappers from an ABI list.
    pub fn from_abi_list(abi_entries: &[ABIListEntry]) -> Self {
        let mut registry = Self::new();
        for entry in abi_entries {
            if entry.custom || entry.discard {
                let wrapper = CustomFunctionWrapper {
                    function_name: entry.name.clone(),
                    arg_count: entry.arg_rules.len(),
                    arg_rules: entry.arg_rules.clone(),
                    ret_label_union: entry.ret_union,
                    is_custom: entry.custom,
                    wrapper_id: registry.wrappers.len() as u32,
                };
                registry.register(wrapper);
            }
        }
        registry
    }
}

// ============================================================================
// Shadow Memory Management for DFSan
// ============================================================================

/// Manages the shadow memory for DataFlowSanitizer.
/// Maps each application memory location to a shadow label.
#[derive(Debug, Clone)]
pub struct ShadowMemory {
    /// Shadow data: shadow_addr → label (sparse map for simplicity).
    shadow_map: HashMap<u64, Label>,
    /// Shadow offset for address translation.
    offset: u64,
    /// Shadow scale (log2 of granularity).
    scale: u8,
}

impl ShadowMemory {
    pub fn new(offset: u64, scale: u8) -> Self {
        Self {
            shadow_map: HashMap::new(),
            offset,
            scale,
        }
    }

    /// Compute shadow address from application address.
    pub fn app_to_shadow(&self, app_addr: u64) -> u64 {
        app_addr
            .wrapping_shl(self.scale as u32)
            .wrapping_add(self.offset)
    }

    /// Get the label for an application address.
    pub fn get_label(&self, app_addr: u64) -> Label {
        let shadow = self.app_to_shadow(app_addr);
        self.shadow_map.get(&shadow).copied().unwrap_or(LABEL_CLEAN)
    }

    /// Set the label for an application address range.
    pub fn set_label(&mut self, app_addr: u64, size: usize, label: Label) {
        let shadow_start = self.app_to_shadow(app_addr);
        let shadow_end = self.app_to_shadow(app_addr + size as u64 - 1);
        for sa in shadow_start..=shadow_end {
            self.shadow_map.insert(sa, label);
        }
    }

    /// Store a label with union semantics (OR with existing).
    pub fn union_label(&mut self, app_addr: u64, size: usize, label: Label) {
        let shadow_start = self.app_to_shadow(app_addr);
        let shadow_end = self.app_to_shadow(app_addr + size as u64 - 1);
        for sa in shadow_start..=shadow_end {
            let existing = self.shadow_map.get(&sa).copied().unwrap_or(LABEL_CLEAN);
            self.shadow_map.insert(sa, existing | label);
        }
    }

    /// Clear the label for an application address range.
    pub fn clear_label(&mut self, app_addr: u64, size: usize) {
        let shadow_start = self.app_to_shadow(app_addr);
        let shadow_end = self.app_to_shadow(app_addr + size as u64 - 1);
        for sa in shadow_start..=shadow_end {
            self.shadow_map.remove(&sa);
        }
    }

    /// Get the label for a range, returning the union of all bytes.
    pub fn get_label_range(&self, app_addr: u64, size: usize) -> Label {
        let shadow_start = self.app_to_shadow(app_addr);
        let shadow_end = self.app_to_shadow(app_addr + size as u64 - 1);
        let mut label = LABEL_CLEAN;
        for sa in shadow_start..=shadow_end {
            if let Some(&l) = self.shadow_map.get(&sa) {
                label |= l;
            }
        }
        label
    }

    /// Number of shadowed bytes.
    pub fn shadowed_bytes(&self) -> usize {
        self.shadow_map.len()
    }
}

impl Default for ShadowMemory {
    fn default() -> Self {
        Self::new(DEFAULT_SHADOW_OFFSET, DEFAULT_SHADOW_SCALE as u8)
    }
}

// ============================================================================
// DFSan Instrumentation Statistics and Reporting
// ============================================================================

/// Statistics about DFSan instrumentation.
#[derive(Debug, Clone, Default)]
pub struct DFSanStats {
    /// Number of load instructions instrumented.
    pub loads_instrumented: usize,
    /// Number of store instructions instrumented.
    pub stores_instrumented: usize,
    /// Number of call instructions instrumented.
    pub calls_instrumented: usize,
    /// Number of functions instrumented.
    pub functions_instrumented: usize,
    /// Number of custom wrapper calls.
    pub custom_wrapper_calls: usize,
    /// Number of unique labels allocated.
    pub labels_allocated: usize,
    /// Number of ABI list entries loaded.
    pub abi_entries_loaded: usize,
}

impl DFSanStats {
    /// Total instrumentations.
    pub fn total(&self) -> usize {
        self.loads_instrumented + self.stores_instrumented + self.calls_instrumented
    }
}

/// Enhanced DFSan that tracks statistics.
#[derive(Debug, Clone)]
pub struct DataFlowSanitizerFull {
    /// The base DFSan instance.
    pub inner: DataFlowSanitizer,
    /// Instrumentation statistics.
    pub stats: DFSanStats,
    /// Union table for fast label combination.
    pub union_table: UnionTable,
    /// Origin tracker.
    pub origins: OriginTracker,
    /// Custom wrapper registry.
    pub wrappers: CustomWrapperRegistry,
    /// Shadow memory manager.
    pub shadow_memory: ShadowMemory,
}

impl DataFlowSanitizerFull {
    pub fn new() -> Self {
        let inner = DataFlowSanitizer::new();
        let shadow_memory = ShadowMemory::new(inner.shadow_offset, inner.shadow_scale as u8);
        Self {
            inner,
            stats: DFSanStats::default(),
            union_table: UnionTable::default(),
            origins: OriginTracker::new(),
            wrappers: CustomWrapperRegistry::new(),
            shadow_memory,
        }
    }

    /// Instrument a full module with all DFSan capabilities.
    pub fn instrument_full(&mut self, module: &mut Module) {
        self.inner.instrument(module);
        // Count instrumentations from the base instrumenter
        self.stats.functions_instrumented = self.inner.instrumented_count;
    }

    /// Load ABI list and build wrappers.
    pub fn load_abi_list(&mut self, contents: &str) {
        let mut parser = ABIListParser::new();
        parser.parse(contents);
        self.stats.abi_entries_loaded = parser.entries.len();
        self.wrappers = CustomWrapperRegistry::from_abi_list(&parser.entries);
        for entry in parser.entries {
            self.inner.abi_list.insert(entry.name.clone(), entry);
        }
    }

    /// Allocate a new label with origin tracking.
    pub fn allocate_label_with_origin(&mut self, origin: &str) -> Label {
        let label = self.origins.register_source(origin);
        self.stats.labels_allocated += 1;
        label
    }

    /// Track a memory store with origin propagation.
    pub fn track_store(&mut self, addr: u64, size: usize, label: Label) {
        self.shadow_memory.set_label(addr, size, label);
        self.stats.stores_instrumented += 1;
    }

    /// Track a memory load with origin query.
    pub fn track_load(&mut self, addr: u64, size: usize) -> Label {
        self.stats.loads_instrumented += 1;
        self.shadow_memory.get_label_range(addr, size)
    }

    /// Union two labels using the fast union table.
    pub fn union_labels(&self, a: Label, b: Label) -> Label {
        self.union_table.union(a, b)
    }

    /// Get origin descriptions for a label.
    pub fn describe_origins(&self, label: Label) -> Vec<String> {
        self.origins.origins_of(label)
    }
}

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

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

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

    #[test]
    fn test_dfsan_new() {
        let dfsan = DataFlowSanitizer::new();
        assert_eq!(dfsan.shadow_offset, DEFAULT_SHADOW_OFFSET);
        assert_eq!(dfsan.shadow_scale, DEFAULT_SHADOW_SCALE);
        assert_eq!(dfsan.instrumented_count, 0);
        assert!(dfsan.abi_list.is_empty());
    }

    #[test]
    fn test_default() {
        let dfsan = DataFlowSanitizer::default();
        assert_eq!(dfsan.shadow_offset, DEFAULT_SHADOW_OFFSET);
    }

    #[test]
    fn test_union_label_clean() {
        let dfsan = DataFlowSanitizer::new();
        assert_eq!(dfsan.create_union_shadow_label(0, 0), 0);
    }

    #[test]
    fn test_union_label_mixed() {
        let dfsan = DataFlowSanitizer::new();
        assert_eq!(dfsan.create_union_shadow_label(0b01, 0b10), 0b11);
    }

    #[test]
    fn test_union_label_both_set() {
        let dfsan = DataFlowSanitizer::new();
        assert_eq!(dfsan.create_union_shadow_label(0b11, 0b10), 0b11);
    }

    #[test]
    fn test_allocate_label() {
        let mut dfsan = DataFlowSanitizer::new();
        let l1 = dfsan.allocate_label();
        assert_eq!(l1, 1);
        let l2 = dfsan.allocate_label();
        assert_eq!(l2, 2);
        let l3 = dfsan.allocate_label();
        assert_eq!(l3, 4);
    }

    #[test]
    fn test_allocate_label_exhaustion() {
        let mut dfsan = DataFlowSanitizer::new();
        for _ in 0..32 {
            let label = dfsan.allocate_label();
            assert!(label != LABEL_ALL || dfsan.next_label_bit >= 32);
        }
        // Next allocation returns LABEL_ALL
        let label = dfsan.allocate_label();
        assert_eq!(label, LABEL_ALL);
    }

    #[test]
    fn test_abi_list_default() {
        let dfsan = DataFlowSanitizer::new();
        let list = dfsan.create_abi_list();
        assert!(!list.is_empty());
        assert!(list.contains(&"memcpy".to_string()));
        assert!(list.contains(&"malloc".to_string()));
        assert!(list.contains(&"strlen".to_string()));
    }

    #[test]
    fn test_register_abi_function() {
        let mut dfsan = DataFlowSanitizer::new();
        let entry = ABIListEntry {
            name: "custom_func".to_string(),
            arg_rules: vec![ArgLabelRule::Propagate],
            ret_union: true,
            discard: false,
            custom: false,
        };
        dfsan.register_abi_function("custom_func", entry);
        assert!(dfsan.is_in_abi_list("custom_func"));
    }

    #[test]
    fn test_not_in_abi_list() {
        let dfsan = DataFlowSanitizer::new();
        assert!(!dfsan.is_in_abi_list("nonexistent_func"));
    }

    #[test]
    fn test_instrument_empty_module() {
        let mut dfsan = DataFlowSanitizer::new();
        let mut module = llvm_native_core::module::Module::new("test");
        let count = dfsan.instrument(&mut module);
        assert_eq!(count, 0);
    }

    #[test]
    fn test_type_size_integers() {
        assert_eq!(Type::i1().size_bytes(), 1);
        assert_eq!(Type::i8().size_bytes(), 1);
        assert_eq!(Type::i32().size_bytes(), 4);
        assert_eq!(Type::i64().size_bytes(), 8);
    }

    #[test]
    fn test_type_size_void() {
        assert_eq!(Type::void().size_bytes(), 0);
        assert!(Type::void().is_void());
    }

    #[test]
    fn test_shadow_alloca_creation() {
        let dfsan = DataFlowSanitizer::new();
        let val = valref(Value::new(Type::i32()));
        let shadow = dfsan.create_shadow_alloca(&val, 4);
        assert_eq!(shadow.borrow().name, "dfsan_label");
    }

    #[test]
    fn test_abi_list_entry_propagate() {
        let entry = ABIListEntry {
            name: "test".to_string(),
            arg_rules: vec![ArgLabelRule::Propagate, ArgLabelRule::Ignore],
            ret_union: true,
            discard: false,
            custom: false,
        };
        assert_eq!(entry.name, "test");
        assert_eq!(entry.arg_rules.len(), 2);
        assert!(entry.ret_union);
        assert!(!entry.custom);
    }

    #[test]
    fn test_abi_list_entry_custom() {
        let entry = ABIListEntry {
            name: "custom".to_string(),
            arg_rules: Vec::new(),
            ret_union: false,
            discard: true,
            custom: true,
        };
        assert!(entry.discard);
        assert!(entry.custom);
    }
}

use std::rc::Rc;