nabla-decompiler 0.1.2

//! Function identification and analysis

use anyhow::Result;
use std::collections::HashSet;
use uuid::Uuid;

use crate::cfg_analyzer::{CfgAnalyzer, CfgAnalyzerRegistry};
use crate::types::{
    Address, BasicBlock, BlockType, Disassembly, Function, 
    Instruction, InstructionGroup
};

/// Identify functions from disassembly using YARA rules for better accuracy
pub fn identify_functions_with_yara(disasm: &Disassembly, arch: &str, format: &str) -> Result<Vec<Function>> {
    let registry = CfgAnalyzerRegistry::new();
    let analyzer = registry.get_analyzer(arch, format)
        .ok_or_else(|| anyhow::anyhow!("No analyzer found for architecture: {} format: {}", arch, format))?;
    
    let mut functions = Vec::new();
    let mut processed_addresses = HashSet::new();
    
    // Find function entry points using YARA rules
    let entry_points = find_function_entry_points_with_yara(disasm, analyzer)?;
    
    for &entry_point in &entry_points {
        if processed_addresses.contains(&entry_point) {
            continue;
        }
        
        if let Ok(function) = analyze_function_from_entry(disasm, entry_point, analyzer) {
            // Mark all addresses in this function as processed
            for &addr in &function.instructions {
                processed_addresses.insert(addr);
            }
            functions.push(function);
        }
    }
    
    Ok(functions)
}

/// Identify functions from disassembly (fallback with default ARM analyzer)
pub fn identify_functions(disasm: &Disassembly) -> Result<Vec<Function>> {
    identify_functions_with_yara(disasm, "arm", "elf")
}

/// Analyze a single function starting from an address with architecture detection
pub fn analyze_single_function_with_arch(disasm: &Disassembly, address: Address, arch: &str, format: &str) -> Result<Function> {
    let registry = CfgAnalyzerRegistry::new();
    let analyzer = registry.get_analyzer(arch, format)
        .ok_or_else(|| anyhow::anyhow!("No analyzer found for architecture: {} format: {}", arch, format))?;
    
    analyze_function_from_entry(disasm, address, analyzer)
}

/// Analyze a single function starting from an address (fallback with default ARM analyzer)
pub fn analyze_single_function(disasm: &Disassembly, address: Address) -> Result<Function> {
    analyze_single_function_with_arch(disasm, address, "arm", "elf")
}

/// Find potential function entry points using YARA rules and architecture-specific analyzer
fn find_function_entry_points_with_yara(disasm: &Disassembly, analyzer: &dyn CfgAnalyzer) -> Result<Vec<Address>> {
    let mut entry_points = Vec::new();
    
    // Use YARA rules to classify symbols and filter for functions only
    let function_symbols = classify_function_symbols_with_yara(&disasm.symbols)?;
    entry_points.extend(function_symbols);
    
    // Find call targets using analyzer
    for instruction in &disasm.instructions {
        if analyzer.classify_instruction(instruction) == InstructionGroup::Call {
            if let Some(target) = analyzer.extract_jump_target(instruction) {
                entry_points.push(target);
            }
        }
    }
    
    // Look for function prologues (common patterns)
    for window in disasm.instructions.windows(2) {
        if is_function_prologue(&window) {
            entry_points.push(window[0].address);
        }
    }
    
    entry_points.sort();
    entry_points.dedup();
    Ok(entry_points)
}

