damn_simple_architecture/assembler/src/tooling/brainf.rs

use std::{fs, path::Path};

use common::prelude::*;

use crate::{
    assembler::{
        Module, Node, Opcode, Symbol, Token, codegen, create_sections, expand_pseudo_ops,
        resolve_symbols,
    },
    node,
};

#[must_use]
pub fn build(src: &Path) -> Vec<Instruction> {
    let src = fs::read_to_string(src).expect("Failed to read source file");
    let mut nodes = parse(&src);

    // we need to expand pseudoinstructions etc now
    nodes = expand_pseudo_ops(nodes, 0).expect("Failed to expand pseudo-operations");

    create_sections(&mut nodes).expect("Failed to create sections");

    for n in &nodes {
        println!("{n}");
    }

    resolve_symbols(&mut nodes).expect("Failed to resolve symbols");

    codegen(nodes).expect("Failed to generate code from nodes")
}

#[must_use]
#[expect(clippy::too_many_lines)]
pub fn parse(src: &str) -> Vec<Node> {
    let stack = Token::Immediate(0x10000);
    let acc = Token::Register(Register::Acc);
    let rga = Token::Register(Register::Rga);

    let bpr = Token::Register(Register::Bpr);
    let spr = Token::Register(Register::Spr);
    let mut nodes = Vec::<Node>::new();

    // Define symbols
    let print_start = Symbol {
        name: "print".to_string(),
        module: Module::Resolved(0),
    };

    let tokens = lex(src);

    // let _id = 0;
    let mut idstack = Vec::<u32>::new();

    nodes.extend(vec![
        // set up a stack
        node!(None, Opcode::Lwi, stack, bpr),
        node!(None, Opcode::Mov, bpr, spr),
        // set up the data pointer
        node!(
            Some(Symbol {
                name: "main".to_string(),
                module: Module::Resolved(0)
            }),
            Opcode::Lwi,
            Token::Immediate(0x30000),
            rga
        ),
    ]);

    for (id, tok) in tokens.iter().enumerate() {
        match tok {
            BfToken::Inc => {
                // inc acc
                nodes.extend(vec![node!(None, Opcode::Inc, acc)]);
            }
            BfToken::Dec => {
                // dec acc
                nodes.extend(vec![node!(None, Opcode::Dec, acc)]);
            }
            BfToken::IncPtr => {
                // stb acc, rga
                // add rga, 4
                // ldb rga, acc
                nodes.extend(vec![
                    node!(None, Opcode::Stw, acc, rga, 0),
                    node!(None, Opcode::AddI, rga, 4, rga),
                    node!(None, Opcode::Ldw, rga, acc, 0),
                ]);
            }
            BfToken::DecPtr => {
                // stb acc, rga
                // sub rga, 4
                // ldb rga, acc
                nodes.extend(vec![
                    node!(None, Opcode::Stw, acc, rga, 0),
                    node!(None, Opcode::SubI, rga, 4, rga),
                    node!(None, Opcode::Ldw, rga, acc, 0),
                ]);
            }
            BfToken::Out => {
                // push rga
                // call print
                // pop zero
                nodes.extend(vec![
                    node!(None, Opcode::Push, acc),
                    node!(None, Opcode::Call, Token::Symbol(print_start.clone())),
                    node!(None, Opcode::Pop, Token::Register(Register::Zero)),
                ]);
            }
            BfToken::In => {
                // Read a byte from input and store it at the current data pointer
                // Assuming we have an input function mapped to a specific memory location or I/O port
                nodes.extend(vec![
                    // Read input (assuming input is mapped to memory address 0x40000)
                    node!(None, Opcode::Ldw, Token::Immediate(0x40000), acc, 0),
                    // Store the input byte at the current data pointer
                ]);
            }
            BfToken::Forward => {
                // Start of loop [
                let loop_start = format!("loop_start_{id}");
                let loop_end = format!("loop_end_{id}");

                // Push the current position for the matching ]
                idstack.push(id as u32);

                // Load current cell value and check if zero
                nodes.extend(vec![
                    // Compare with zero
                    node!(None, Opcode::Cmp, acc, Token::Register(Register::Zero)),
                    // If zero, jump to end of loop
                    node!(
                        None,
                        Opcode::Jeq,
                        Token::Symbol(Symbol {
                            name: loop_end,
                            module: Module::Resolved(0),
                        }),
                        Token::Register(Register::Zero)
                    ),
                ]);

                // Add label for loop start
                nodes.push(node!(
                    Some(Symbol {
                        name: loop_start,
                        module: Module::Resolved(0),
                    }),
                    Opcode::Nop
                ));
            }
            BfToken::Back => {
                // End of loop ]
                if let Some(start_id) = idstack.pop() {
                    let loop_start = format!("loop_start_{start_id}");
                    let loop_end = format!("loop_end_{start_id}");

                    // Jump back to the start of the loop
                    nodes.extend(vec![
                        // Compare with zero
                        node!(None, Opcode::Cmp, acc, Token::Register(Register::Zero)),
                        // If not zero, jump back to start of loop
                        node!(
                            None,
                            Opcode::Jne,
                            Token::Symbol(Symbol {
                                name: loop_start,
                                module: Module::Resolved(0),
                            }),
                            Token::Register(Register::Zero)
                        ),
                        // Add label for loop end
                        node!(
                            Some(Symbol {
                                name: loop_end,
                                module: Module::Resolved(0),
                            }),
                            Opcode::Nop
                        ),
                    ]);
                } else {
                    // Unmatched closing bracket - could add error handling here
                    eprintln!("Warning: Unmatched ']' at position {id}");
                }
            }
        }
    }

    nodes.push(node!(None, Opcode::Hlt));

    insert_lib(&mut nodes);

    nodes
}

