nabla-decompiler 0.1.2

//! Control Flow Graph generation using Petgraph

use anyhow::Result;
use petgraph::Graph;
use std::collections::HashMap;

use crate::cfg_analyzer::{CfgAnalyzer, CfgAnalyzerRegistry};
use crate::types::{
    Address, BasicBlockNode, CallEdge, CallGraph, CallType, ControlFlowEdge, 
    ControlFlowGraph, ControlFlowType, Disassembly, Function, FunctionNode,
    InstructionGroup
};

/// Build call graph from functions
pub fn build_call_graph(functions: &[Function]) -> Result<CallGraph> {
    let mut graph = Graph::new();
    let mut addr_to_node = HashMap::new();
    
    // Add function nodes
    for function in functions {
        let node = graph.add_node(FunctionNode {
            address: function.address,
            name: function.name.clone(),
            size: function.size,
        });
        addr_to_node.insert(function.address, node);
    }
    
    // Add call edges
    for function in functions {
        if let Some(&caller_node) = addr_to_node.get(&function.address) {
            for &called_addr in &function.calls {
                if let Some(&callee_node) = addr_to_node.get(&called_addr) {
                    graph.add_edge(caller_node, callee_node, CallEdge {
                        call_type: CallType::Direct, // Simplified for now
                        call_site: function.address, // Simplified
                    });
                }
            }
        }
    }
    
    Ok(graph)
}

/// Build CFGs for all functions with architecture detection
pub fn build_function_cfgs(functions: &[Function], disasm: &Disassembly, arch: &str, format: &str) -> Result<HashMap<Address, ControlFlowGraph>> {
    let registry = CfgAnalyzerRegistry::new();
    let analyzer = registry.get_analyzer(arch, format)
        .ok_or_else(|| {
            let available = registry.list_analyzers();
            tracing::warn!("No CFG analyzer found for architecture: '{}' format: '{}'. Available analyzers: {:?}", arch, format, available);
            anyhow::anyhow!("No CFG analyzer found for architecture: {} format: {}", arch, format)
        })?;
    
    tracing::info!("Using '{}' analyzer for architecture: '{}' format: '{}'", analyzer.name(), arch, format);
    let mut cfgs = HashMap::new();
    
    for function in functions {
        let cfg = build_function_cfg(function, disasm, analyzer)?;
        cfgs.insert(function.address, cfg);
    }
    
    Ok(cfgs)
}