/// Analyze function starting from entry point using architecture-specific analyzer
fn analyze_function_from_entry(disasm: &Disassembly, entry_point: Address, analyzer: &dyn CfgAnalyzer) -> Result<Function> {
    let mut instructions = Vec::new();
    let mut exit_points = Vec::new();
    let mut calls = Vec::new();
    let _current_addr = entry_point;
    
    // Simple linear sweep to find function boundaries
    let mut visited = HashSet::new();
    let mut to_process = vec![entry_point];
    
    while let Some(addr) = to_process.pop() {
        if visited.contains(&addr) {
            continue;
        }
        visited.insert(addr);
        
        if let Some(instruction) = find_instruction_at_address(disasm, addr) {
            instructions.push(instruction.address);
            
            let instruction_group = analyzer.classify_instruction(instruction);
            match instruction_group {
                InstructionGroup::Return => {
                    exit_points.push(instruction.address);
                    // Don't follow returns
                }
                InstructionGroup::Call => {
                    if let Some(target) = analyzer.extract_jump_target(instruction) {
                        calls.push(target);
                    }
                    // Continue to next instruction after call
                    if let Some(next_addr) = analyzer.fall_through_address(instruction) {
                        to_process.push(next_addr);
                    }
                }
                InstructionGroup::Jump => {
                    if let Some(target) = analyzer.extract_jump_target(instruction) {
                        to_process.push(target);
                    }
                    // For conditional jumps, also follow fall-through
                    if analyzer.is_conditional(instruction) {
                        if let Some(next_addr) = analyzer.fall_through_address(instruction) {
                            to_process.push(next_addr);
                        }
                    }
                }
                _ => {
                    // Continue to next instruction
                    if let Some(next_addr) = analyzer.fall_through_address(instruction) {
                        to_process.push(next_addr);
                    }
                }
            }
        } else {
            // No instruction found, probably end of function
            break;
        }
        
        // Limit function size to prevent runaway analysis
        if instructions.len() > 1000 {
            break;
        }
    }
    
    instructions.sort();
    
    // Build basic blocks with analyzer
    let basic_blocks = build_basic_blocks_with_analyzer(&instructions, disasm, analyzer)?;
    
    let function_size = if let (Some(&first), Some(&last)) = (instructions.first(), instructions.last()) {
        last - first + 4 // Approximate size
    } else {
        0
    };
    
    // Try to get function name from symbols with fuzzy lookup
    tracing::debug!("Looking for symbol at function address 0x{:x}, total symbols: {}", entry_point, disasm.symbols.len());
    if disasm.symbols.len() > 0 && disasm.symbols.len() < 20 {
        for (&addr, name) in &disasm.symbols {
            tracing::debug!("  Available symbol: '{}' at 0x{:x}", name, addr);
        }
    }
    let name = find_closest_symbol(&disasm.symbols, entry_point);
    
    Ok(Function {
        address: entry_point,
        name,
        size: function_size,
        instructions,
        basic_blocks,
        entry_point,
        exit_points,
        calls,
        called_by: Vec::new(), // Will be filled in later pass
    })
}

