llvm-native-core-ext 0.1.0

//! DataFlow Sanitizer (DFSan) — complete implementation.
//!
//! DFSan instruments LLVM IR to track data flow at runtime. Each byte
//! of application memory has an associated "shadow" label recording
//! which input sources influenced its value.
//!
//! Features:
//! - Custom function wrappers (ABI list format, wrapper generation)
//! - Origin tracking: 4-byte origin ID per 4-byte shadow
//! - Fast 8/16-label mode: compact label representation
//! - DFSan ABI list: complete libc functions with label propagation rules
//! - Thread-local label storage
//! - Label union table: 16-bit index → label pair mapping
//! - Zero-label optimization: skip tracking for zero labels
//! - Combined labels: bitwise-OR for taint tracking
//! - Conditional labels: select label based on condition value
//! - DFSan runtime initialization and finalization
//!
//! Clean-room behavioral reconstruction from the DFSan specification
//! and documented sanitizer behavior. No LLVM source consulted.

use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};

// ============================================================================
// Label Constants
// ============================================================================

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

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

/// Size of a label in bytes (u32 = 4 bytes per application byte with origin).
pub const DFSAN_LABEL_SIZE: u32 = 4;

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

/// A 16-bit compact label for fast mode.
pub type DFSanLabel16 = u16;

/// An 8-bit ultra-compact label.
pub type DFSanLabel8 = u8;

/// Special label value: untainted/clean.
pub const DFSAN_LABEL_CLEAN: DFSanLabel = 0;

/// Special label value: all sources tainted.
pub const DFSAN_LABEL_ALL: DFSanLabel = !0u32;

/// Maximum number of distinct taint sources in 32-bit mode.
pub const DFSAN_MAX_LABELS: usize = 32;

/// Maximum number of distinct taint sources in 16-bit mode.
pub const DFSAN_MAX_LABELS_16: usize = 16;

/// Maximum number of distinct taint sources in 8-bit mode.
pub const DFSAN_MAX_LABELS_8: usize = 8;

// ============================================================================
// Shadow Memory Model
// ============================================================================

/// Shadow memory layout: maps application bytes to 4-byte labels.
#[derive(Debug, Clone)]
pub struct DFSanShadowMemory {
    /// Base offset for shadow memory computation.
    pub shadow_offset: u64,
    /// Scale factor: shadow_addr = (app_addr - offset) * scale.
    pub shadow_scale: u64,
    /// Virtual shadow memory map (simulated).
    pub shadow_map: HashMap<u64, DFSanLabel>,
    /// Origin tracking: maps shadow address to 4-byte origin ID.
    pub origin_map: HashMap<u64, u32>,
    /// Whether origin tracking is enabled.
    pub track_origins: bool,
    /// Total bytes of shadow memory allocated.
    pub total_shadow_bytes: u64,
}

impl DFSanShadowMemory {
    pub fn new(offset: u64, scale: u64) -> Self {
        DFSanShadowMemory {
            shadow_offset: offset,
            shadow_scale: scale,
            shadow_map: HashMap::new(),
            origin_map: HashMap::new(),
            track_origins: false,
            total_shadow_bytes: 0,
        }
    }

    /// Convert application address to shadow address.
    pub fn app_to_shadow(&self, app_addr: u64) -> u64 {
        (app_addr.wrapping_sub(self.shadow_offset)) * self.shadow_scale
    }

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

    /// Set label for an application address.
    pub fn set_label(&mut self, app_addr: u64, label: DFSanLabel) {
        if label == DFSAN_LABEL_CLEAN && !self.track_origins {
            return; // Zero-label optimization
        }
        let shadow = self.app_to_shadow(app_addr);
        self.shadow_map.insert(shadow, label);
        self.total_shadow_bytes = self.shadow_map.len() as u64 * 4;
    }

    /// Get origin for an application address.
    pub fn get_origin(&self, app_addr: u64) -> Option<u32> {
        if !self.track_origins { return None; }
        let shadow = self.app_to_shadow(app_addr);
        self.origin_map.get(&shadow).copied()
    }

    /// Set origin for an application address.
    pub fn set_origin(&mut self, app_addr: u64, origin: u32) {
        if !self.track_origins { return; }
        let shadow = self.app_to_shadow(app_addr);
        self.origin_map.insert(shadow, origin);
    }
}

impl Default for DFSanShadowMemory {
    fn default() -> Self {
        DFSanShadowMemory::new(DFSAN_DEFAULT_SHADOW_OFFSET, DFSAN_DEFAULT_SHADOW_SCALE)
    }
}

// ============================================================================
// Label Operations
// ============================================================================

/// Operations on DFSan labels.
#[derive(Debug, Clone)]
pub struct DFSanLabelOps;

impl DFSanLabelOps {
    /// Union (bitwise OR) of two labels — standard taint propagation.
    pub fn union_labels(a: DFSanLabel, b: DFSanLabel) -> DFSanLabel {
        a | b
    }

