rubbler 0.1.2

use crate::{
    expression::{ExprType, Expression},
    generator::Code,
    stanalyzer::Context,
};

pub fn is_directive(string: &str) -> bool {
    match string {
        ".align" | ".comm" | ".common" | ".section" | ".text" | ".data" | ".rodata" | ".bss"
        | ".string" | ".asciz" | ".equ" | ".byte" | ".2byte" | ".half" | ".short" | ".4byte"
        | ".word" | ".long" | ".8byte" | ".dword" | ".quad" | ".p2align" => true,
        _ => false,
    }
}

pub fn execute_directive(
    dir: &str,
    args: &Vec<Expression>,
    ctx: &mut Context,
) -> Result<(), String> {
    match dir {
        ".align" | ".p2align" => exec_align(args, ctx),
        ".comm" | ".common" => exec_comm(args, ctx),
        ".section" => exec_section(args, ctx),
        ".equ" => exec_equ(args, ctx),
        ".byte" => exec_byte(args, ctx),
        _ => return Err(format!("`{}` directive not yet supported.", dir)),
    }
    .map_err(|e| format!("{}: {}", dir, e))
}

fn exec_align(args: &Vec<Expression>, ctx: &mut Context) -> Result<(), String> {
    let args: Vec<&ExprType> = args.iter().map(|arg| arg.get_type()).collect();
    if let [ExprType::NumberLiteral(n)] = args[..] {
        let alignment = 1usize << *n;
        ctx.align_addr(alignment);
        Ok(())
    } else {
        Err("arguments should be: `integer`".to_string())
    }
}

fn exec_comm(args: &Vec<Expression>, ctx: &mut Context) -> Result<(), String> {
    let args: Vec<&ExprType> = args.iter().map(|arg| arg.get_type()).collect();
    if let [ExprType::Symbol(sym), ExprType::NumberLiteral(size), ExprType::NumberLiteral(align)] =
        &args[..]
    {
        ctx.to_bss();
        let align: usize = (*align)
            .try_into()
            .map_err(|_| "`align` argument must be a positive integer".to_string())?;
        ctx.align_addr(align);
        ctx.emit_sym(sym.to_string());
        let size: usize = (*size)
            .try_into()
            .map_err(|_| "`size` argument must be a positive integer".to_string())?;
        ctx.inc_addr(size);
        Ok(())
    } else {
        Err("arguments should be: `symbol`, `size`, `align`".to_string())
    }
}

fn exec_section(args: &Vec<Expression>, ctx: &mut Context) -> Result<(), String> {
    let args: Vec<&ExprType> = args.iter().map(|arg| arg.get_type()).collect();
    if args.len() == 0 {
        ctx.to_text();
        Ok(())
    } else if let [ExprType::Symbol(sym)] = &args[..] {
        if sym == ".text" {
            ctx.to_text()
        } else if sym == ".data" {
            ctx.to_data()
        } else if sym == ".rodata" {
            ctx.to_rodata()
        } else if sym == ".bss" {
            ctx.to_bss()
        } else {
            return Err("arguments should be: {`.text`, `.data`, `.rodata`, `.bss`}.".to_string());
        }
        Ok(())
    } else {
        Err("arguments should be: {`.text`, `.data`, `.rodata`, `.bss`}.".to_string())
    }
}

fn exec_equ(args: &Vec<Expression>, ctx: &mut Context) -> Result<(), String> {
    let args: Vec<&ExprType> = args.iter().map(|arg| arg.get_type()).collect();
    if let [ExprType::Symbol(name), ExprType::NumberLiteral(value)] = &args[..] {
        ctx.emit_const(name.to_string(), *value);
        Ok(())
    } else {
        Err("arguments should be: `name`, `value`.".to_string())
    }
}

fn exec_byte(args: &Vec<Expression>, ctx: &mut Context) -> Result<(), String> {
    let args: Vec<&ExprType> = args.iter().map(|arg| arg.get_type()).collect();
    let mut n_bytes = 0;
    for arg in args {
        match arg {
            ExprType::NumberLiteral(_) => n_bytes += 1,
            _ => return Err("arguments should be: number[, number]*.".to_string()),
        }
    }
    ctx.inc_addr(n_bytes);
    Ok(())
}

pub fn generate_directive(
    dir: &str,
    args: &Vec<Expression>,
    code: &mut Code,
) -> Result<(), String> {
    match dir {
        ".align" | ".p2align" => generate_align(args, code),
        ".comm" | ".common" => generate_comm(args, code),
        ".section" => generate_section(args, code),
        ".equ" => Ok(()),
        ".byte" => generate_bytes(args, code),
        _ => return Err(format!("{} directive not yet unsupported.", dir)),
    }
    .map_err(|e| format!("{}: {}", dir, e))
}

fn generate_align(args: &Vec<Expression>, code: &mut Code) -> Result<(), String> {
    let args: Vec<&ExprType> = args.iter().map(|arg| arg.get_type()).collect();
    if let [ExprType::NumberLiteral(n)] = args[..] {
        let alignment = 1usize << n;
        code.align_addr(alignment);
        Ok(())
    } else {
        Err("arguments should be: `integer`".to_string())
    }
}

fn generate_comm(_args: &Vec<Expression>, code: &mut Code) -> Result<(), String> {
    code.to_bss();
    Ok(())
}

fn generate_section(args: &Vec<Expression>, code: &mut Code) -> Result<(), String> {
    let args: Vec<&ExprType> = args.iter().map(|arg| arg.get_type()).collect();
    if args.len() == 0 {
        code.to_text();
        Ok(())
    } else if let [ExprType::Symbol(sym)] = &args[..] {
        if sym == ".text" {
            code.to_text()
        } else if sym == ".data" {
            code.to_data()
        } else if sym == ".rodata" {
            code.to_rodata()
        } else if sym == ".bss" {
            code.to_bss()
        } else {
            return Err("arguments should be: {`.text`, `.data`, `.rodata`, `.bss`}.".to_string());
        }
        Ok(())
    } else {
        Err("arguments should be: {`.text`, `.data`, `.rodata`, `.bss`}.".to_string())
    }
}

fn generate_bytes(args: &Vec<Expression>, code: &mut Code) -> Result<(), String> {
    let args: Vec<&ExprType> = args.iter().map(|arg| arg.get_type()).collect();
    let mut bytes = vec![];
    for arg in args {
        match arg {
            ExprType::NumberLiteral(n) => {
                let n: u8 = (*n as u32)
                    .try_into()
                    .map_err(|_| format!("argument does not fit inside a single byte."))?;
                bytes.push(n);
            }
            _ => return Err("arguments should be: number[, number]*.".to_string()),
        }
    }
    code.append_bytes(bytes)?;
    Ok(())
}