/// Build basic blocks for a function using architecture-specific analyzer
fn build_basic_blocks_with_analyzer(instructions: &[Address], disasm: &Disassembly, analyzer: &dyn CfgAnalyzer) -> Result<Vec<BasicBlock>> {
    if instructions.is_empty() {
        tracing::debug!("No instructions provided for basic block analysis");
        return Ok(Vec::new());
    }
    
    tracing::debug!("Building basic blocks from {} instructions", instructions.len());
    
    // Find basic block boundaries
    let mut block_starts = HashSet::new();
    block_starts.insert(instructions[0]); // Function entry
    
    tracing::debug!("Initial block start at 0x{:x}", instructions[0]);
    
    // Add targets of jumps and branches using analyzer
    let mut jump_targets_found = 0;
    let mut valid_instruction_count = 0;
    
    for &addr in instructions {
        if let Some(instruction) = find_instruction_at_address(disasm, addr) {
            let instruction_group = analyzer.classify_instruction(instruction);
            tracing::debug!("Analyzing instruction at 0x{:x}: {} {} (group: {:?})", 
                      addr, instruction.mnemonic, instruction.operands, instruction_group);
            
            // Count valid instructions (filter out data/invalid instructions)
            if !instruction.mnemonic.is_empty() && instruction.mnemonic != "???" {
                valid_instruction_count += 1;
            }
            
            match instruction_group {
                InstructionGroup::Jump | InstructionGroup::Call => {
                    // Next instruction after jump/call is a block start
                    if let Some(next_addr) = analyzer.fall_through_address(instruction) {
                        if instructions.contains(&next_addr) {
                            tracing::debug!("  Adding fall-through block start at 0x{:x}", next_addr);
                            block_starts.insert(next_addr);
                        }
                    }
                    
                    // Jump target is a block start
                    if let Some(target) = analyzer.extract_jump_target(instruction) {
                        if instructions.contains(&target) {
                            tracing::debug!("  Adding jump target block start at 0x{:x}", target);
                            block_starts.insert(target);
                            jump_targets_found += 1;
                        } else {
                            tracing::debug!("  Jump target 0x{:x} is outside function", target);
                        }
                    }
                }
                InstructionGroup::Return => {
                    // Next instruction after return is a block start (if it exists)
                    if let Some(next_addr) = analyzer.fall_through_address(instruction) {
                        if instructions.contains(&next_addr) {
                            tracing::debug!("  Adding post-return block start at 0x{:x}", next_addr);
                            block_starts.insert(next_addr);
                        }
                    }
                }
                _ => {}
            }
        }
    }
    
    // If we have a very long sequence of instructions with no jumps detected,
    // it might be data misinterpreted as instructions. Split it artificially for better visualization.
    if jump_targets_found == 0 && valid_instruction_count > 20 {
        tracing::debug!("Long instruction sequence detected ({} instructions), adding artificial splits", valid_instruction_count);
        
        // Split every 10-15 instructions for better CFG visualization
        let chunk_size = 12;
        for (i, &addr) in instructions.iter().enumerate() {
            if i > 0 && i % chunk_size == 0 && i < instructions.len() - 1 {
                tracing::debug!("  Adding artificial block split at 0x{:x} (position {})", addr, i);
                block_starts.insert(addr);
                jump_targets_found += 1; // Count as found to avoid warnings
            }
        }
    }
    
    tracing::debug!("Found {} jump targets, total block starts: {}", jump_targets_found, block_starts.len());
    
    let mut block_starts: Vec<Address> = block_starts.into_iter().collect();
    block_starts.sort();
    
    tracing::debug!("Block starts: {:?}", block_starts.iter().map(|&addr| format!("0x{:x}", addr)).collect::<Vec<_>>());
    
    // Create basic blocks
    let mut basic_blocks = Vec::new();
    
    for i in 0..block_starts.len() {
        let start_addr = block_starts[i];
        let end_addr = if i + 1 < block_starts.len() {
            block_starts[i + 1]
        } else {
            instructions.last().copied().unwrap_or(start_addr) + 4
        };
        
        let block_instructions: Vec<Address> = instructions.iter()
            .filter(|&&addr| addr >= start_addr && addr < end_addr)
            .copied()
            .collect();
        
        if !block_instructions.is_empty() {
            let block_type = if i == 0 {
                BlockType::Entry
            } else if i == block_starts.len() - 1 {
                BlockType::Exit
            } else {
                BlockType::Normal
            };
            
            tracing::debug!("Created basic block {}: 0x{:x}-0x{:x} ({} instructions, type: {:?})", 
                      i, start_addr, end_addr, block_instructions.len(), block_type);
            
            basic_blocks.push(BasicBlock {
                id: Uuid::new_v4(),
                start_address: start_addr,
                end_address: end_addr,
                instructions: block_instructions,
                predecessors: Vec::new(), // Will be filled in later
                successors: Vec::new(),   // Will be filled in later
                block_type,
            });
        }
    }
    
    tracing::debug!("Final result: {} basic blocks created", basic_blocks.len());
    
    Ok(basic_blocks)
}


/// Find instruction at specific address
fn find_instruction_at_address(disasm: &Disassembly, address: Address) -> Option<&Instruction> {
    disasm.instructions.iter().find(|insn| insn.address == address)
}