/// Build control flow graph for a single function using architecture-specific analyzer
pub fn build_function_cfg(function: &Function, disasm: &Disassembly, analyzer: &dyn CfgAnalyzer) -> Result<ControlFlowGraph> {
    let mut graph = Graph::new();
    let mut block_to_node = HashMap::new();
    
    // Debug logging - more comprehensive diagnostics
    tracing::debug!("Building CFG for function at 0x{:x} using {} analyzer", function.address, analyzer.name());
    tracing::debug!("Function has {} basic blocks and {} total instructions", function.basic_blocks.len(), function.instructions.len());
    
    // Log all basic blocks to understand function structure
    for (i, block) in function.basic_blocks.iter().enumerate() {
        tracing::debug!("  Block {}: 0x{:x}-0x{:x} ({} instructions) - Type: {:?}", 
                   i, block.start_address, block.end_address, block.instructions.len(), block.block_type);
        
        // Log first and last instruction of each block
        if let Some(&first_addr) = block.instructions.first() {
            if let Some(first_instr) = disasm.instructions.iter().find(|i| i.address == first_addr) {
                tracing::debug!("    First: 0x{:x}: {} {} (group: {:?})", 
                          first_addr, first_instr.mnemonic, first_instr.operands, first_instr.group);
            }
        }
        if let Some(&last_addr) = block.instructions.last() {
            if let Some(last_instr) = disasm.instructions.iter().find(|i| i.address == last_addr) {
                tracing::debug!("    Last: 0x{:x}: {} {} (group: {:?})", 
                          last_addr, last_instr.mnemonic, last_instr.operands, last_instr.group);
            }
        }
    }
    
    // Add basic block nodes
    for block in &function.basic_blocks {
        let node = graph.add_node(BasicBlockNode {
            id: block.id,
            start_address: block.start_address,
            end_address: block.end_address,
            instruction_count: block.instructions.len(),
            block_type: block.block_type.clone(),
        });
        block_to_node.insert(block.id, node);
        tracing::debug!("Added block: 0x{:x}-0x{:x} ({} instructions)", 
                   block.start_address, block.end_address, block.instructions.len());
    }
    
    // Add control flow edges using architecture-specific analyzer
    let mut edges_added = 0;
    for block in &function.basic_blocks {
        if let Some(&source_node) = block_to_node.get(&block.id) {
            // Get the last instruction of this block
            if let Some(&last_addr) = block.instructions.last() {
                if let Some(instruction) = disasm.instructions.iter().find(|i| i.address == last_addr) {
                    let instruction_group = analyzer.classify_instruction(instruction);
                    
                    tracing::debug!("Processing block ending at 0x{:x}: {} {} (classified as {:?})", 
                              last_addr, instruction.mnemonic, instruction.operands, instruction_group);
                    
                    match instruction_group {
                        InstructionGroup::Jump => {
                            // Find target block and add edge using analyzer
                            if let Some(target_addr) = analyzer.extract_jump_target(instruction) {
                                tracing::debug!("Jump target extracted: 0x{:x}", target_addr);
                                if let Some(target_block) = find_block_containing_address(&function.basic_blocks, target_addr) {
                                    if let Some(&target_node) = block_to_node.get(&target_block.id) {
                                        let edge_type = if analyzer.is_conditional(instruction) {
                                            ControlFlowType::ConditionalTrue
                                        } else {
                                            ControlFlowType::Unconditional
                                        };
                                        
                                        graph.add_edge(source_node, target_node, ControlFlowEdge {
                                            edge_type: edge_type.clone(),
                                            condition: analyzer.extract_condition(instruction),
                                        });
                                        edges_added += 1;
                                        tracing::debug!("Added {:?} edge: 0x{:x} -> 0x{:x}", 
                                                  edge_type, block.start_address, target_addr);
                                    }
                                } else {
                                    tracing::debug!("No block found containing jump target 0x{:x}", target_addr);
                                }
                                
                                // For conditional jumps, also add fall-through edge
                                if analyzer.is_conditional(instruction) {
                                    if let Some(next_addr) = analyzer.fall_through_address(instruction) {
                                        tracing::debug!("Fall-through address: 0x{:x}", next_addr);
                                        if let Some(next_block) = find_block_containing_address(&function.basic_blocks, next_addr) {
                                            if let Some(&next_node) = block_to_node.get(&next_block.id) {
                                                graph.add_edge(source_node, next_node, ControlFlowEdge {
                                                    edge_type: ControlFlowType::ConditionalFalse,
                                                    condition: analyzer.negate_condition(&analyzer.extract_condition(instruction)),
                                                });
                                                edges_added += 1;
                                                tracing::debug!("Added ConditionalFalse edge: 0x{:x} -> 0x{:x}", 
                                                          block.start_address, next_addr);
                                            }
                                        }
                                    }
                                }
                            } else {
                                tracing::debug!("No jump target extracted from: {} {}", instruction.mnemonic, instruction.operands);
                            }
                        }
                        InstructionGroup::Return => {
                            tracing::debug!("Return instruction, no outgoing edges");
                        }
                        _ => {
                            // Fall-through to next block
                            if let Some(next_addr) = analyzer.fall_through_address(instruction) {
                                tracing::debug!("Fall-through address: 0x{:x}", next_addr);
                                if let Some(next_block) = find_block_containing_address(&function.basic_blocks, next_addr) {
                                    if let Some(&next_node) = block_to_node.get(&next_block.id) {
                                        graph.add_edge(source_node, next_node, ControlFlowEdge {
                                            edge_type: ControlFlowType::FallThrough,
                                            condition: None,
                                        });
                                        edges_added += 1;
                                        tracing::debug!("Added FallThrough edge: 0x{:x} -> 0x{:x}", 
                                                  block.start_address, next_addr);
                                    }
                                } else {
                                    tracing::debug!("No block found for fall-through address 0x{:x}", next_addr);
                                }
                            } else {
                                tracing::debug!("No fall-through address for instruction: {} {}", 
                                          instruction.mnemonic, instruction.operands);
                            }
                        }
                    }
                } else {
                    tracing::debug!("No instruction found at address 0x{:x}", last_addr);
                }
            } else {
                tracing::debug!("Block has no instructions");
            }
        }
    }
    
    // If no edges were added but we have multiple blocks, add fall-through edges
    // This handles cases where blocks were split artificially or instruction classification failed
    if edges_added == 0 && function.basic_blocks.len() > 1 {
        tracing::debug!("No edges detected, adding fall-through edges between consecutive blocks");
        
        for i in 0..function.basic_blocks.len() - 1 {
            let current_block = &function.basic_blocks[i];
            let next_block = &function.basic_blocks[i + 1];
            
            if let (Some(&current_node), Some(&next_node)) = (
                block_to_node.get(&current_block.id),
                block_to_node.get(&next_block.id)
            ) {
                // Check if the blocks are consecutive in memory
                if current_block.end_address == next_block.start_address ||
                   (current_block.end_address <= next_block.start_address && 
                    next_block.start_address - current_block.end_address <= 8) { // Small gap tolerance
                    
                    graph.add_edge(current_node, next_node, ControlFlowEdge {
                        edge_type: ControlFlowType::FallThrough,
                        condition: None,
                    });
                    edges_added += 1;
                    tracing::debug!("Added consecutive fall-through edge: 0x{:x} -> 0x{:x}", 
                              current_block.start_address, next_block.start_address);
                }
            }
        }
    }
    
    tracing::debug!("CFG construction complete: {} edges added", edges_added);
    
    if edges_added == 0 {
        tracing::warn!("No control flow edges detected! This may indicate:");
        tracing::warn!("- Incorrect instruction classification");
        tracing::warn!("- Jump target extraction failures");
        tracing::warn!("- Missing basic blocks");
        tracing::warn!("Function has {} basic blocks:", function.basic_blocks.len());
        for (i, block) in function.basic_blocks.iter().enumerate() {
            tracing::warn!("  Block {}: 0x{:x}-0x{:x} ({} instructions)", 
                     i, block.start_address, block.end_address, block.instructions.len());
            if let Some(&last_addr) = block.instructions.last() {
                if let Some(instruction) = disasm.instructions.iter().find(|i| i.address == last_addr) {
                    tracing::warn!("    Last instruction: {} {} (group: {:?})", 
                             instruction.mnemonic, instruction.operands, instruction.group);
                }
            }
        }
    }
    
    Ok(graph)
}


/// Build control flow graph for a single function with architecture detection
pub fn build_function_cfg_with_arch(function: &Function, disasm: &Disassembly, arch: &str, format: &str) -> Result<ControlFlowGraph> {
    let registry = CfgAnalyzerRegistry::new();
    let analyzer = registry.get_analyzer(arch, format)
        .ok_or_else(|| anyhow::anyhow!("No CFG analyzer found for architecture: {} format: {}", arch, format))?;
    
    build_function_cfg(function, disasm, analyzer)
}

/// Get available CFG analyzers for debugging and info purposes  
pub fn get_available_analyzers() -> Vec<&'static str> {
    let registry = CfgAnalyzerRegistry::new();
    registry.list_analyzers()
}

/// Get the analyzer name that would be used for given architecture and format
pub fn get_analyzer_for_binary(arch: &str, format: &str) -> Option<&'static str> {
    let registry = CfgAnalyzerRegistry::new();
    registry.get_analyzer(arch, format).map(|a| a.name())
}

/// Find basic block containing an address
fn find_block_containing_address(blocks: &[crate::types::BasicBlock], address: Address) -> Option<&crate::types::BasicBlock> {
    blocks.iter().find(|block| {
        address >= block.start_address && address < block.end_address
    })
}