lamina-mir 0.2.1

//! Module representation in LUMIR.
//!
//! A module is a collection of functions and global data. Modules are the
//! top-level unit of organization in LUMIR and can be compiled independently.
use lamina_platform::PointerWidth;

pub use lamina_platform::TlsModel;

use crate::{
    function::{Function, SymbolVisibility},
    types::MirType,
    verify::{function_errors, verify_module_in_pipeline},
};
use std::{
    collections::{BTreeMap, HashMap, HashSet},
    fmt,
    num::NonZeroU32,
};

/// Global variable declaration
#[derive(Debug, Clone, PartialEq)]
pub struct Global {
    /// Global variable name
    pub name: String,

    /// Type of the global
    pub ty: MirType,

    /// Whether this is mutable
    pub mutable: bool,

    /// Initial value (as raw bytes)
    pub initializer: Option<Vec<u8>>,

    /// Whether the symbol is visible to other modules
    pub exported: bool,

    /// Symbol visibility requested by `@internal`, `@private`, `@hidden` or `@protected`
    pub visibility: Option<SymbolVisibility>,

    /// Whether the symbol is a weak definition
    pub weak: bool,

    /// Bytes the symbol occupies. Larger than `ty` for an aggregate, whose element
    /// type is all MIR can name.
    pub size: usize,

    /// Alignment the symbol needs, in bytes.
    pub align: usize,

    /// The thread-local model this symbol is placed under, or `None` for ordinary
    /// static storage.
    pub thread_local: Option<TlsModel>,

    /// Addresses the initializer bytes cannot hold, because a symbol's address is
    /// not known until link time. Each one covers `size` bytes at `offset`.
    pub relocations: Vec<GlobalRelocation>,
}

/// One symbol address inside a global's initializer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GlobalRelocation {
    /// Byte offset into the initializer.
    pub offset: usize,
    /// Symbol whose address is written, before mangling.
    pub symbol: String,
    /// Added to that address.
    pub addend: i64,
    /// Width of the slot, in bytes.
    pub size: usize,
}

impl Global {
    pub fn new(name: impl Into<String>, ty: MirType, pointer: PointerWidth) -> Self {
        Self {
            name: name.into(),
            ty,
            mutable: true,
            initializer: None,
            exported: false,
            visibility: None,
            weak: false,
            size: ty.size_bytes(pointer),
            align: ty.alignment(pointer),
            thread_local: None,
            relocations: Vec::new(),
        }
    }

    /// Place the global in thread-local storage under `model`.
    pub fn thread_local(mut self, model: TlsModel) -> Self {
        self.thread_local = Some(model);
        self
    }

    /// Give the global a layout of its own, for an aggregate that `ty` cannot describe.
    pub fn with_layout(mut self, size: usize, align: usize) -> Self {
        self.size = size;
        self.align = align;
        self
    }

    pub fn exported(mut self) -> Self {
        self.exported = true;
        self
    }

    pub fn immutable(mut self) -> Self {
        self.mutable = false;
        self
    }

    pub fn with_initializer(mut self, data: Vec<u8>) -> Self {
        self.initializer = Some(data);
        self
    }
}

/// How generated code addresses symbols, set by the `@pic` and `@pie` module annotations.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RelocationModel {
    /// Non-PIC code: symbols are addressed directly.
    #[default]
    Static,
    /// Position-independent code: external and exported symbols are reached through the GOT.
    Pic,
    /// Position-independent executable: external symbols are reached through the GOT.
    Pie,
}

/// LUMIR module
#[derive(Debug, Clone, PartialEq)]
pub struct Module {
    /// Module name
    pub name: String,

    /// Functions in this module
    pub functions: HashMap<String, Function>,

    /// Global variables
    pub globals: HashMap<String, Global>,

    /// Names of external functions (functions with external linkage)
    pub external_functions: HashSet<String>,

    /// Names of functions exported with global linkage
    pub exported_functions: HashSet<String>,

    /// Relocation model requested by the module annotations
    pub relocation: RelocationModel,

    pub source_files: BTreeMap<NonZeroU32, String>,
}