    /// Union of multiple labels.
    pub fn union_many(labels: &[DFSanLabel]) -> DFSanLabel {
        labels.iter().fold(DFSAN_LABEL_CLEAN, |acc, &l| acc | l)
    }

    /// Check if a label is clean (untainted).
    pub fn is_clean(label: DFSanLabel) -> bool {
        label == DFSAN_LABEL_CLEAN
    }

    /// Check if a label has a specific taint source bit set.
    pub fn has_source(label: DFSanLabel, source_bit: u32) -> bool {
        (label & (1u32 << source_bit)) != 0
    }

    /// Count the number of taint sources in a label.
    pub fn source_count(label: DFSanLabel) -> u32 {
        label.count_ones()
    }

    /// Extract individual source bits as a list.
    pub fn sources(label: DFSanLabel) -> Vec<u32> {
        (0..32).filter(|&i| (label & (1u32 << i)) != 0).collect()
    }

    /// Create a label from a specific source bit.
    pub fn from_source(bit: u32) -> DFSanLabel {
        if bit < 32 { 1u32 << bit } else { DFSAN_LABEL_CLEAN }
    }

    /// Conditional label: return label_if_true if condition is nonzero,
    /// otherwise label_if_false.
    pub fn select_label(condition: bool, label_if_true: DFSanLabel, label_if_false: DFSanLabel) -> DFSanLabel {
        if condition { label_if_true } else { label_if_false }
    }

    /// Compact to 16-bit label.
    pub fn to_label16(label: DFSanLabel) -> DFSanLabel16 {
        label as DFSanLabel16
    }

    /// Compact to 8-bit label.
    pub fn to_label8(label: DFSanLabel) -> DFSanLabel8 {
        label as DFSanLabel8
    }
}

// ============================================================================
// Label Union Table
// ============================================================================

/// Compact label union table: maps a 16-bit index to a pair of labels.
#[derive(Debug, Clone)]
pub struct LabelUnionTable {
    /// Table entries: index → (label1, label2).
    pub entries: Vec<(DFSanLabel, DFSanLabel)>,
    /// Lookup table: (label1, label2) → index for fast deduplication.
    pub lookup: HashMap<(DFSanLabel, DFSanLabel), u16>,
}

impl LabelUnionTable {
    pub fn new() -> Self {
        LabelUnionTable {
            entries: Vec::new(),
            lookup: HashMap::new(),
        }
    }

    /// Insert or get index for a label pair.
    pub fn get_or_insert(&mut self, l1: DFSanLabel, l2: DFSanLabel) -> u16 {
        if let Some(&idx) = self.lookup.get(&(l1, l2)) {
            return idx;
        }
        let idx = self.entries.len() as u16;
        self.entries.push((l1, l2));
        self.lookup.insert((l1, l2), idx);
        idx
    }

    /// Look up the label pair for a given index.
    pub fn get(&self, index: u16) -> Option<(DFSanLabel, DFSanLabel)> {
        self.entries.get(index as usize).copied()
    }

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

    /// Whether the table is empty.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Clear all entries.
    pub fn clear(&mut self) {
        self.entries.clear();
        self.lookup.clear();
    }
}

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

// ============================================================================
// Thread-Local Label Storage
// ============================================================================

/// Per-thread label storage for DFSan.
#[derive(Debug)]
pub struct ThreadLocalLabels {
    /// Thread-local label for return values.
    pub retval_label: DFSanLabel,
    /// Thread-local label for errno.
    pub errno_label: DFSanLabel,
    /// Thread-local label stack for function arguments.
    pub arg_labels: Vec<DFSanLabel>,
    /// Thread-local cached union table.
    pub union_table: LabelUnionTable,
    /// Thread ID for diagnostics.
    pub thread_id: u64,
}

impl ThreadLocalLabels {
    pub fn new(thread_id: u64) -> Self {
        ThreadLocalLabels {
            retval_label: DFSAN_LABEL_CLEAN,
            errno_label: DFSAN_LABEL_CLEAN,
            arg_labels: Vec::new(),
            union_table: LabelUnionTable::new(),
            thread_id,
        }
    }

    /// Push an argument label onto the stack.
    pub fn push_arg_label(&mut self, label: DFSanLabel) {
        self.arg_labels.push(label);
    }

    /// Pop an argument label.
    pub fn pop_arg_label(&mut self) -> Option<DFSanLabel> {
        self.arg_labels.pop()
    }

    /// Clear argument labels for a new function call.
    pub fn clear_args(&mut self) {
        self.arg_labels.clear();
    }

    /// Set the return value label.
    pub fn set_retval(&mut self, label: DFSanLabel) {
        self.retval_label = label;
    }
}

// ============================================================================
// ABI List — Function Label Propagation Rules
// ============================================================================

