gaia-assembler 0.1.1

//! Native x86_64 backend compiler for Windows
//! This backend generates native x86_64 machine code wrapped in a PE container

#[cfg(feature = "x86_64-assembler")]
use crate::{
    backends::{
        windows::emitter::{RelocationKind, X64Emitter},
        Backend, GeneratedFiles,
    },
    config::GaiaConfig,
    instruction::{CoreInstruction, GaiaInstruction, ManagedInstruction},
    program::{GaiaConstant, GaiaFunction, GaiaModule},
    types::GaiaType,
};
use gaia_types::{
    helpers::{AbiCompatible, ApiCompatible, Architecture, ArtifactType, CompilationTarget},
    GaiaError, Result,
};

#[cfg(feature = "pe-assembler")]
use pe_assembler::formats::exe::writer::ExeWriter;
use std::collections::HashMap;

#[cfg(feature = "x86_64-assembler")]
use x86_64_assembler::{
    encoder::InstructionEncoder,
    instruction::{Instruction, Operand, Register},
};

/// Native x86_64 Backend implementation for Windows
#[cfg(feature = "x86_64-assembler")]
pub struct X64Backend {
    encoder: InstructionEncoder,
}

#[cfg(feature = "x86_64-assembler")]
impl X64Backend {
    /// Create a new X64Backend instance
    pub fn new() -> Self {
        Self {
            encoder: InstructionEncoder::new(Architecture::X86_64),
        }
    }
}

#[cfg(feature = "x86_64-assembler")]
impl Default for X64Backend {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(feature = "x86_64-assembler")]
impl Backend for X64Backend {
    fn name(&self) -> &'static str {
        "Windows (Native x86_64)"
    }

    fn primary_target(&self) -> CompilationTarget {
        CompilationTarget { build: Architecture::X86_64, host: AbiCompatible::PE, target: ApiCompatible::MicrosoftVisualC }
    }

    fn artifact_type(&self) -> ArtifactType {
        ArtifactType::Executable
    }

    fn match_score(&self, target: &CompilationTarget) -> f32 {
        if target.build == Architecture::X86_64 && target.host == AbiCompatible::PE {
            if target.target == ApiCompatible::MicrosoftVisualC {
                return 100.0;
            }
            return 80.0;
        }
        0.0
    }

    fn generate(&self, program: &GaiaModule, config: &GaiaConfig) -> Result<GeneratedFiles> {
        let mut emitter = X64Emitter::new(program);
        emitter.emit()?;
        let (instructions, relocations, _rdata) = emitter.take_result();

        let mut code = Vec::new();
        let mut labels = HashMap::new();
        let mut inst_offsets = Vec::new();

        // Pass 1: Encode and find labels
        for inst in &instructions {
            inst_offsets.push(code.len());
            match inst {
                Instruction::Label(name) => {
                    labels.insert(name.clone(), code.len());
                }
                _ => {
                    let bytes = self.encode_inst(inst)?;
                    code.extend_from_slice(&bytes);
                }
            }
        }

        // Pass 2: Apply relocations (Simplified for now)
        let mut exit_pos = 0;
        let mut external_calls: HashMap<String, Vec<usize>> = HashMap::new();
        let string_patches: Vec<(usize, usize)> = Vec::new();

        for reloc in &relocations {
            let inst_pos = inst_offsets[reloc.instruction_index];
            match reloc.kind {
                RelocationKind::Relative32 => {
                    if let Some(&target_offset) = labels.get(&reloc.target) {
                        let rel = (target_offset as i32) - (inst_pos as i32 + 5);
                        code[inst_pos + 1..inst_pos + 5].copy_from_slice(&rel.to_le_bytes());
                    }
                }
                RelocationKind::RipRelative => {
                    if reloc.target == "ExitProcess" {
                        exit_pos = inst_pos;
                    }
                    external_calls.entry(reloc.target.clone()).or_default().push(inst_pos);
                }
                _ => {}
            }
        }

        let mut files = HashMap::new();
        files.insert("main.exe".to_string(), code.clone());

        #[cfg(feature = "pe-assembler")]
        {
            let rdata = &[];
            let pe_bytes = self.create_pe_exe(&code, program, exit_pos, &external_calls, rdata, &string_patches)?;
            files.insert("main.exe".to_string(), pe_bytes);
        }

        Ok(GeneratedFiles { artifact_type: self.artifact_type(), files, custom: None, diagnostics: vec![] })
    }
}

#[cfg(feature = "x86_64-assembler")]
impl X64Backend {
    fn encode_inst(&self, inst: &Instruction) -> Result<Vec<u8>> {
        self.encoder.encode(inst)
    }

    #[cfg(feature = "pe-assembler")]
    fn create_pe_exe(
        &self,
        code: &[u8],
        program: &GaiaModule,
        exit_pos: usize,
        external_calls: &HashMap<String, Vec<usize>>,
        rdata: &[u8],
        string_patches: &[(usize, usize)],
    ) -> Result<Vec<u8>> {
        use pe_assembler::formats::exe::writer::ExeWriter;
        use pe_assembler::helpers::PeWriter;
        use pe_assembler::types::{ImportEntry, ImportTable, PeProgram};
        use std::io::Cursor;
        
        // 创建 PE 程序
        let mut pe_program = PeProgram::create_executable(code.to_vec());
        
        // 添加导入表（如果有外部调用）
        if !external_calls.is_empty() {
            let mut imports = ImportTable::default();
            
            // 处理 kernel32.dll 导入
            if external_calls.contains_key("ExitProcess") {
                let kernel32_entry = ImportEntry {
                    dll_name: "kernel32.dll".to_string(),
                    functions: vec!["ExitProcess".to_string()],
                };
                imports.entries.push(kernel32_entry);
            }
            
            // 处理其他 DLL 导入
            if external_calls.contains_key("printf") {
                let msvcrt_entry = ImportEntry {
                    dll_name: "msvcrt.dll".to_string(),
                    functions: vec!["printf".to_string()],
                };
                imports.entries.push(msvcrt_entry);
            }
            
            if !imports.entries.is_empty() {
                pe_program = pe_program.with_imports(imports);
            }
        }
        
        // 写入 PE 文件
        let mut buffer = Vec::new();
        let mut cursor = Cursor::new(&mut buffer);
        let mut writer = ExeWriter::new(cursor);
        writer.write_program(&pe_program)?;
        Ok(buffer)
    }
}