nabla-decompiler 0.1.2

//! Architecture-specific CFG analysis traits

use crate::types::{Address, Instruction, InstructionGroup};

/// Trait for architecture-specific control flow graph analysis
pub trait CfgAnalyzer: Send + Sync {
    /// Get the name of this analyzer (e.g., "ARM64", "x86_64", "Intel8051")
    fn name(&self) -> &'static str;
    
    /// Check if this analyzer can handle the given architecture
    fn can_analyze(&self, arch: &str, format: &str) -> bool;
    
    /// Extract jump target address from instruction
    fn extract_jump_target(&self, instruction: &Instruction) -> Option<Address>;
    
    /// Check if instruction is conditional
    fn is_conditional(&self, instruction: &Instruction) -> bool;
    
    /// Extract condition string from conditional instruction
    fn extract_condition(&self, instruction: &Instruction) -> Option<String>;
    
    /// Negate a condition string
    fn negate_condition(&self, condition: &Option<String>) -> Option<String>;
    
    /// Classify instruction for control flow purposes
    fn classify_instruction(&self, instruction: &Instruction) -> InstructionGroup;
    
    /// Calculate fall-through address (next instruction)
    fn fall_through_address(&self, instruction: &Instruction) -> Option<Address> {
        Some(instruction.address + instruction.size as u64)
    }
}

/// ARM architecture CFG analyzer
pub struct ArmCfgAnalyzer;

impl CfgAnalyzer for ArmCfgAnalyzer {
    fn name(&self) -> &'static str {
        "ARM"
    }
    
    fn can_analyze(&self, arch: &str, _format: &str) -> bool {
        arch.to_lowercase().contains("arm") || 
        arch.to_lowercase().contains("aarch64")
    }
    
    fn extract_jump_target(&self, instruction: &Instruction) -> Option<Address> {
        let operands = &instruction.operands;
        
        // Direct address: b #0x1234 or b 0x1234
        if let Some(hex_start) = operands.find("0x") {
            let hex_part = &operands[hex_start + 2..];
            if let Some(space_pos) = hex_part.find(' ') {
                if let Ok(addr) = u64::from_str_radix(&hex_part[..space_pos], 16) {
                    return Some(addr);
                }
            } else if let Ok(addr) = u64::from_str_radix(hex_part, 16) {
                return Some(addr);
            }
        }
        
        // Immediate value: b #1234
        if let Some(hash_pos) = operands.find('#') {
            let num_part = &operands[hash_pos + 1..];
            if let Some(space_pos) = num_part.find(' ') {
                if let Ok(addr) = num_part[..space_pos].parse::<u64>() {
                    return Some(addr);
                }
            } else if let Ok(addr) = num_part.parse::<u64>() {
                return Some(addr);
            }
        }
        
        None
    }
    
    fn is_conditional(&self, instruction: &Instruction) -> bool {
        let mnemonic = instruction.mnemonic.to_lowercase();
        
        // ARM conditional suffixes
        mnemonic.ends_with("eq") || mnemonic.ends_with("ne") || 
        mnemonic.ends_with("lt") || mnemonic.ends_with("le") ||
        mnemonic.ends_with("gt") || mnemonic.ends_with("ge") ||
        mnemonic.ends_with("cs") || mnemonic.ends_with("cc") ||
        mnemonic.ends_with("mi") || mnemonic.ends_with("pl") ||
        mnemonic.ends_with("vs") || mnemonic.ends_with("vc") ||
        mnemonic.ends_with("hi") || mnemonic.ends_with("ls") ||
        // Common conditional branches
        (mnemonic.starts_with("b") && mnemonic.len() > 1 && mnemonic != "bl" && mnemonic != "blx")
    }
    
    fn extract_condition(&self, instruction: &Instruction) -> Option<String> {
        let mnemonic = instruction.mnemonic.to_lowercase();
        
        if mnemonic.ends_with("eq") {
            Some("== 0".to_string())
        } else if mnemonic.ends_with("ne") {
            Some("!= 0".to_string())
        } else if mnemonic.ends_with("lt") {
            Some("< 0".to_string())
        } else if mnemonic.ends_with("le") {
            Some("<= 0".to_string())
        } else if mnemonic.ends_with("gt") {
            Some("> 0".to_string())
        } else if mnemonic.ends_with("ge") {
            Some(">= 0".to_string())
        } else if mnemonic.ends_with("cs") {
            Some("carry set".to_string())
        } else if mnemonic.ends_with("cc") {
            Some("carry clear".to_string())
        } else if mnemonic.ends_with("mi") {
            Some("< 0".to_string())
        } else if mnemonic.ends_with("pl") {
            Some(">= 0".to_string())
        } else if mnemonic.ends_with("vs") {
            Some("overflow".to_string())
        } else if mnemonic.ends_with("vc") {
            Some("no overflow".to_string())
        } else if mnemonic.ends_with("hi") {
            Some("unsigned >".to_string())
        } else if mnemonic.ends_with("ls") {
            Some("unsigned <=".to_string())
        } else {
            None
        }
    }
    
    fn negate_condition(&self, condition: &Option<String>) -> Option<String> {
        condition.as_ref().map(|c| {
            match c.as_str() {
                "== 0" => "!= 0".to_string(),
                "!= 0" => "== 0".to_string(),
                "< 0" => ">= 0".to_string(),
                "<= 0" => "> 0".to_string(),
                "> 0" => "<= 0".to_string(),
                ">= 0" => "< 0".to_string(),
                "carry set" => "carry clear".to_string(),
                "carry clear" => "carry set".to_string(),
                "overflow" => "no overflow".to_string(),
                "no overflow" => "overflow".to_string(),
                "unsigned >" => "unsigned <=".to_string(),
                "unsigned <=" => "unsigned >".to_string(),
                _ => format!("!({c})"),
            }
        })
    }
    
    fn classify_instruction(&self, instruction: &Instruction) -> InstructionGroup {
        // Use the existing classification from the instruction
        instruction.group.clone()
    }
}