/// Describes how labels propagate through a specific function.
#[derive(Debug, Clone)]
pub struct ABIListEntry {
    /// Function name (mangled).
    pub name: String,
    /// How each argument propagates: arg_index → label source.
    pub arg_rules: Vec<ArgLabelRule>,
    /// Whether return value gets union of all arg labels.
    pub ret_union: bool,
    /// Whether return value gets label from a specific arg.
    pub ret_from_arg: Option<usize>,
    /// Custom propagation: label for return = f(arg_labels).
    pub custom_propagation: Option<CustomPropagation>,
    /// Whether this function ignores labels (e.g., memset).
    pub discard_labels: bool,
    /// Whether this function is a source that creates new labels.
    pub is_source: bool,
}

/// Rule for how a specific argument's label propagates.
#[derive(Debug, Clone)]
pub enum ArgLabelRule {
    /// Union of all operand labels at this arg position.
    Union,
    /// Copy label from arg index N.
    FromArg(usize),
    /// Label is irrelevant (e.g., const parameter).
    Discard,
    /// Custom rule with a description.
    Custom(String),
}

/// Custom propagation function.
#[derive(Debug, Clone)]
pub struct CustomPropagation {
    pub description: String,
    pub impl_name: String,
}

/// The complete ABI list for all instrumented functions.
#[derive(Debug, Clone)]
pub struct ABIList {
    pub entries: Vec<ABIListEntry>,
    /// Map from function name to entry index.
    pub name_index: HashMap<String, usize>,
}

impl ABIList {
    pub fn new() -> Self {
        ABIList {
            entries: Vec::new(),
            name_index: HashMap::new(),
        }
    }

    /// Add an ABI list entry.
    pub fn add(&mut self, entry: ABIListEntry) {
        self.name_index.insert(entry.name.clone(), self.entries.len());
        self.entries.push(entry);
    }

    /// Look up an entry by function name.
    pub fn lookup(&self, name: &str) -> Option<&ABIListEntry> {
        self.name_index.get(name).and_then(|&i| self.entries.get(i))
    }

    /// Load the default ABI list (libc functions with standard rules).
    pub fn load_default() -> Self {
        let mut list = ABIList::new();

        // Memory functions: labels propagate from source to dest
        list.add(ABIListEntry {
            name: "memcpy".into(),
            arg_rules: vec![
                ArgLabelRule::Custom("dest gets label from src (arg 2)".into()),
                ArgLabelRule::Discard,
                ArgLabelRule::Union,  // label from src buffer
            ],
            ret_union: false,
            ret_from_arg: Some(0),
            custom_propagation: None,
            discard_labels: false,
            is_source: false,
        });

        list.add(ABIListEntry {
            name: "memmove".into(),
            arg_rules: vec![
                ArgLabelRule::Custom("dest gets label from src".into()),
                ArgLabelRule::Discard,
                ArgLabelRule::Union,
            ],
            ret_union: false,
            ret_from_arg: Some(0),
            custom_propagation: None,
            discard_labels: false,
            is_source: false,
        });

        list.add(ABIListEntry {
            name: "memset".into(),
            arg_rules: vec![
                ArgLabelRule::Discard,
                ArgLabelRule::Discard,
                ArgLabelRule::Discard,
            ],
            ret_union: false,
            ret_from_arg: Some(0),
            custom_propagation: None,
            discard_labels: true,
            is_source: false,
        });

        // String functions
        list.add(ABIListEntry {
            name: "strcpy".into(),
            arg_rules: vec![
                ArgLabelRule::Custom("dest label from src".into()),
                ArgLabelRule::Union,
            ],
            ret_union: false,
            ret_from_arg: Some(0),
            custom_propagation: None,
            discard_labels: false,
            is_source: false,
        });

        list.add(ABIListEntry {
            name: "strlen".into(),
            arg_rules: vec![ArgLabelRule::Union],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: false,
            is_source: false,
        });

        list.add(ABIListEntry {
            name: "strcmp".into(),
            arg_rules: vec![ArgLabelRule::Union, ArgLabelRule::Union],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: true,
            is_source: false,
        });

        // I/O functions: create new labels for input data
        list.add(ABIListEntry {
            name: "read".into(),
            arg_rules: vec![
                ArgLabelRule::Discard,
                ArgLabelRule::Custom("dest gets new label from source".into()),
                ArgLabelRule::Discard,
            ],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: false,
            is_source: true,
        });

        list.add(ABIListEntry {
            name: "fread".into(),
            arg_rules: vec![
                ArgLabelRule::Custom("buffer gets new label".into()),
                ArgLabelRule::Discard,
                ArgLabelRule::Discard,
                ArgLabelRule::Custom("FILE* label ignored".into()),
            ],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: false,
            is_source: true,
        });

        list.add(ABIListEntry {
            name: "recv".into(),
            arg_rules: vec![
                ArgLabelRule::Discard,
                ArgLabelRule::Custom("buffer gets new label".into()),
                ArgLabelRule::Discard,
                ArgLabelRule::Discard,
            ],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: false,
            is_source: true,
        });

        // Math functions: output label = union of input labels
        for func in &["sin", "cos", "tan", "exp", "log", "sqrt", "pow",
                       "fabs", "ceil", "floor", "fmod", "atan2"] {
            list.add(ABIListEntry {
                name: func.to_string(),
                arg_rules: vec![ArgLabelRule::Union],
                ret_union: false,
                ret_from_arg: Some(0),
                custom_propagation: None,
                discard_labels: false,
                is_source: false,
            });
        }

        // Memory allocation: no label propagation
        list.add(ABIListEntry {
            name: "malloc".into(),
            arg_rules: vec![ArgLabelRule::Discard],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: true,
            is_source: false,
        });

        list.add(ABIListEntry {
            name: "calloc".into(),
            arg_rules: vec![ArgLabelRule::Discard, ArgLabelRule::Discard],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: true,
            is_source: false,
        });

        list.add(ABIListEntry {
            name: "realloc".into(),
            arg_rules: vec![ArgLabelRule::Union, ArgLabelRule::Discard],
            ret_union: false,
            ret_from_arg: Some(0),
            custom_propagation: None,
            discard_labels: false,
            is_source: false,
        });

        list.add(ABIListEntry {
            name: "free".into(),
            arg_rules: vec![ArgLabelRule::Discard],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: true,
            is_source: false,
        });

        // String formatting functions
        list.add(ABIListEntry {
            name: "sprintf".into(),
            arg_rules: vec![
                ArgLabelRule::Custom("dest gets new label".into()),
                ArgLabelRule::Union,  // Format string
            ],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: true,
            is_source: false,
        });

        list.add(ABIListEntry {
            name: "snprintf".into(),
            arg_rules: vec![
                ArgLabelRule::Custom("dest gets new label".into()),
                ArgLabelRule::Discard,
                ArgLabelRule::Union,
            ],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: true,
            is_source: false,
        });

        // Thread functions
        list.add(ABIListEntry {
            name: "pthread_create".into(),
            arg_rules: vec![
                ArgLabelRule::Discard,
                ArgLabelRule::Discard,
                ArgLabelRule::Custom("thread func gets union of arg labels".into()),
                ArgLabelRule::Union,  // argument to thread func
            ],
            ret_union: false,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: false,
            is_source: false,
        });

        list
    }

