rustcc 0.1.1

//! # RustCC Code Generator
//! Generates x86-64 assembly code from abstract syntax trees.

use std::fs::File;
use std::io::{Result, Write};
use crate::ast::{ASTNode, NodeType};

/// A code generator that produces x86-64 assembly code from ASTs.
/// The generated code is written to an output file.
pub struct CodeGenerator {
    out_file: File,
    free_regs: [bool; 4],
    reg_names: [&'static str; 4],
}

impl CodeGenerator {
    /// Creates a new code generator.
    /// 
    /// # Arguments
    /// * `out_path` - Path to the output assembly file
    pub fn new(out_path: &str) -> Result<Self> {
        let out_file = File::create(out_path)?;
        Ok(CodeGenerator {
            out_file,
            free_regs: [true; 4],
            reg_names: ["%r8", "%r9", "%r10", "%r11"],
        })
    }

    /// Frees all allocated registers.
    fn free_all_registers(&mut self) {
        self.free_regs = [true; 4];
    }

    /// Allocates a free register.
    /// 
    /// # Returns
    /// Index of the allocated register
    fn alloc_register(&mut self) -> usize {
        for (i, free) in self.free_regs.iter().enumerate() {
            if *free {
                self.free_regs[i] = false;
                return i;
            }
        }
        panic!("Out of registers!");
    }

    /// Frees a previously allocated register.
    /// 
    /// # Arguments
    /// * `reg` - Index of the register to free
    fn free_register(&mut self, reg: usize) {
        if self.free_regs[reg] {
            panic!("Error trying to free register {}", reg);
        }
        self.free_regs[reg] = true;
    }

    /// Writes the assembly preamble to the output file.
    fn write_preamble(&mut self) -> Result<()> {
        self.free_all_registers();
        writeln!(
            &mut self.out_file,
            r#"section .data
LC0 db 'Result: %d', 0
LC1 db 13, 10, 0 ; Carriage return and line feed

section .text

extern printf

global main
main:
	push rbp
	mov rbp, rsp
	sub rsp, 40 ; 16-byte alignment for Windows x64
"#
        )
    }

    /// Writes the assembly postamble to the output file.
    fn write_postamble(&mut self) -> Result<()> {
        writeln!(
            &mut self.out_file,
            r#"	mov eax, 0
	pop rbp
	ret

; Function to print a number
print_number:
	push rbp
	mov rbp, rsp
	sub rsp, 40 ; 32 bytes shadow space + 8 bytes for alignment
	
	; Convert number to string (RCX contains the number)
	mov eax, ecx
	mov rdi, rsp
	call int_to_str
	
	; Print the string (use RCX for first parameter per Windows x64 calling convention)
	mov rcx, rsp
	mov eax, 0
	call printf
	
	leave
	ret

; Function to convert integer to string
int_to_str:
	push rbx ; Save non-volatile register
	mov rsi, rdi
	mov r11, 0
	mov ebx, 10
	
	; Handle zero
	cmp eax, 0
	jne .loop
	mov byte [rsi], '0'
	inc rsi
	jmp .end
	
.loop:
	cmp eax, 0
	je .reverse
	mov edx, 0
	div ebx
	add dl, '0'
	mov [rsi], dl
	inc rsi
	inc r11
	jmp .loop
	
.reverse:
	mov rdi, rsi
	sub rdi, r11
	dec rsi
	
.reverse_loop:
	cmp rdi, rsi
	jge .end
	mov al, [rdi]
	mov ah, [rsi]
	mov [rdi], ah
	mov [rsi], al
	inc rdi
	dec rsi
	jmp .reverse_loop
	
.end:
	mov byte [rsi], 0
	pop rbx ; Restore non-volatile register
	ret
"#
        )
    }

    /// Loads an integer value into a register.
    /// 
    /// # Arguments
    /// * `value` - The integer value to load
    /// 
    /// # Returns
    /// Index of the register containing the value
    fn load_int(&mut self, value: i32) -> Result<usize> {
        let reg = self.alloc_register();
        writeln!(&mut self.out_file, "\tmov rax, {}", value)?;
        writeln!(&mut self.out_file, "\tmov {}, rax", self.reg_names[reg].replace('%', ""))?;
            Ok(reg)
    }

    /// Adds the values in two registers and returns the result register.
    /// 
    /// # Arguments
    /// * `r1` - Index of the first register
    /// * `r2` - Index of the second register
    /// 
    /// # Returns
    /// Index of the register containing the result
    fn add(&mut self, r1: usize, r2: usize) -> Result<usize> {
        writeln!(&mut self.out_file, "\tadd {}, {}", self.reg_names[r2].replace('%', ""), self.reg_names[r1].replace('%', ""))?;
        self.free_register(r1);
        Ok(r2)
    }