/// x86/x64 architecture CFG analyzer
pub struct X86CfgAnalyzer;

impl CfgAnalyzer for X86CfgAnalyzer {
    fn name(&self) -> &'static str {
        "x86_64"
    }
    
    fn can_analyze(&self, arch: &str, format: &str) -> bool {
        let arch_lower = arch.to_lowercase();
        let format_lower = format.to_lowercase();
        
        // Handle all x86 variants and formats
        (arch_lower.contains("x86") || 
         arch_lower.contains("x64") ||
         arch_lower.contains("amd64") ||
         arch_lower.contains("i386") ||
         arch_lower.contains("x86_64")) &&
        // Support various binary formats
        (format_lower.contains("elf") ||
         format_lower.contains("pe") ||
         format_lower.contains("macho") ||
         format_lower.contains("mach-o") ||
         format_lower.contains("coff"))
    }
    
    fn extract_jump_target(&self, instruction: &Instruction) -> Option<Address> {
        let operands = &instruction.operands;
        let mnemonic = instruction.mnemonic.to_lowercase();
        
        // Enhanced debug logging for MachO analysis
        // tracing::debug!("Extracting jump target from: '{}' operands: '{}'", mnemonic, operands);
        
        // Direct absolute jump: jmp 0x401000, call 0x401000
        if let Some(hex_start) = operands.find("0x") {
            let hex_part = &operands[hex_start + 2..];
            // Handle various formats: 0x401000, 0x401000h, 0x401000 <symbol>, etc.
            let end_pos = hex_part.find(' ')
                .or_else(|| hex_part.find('<'))
                .or_else(|| hex_part.find('h'))
                .or_else(|| hex_part.find(','))
                .unwrap_or(hex_part.len());
            
            if let Ok(addr) = u64::from_str_radix(&hex_part[..end_pos], 16) {
                tracing::debug!("Extracted hex address: 0x{:x}", addr);
                return Some(addr);
            }
        }
        
        // Handle relative jumps (e.g., "jmp +0x42", "jmp -0x10")
        if let Some(plus_pos) = operands.find("+0x") {
            let hex_part = &operands[plus_pos + 3..];
            let end_pos = hex_part.find(' ').unwrap_or(hex_part.len());
            if let Ok(offset) = i64::from_str_radix(&hex_part[..end_pos], 16) {
                let target = (instruction.address as i64) + (instruction.size as i64) + offset;
                if target >= 0 {
                    return Some(target as u64);
                }
            }
        }
        
        if let Some(minus_pos) = operands.find("-0x") {
            let hex_part = &operands[minus_pos + 3..];
            let end_pos = hex_part.find(' ').unwrap_or(hex_part.len());
            if let Ok(offset) = i64::from_str_radix(&hex_part[..end_pos], 16) {
                let target = (instruction.address as i64) + (instruction.size as i64) - offset;
                if target >= 0 {
                    return Some(target as u64);
                }
            }
        }
        
        // Handle plain hex values (e.g., "call 401000h")
        if operands.ends_with('h') && operands.len() > 1 {
            let hex_part = &operands[..operands.len() - 1];
            if let Ok(addr) = u64::from_str_radix(hex_part, 16) {
                tracing::debug!("Extracted hex address with 'h' suffix: 0x{:x}", addr);
                return Some(addr);
            }
        }
        
        // MachO-specific formats: Handle RIP-relative addressing like "jmp 0x100003f60"
        // or "jmp [rip + offset]" style addressing common in x86_64 MachO
        if operands.contains("rip") {
            // Parse RIP-relative: "0x12345678 # 0x100003f60" or similar
            if let Some(hash_pos) = operands.find('#') {
                let addr_part = &operands[hash_pos + 1..].trim();
                if addr_part.starts_with("0x") {
                    let hex_part = &addr_part[2..];
                    let end_pos = hex_part.find(' ').unwrap_or(hex_part.len());
                    if let Ok(addr) = u64::from_str_radix(&hex_part[..end_pos], 16) {
                        tracing::debug!("Extracted RIP-relative address: 0x{:x}", addr);
                        return Some(addr);
                    }
                }
            }
        }
        
        // Handle bare hex addresses (no prefix): "jmp 401000" 
        if mnemonic.starts_with('j') || mnemonic == "call" {
            let trimmed = operands.trim();
            // Try parsing as hex first (common in disassemblers)
            if trimmed.chars().all(|c| c.is_ascii_hexdigit()) && trimmed.len() >= 4 {
                if let Ok(addr) = u64::from_str_radix(trimmed, 16) {
                    tracing::debug!("Extracted bare hex address: 0x{:x}", addr);
                    return Some(addr);
                }
            }
            
            // Simple positive decimal: jmp 1234
            if let Ok(offset) = trimmed.parse::<i32>() {
                let target = (instruction.address as i64) + (instruction.size as i64) + (offset as i64);
                if target >= 0 {
                    tracing::debug!("Calculated relative target: 0x{:x}", target as u64);
                    return Some(target as u64);
                }
            }
        }
        
        // Enhanced parsing for various operand formats
        // Handle signed decimal offsets like "jmp +42", "jmp -10" (without 0x)
        if (mnemonic.starts_with('j') || mnemonic == "call") && !operands.is_empty() {
            let trimmed = operands.trim();
            
            // Relative offset with + or - sign
            if trimmed.starts_with('+') || trimmed.starts_with('-') {
                if let Ok(offset) = trimmed.parse::<i32>() {
                    let target = (instruction.address as i64) + (instruction.size as i64) + (offset as i64);
                    if target >= 0 {
                        tracing::debug!("Calculated signed relative target: 0x{:x}", target as u64);
                        return Some(target as u64);
                    }
                }
            }
            
            // Try to parse operands like "jmp loc_100001234" or "jmp sub_100001234"
            if let Some(underscore_pos) = trimmed.rfind('_') {
                let addr_part = &trimmed[underscore_pos + 1..];
                if let Ok(addr) = u64::from_str_radix(addr_part, 16) {
                    tracing::debug!("Extracted address from symbol: 0x{:x}", addr);
                    return Some(addr);
                }
            }
            
            // Try parsing the entire operand as hex (might be missing 0x prefix)
            if trimmed.len() >= 6 && trimmed.len() <= 16 && trimmed.chars().all(|c| c.is_ascii_hexdigit()) {
                if let Ok(addr) = u64::from_str_radix(trimmed, 16) {
                    tracing::debug!("Extracted hex address (no prefix): 0x{:x}", addr);
                    return Some(addr);
                }
            }
        }
        
        tracing::debug!("No jump target extracted from operands: '{}'", operands);
        None
    }
    
    fn is_conditional(&self, instruction: &Instruction) -> bool {
        let mnemonic = instruction.mnemonic.to_lowercase();
        
        // x86 conditional jumps (including all variants)
        matches!(mnemonic.as_str(),
            // Zero/Equal conditions
            "je" | "jz" | "jne" | "jnz" | 
            // Signed comparisons
            "jl" | "jnge" | "jle" | "jng" | "jg" | "jnle" | "jge" | "jnl" |
            // Unsigned comparisons  
            "ja" | "jnbe" | "jae" | "jnb" | "jb" | "jnae" | "jbe" | "jna" |
            // Carry flag
            "jc" | "jnc" |
            // Overflow flag
            "jo" | "jno" |
            // Sign flag
            "js" | "jns" |
            // Parity flag
            "jp" | "jpe" | "jnp" | "jpo" |
            // Loop instructions (also conditional)
            "loop" | "loope" | "loopz" | "loopne" | "loopnz" |
            // Conditional moves (not jumps but similar logic)
            "jecxz" | "jrcxz"
        )
    }
    
    fn extract_condition(&self, instruction: &Instruction) -> Option<String> {
        let mnemonic = instruction.mnemonic.to_lowercase();
        
        match mnemonic.as_str() {
            "je" | "jz" => Some("== 0".to_string()),
            "jne" | "jnz" => Some("!= 0".to_string()),
            "jl" | "jnge" => Some("< 0".to_string()),
            "jle" | "jng" => Some("<= 0".to_string()),
            "jg" | "jnle" => Some("> 0".to_string()),
            "jge" | "jnl" => Some(">= 0".to_string()),
            "ja" | "jnbe" => Some("unsigned >".to_string()),
            "jae" | "jnb" => Some("unsigned >=".to_string()),
            "jb" | "jnae" => Some("unsigned <".to_string()),
            "jbe" | "jna" => Some("unsigned <=".to_string()),
            "jc" => Some("carry set".to_string()),
            "jnc" => Some("carry clear".to_string()),
            "jo" => Some("overflow".to_string()),
            "jno" => Some("no overflow".to_string()),
            "js" => Some("sign set".to_string()),
            "jns" => Some("sign clear".to_string()),
            "jp" | "jpe" => Some("parity even".to_string()),
            "jnp" | "jpo" => Some("parity odd".to_string()),
            _ => None,
        }
    }
    
    fn negate_condition(&self, condition: &Option<String>) -> Option<String> {
        condition.as_ref().map(|c| {
            match c.as_str() {
                "== 0" => "!= 0".to_string(),
                "!= 0" => "== 0".to_string(),
                "< 0" => ">= 0".to_string(),
                "<= 0" => "> 0".to_string(),
                "> 0" => "<= 0".to_string(),
                ">= 0" => "< 0".to_string(),
                "unsigned >" => "unsigned <=".to_string(),
                "unsigned >=" => "unsigned <".to_string(),
                "unsigned <" => "unsigned >=".to_string(),
                "unsigned <=" => "unsigned >".to_string(),
                "carry set" => "carry clear".to_string(),
                "carry clear" => "carry set".to_string(),
                "overflow" => "no overflow".to_string(),
                "no overflow" => "overflow".to_string(),
                "sign set" => "sign clear".to_string(),
                "sign clear" => "sign set".to_string(),
                "parity even" => "parity odd".to_string(),
                "parity odd" => "parity even".to_string(),
                _ => format!("!({c})"),
            }
        })
    }
    
    fn classify_instruction(&self, instruction: &Instruction) -> InstructionGroup {
        let mnemonic = instruction.mnemonic.to_lowercase();
        
        tracing::debug!("Classifying instruction: '{}' (original group: {:?})", mnemonic, instruction.group);
        
        // Override instruction classification if needed for x86 architecture
        let result = match mnemonic.as_str() {
            // Jump instructions - be more aggressive in matching patterns
            "jmp" | "jmpq" => InstructionGroup::Jump,
            "je" | "jz" | "jne" | "jnz" | "jl" | "jnge" | "jle" | "jng" |
            "jg" | "jnle" | "jge" | "jnl" | "ja" | "jnbe" | "jae" | "jnb" |
            "jb" | "jnae" | "jbe" | "jna" | "jc" | "jnc" | "jo" | "jno" |
            "js" | "jns" | "jp" | "jpe" | "jnp" | "jpo" | "loop" | 
            "loope" | "loopz" | "loopne" | "loopnz" | "jecxz" | "jrcxz" => InstructionGroup::Jump,
            
            // Call instructions
            "call" | "callq" => InstructionGroup::Call,
            
            // Return instructions  
            "ret" | "retq" | "retn" | "retf" | "iret" | "iretd" | "iretq" => InstructionGroup::Return,
            
            // Move instructions
            "mov" | "movq" | "movl" | "movw" | "movb" | "movzx" | "movsx" | "movsxd" => InstructionGroup::Move,
            
            // Arithmetic
            "add" | "sub" | "mul" | "div" | "imul" | "idiv" | "inc" | "dec" |
            "addq" | "subq" | "mulq" | "divq" | "incq" | "decq" => InstructionGroup::Arithmetic,
            
            // Logical
            "and" | "or" | "xor" | "not" | "shl" | "shr" | "sal" | "sar" |
            "andq" | "orq" | "xorq" | "notq" | "shlq" | "shrq" => InstructionGroup::Logical,
            
            // Compare
            "cmp" | "test" | "cmpq" | "testq" => InstructionGroup::Compare,
            
            // Load/Store
            "push" | "pop" | "pushq" | "popq" | "lea" | "leaq" => InstructionGroup::Load,
            
            // No operation
            "nop" => InstructionGroup::Nop,
            
            // Fall back to original classification
            _ => instruction.group.clone(),
        };
        
       // tracing::trace!("Classified '{}' as {:?}", mnemonic, result);
        result
    }
}