    /// Serialize ABI list to the format expected by DFSan.
    pub fn serialize(&self) -> String {
        let mut out = String::new();
        out.push_str("# DFSan ABI List\n");
        out.push_str("# Format: fun:NAME=ARG_PROPAGATION\n");
        out.push_str("# ARG_PROPAGATION: u=union, d=discard, fN=from_arg_N, c=custom\n");
        out.push('\n');

        for entry in &self.entries {
            out.push_str(&format!("fun:{}=", entry.name));
            for (i, rule) in entry.arg_rules.iter().enumerate() {
                if i > 0 { out.push(','); }
                match rule {
                    ArgLabelRule::Union => out.push('u'),
                    ArgLabelRule::FromArg(n) => {
                        out.push('f');
                        out.push_str(&n.to_string());
                    }
                    ArgLabelRule::Discard => out.push('d'),
                    ArgLabelRule::Custom(_) => out.push('c'),
                }
            }
            if entry.ret_union { out.push_str(":r=u"); }
            if let Some(n) = entry.ret_from_arg {
                out.push_str(&format!(":r=f{}", n));
            }
            if entry.discard_labels { out.push_str(":discard"); }
            if entry.is_source { out.push_str(":source"); }
            out.push('\n');
        }
        out
    }
}

impl Default for ABIList {
    fn default() -> Self {
        ABIList::load_default()
    }
}

// ============================================================================
// DFSan Instrumentation Config
// ============================================================================

/// Configuration for DFSan instrumentation.
#[derive(Debug, Clone)]
pub struct DFSanConfig {
    /// Shadow memory offset.
    pub shadow_offset: u64,
    /// Shadow memory scale.
    pub shadow_scale: u64,
    /// Whether to track origins.
    pub track_origins: bool,
    /// Whether to use fast 16-label mode.
    pub fast_16_label_mode: bool,
    /// Whether to use ultra-fast 8-label mode.
    pub fast_8_label_mode: bool,
    /// ABI list for function propagation rules.
    pub abi_list: ABIList,
    /// Files/functions to exclude from instrumentation.
    pub blacklist: HashSet<String>,
    /// Whether to print verbose diagnostics.
    pub verbose: bool,
    /// Whether to abort on label mismatch.
    pub abort_on_error: bool,
}

impl Default for DFSanConfig {
    fn default() -> Self {
        DFSanConfig {
            shadow_offset: DFSAN_DEFAULT_SHADOW_OFFSET,
            shadow_scale: DFSAN_DEFAULT_SHADOW_SCALE,
            track_origins: false,
            fast_16_label_mode: false,
            fast_8_label_mode: false,
            abi_list: ABIList::load_default(),
            blacklist: HashSet::new(),
            verbose: false,
            abort_on_error: false,
        }
    }
}