    /// Subtracts the value in the second register from the first and returns the result register.
    /// 
    /// # Arguments
    /// * `r1` - Index of the first register
    /// * `r2` - Index of the second register
    /// 
    /// # Returns
    /// Index of the register containing the result
    fn subtract(&mut self, r1: usize, r2: usize) -> Result<usize> {
        writeln!(&mut self.out_file, "\tsub {}, {}", self.reg_names[r1].replace('%', ""), self.reg_names[r2].replace('%', ""))?;
        self.free_register(r2);
        Ok(r1)
    }

    /// Multiplies the values in two registers and returns the result register.
    /// 
    /// # Arguments
    /// * `r1` - Index of the first register
    /// * `r2` - Index of the second register
    /// 
    /// # Returns
    /// Index of the register containing the result
    fn multiply(&mut self, r1: usize, r2: usize) -> Result<usize> {
        writeln!(&mut self.out_file, "\timul {}, {}", self.reg_names[r1].replace('%', ""), self.reg_names[r2].replace('%', ""))?;
        self.free_register(r2);
        Ok(r1)
    }

    /// Divides the value in the first register by the second and returns the result register.
    /// 
    /// # Arguments
    /// * `r1` - Index of the first register
    /// * `r2` - Index of the second register
    /// 
    /// # Returns
    /// Index of the register containing the result
    fn divide(&mut self, r1: usize, r2: usize) -> Result<usize> {
        writeln!(&mut self.out_file, "\tmov rax, {}", self.reg_names[r1].replace('%', ""))?;
        writeln!(&mut self.out_file, "\tcqo")?;
        writeln!(&mut self.out_file, "\tidiv {}", self.reg_names[r2].replace('%', ""))?;
        writeln!(&mut self.out_file, "\tmov {}, rax", self.reg_names[r1].replace('%', ""))?;
        self.free_register(r2);
        Ok(r1)
    }

    /// Generates code to print an integer from a register.
    /// 
    /// # Arguments
    /// * `reg` - Index of the register containing the value to print
    fn print_int(&mut self, reg: usize) -> Result<()> {
        // Print the result using printf with format string
        writeln!(&mut self.out_file, "\tmov rcx, LC0")?;
        writeln!(&mut self.out_file, "\tmov rdx, {}", self.reg_names[reg].replace('%', ""))?;
        writeln!(&mut self.out_file, "\tmov eax, 0")?;
        writeln!(&mut self.out_file, "\tcall printf")?;
        
        // Print newline
        writeln!(&mut self.out_file, "\tmov rcx, LC1")?;
        writeln!(&mut self.out_file, "\tmov eax, 0")?;
        writeln!(&mut self.out_file, "\tcall printf")?;
        
        self.free_register(reg);
        Ok(())
    }

    /// Recursively generates code for an AST node.
    /// 
    /// # Arguments
    /// * `node` - The AST node to generate code for
    /// 
    /// # Returns
    /// Index of the register containing the result
    fn gen_ast(&mut self, node: &ASTNode) -> Result<usize> {
        match node.op {
            NodeType::Add => {
                let left_reg = self.gen_ast(node.left.as_ref().expect("Missing left operand"))?;
                let right_reg = self.gen_ast(node.right.as_ref().expect("Missing right operand"))?;
                self.add(left_reg, right_reg)
            },
            NodeType::Subtract => {
                let left_reg = self.gen_ast(node.left.as_ref().expect("Missing left operand"))?;
                let right_reg = self.gen_ast(node.right.as_ref().expect("Missing right operand"))?;
                self.subtract(left_reg, right_reg)
            },
            NodeType::Multiply => {
                let left_reg = self.gen_ast(node.left.as_ref().expect("Missing left operand"))?;
                let right_reg = self.gen_ast(node.right.as_ref().expect("Missing right operand"))?;
                self.multiply(left_reg, right_reg)
            },
            NodeType::Divide => {
                let left_reg = self.gen_ast(node.left.as_ref().expect("Missing left operand"))?;
                let right_reg = self.gen_ast(node.right.as_ref().expect("Missing right operand"))?;
                self.divide(left_reg, right_reg)
            },
            NodeType::IntLit => {
                self.load_int(node.int_value.expect("Missing integer value"))
            },
        }
    }

    /// Generates assembly code for the given AST and writes it to the output file.
    /// 
    /// # Arguments
    /// * `ast` - The AST to generate code from
    pub fn generate_code(&mut self, ast: &ASTNode) -> Result<()> {
        self.write_preamble()?;
        let reg = self.gen_ast(ast)?;
        self.print_int(reg)?;
        self.write_postamble()
    }
}