damn_simple_architecture/assembler/src/assembler/assembler.rs

use std::{
    collections::HashSet,
    fs,
    path::{self, Path, PathBuf},
    sync::{Arc, Mutex},
    thread::{self, JoinHandle},
};

use crate::assembler::{AssembleError, Token, expand_pseudo_ops, lexer, quick_hash};
use crate::assembler::{Node, Parser, resolve_dependencies};
// use crate::util::logging::Logger;

// pub fn new_assemble(path: &Path) {
//     let program = Program::new();
//     let program_ref = ProgramRef::new(program);

//     let task = Module::build(path.to_path_buf(), program_ref.clone());
//     program_ref.add_task(task);

//     // wait on all tasks to finish
//     for task in program_ref.get_tasks() {
//         let module = task.module.join().unwrap();
//         program_ref.add_module(module);
//     }
// }

pub struct Program {
    pub main_path: PathBuf,
    registry: HashSet<u64>,
    modules: Vec<Module>,
    tasks: Vec<Task>,
    logger: Logger,
}

impl Program {
    #[must_use]
    pub fn new() -> Self {
        Self {
            registry: HashSet::new(),
            modules: Vec::new(),
            tasks: Vec::new(),
            main_path: PathBuf::new(),
            logger: Logger::new(),
        }
    }

    pub fn add_task(&mut self, task: Task) {
        self.tasks.push(task);
    }
}

impl Default for Program {
    fn default() -> Self {
        Self::new()
    }
}

pub struct ProgramRef {
    program: Arc<Mutex<Program>>,
}

impl ProgramRef {
    #[must_use]
    pub fn new(program: Program) -> Self {
        Self {
            program: Arc::new(Mutex::new(program)),
        }
    }

    pub fn register(&self, path: &Path) {
        self.program
            .lock()
            .expect("Failed to acquire program lock")
            .registry
            .insert(quick_hash(path));
    }

    #[must_use]
    pub fn is_registered(&self, path: &Path) -> bool {
        self.program
            .lock()
            .expect("Failed to acquire program lock")
            .registry
            .contains(&quick_hash(path))
    }

    // pub fn get_tasks(&self) -> Vec<&Task> {
    //     self.program.lock().unwrap().tasks.iter().collect()
    // }

    pub fn add_task(&self, task: Task) {
        self.program
            .lock()
            .expect("Failed to acquire program lock")
            .add_task(task);
    }

    pub fn add_module(&self, module: Module) {
        self.program
            .lock()
            .expect("Failed to acquire program lock")
            .modules
            .push(module);
    }

    pub fn log(&self, message: &str) {
        self.program
            .lock()
            .expect("Failed to acquire program lock")
            .logger
            .log(message);
    }
}

impl Clone for ProgramRef {
    fn clone(&self) -> Self {
        Self {
            program: self.program.clone(),
        }
    }
}

pub struct Module {
    pub path: PathBuf,
    pub hash: u64,
    pub nodes: Vec<Node>,
    program: ProgramRef,
}

impl Module {
    #[must_use]
    pub const fn new(
        path: PathBuf,
        hash: u64,
        nodes: Vec<Node>,
        program: ProgramRef,
    ) -> Self {
        Self {
            path,
            hash,
            nodes,
            program,
        }
    }

    pub fn build(path: PathBuf, program: ProgramRef) -> Result<Task, AssembleError> {
        // Spawn a thread that creates the main function and executes the lexer and
        // parser.
        let handle = thread::spawn(move || {
            let mut module =
                Self::new(path.clone(), quick_hash(&path), Vec::new(), program.clone());

            match module.lex() {
                Ok(tokens) => {
                    module.parse(tokens);
                    module.expand();
                    module.prepare_dependencies();
                    module
                }
                Err(why) => {
                    eprintln!(
                        "Error building program at path `{}`: {why}",
                        path.display()
                    );

                    // TODO: Find a way to make this work without panicking.
                    unreachable!()
                }
            }
        });

        Ok(Task { module: handle })
    }

    fn lex(&self) -> Result<Vec<Token>, AssembleError> {
        if let Ok(path) = self.path.canonicalize() {
            self.program.log(&format!(
                "{:20} {:20} [{}]",
                "Building",
                self.get_filename(),
                path.display()
            ));
        }

        let src = fs::read_to_string(&self.path)
            .map_err(|_| AssembleError::InvalidFile(self.path.clone()))?;

        let file_hash = quick_hash(&self.path);

        self.program
            .log(&format!("{:20} {:20}", "Tokenising", self.get_filename()));

        lexer::lexer(src, file_hash)
    }

    fn parse(&mut self, tokens: Vec<Token>) -> Result<(), AssembleError> {
        self.program
            .log(&format!("{:20} {:20}", "Parsing", self.get_filename()));

        let parsed = Parser::parse_nodes(tokens)?;
        self.nodes = parsed;

        Ok(())
    }

    fn expand(&mut self) -> Result<(), AssembleError> {
        self.program
            .log(&format!("{:20} {:20}", "Expanding", self.get_filename()));

        let expanded = expand_pseudo_ops(self.nodes.clone(), self.hash)?;
        self.nodes = expanded;

        Ok(())
    }

    fn prepare_dependencies(&self) -> Result<(), AssembleError> {
        let nodes = resolve_dependencies(
            self.nodes.clone(),
            self.path.parent().expect("File should have a parent path!"),
        )?;

        let dependencies = Parser::get_dependencies(&nodes, &self.path)?;

        for dep in dependencies {
            if self.program.is_registered(&dep) {
                // we have already built this module!
                continue;
            }

            self.program.register(&dep);

            // create new module
            // add the task to the program

            match Self::build(dep, self.program.clone()) {
                Ok(task) => self.program.add_task(task),
                Err(why) => {
                    eprintln!("Error building program: {why}");
                }
            }
        }

        Ok(())
    }

    /// Gets the filename from a [`PathBuf`].
    fn get_filename(&self) -> &str {
        self.path
            .file_name()
            .and_then(|f| f.to_str())
            .unwrap_or_default()
    }

    /// Gets the parent filepath from a [`PathBuf`].
    fn get_parent(&self) -> &str {
        self.path
            .parent()
            .and_then(|f| f.to_str())
            .unwrap_or_default()
    }
}

pub struct Task {
    module: JoinHandle<Module>,
}