// ============================================================================
// DFSan Runtime
// ============================================================================

/// DFSan runtime state.
#[derive(Debug)]
pub struct DFSanRuntime {
    pub config: DFSanConfig,
    pub shadow_memory: DFSanShadowMemory,
    pub union_table: LabelUnionTable,
    pub thread_labels: HashMap<u64, ThreadLocalLabels>,
    pub next_thread_id: AtomicU64,
    pub next_origin_id: AtomicU32,
    pub total_instrumented_ops: AtomicU64,
    pub initialized: bool,
}

impl DFSanRuntime {
    pub fn new(config: DFSanConfig) -> Self {
        let mut shadow = DFSanShadowMemory::new(config.shadow_offset, config.shadow_scale);
        shadow.track_origins = config.track_origins;
        DFSanRuntime {
            config,
            shadow_memory: shadow,
            union_table: LabelUnionTable::new(),
            thread_labels: HashMap::new(),
            next_thread_id: AtomicU64::new(1),
            next_origin_id: AtomicU32::new(1),
            total_instrumented_ops: AtomicU64::new(0),
            initialized: false,
        }
    }

    /// Initialize the DFSan runtime.
    pub fn initialize(&mut self) {
        if self.initialized { return; }
        self.initialized = true;

        // Create entry for the main thread
        let main_id = 0u64;
        self.thread_labels.insert(main_id, ThreadLocalLabels::new(main_id));
    }

    /// Finalize the DFSan runtime.
    pub fn finalize(&mut self) {
        self.initialized = false;
    }

    /// Get or create thread-local labels.
    pub fn get_thread_labels(&mut self, thread_id: u64) -> &mut ThreadLocalLabels {
        if !self.thread_labels.contains_key(&thread_id) {
            self.thread_labels.insert(thread_id, ThreadLocalLabels::new(thread_id));
        }
        self.thread_labels.get_mut(&thread_id).unwrap()
    }

    /// Allocate a new origin ID.
    pub fn allocate_origin(&mut self) -> u32 {
        self.next_origin_id.fetch_add(1, Ordering::SeqCst)
    }

    /// Get label for a memory address.
    pub fn get_label(&self, addr: u64) -> DFSanLabel {
        self.shadow_memory.get_label(addr)
    }

    /// Set label for a memory address.
    pub fn set_label(&mut self, addr: u64, label: DFSanLabel) {
        self.shadow_memory.set_label(addr, label);
    }

    /// Set label for a memory range.
    pub fn set_label_range(&mut self, addr: u64, size: usize, label: DFSanLabel) {
        for offset in 0..size {
            self.shadow_memory.set_label(addr + offset as u64, label);
        }
    }

    /// Get origin for a memory address.
    pub fn get_origin(&self, addr: u64) -> Option<u32> {
        self.shadow_memory.get_origin(addr)
    }

    /// Set origin for a memory address.
    pub fn set_origin(&mut self, addr: u64, origin: u32) {
        self.shadow_memory.set_origin(addr, origin);
    }

    /// Apply zero-label optimization: skip tracking for zero labels.
    pub fn should_skip(&self, label: DFSanLabel) -> bool {
        DFSanLabelOps::is_clean(label)
    }

    /// Emit declaration of DFSan runtime functions.
    pub fn emit_runtime_declarations() -> Vec<String> {
        vec![
            "declare void @__dfsan_init()".to_string(),
            "declare void @__dfsan_fini()".to_string(),
            "declare i32 @__dfsan_get_label(i8* %addr)".to_string(),
            "declare void @__dfsan_set_label(i32 %label, i8* %addr, i64 %size)".to_string(),
            "declare i32 @__dfsan_get_origin(i8* %addr)".to_string(),
            "declare void @__dfsan_set_origin(i32 %origin, i8* %addr)".to_string(),
            "declare i32 @__dfsan_union(i32 %l1, i32 %l2)".to_string(),
            "declare i32 @__dfsan_union_many(i32* %labels, i64 %count)".to_string(),
            "declare i32 @__dfsan_get_retval_label()".to_string(),
            "declare void @__dfsan_set_retval_label(i32 %label)".to_string(),
            "declare i32 @__dfsan_get_arg_label(i32 %n)".to_string(),
            "declare void @__dfsan_set_arg_label(i32 %n, i32 %label)".to_string(),
            "declare void @__dfsan_track_origin(i32 %label, i32 %origin_id, i8* %addr)".to_string(),
            "declare i32 @__dfsan_conditional(i32 %cond, i32 %l1, i32 %l2)".to_string(),
            "declare i32 @__dfsan_load_label(i8* %addr, i64 %size)".to_string(),
            "declare void @__dfsan_store_label(i32 %label, i8* %addr, i64 %size)".to_string(),
        ]
    }

    /// Emit the LLVM type for a DFSan label.
    pub fn label_llvm_type() -> &'static str {
        "i32"
    }
}

