damn_simple_architecture/assembler/src/expand.rs

use common::prelude::Register;

use crate::{
    AssembleError, expect_token, expect_type,
    model::{Node, Opcode, Token},
    node,
};

pub fn expand_pseudo_ops(
    mut nodes: Vec<Node>,
    module: u64,
) -> Result<Vec<Node>, AssembleError> {
    let mut result = Vec::<Node>::with_capacity(nodes.len());

    for node in nodes.iter_mut() {
        if try_expand(node.clone(), &mut result, module).is_err() {
            result.push(node.clone());
        }
    }

    Ok(result)
}

fn try_expand(
    node: Node,
    result: &mut Vec<Node>,
    _module: u64,
) -> Result<(), AssembleError> {
    match node.opcode() {
        Opcode::Push => expand_push(node.clone(), result)?,
        Opcode::Pop => expand_pop(node.clone(), result)?,
        Opcode::Ldb | Opcode::Ldbs | Opcode::Ldh | Opcode::Ldhs | Opcode::Ldw => {
            expand_ldx(node.clone(), result)?
        }
        Opcode::Stb | Opcode::Sth | Opcode::Stw => expand_stx(node.clone(), result)?,

        Opcode::Lwi => expand_lwi(node.clone(), result)?,
        Opcode::Resb | Opcode::Resh | Opcode::Resw => expand_resx(node.clone(), result)?,
        Opcode::Db | Opcode::Dh | Opcode::Dw => expand_dx(node.clone(), result)?,
        _ => result.push(node.clone()),
    };
    Ok(())
}

fn expand_push(current: Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let label = current.label();
    let reg = expect_type!(current.arg(0).unwrap(), Register)?;

    nodes.extend(vec![
        node!(
            label,
            Opcode::Isub,
            Token::Register(Register::Spr),
            Token::Immediate(4),
            Token::Register(Register::Spr)
        ),
        node!(
            None,
            Opcode::Stw,
            reg,
            Token::Register(Register::Spr),
            Token::Immediate(0)
        ),
    ]);

    Ok(())
}

fn expand_pop(current: Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let label = current.label();
    let reg = expect_type!(current.arg(0).unwrap(), Register)?;

    nodes.extend(vec![
        node!(
            label,
            Opcode::Ldw,
            Token::Register(Register::Spr),
            reg,
            Token::Immediate(0)
        ),
        node!(
            None,
            Opcode::Iadd,
            Token::Register(Register::Spr),
            Token::Immediate(4),
            Token::Register(Register::Spr)
        ),
    ]);

    Ok(())
}

fn expand_ldx(current: Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let opcode = current.opcode();
    let name = expect_type!(current.arg(0).unwrap(), Symbol)?;
    let reg = expect_type!(current.arg(1).unwrap(), Register)?;
    let offset = expect_type!(current.arg(2).unwrap(), Immediate)?;

    nodes.extend(vec![
        node!(current.label(), Opcode::Lli, name.clone(), reg.clone()),
        node!(None, Opcode::Lui, name.clone(), reg.clone()),
        node!(None, opcode, reg.clone(), reg, offset),
    ]);

    Ok(())
}

fn expand_stx(current: Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let opcode = current.opcode();
    let base = expect_type!(current.arg(0).unwrap(), Register)?;
    let dest = expect_type!(current.arg(1).unwrap(), Symbol)?;
    let offset = expect_type!(current.arg(2).unwrap(), Immediate)?;
    nodes.extend(vec![
        node!(
            current.label(),
            Opcode::Lli,
            dest.clone(),
            Token::Register(Register::Rgf)
        ),
        node!(
            None,
            Opcode::Lui,
            dest.clone(),
            Token::Register(Register::Rgf)
        ),
        node!(None, opcode, base, Token::Register(Register::Rgf), offset),
    ]);

    Ok(())
}

fn expand_lwi(current: Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let val = expect_type!(current.arg(0).unwrap(), Symbol, Immediate)?;
    let reg = expect_type!(current.arg(1).unwrap(), Register)?;

    nodes.extend(vec![
        node!(current.label(), Opcode::Lli, val.clone(), reg.clone()),
        node!(None, Opcode::Lui, val.clone(), reg.clone()),
    ]);

    Ok(())
}

fn expand_resx(current: Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let region_label = expect_token!(current.arg(0).unwrap(), Symbol)?;
    let size = expect_token!(current.arg(1).unwrap(), Immediate)?;

    let units_per = match current.opcode() {
        Opcode::Resb => 4,
        Opcode::Resh => 2,
        Opcode::Resw => 1,
        _ => unreachable!(),
    };

    let mut buffer = vec![];
    // push the inital node with the label
    for _ in 0..size.div_ceil(units_per) {
        // push the rest of the nodes
        buffer.push(node!(None, Opcode::Data, Token::Immediate(0)));
    }
    buffer[0].symbol = Some(region_label);
    nodes.extend(buffer);

    Ok(())
}

fn expand_dx(current: Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let region_label = expect_token!(current.arg(0).unwrap(), Symbol)?;
    let size = match current.opcode() {
        Opcode::Db => 4,
        Opcode::Dh => 2,
        Opcode::Dw => 1,
        _ => unreachable!(),
    };

    let mut buffer = vec![];

    let mut args = current.args();
    let _label = args.remove(0);

    for word in process_dx_data(args, size)? {
        buffer.push(node!(None, Opcode::Data, Token::Immediate(word)));
    }
    buffer[0].symbol = Some(region_label);

    nodes.extend(buffer);
    Ok(())
}

fn process_dx_data(args: Vec<Token>, size: usize) -> Result<Vec<u32>, AssembleError> {
    assert!(matches!(size, 1 | 2 | 4));

    let mut buffer = Vec::<u8>::new();

    // Process each token
    for token in args {
        match token {
            Token::StringLit(mut s) => {
                s.push('\0');
                // Split string into chars and write as bytes
                for ch in s.chars() {
                    // Convert char to bytes (UTF-8 encoding)
                    let mut char_buf = [0u8; 4];
                    let char_bytes = ch.encode_utf8(&mut char_buf);
                    buffer.extend_from_slice(char_bytes.as_bytes());
                }
            }
            Token::Immediate(value) => {
                // Split u32 into bytes (little-endian)
                buffer.extend_from_slice(&value.to_be_bytes());
            }
            _ => {
                return Err(AssembleError::Generic);
            }
        }

        // Pad buffer to alignment boundary with zeros
        let remainder = buffer.len() % size;
        if remainder != 0 {
            let padding = size - remainder;
            buffer.resize(buffer.len() + padding, 0);
        }
    }

    // Convert byte buffer to u32 chunks
    // Pad final buffer to u32 boundary if needed
    let remainder = buffer.len() % 4;
    if remainder != 0 {
        let padding = 4 - remainder;
        buffer.resize(buffer.len() + padding, 0);
    }

    // Convert bytes to u32s efficiently using chunks_exact
    let result = buffer
        .chunks_exact(4)
        .map(|chunk| {
            // Convert 4 bytes to u32 (little-endian)
            u32::from_be_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])
        })
        .collect();

    Ok(result)
}