/// Check if instructions form a function prologue
fn is_function_prologue(instructions: &[Instruction]) -> bool {
    if instructions.len() < 2 {
        return false;
    }
    
    let first = &instructions[0];
    let second = &instructions[1];
    
    // x86/x64 function prologue patterns
    if (first.mnemonic == "push" && first.operands.contains("ebp")) &&
       (second.mnemonic == "mov" && second.operands.contains("ebp") && second.operands.contains("esp")) {
        return true; // push ebp; mov ebp, esp (x86)
    }
    
    if (first.mnemonic == "push" && first.operands.contains("rbp")) &&
       (second.mnemonic == "mov" && second.operands.contains("rbp") && second.operands.contains("rsp")) {
        return true; // push rbp; mov rbp, rsp (x64)
    }
    
    // Common x86 stack allocation
    if first.mnemonic == "sub" && first.operands.contains("esp") {
        return true; // sub esp, imm
    }
    
    if first.mnemonic == "sub" && first.operands.contains("rsp") {
        return true; // sub rsp, imm
    }
    
    // Single push ebp/rbp (common in optimized code)
    if first.mnemonic == "push" && (first.operands.contains("ebp") || first.operands.contains("rbp")) {
        return true;
    }
    
    // Common ARM function prologue patterns
    // push {r4, r5, r6, r7, lr} or similar
    (first.mnemonic == "push" && first.operands.contains("lr")) ||
    // stmfd sp!, {r4, r5, r6, r7, lr}
    (first.mnemonic == "stmfd" && first.operands.contains("sp!") && first.operands.contains("lr")) ||
    // mov r7, sp (frame pointer setup)
    (first.mnemonic == "mov" && first.operands.contains("r7") && first.operands.contains("sp"))
}

/// Find the closest symbol to a given address (within reasonable range)
fn find_closest_symbol(symbols: &std::collections::HashMap<Address, String>, target_address: Address) -> Option<String> {
    // First try exact match
    if let Some(name) = symbols.get(&target_address) {
        return Some(name.clone());
    }
    
    // If no exact match, find the closest symbol within a reasonable range
    // This accounts for function alignment, padding, and address calculations
    let mut closest_distance = u64::MAX;
    let mut closest_symbol = None;
    
    for (&symbol_addr, symbol_name) in symbols {
        let distance = if symbol_addr <= target_address {
            target_address - symbol_addr
        } else {
            symbol_addr - target_address
        };
        
        // Consider symbols within 256 bytes of the target (both before and after)
        // This handles cases where symbol addresses may be calculated differently
        if distance <= 256 && distance < closest_distance {
            closest_distance = distance;
            closest_symbol = Some(symbol_name.clone());
            
            // Log the match for debugging
            tracing::debug!("Found symbol '{}' at 0x{:x} for function at 0x{:x} (distance: {})", 
                symbol_name, symbol_addr, target_address, distance);
        }
    }
    
    if closest_symbol.is_none() {
        tracing::debug!("No symbol found for function at 0x{:x}", target_address);
    }
    
    closest_symbol
}

/// Classify symbols using YARA rules to identify which ones are actually functions
fn classify_function_symbols_with_yara(symbols: &std::collections::HashMap<Address, String>) -> Result<Vec<Address>> {
    let mut function_addresses = Vec::new();
    
    // First use basic heuristics to filter obvious function symbols
    for (&addr, name) in symbols {
        if is_likely_function_symbol(name) {
            function_addresses.push(addr);
        }
    }
    
    tracing::debug!("Basic heuristics identified {} potential functions from {} symbols", 
        function_addresses.len(), symbols.len());
    
    Ok(function_addresses)
}

/// Check if a symbol name is likely to be a function
fn is_likely_function_symbol(name: &str) -> bool {
    // C++ mangled functions
    if name.starts_with("_Z") {
        return true;
    }
    
    // Function-like names
    if name.contains("()") || name.ends_with("()") {
        return true;
    }
    
    // Common function names
    let function_keywords = ["main", "init", "start", "setup", "run", "execute", 
                           "create", "destroy", "handle", "process", "update"];
    if function_keywords.iter().any(|&keyword| name.to_lowercase().contains(keyword)) {
        return true;
    }
    
    // Avoid obvious data symbols
    let data_keywords = ["_data", "_rodata", "_bss", "String", "string", 
                        "variable", "const", "static", "_var", "__"];
    if data_keywords.iter().any(|&keyword| name.to_lowercase().contains(keyword)) {
        return false;
    }
    
    // If it contains common programming patterns, likely a function
    if name.contains("::") || name.contains("get") || name.contains("set") {
        return true;
    }
    
    // Default: if it's not obviously data, consider it a potential function
    // This is conservative but better than including everything
    !name.chars().any(|c| c.is_ascii_digit()) || name.len() > 8
}