impl Module {
    /// Create a new module with the given name
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            functions: HashMap::new(),
            globals: HashMap::new(),
            external_functions: HashSet::new(),
            exported_functions: HashSet::new(),
            relocation: RelocationModel::Static,
            source_files: BTreeMap::new(),
        }
    }

    /// Add a function to this module
    pub fn add_function(&mut self, func: Function) {
        let name = func.sig.name.clone();
        self.functions.insert(name, func);
    }

    /// Mark a function as external
    pub fn mark_external(&mut self, name: impl Into<String>) {
        self.external_functions.insert(name.into());
    }

    /// Check if a function is external
    pub fn is_external(&self, name: &str) -> bool {
        self.external_functions.contains(name)
    }

    /// Mark a function as exported
    pub fn mark_exported(&mut self, name: impl Into<String>) {
        self.exported_functions.insert(name.into());
    }

    /// Check if a function is exported
    pub fn is_exported(&self, name: &str) -> bool {
        self.exported_functions.contains(name)
    }

    /// Add a global variable to this module
    pub fn add_global(&mut self, global: Global) {
        let name = global.name.clone();
        self.globals.insert(name, global);
    }

    /// Get a function by name
    pub fn get_function(&self, name: &str) -> Option<&Function> {
        self.functions.get(name)
    }

    /// Get a mutable reference to a function by name
    pub fn get_function_mut(&mut self, name: &str) -> Option<&mut Function> {
        self.functions.get_mut(name)
    }

    /// Get a global by name
    pub fn get_global(&self, name: &str) -> Option<&Global> {
        self.globals.get(name)
    }

    /// Get all function names
    pub fn function_names(&self) -> Vec<&str> {
        self.functions.keys().map(String::as_str).collect()
    }

    /// Get all global names
    pub fn global_names(&self) -> Vec<&str> {
        self.globals.keys().map(String::as_str).collect()
    }

    /// Total number of instructions across all functions
    pub fn instruction_count(&self) -> usize {
        self.functions
            .values()
            .map(Function::instruction_count)
            .sum()
    }

    /// Validate the entire module
    pub fn validate(&self) -> Result<(), Vec<String>> {
        let mut errors = Vec::new();

        for func in self.functions.values() {
            errors.extend(function_errors(func, self));
        }

        errors.extend(self.module_errors());

        if errors.is_empty() {
            Ok(())
        } else {
            Err(errors)
        }
    }

    /// Validate a module at compiler pipeline boundaries.
    #[doc(hidden)]
    pub fn validate_in_pipeline(&self) -> Result<(), Vec<String>> {
        let mut errors = match verify_module_in_pipeline(self) {
            Ok(()) => Vec::new(),
            Err(errors) => errors,
        };
        errors.extend(self.module_errors());

        if errors.is_empty() {
            Ok(())
        } else {
            Err(errors)
        }
    }

    fn module_errors(&self) -> Vec<String> {
        let mut errors = Vec::new();

        // Check for duplicate function/global names
        for func_name in self.functions.keys() {
            if self.globals.contains_key(func_name) {
                errors.push(format!(
                    "Name '{func_name}' used for both function and global"
                ));
            }
        }

        errors
    }
}

impl fmt::Display for Module {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Functions
        let mut names: Vec<_> = self.functions.keys().collect();
        names.sort();
        for (i, name) in names.iter().enumerate() {
            if let Some(func) = self.functions.get(*name) {
                writeln!(f, "{func}")?;
                if i < names.len() - 1 {
                    writeln!(f)?;
                }
            }
        }
        Ok(())
    }
}

/// Builder for constructing modules
pub struct ModuleBuilder {
    module: Module,
}

impl ModuleBuilder {
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            module: Module::new(name),
        }
    }

    pub fn function(mut self, func: Function) -> Self {
        self.module.add_function(func);
        self
    }

    pub fn global(mut self, global: Global) -> Self {
        self.module.add_global(global);
        self
    }

    /// Returns the module without running `Module::validate`; validate separately if the
    /// inputs are untrusted.
    pub fn build(self) -> Module {
        self.module
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        block::Block,
        function::{Function, Parameter, Signature},
        instruction::{Instruction, Operand},
        register::{Register, VirtualReg},
        types::{MirType, ScalarType},
    };

    #[test]
    fn test_module_add_function() {
        let mut module = Module::new("test");
        let func = Function::new(Signature::new("test_func"));

        module.add_function(func);
        assert_eq!(module.functions.len(), 1);
        assert!(module.get_function("test_func").is_some());
    }

    #[test]
    fn test_module_add_global() {
        let mut module = Module::new("test");
        let global = Global::new(
            "my_global",
            MirType::Scalar(ScalarType::I32),
            PointerWidth::Bits64,
        );

        module.add_global(global);
        assert_eq!(module.globals.len(), 1);
        assert!(module.get_global("my_global").is_some());
    }

    #[test]
    fn test_module_builder() {
        let func = Function::new(Signature::new("main"));
        let global = Global::new(
            "counter",
            MirType::Scalar(ScalarType::I64),
            PointerWidth::Bits64,
        );

        let module = ModuleBuilder::new("my_module")
            .function(func)
            .global(global)
            .build();

        assert_eq!(module.name, "my_module");
        assert_eq!(module.functions.len(), 1);
        assert_eq!(module.globals.len(), 1);
    }

    #[test]
    fn test_validate_rejects_tailcall_arity_mismatch() {
        let i64_ty = MirType::Scalar(ScalarType::I64);
        let mut sink = Function::new(Signature::new("sink").with_return(i64_ty));
        let mut sink_entry = Block::new("entry");
        sink_entry.push(Instruction::Ret {
            value: Some(Operand::Register(Register::Virtual(VirtualReg::gpr(0)))),
        });
        sink.add_block(sink_entry);

        let mut bad = Function::new(Signature::new("bad").with_return(i64_ty).with_params(vec![
            Parameter::new(Register::Virtual(VirtualReg::gpr(0)), i64_ty),
        ]));
        let mut bad_entry = Block::new("entry");
        bad_entry.push(Instruction::TailCall {
            name: "sink".to_string(),
            args: vec![Operand::Register(Register::Virtual(VirtualReg::gpr(0)))],
        });
        bad.add_block(bad_entry);

        let mut module = Module::new("m");
        module.add_function(sink);
        module.add_function(bad);

        let errs = module.validate().expect_err("arity mismatch");
        assert!(
            errs.iter()
                .any(|e| e.contains("tailcall") && e.contains("sink")),
            "{errs:?}"
        );
    }
}