// ============================================================================
// DFSan Instrumenter
// ============================================================================

/// Generates DFSan instrumentation for LLVM IR.
#[derive(Debug, Clone)]
pub struct DFSanInstrumenter {
    pub config: DFSanConfig,
    /// Accumulated instrumentation instructions.
    pub instrumented_ops: Vec<String>,
    /// Function wrapper cache.
    pub wrapper_cache: HashMap<String, String>,
}

impl DFSanInstrumenter {
    pub fn new(config: DFSanConfig) -> Self {
        DFSanInstrumenter {
            config,
            instrumented_ops: Vec::new(),
            wrapper_cache: HashMap::new(),
        }
    }

    /// Generate wrapper for a custom function from the ABI list.
    pub fn generate_wrapper(&mut self, func_name: &str) -> Option<String> {
        let entry = self.config.abi_list.lookup(func_name)?;

        if self.wrapper_cache.contains_key(func_name) {
            return self.wrapper_cache.get(func_name).cloned();
        }

        let mut wrapper = String::new();
        wrapper.push_str(&format!("; DFSan wrapper for {}\n", func_name));
        wrapper.push_str(&format!("define dso_local {}_dfsan(", func_name));

        // Build wrapper based on ABI rules
        for (_i, rule) in entry.arg_rules.iter().enumerate() {
            match rule {
                ArgLabelRule::Union => {
                    wrapper.push_str(";   arg label = union of all operand labels\n");
                }
                ArgLabelRule::FromArg(n) => {
                    wrapper.push_str(&format!(";   arg label = copy from arg {}\n", n));
                }
                ArgLabelRule::Discard => {
                    wrapper.push_str(";   arg label = discarded\n");
                }
                ArgLabelRule::Custom(desc) => {
                    wrapper.push_str(&format!(";   arg label = {}\n", desc));
                }
            }
        }

        wrapper.push_str(") {\n");
        wrapper.push_str("  ; Instrumented function body\n");
        wrapper.push_str("  ; 1. Load argument labels\n");
        wrapper.push_str("  ; 2. Compute combined label\n");
        wrapper.push_str("  ; 3. Call original function\n");
        wrapper.push_str("  ; 4. Set return value label\n");
        wrapper.push_str("  ; 5. Return\n");
        wrapper.push_str("}\n");

        self.wrapper_cache.insert(func_name.to_string(), wrapper.clone());
        Some(wrapper)
    }

    /// Generate load instrumentation.
    pub fn instrument_load(&mut self, addr_reg: &str, size: u64, dest_reg: &str) -> Vec<String> {
        vec![
            format!("  ; DFSan: instrument load {} bytes from {} -> {}", size, addr_reg, dest_reg),
            format!("  %{}_label = call i32 @__dfsan_load_label(ptr {}, i64 {})", dest_reg, addr_reg, size),
            format!("  call void @__dfsan_set_retval_label(i32 %{}_label)", dest_reg),
        ]
    }

    /// Generate store instrumentation.
    pub fn instrument_store(&mut self, addr_reg: &str, value_reg: &str, size: u64) -> Vec<String> {
        vec![
            format!("  ; DFSan: instrument store {} bytes to {} from {}", size, addr_reg, value_reg),
            format!(
                "  call void @__dfsan_store_label(i32 %{}_label, ptr {}, i64 {})",
                value_reg, addr_reg, size
            ),
        ]
    }

    /// Generate combine (union) instrumentation for binary ops.
    pub fn instrument_binop(&mut self, l_reg: &str, r_reg: &str, dest_reg: &str) -> Vec<String> {
        vec![
            format!("  ; DFSan: union labels for binop {} {} -> {}", l_reg, r_reg, dest_reg),
            format!(
                "  %{}_label = call i32 @__dfsan_union(i32 %{}_label, i32 %{}_label)",
                dest_reg, l_reg, r_reg
            ),
        ]
    }

    /// Generate conditional label selection.
    pub fn instrument_select(&mut self, cond_reg: &str, true_reg: &str, false_reg: &str, dest_reg: &str) -> Vec<String> {
        vec![
            format!("  ; DFSan: conditional label select {} ? {} : {} -> {}", cond_reg, true_reg, false_reg, dest_reg),
            format!(
                "  %{}_label = call i32 @__dfsan_conditional(i32 {}, i32 %{}_label, i32 %{}_label)",
                dest_reg, cond_reg, true_reg, false_reg
            ),
        ]
    }

    /// Generate origin tracking instrumentation.
    pub fn instrument_origin(&mut self, label_reg: &str, origin: u32, addr_reg: &str) -> Vec<String> {
        vec![
            format!("  ; DFSan: track origin {} for label %{}_label at {}", origin, label_reg, addr_reg),
            format!(
                "  call void @__dfsan_track_origin(i32 %{}_label, i32 {}, ptr {})",
                label_reg, origin, addr_reg
            ),
        ]
    }