fn insert_lib(nodes: &mut Vec<Node>) {
    let bpr = Token::Register(Register::Bpr);
    let spr = Token::Register(Register::Spr);
    let rg0 = Token::Register(Register::Rg0);
    let rg1 = Token::Register(Register::Rg1);

    let print_start = Symbol {
        name: "print".to_string(),
        module: Module::Resolved(0),
    };
    let current = Symbol {
        name: "current".to_string(),
        module: Module::Resolved(0),
    };
    // set up the program framework.
    nodes.extend(vec![
        // set display to 0x20000
        node!(
            None,
            Opcode::Dw,
            Token::Symbol(current.clone()),
            Token::Immediate(0x20000)
        ),
        // print function
        // initialisation
        node!(Some(print_start), Opcode::Push, bpr),
        node!(None, Opcode::Mov, spr, bpr),
        // function body
        node!(None, Opcode::Ldw, bpr, rg0, Token::Immediate(8)),
        node!(
            None,
            Opcode::Ldw,
            Token::Symbol(current.clone()), // Load address of current
            rg1,
            Token::Immediate(0)
        ),
        node!(None, Opcode::Stb, rg0, rg1, Token::Immediate(0)),
        node!(None, Opcode::AddI, rg1, Token::Immediate(1), rg1),
        // function return according to spec.
        node!(
            None,
            Opcode::Stw,
            rg1,
            Token::Symbol(current), // Store back to current
            Token::Immediate(0)
        ),
        node!(None, Opcode::Mov, bpr, spr),
        node!(None, Opcode::Pop, bpr),
        node!(None, Opcode::Return),
    ]);
}

enum BfToken {
    Inc,
    Dec,
    IncPtr,
    DecPtr,
    Out,
    In,
    Forward,
    Back,
}

fn lex(src: &str) -> Vec<BfToken> {
    src.chars()
        .filter_map(|c| match c {
            '+' => Some(BfToken::Inc),
            '-' => Some(BfToken::Dec),
            '>' => Some(BfToken::IncPtr),
            '<' => Some(BfToken::DecPtr),
            '.' => Some(BfToken::Out),
            ',' => Some(BfToken::In),
            '[' => Some(BfToken::Forward),
            ']' => Some(BfToken::Back),
            _ => None,
        })
        .collect()
}

fn _create_symbol(id: u32) -> Symbol {
    Symbol {
        name: format!("label_{id}"),
        module: Module::Resolved(0),
    }
}