damn_simple_architecture/assembler/src/assembler/codegen.rs

use common::{args, prelude::*};

use crate::assembler::model::{Node, Opcode};
use crate::{assembler::AssembleError, expect_token};

fn log(message: &str) {
    println!("\x1b[32mINFO:\x1b[0m {message}");
}

pub fn codegen(nodes: Vec<Node>) -> Result<Vec<Instruction>, AssembleError> {
    let mut instructions = vec![];

    for node in nodes {
        instructions.push(build_instruction(&node)?);
    }

    println!("------------------------");
    log("Compilation Success ✅");

    Ok(instructions)
}

fn build_instruction(node: &Node) -> Result<Instruction, AssembleError> {
    let opcode = node.opcode();
    let args = node.args();

    match opcode {
        Opcode::Nop => Ok(Instruction::Nop),
        Opcode::Mov | Opcode::Movs => build_mov_instruction(opcode, &args),
        Opcode::Ldb
        | Opcode::Ldw
        | Opcode::Ldh
        | Opcode::Ldbs
        | Opcode::Ldhs
        | Opcode::Stb
        | Opcode::Stw
        | Opcode::Sth => build_memory_instruction(opcode, &args),
        Opcode::Lli | Opcode::Lui => build_load_immediate_instruction(opcode, &args),
        Opcode::Jmp
        | Opcode::Jeq
        | Opcode::Jne
        | Opcode::Jgt
        | Opcode::Jge
        | Opcode::Jlt
        | Opcode::Jle => build_jump_instruction(opcode, &args),
        Opcode::Cmp => build_compare_instruction(&args),
        Opcode::Inc | Opcode::Dec => build_inc_dec_instruction(opcode, &args),
        Opcode::Shl | Opcode::Shr => build_shift_instruction(opcode, &args),
        Opcode::Add
        | Opcode::Sub
        | Opcode::And
        | Opcode::Or
        | Opcode::Xor
        | Opcode::Nand
        | Opcode::Nor
        | Opcode::Xnor => build_arithmetic_instruction(opcode, &args),
        Opcode::AddI | Opcode::SubI => {
            build_arithmetic_immediate_instruction(opcode, &args)
        }
        Opcode::Not => build_not_instruction(&args),
        Opcode::Int => build_interrupt_instruction(&args),
        Opcode::Irt => Ok(Instruction::IntReturn),
        Opcode::Hlt => Ok(Instruction::Halt),
        Opcode::Data => build_data_instruction(&args),
        Opcode::Segment => build_segment_instruction(&args),
        // These pseudo-instructions should have been expanded!
        Opcode::Db
        | Opcode::Dh
        | Opcode::Dw
        | Opcode::Resb
        | Opcode::Resh
        | Opcode::Resw
        | Opcode::Push
        | Opcode::Pop
        | Opcode::Lwi
        | Opcode::Include
        | Opcode::Call
        | Opcode::Return
        | Opcode::Pusha
        | Opcode::Popa => Err(AssembleError::InvalidArg),
    }
}