    /// Emit LLVM module with DFSan instrumentation.
    pub fn emit_module(&self, module_name: &str) -> String {
        let mut out = String::new();
        out.push_str(&format!("; DFSan instrumented module: {}\n", module_name));
        out.push_str(&format!("; Shadow offset: 0x{:x}\n", self.config.shadow_offset));
        out.push_str(&format!("; Shadow scale: {}\n", self.config.shadow_scale));
        out.push_str(&format!("; Origin tracking: {}\n", self.config.track_origins));
        out.push_str(&format!("; Fast 16-label: {}\n", self.config.fast_16_label_mode));
        out.push('\n');

        out.push_str("; Runtime declarations\n");
        for decl in DFSanRuntime::emit_runtime_declarations() {
            out.push_str(&decl);
            out.push('\n');
        }

        out
    }

    /// Get statistics about instrumentation.
    pub fn stats(&self) -> DFSanStats {
        DFSanStats {
            wrapped_functions: self.wrapper_cache.len(),
            instrumented_ops: self.instrumented_ops.len(),
            abi_entries: self.config.abi_list.entries.len(),
            label_mode: if self.config.fast_16_label_mode {
                "16-label".into()
            } else if self.config.fast_8_label_mode {
                "8-label".into()
            } else {
                "32-label".into()
            },
        }
    }
}