/// Intel 8051 microcontroller CFG analyzer
pub struct Intel8051CfgAnalyzer;

impl CfgAnalyzer for Intel8051CfgAnalyzer {
    fn name(&self) -> &'static str {
        "Intel8051"
    }
    
    fn can_analyze(&self, arch: &str, format: &str) -> bool {
        let arch_lower = arch.to_lowercase();
        let format_lower = format.to_lowercase();
        
        arch_lower.contains("8051") || 
        format_lower.contains("intel_hex") ||
        format_lower.contains("ihex")
    }
    
    fn extract_jump_target(&self, instruction: &Instruction) -> Option<Address> {
        let operands = &instruction.operands;
        
        // 8051 absolute jump: LJMP 1234h or SJMP rel_addr
        if let Some(hex_end) = operands.find('h') {
            if let Ok(addr) = u16::from_str_radix(&operands[..hex_end], 16) {
                return Some(addr as u64);
            }
        }
        
        // Hex format: LJMP 0x1234
        if let Some(hex_start) = operands.find("0x") {
            let hex_part = &operands[hex_start + 2..];
            if let Ok(addr) = u16::from_str_radix(hex_part, 16) {
                return Some(addr as u64);
            }
        }
        
        None
    }
    
    fn is_conditional(&self, instruction: &Instruction) -> bool {
        let mnemonic = instruction.mnemonic.to_uppercase();
        
        // 8051 conditional jumps
        matches!(mnemonic.as_str(),
            "JZ" | "JNZ" | "JC" | "JNC" | "JB" | "JNB" | "JBC" |
            "CJNE" | "DJNZ"
        )
    }
    
    fn extract_condition(&self, instruction: &Instruction) -> Option<String> {
        let mnemonic = instruction.mnemonic.to_uppercase();
        
        match mnemonic.as_str() {
            "JZ" => Some("A == 0".to_string()),
            "JNZ" => Some("A != 0".to_string()),
            "JC" => Some("carry set".to_string()),
            "JNC" => Some("carry clear".to_string()),
            "JB" => Some("bit set".to_string()),
            "JNB" => Some("bit clear".to_string()),
            "JBC" => Some("bit set and clear".to_string()),
            "CJNE" => Some("not equal".to_string()),
            "DJNZ" => Some("decrement and not zero".to_string()),
            _ => None,
        }
    }
    
    fn negate_condition(&self, condition: &Option<String>) -> Option<String> {
        condition.as_ref().map(|c| {
            match c.as_str() {
                "A == 0" => "A != 0".to_string(),
                "A != 0" => "A == 0".to_string(),
                "carry set" => "carry clear".to_string(),
                "carry clear" => "carry set".to_string(),
                "bit set" => "bit clear".to_string(),
                "bit clear" => "bit set".to_string(),
                "not equal" => "equal".to_string(),
                "decrement and not zero" => "decrement and zero".to_string(),
                _ => format!("!({c})"),
            }
        })
    }
    
    fn classify_instruction(&self, instruction: &Instruction) -> InstructionGroup {
        instruction.group.clone()
    }
}

/// Registry for managing different CFG analyzers
pub struct CfgAnalyzerRegistry {
    analyzers: Vec<Box<dyn CfgAnalyzer>>,
}

impl CfgAnalyzerRegistry {
    pub fn new() -> Self {
        let mut registry = Self {
            analyzers: Vec::new(),
        };
        
        // Register default analyzers
        registry.register(Box::new(ArmCfgAnalyzer));
        registry.register(Box::new(X86CfgAnalyzer));
        registry.register(Box::new(Intel8051CfgAnalyzer));
        
        registry
    }
    
    pub fn register(&mut self, analyzer: Box<dyn CfgAnalyzer>) {
        self.analyzers.push(analyzer);
    }
    
    pub fn get_analyzer(&self, arch: &str, format: &str) -> Option<&dyn CfgAnalyzer> {
        self.analyzers
            .iter()
            .find(|analyzer| analyzer.can_analyze(arch, format))
            .map(|analyzer| analyzer.as_ref())
    }
    
    pub fn list_analyzers(&self) -> Vec<&'static str> {
        self.analyzers.iter().map(|a| a.name()).collect()
    }
}

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