damn_simple_architecture/assembler/src/resolver.rs

use std::{collections::HashMap, path::PathBuf};

use common::prelude::Register;

use crate::{
    AssembleError,
    model::{Module, Node, Opcode, Symbol, Token},
    node, quick_hash,
};

pub fn resolve_symbols(nodes: &mut [Node]) -> Result<(), AssembleError> {
    let symbol_table = generate_symbol_table(nodes)?;

    for node in nodes.iter_mut() {
        match node.opcode() {
            Opcode::Lli => {
                if let Token::Symbol(symbol) =
                    node.arg(0).expect("The spanish Inquisition")
                {
                    if let Some(address) = symbol_table.get(&symbol) {
                        node.tokens[0] = Token::Immediate(*address);
                    } else {
                        return Err(AssembleError::UndefinedSymbol(symbol.clone()));
                    }
                }
            }
            Opcode::Lui => {
                if let Token::Symbol(symbol) =
                    node.arg(0).expect("The spanish Inquisition")
                {
                    if let Some(address) = symbol_table.get(&symbol) {
                        node.tokens[0] = Token::Immediate(*address);
                    } else {
                        return Err(AssembleError::UndefinedSymbol(symbol.clone()));
                    }
                }
            }
            Opcode::Jmp
            | Opcode::Jeq
            | Opcode::Jne
            | Opcode::Jgt
            | Opcode::Jge
            | Opcode::Jlt
            | Opcode::Jle => {
                if let Token::Symbol(symbol) =
                    node.arg(0).expect("The spanish Inquisition")
                {
                    if let Some(address) = symbol_table.get(&symbol) {
                        node.tokens[0] = Token::Immediate(*address);
                    } else {
                        return Err(AssembleError::UndefinedSymbol(symbol.clone()));
                    }
                }
            }
            _ => (),
        }
    }

    Ok(())
}

fn generate_symbol_table(nodes: &[Node]) -> Result<HashMap<Symbol, u32>, AssembleError> {
    let mut table = HashMap::new();

    for (i, node) in nodes.iter().enumerate() {
        if let Some(symbol) = node.label() {
            table.insert(symbol, 4 * i as u32);
        }
    }

    Ok(table)
}

pub fn resolve_dependencies(mut nodes: Vec<Node>) -> Result<Vec<Node>, AssembleError> {
    // First we get a list of imports.
    let mut dependencies = Vec::new();
    for node in &nodes {
        if let Opcode::Include = node.opcode() {
            // we want the path, and the name
            let name = if let Token::Symbol(name) =
                node.arg(0).expect("The spanish Inquisition")
            {
                name.name.clone()
            } else {
                unreachable!()
            }; //node.2.get(0).expect("The spanish Inquisition") 
            let path = if let Token::StringLit(path) =
                node.arg(1).expect("The spanish Inquisition")
            {
                path
            } else {
                unreachable!()
            };
            let hash = quick_hash(
                &PathBuf::from(path)
                    .canonicalize()
                    .expect("ERROR: Invalid import path."),
            );

            dependencies.push((name, hash));
        }
    }

    let mut changes = Vec::<(u32, u32, Symbol)>::new();
    // now we resolve the symbols on all the nodes
    // we need to check all operands for unresolved signals
    for (i, node) in nodes.clone().iter().enumerate() {
        let Node {
            tokens: operands, ..
        } = node;
        for (j, token) in operands.iter().enumerate() {
            if let Token::Symbol(symbol) = token {
                for d in &dependencies {
                    if let Module::Unresolved(name) = symbol.module.clone() {
                        if name != d.0 {
                            continue;
                        }

                        let symbol = Symbol {
                            name: symbol.name.clone(),
                            module: Module::Resolved(d.1),
                        };
                        changes.push((i as u32, j as u32, symbol));
                    }
                }
            }
        }
    }

    for (i, j, symbol) in changes {
        nodes[i as usize].tokens[j as usize] = Token::Symbol(symbol);
    }

    Ok(nodes)
}

pub fn create_sections(nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let mut res = Vec::<Node>::with_capacity(nodes.len());

    res.push(node!(None, Opcode::Segment, Token::Immediate(0)));

    for n in nodes.iter() {
        if n.opcode() == Opcode::Data {
            res.push(n.clone());
        }
    }

    let start = res.len() + 1;
    res.insert(
        0,
        node!(
            None,
            Opcode::Jmp,
            Token::Immediate(start as u32 * 4),
            Token::Register(Register::Zero)
        ),
    );
    for n in nodes.iter() {
        if !matches!(n.opcode(), Opcode::Data | Opcode::Include) {
            res.push(n.clone());
        }
    }

    *nodes = res;

    Ok(())
}