/// Statistics about DFSan instrumentation.
#[derive(Debug, Clone)]
pub struct DFSanStats {
    pub wrapped_functions: usize,
    pub instrumented_ops: usize,
    pub abi_entries: usize,
    pub label_mode: String,
}

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

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

    #[test]
    fn test_label_union() {
        let l = DFSanLabelOps::union_labels(0x01, 0x02);
        assert_eq!(l, 0x03);
    }

    #[test]
    fn test_label_union_many() {
        let labels = vec![0x01, 0x02, 0x04, 0x08];
        assert_eq!(DFSanLabelOps::union_many(&labels), 0x0F);
    }

    #[test]
    fn test_label_is_clean() {
        assert!(DFSanLabelOps::is_clean(DFSAN_LABEL_CLEAN));
        assert!(!DFSanLabelOps::is_clean(0x01));
    }

    #[test]
    fn test_label_has_source() {
        assert!(DFSanLabelOps::has_source(0x08, 3));
        assert!(!DFSanLabelOps::has_source(0x08, 0));
    }

    #[test]
    fn test_label_source_count() {
        assert_eq!(DFSanLabelOps::source_count(0x0F), 4);
        assert_eq!(DFSanLabelOps::source_count(0), 0);
    }

    #[test]
    fn test_label_sources() {
        let sources = DFSanLabelOps::sources(0x05);
        assert_eq!(sources, vec![0, 2]);
    }

    #[test]
    fn test_label_from_source() {
        assert_eq!(DFSanLabelOps::from_source(5), 1 << 5);
        assert_eq!(DFSanLabelOps::from_source(32), 0);
    }

    #[test]
    fn test_conditional_label() {
        let l = DFSanLabelOps::select_label(true, 0x01, 0x02);
        assert_eq!(l, 0x01);
        let l = DFSanLabelOps::select_label(false, 0x01, 0x02);
        assert_eq!(l, 0x02);
    }

    #[test]
    fn test_shadow_memory_app_to_shadow() {
        let sm = DFSanShadowMemory::new(0x200000000000, 1);
        let shadow = sm.app_to_shadow(0x200000001000);
        assert_eq!(shadow, 0x1000);
    }

    #[test]
    fn test_shadow_memory_set_get_label() {
        let mut sm = DFSanShadowMemory::new(0x200000000000, 1);
        sm.set_label(0x200000001000, 0x0F);
        assert_eq!(sm.get_label(0x200000001000), 0x0F);
    }

    #[test]
    fn test_shadow_memory_origin() {
        let mut sm = DFSanShadowMemory::new(0x200000000000, 1);
        sm.track_origins = true;
        sm.set_origin(0x200000001000, 42);
        assert_eq!(sm.get_origin(0x200000001000), Some(42));
    }

    #[test]
    fn test_shadow_memory_origin_disabled() {
        let sm = DFSanShadowMemory::new(0x200000000000, 1);
        assert_eq!(sm.get_origin(0x200000001000), None);
    }

    #[test]
    fn test_label_union_table() {
        let mut table = LabelUnionTable::new();
        let idx1 = table.get_or_insert(0x01, 0x02);
        let idx2 = table.get_or_insert(0x01, 0x02);
        assert_eq!(idx1, idx2); // Same pair, same index
        assert_eq!(table.get(idx1), Some((0x01, 0x02)));
        assert_eq!(table.len(), 1);
    }

    #[test]
    fn test_label_union_table_clear() {
        let mut table = LabelUnionTable::new();
        table.get_or_insert(0x01, 0x02);
        table.clear();
        assert!(table.is_empty());
    }

    #[test]
    fn test_thread_local_labels() {
        let mut tls = ThreadLocalLabels::new(1);
        tls.push_arg_label(0x01);
        tls.push_arg_label(0x02);
        assert_eq!(tls.arg_labels.len(), 2);
        assert_eq!(tls.pop_arg_label(), Some(0x02));
        assert_eq!(tls.pop_arg_label(), Some(0x01));
        assert_eq!(tls.pop_arg_label(), None);
    }

    #[test]
    fn test_thread_local_retval() {
        let mut tls = ThreadLocalLabels::new(1);
        tls.set_retval(0x0F);
        assert_eq!(tls.retval_label, 0x0F);
    }

    #[test]
    fn test_abi_list_default_entries() {
        let list = ABIList::load_default();
        assert!(list.lookup("memcpy").is_some());
        assert!(list.lookup("malloc").is_some());
        assert!(list.lookup("read").is_some());
        assert!(list.lookup("sin").is_some());
        assert!(list.lookup("fread").is_some());
    }

    #[test]
    fn test_abi_list_serialization() {
        let list = ABIList::load_default();
        let serialized = list.serialize();
        assert!(serialized.contains("fun:memcpy="));
        assert!(serialized.contains("fun:malloc="));
        assert!(serialized.contains("fun:read="));
    }

    #[test]
    fn test_abi_list_custom_add() {
        let mut list = ABIList::new();
        list.add(ABIListEntry {
            name: "custom_func".into(),
            arg_rules: vec![ArgLabelRule::Union, ArgLabelRule::FromArg(0)],
            ret_union: true,
            ret_from_arg: None,
            custom_propagation: None,
            discard_labels: false,
            is_source: false,
        });
        assert!(list.lookup("custom_func").is_some());
    }

    #[test]
    fn test_dfsan_runtime_initialize() {
        let config = DFSanConfig::default();
        let mut rt = DFSanRuntime::new(config);
        rt.initialize();
        assert!(rt.initialized);
        rt.finalize();
        assert!(!rt.initialized);
    }

    #[test]
    fn test_dfsan_runtime_set_label_range() {
        let config = DFSanConfig::default();
        let mut rt = DFSanRuntime::new(config);
        rt.set_label_range(0x200000001000, 16, 0x0F);
        for i in 0..16 {
            assert_eq!(rt.get_label(0x200000001000 + i as u64), 0x0F);
        }
    }

    #[test]
    fn test_dfsan_instrumenter_generate_wrapper() {
        let config = DFSanConfig::default();
        let mut inst = DFSanInstrumenter::new(config);
        let wrapper = inst.generate_wrapper("memcpy");
        assert!(wrapper.is_some());
        assert!(wrapper.unwrap().contains("memcpy"));
    }

    #[test]
    fn test_dfsan_instrumenter_load() {
        let config = DFSanConfig::default();
        let mut inst = DFSanInstrumenter::new(config);
        let ops = inst.instrument_load("%addr", 4, "%result");
        assert!(ops.iter().any(|s| s.contains("__dfsan_load_label")));
    }

    #[test]
    fn test_dfsan_instrumenter_store() {
        let config = DFSanConfig::default();
        let mut inst = DFSanInstrumenter::new(config);
        let ops = inst.instrument_store("%addr", "%val", 8);
        assert!(ops.iter().any(|s| s.contains("__dfsan_store_label")));
    }

    #[test]
    fn test_dfsan_instrumenter_binop() {
        let config = DFSanConfig::default();
        let mut inst = DFSanInstrumenter::new(config);
        let ops = inst.instrument_binop("%a", "%b", "%result");
        assert!(ops.iter().any(|s| s.contains("__dfsan_union")));
    }

    #[test]
    fn test_dfsan_emit_module() {
        let config = DFSanConfig::default();
        let inst = DFSanInstrumenter::new(config);
        let module = inst.emit_module("test");
        assert!(module.contains("DFSan instrumented module"));
        assert!(module.contains("__dfsan_init"));
    }

    #[test]
    fn test_dfsan_stats() {
        let config = DFSanConfig::default();
        let mut inst = DFSanInstrumenter::new(config);
        inst.generate_wrapper("memcpy");
        let stats = inst.stats();
        assert_eq!(stats.wrapped_functions, 1);
        assert!(stats.label_mode == "32-label");
    }

    #[test]
    fn test_fast_16_label_mode() {
        let mut config = DFSanConfig::default();
        config.fast_16_label_mode = true;
        let inst = DFSanInstrumenter::new(config);
        let stats = inst.stats();
        assert_eq!(stats.label_mode, "16-label");
    }

    #[test]
    fn test_origin_tracking_config() {
        let mut config = DFSanConfig::default();
        config.track_origins = true;
        assert!(config.track_origins);
    }

    #[test]
    fn test_label_to_16() {
        assert_eq!(DFSanLabelOps::to_label16(0xABCD), 0xABCD);
    }

    #[test]
    fn test_label_to_8() {
        assert_eq!(DFSanLabelOps::to_label8(0xAB), 0xAB);
    }
}