fn build_mov_instruction(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(src_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(dest_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let src = expect_token!(src_token, Register)?;
    let dest = expect_token!(dest_token, Register)?;

    match opcode {
        Opcode::Mov => Ok(Instruction::Mov(args!(R, sr1: src, dr: dest))),
        Opcode::Movs => Ok(Instruction::MovSigned(args!(R, sr1: src, dr: dest))),
        _ => unreachable!(),
    }
}

fn build_memory_instruction(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(base_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(dest_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };
    let Some(offset_token) = args.get(2) else {
        return Err(AssembleError::MissingArgument(2));
    };

    let base = expect_token!(base_token, Register)?;
    let dest = expect_token!(dest_token, Register)?;
    let offset = expect_token!(offset_token, Immediate)?;
    let instruction_args = args!(I, immediate: offset as u16, r1: base, r2: dest);

    match opcode {
        Opcode::Ldb => Ok(Instruction::LoadByte(instruction_args)),
        Opcode::Ldw => Ok(Instruction::LoadWord(instruction_args)),
        Opcode::Ldh => Ok(Instruction::LoadHalfword(instruction_args)),
        Opcode::Ldbs => Ok(Instruction::LoadByteSigned(instruction_args)),
        Opcode::Ldhs => Ok(Instruction::LoadHalfwordSigned(instruction_args)),
        Opcode::Stb => Ok(Instruction::StoreByte(instruction_args)),
        Opcode::Stw => Ok(Instruction::StoreWord(instruction_args)),
        Opcode::Sth => Ok(Instruction::StoreHalfword(instruction_args)),
        _ => unreachable!(),
    }
}

fn build_load_immediate_instruction(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(value_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(dest_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let value = expect_token!(value_token, Immediate)?;
    let dest = expect_token!(dest_token, Register)?;

    match opcode {
        Opcode::Lli => {
            let instruction_args = args!(I, immediate: value as u16, r1: dest);
            Ok(Instruction::LoadLowerImmediate(instruction_args))
        }
        Opcode::Lui => {
            let upper_value = value >> 16;
            let instruction_args = args!(I, immediate: upper_value as u16, r1: dest);
            Ok(Instruction::LoadUpperImmediate(instruction_args))
        }
        _ => unreachable!(),
    }
}

fn build_jump_instruction(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(address_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(offset_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let address = expect_token!(address_token, Immediate)?;
    let offset = expect_token!(offset_token, Register)?;
    let instruction_args = args!(I, immediate: address as u16, r1: offset);

    match opcode {
        Opcode::Jmp => Ok(Instruction::Jump(instruction_args)),
        Opcode::Jeq => Ok(Instruction::JumpEq(instruction_args)),
        Opcode::Jne => Ok(Instruction::JumpNeq(instruction_args)),
        Opcode::Jgt => Ok(Instruction::JumpGt(instruction_args)),
        Opcode::Jge => Ok(Instruction::JumpGe(instruction_args)),
        Opcode::Jlt => Ok(Instruction::JumpLt(instruction_args)),
        Opcode::Jle => Ok(Instruction::JumpLe(instruction_args)),
        _ => unreachable!(),
    }
}

fn build_compare_instruction(
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(left_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(right_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let left = expect_token!(left_token, Register)?;
    let right = expect_token!(right_token, Register)?;
    Ok(Instruction::Compare(args!(R, sr1: left, sr2: right)))
}

fn build_inc_dec_instruction(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(reg_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };

    let reg = expect_token!(reg_token, Register)?;
    match opcode {
        Opcode::Inc => Ok(Instruction::Increment(args!(R, sr1: reg))),
        Opcode::Dec => Ok(Instruction::Decrement(args!(R, sr1: reg))),
        _ => unreachable!(),
    }
}

fn build_shift_instruction(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(reg_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(amount_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let reg = expect_token!(reg_token, Register)?;
    let amount = expect_token!(amount_token, Immediate)? as u8;

    match opcode {
        Opcode::Shl => Ok(Instruction::ShiftLeft(args!(R, sr1: reg, shamt: amount))),
        Opcode::Shr => Ok(Instruction::ShiftRight(args!(R, sr1: reg, shamt: amount))),
        _ => unreachable!(),
    }
}

fn build_arithmetic_instruction(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(left_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(right_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };
    let Some(dest_token) = args.get(2) else {
        return Err(AssembleError::MissingArgument(2));
    };

    let left = expect_token!(left_token, Register)?;
    let right = expect_token!(right_token, Register)?;
    let dest = expect_token!(dest_token, Register)?;
    let instruction_args = args!(R, sr1: left, sr2: right, dr: dest);

    match opcode {
        Opcode::Add => Ok(Instruction::Add(instruction_args)),
        Opcode::Sub => Ok(Instruction::Sub(instruction_args)),
        Opcode::And => Ok(Instruction::And(instruction_args)),
        Opcode::Or => Ok(Instruction::Or(instruction_args)),
        Opcode::Xor => Ok(Instruction::Xor(instruction_args)),
        Opcode::Nand => Ok(Instruction::Nand(instruction_args)),
        Opcode::Nor => Ok(Instruction::Nor(instruction_args)),
        Opcode::Xnor => Ok(Instruction::Xnor(instruction_args)),
        _ => unreachable!(),
    }
}

fn build_arithmetic_immediate_instruction(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(reg_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(immediate_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };
    let Some(dest_token) = args.get(2) else {
        return Err(AssembleError::MissingArgument(2));
    };

    let reg = expect_token!(reg_token, Register)?;
    let immediate = expect_token!(immediate_token, Immediate)? as u16;
    let dest = expect_token!(dest_token, Register)?;
    let instruction_args = args!(I, immediate: immediate, r1: reg, r2: dest);

    match opcode {
        Opcode::AddI => Ok(Instruction::AddImmediate(instruction_args)),
        Opcode::SubI => Ok(Instruction::SubImmediate(instruction_args)),
        _ => unreachable!(),
    }
}

fn build_not_instruction(
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(reg_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(dest_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let reg = expect_token!(reg_token, Register)?;
    let dest = expect_token!(dest_token, Register)?;
    Ok(Instruction::Not(args!(R, sr1: reg, dr: dest)))
}

fn build_interrupt_instruction(
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(code_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };

    let code = expect_token!(code_token, Immediate)? as u8;
    Ok(Instruction::Interrupt(Interrupt::Software(code)))
}

fn build_data_instruction(
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(immediate_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };

    let immediate = expect_token!(immediate_token, Immediate)?;
    Ok(Instruction::Data(immediate))
}

fn build_segment_instruction(
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(immediate_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };

    let immediate = expect_token!(immediate_token, Immediate)?;
    Ok(Instruction::Segment(immediate))
}