lamina_mir/
module.rs

1//! Module representation in LUMIR.
2//!
3//! A module is a collection of functions and global data. Modules are the
4//! top-level unit of organization in LUMIR and can be compiled independently.
5use super::function::Function;
6use super::instruction::Instruction;
7use super::types::MirType;
8use std::collections::{HashMap, HashSet};
9use std::fmt;
10
11/// Global variable declaration
12#[derive(Debug, Clone, PartialEq)]
13pub struct Global {
14    /// Global variable name
15    pub name: String,
16
17    /// Type of the global
18    pub ty: MirType,
19
20    /// Whether this is mutable
21    pub mutable: bool,
22
23    /// Initial value (as raw bytes)
24    pub initializer: Option<Vec<u8>>,
25}
26
27impl Global {
28    pub fn new(name: impl Into<String>, ty: MirType) -> Self {
29        Self {
30            name: name.into(),
31            ty,
32            mutable: true,
33            initializer: None,
34        }
35    }
36
37    pub fn immutable(mut self) -> Self {
38        self.mutable = false;
39        self
40    }
41
42    pub fn with_initializer(mut self, data: Vec<u8>) -> Self {
43        self.initializer = Some(data);
44        self
45    }
46}
47
48/// LUMIR module
49#[derive(Debug, Clone, PartialEq)]
50pub struct Module {
51    /// Module name
52    pub name: String,
53
54    /// Functions in this module
55    pub functions: HashMap<String, Function>,
56
57    /// Global variables
58    pub globals: HashMap<String, Global>,
59
60    /// Names of external functions (functions with external linkage)
61    pub external_functions: HashSet<String>,
62}
63
64impl Module {
65    /// Create a new module with the given name
66    pub fn new(name: impl Into<String>) -> Self {
67        Self {
68            name: name.into(),
69            functions: HashMap::new(),
70            globals: HashMap::new(),
71            external_functions: HashSet::new(),
72        }
73    }
74
75    /// Add a function to this module
76    pub fn add_function(&mut self, func: Function) {
77        let name = func.sig.name.clone();
78        self.functions.insert(name, func);
79    }
80
81    /// Mark a function as external
82    pub fn mark_external(&mut self, name: impl Into<String>) {
83        self.external_functions.insert(name.into());
84    }
85
86    /// Check if a function is external
87    pub fn is_external(&self, name: &str) -> bool {
88        self.external_functions.contains(name)
89    }
90
91    /// Add a global variable to this module
92    pub fn add_global(&mut self, global: Global) {
93        let name = global.name.clone();
94        self.globals.insert(name, global);
95    }
96
97    /// Get a function by name
98    pub fn get_function(&self, name: &str) -> Option<&Function> {
99        self.functions.get(name)
100    }
101
102    /// Get a mutable reference to a function by name
103    pub fn get_function_mut(&mut self, name: &str) -> Option<&mut Function> {
104        self.functions.get_mut(name)
105    }
106
107    /// Get a global by name
108    pub fn get_global(&self, name: &str) -> Option<&Global> {
109        self.globals.get(name)
110    }
111
112    /// Get all function names
113    pub fn function_names(&self) -> Vec<&str> {
114        self.functions.keys().map(|s| s.as_str()).collect()
115    }
116
117    /// Get all global names
118    pub fn global_names(&self) -> Vec<&str> {
119        self.globals.keys().map(|s| s.as_str()).collect()
120    }
121
122    /// Total number of instructions across all functions
123    pub fn instruction_count(&self) -> usize {
124        self.functions.values().map(|f| f.instruction_count()).sum()
125    }
126
127    /// Validate the entire module
128    pub fn validate(&self) -> Result<(), Vec<String>> {
129        let mut errors = Vec::new();
130
131        // Validate all functions
132        for (name, func) in &self.functions {
133            if let Err(e) = func.validate() {
134                errors.push(format!("Function '{}': {}", name, e));
135            }
136        }
137
138        // Check for duplicate function/global names
139        for func_name in self.functions.keys() {
140            if self.globals.contains_key(func_name) {
141                errors.push(format!(
142                    "Name '{}' used for both function and global",
143                    func_name
144                ));
145            }
146        }
147
148        errors.extend(self.validate_internal_call_arities());
149
150        if errors.is_empty() {
151            Ok(())
152        } else {
153            Err(errors)
154        }
155    }
156
157    /// [`Call`](Instruction::Call) / [`TailCall`](Instruction::TailCall) arity vs in-module callee signatures.
158    fn validate_internal_call_arities(&self) -> Vec<String> {
159        let mut out = Vec::new();
160        for (caller_name, caller) in &self.functions {
161            for block in &caller.blocks {
162                for inst in &block.instructions {
163                    let (op, callee_name, argc) = match inst {
164                        Instruction::Call { name, args, .. } => ("call", name.as_str(), args.len()),
165                        Instruction::TailCall { name, args } => {
166                            ("tailcall", name.as_str(), args.len())
167                        }
168                        _ => continue,
169                    };
170                    if let Some(callee) = self.functions.get(callee_name) {
171                        let expected = callee.sig.params.len();
172                        if argc != expected {
173                            out.push(format!(
174                                "Function '{}' {} '{}' with {} args but callee expects {}",
175                                caller_name, op, callee_name, argc, expected
176                            ));
177                        }
178                    }
179                }
180            }
181        }
182        out
183    }
184}
185
186impl fmt::Display for Module {
187    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188        // Functions
189        let mut names: Vec<_> = self.functions.keys().collect();
190        names.sort();
191        for (i, name) in names.iter().enumerate() {
192            if let Some(func) = self.functions.get(*name) {
193                writeln!(f, "{}", func)?;
194                if i < names.len() - 1 {
195                    writeln!(f)?;
196                }
197            }
198        }
199        Ok(())
200    }
201}
202
203/// Builder for constructing modules
204pub struct ModuleBuilder {
205    module: Module,
206}
207
208impl ModuleBuilder {
209    pub fn new(name: impl Into<String>) -> Self {
210        Self {
211            module: Module::new(name),
212        }
213    }
214
215    pub fn function(mut self, func: Function) -> Self {
216        self.module.add_function(func);
217        self
218    }
219
220    pub fn global(mut self, global: Global) -> Self {
221        self.module.add_global(global);
222        self
223    }
224
225    pub fn build(self) -> Module {
226        self.module
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233    use crate::block::Block;
234    use crate::function::{Function, Parameter, Signature};
235    use crate::instruction::{Instruction, Operand};
236    use crate::register::{Register, VirtualReg};
237    use crate::types::{MirType, ScalarType};
238
239    #[test]
240    fn test_module_creation() {
241        let module = Module::new("test_module");
242        assert_eq!(module.name, "test_module");
243        assert!(module.functions.is_empty());
244        assert!(module.globals.is_empty());
245    }
246
247    #[test]
248    fn test_module_add_function() {
249        let mut module = Module::new("test");
250        let func = Function::new(Signature::new("test_func"));
251
252        module.add_function(func);
253        assert_eq!(module.functions.len(), 1);
254        assert!(module.get_function("test_func").is_some());
255    }
256
257    #[test]
258    fn test_module_add_global() {
259        let mut module = Module::new("test");
260        let global = Global::new("my_global", MirType::Scalar(ScalarType::I32));
261
262        module.add_global(global);
263        assert_eq!(module.globals.len(), 1);
264        assert!(module.get_global("my_global").is_some());
265    }
266
267    #[test]
268    fn test_module_builder() {
269        let func = Function::new(Signature::new("main"));
270        let global = Global::new("counter", MirType::Scalar(ScalarType::I64));
271
272        let module = ModuleBuilder::new("my_module")
273            .function(func)
274            .global(global)
275            .build();
276
277        assert_eq!(module.name, "my_module");
278        assert_eq!(module.functions.len(), 1);
279        assert_eq!(module.globals.len(), 1);
280    }
281
282    #[test]
283    fn test_validate_rejects_tailcall_arity_mismatch() {
284        let i64_ty = MirType::Scalar(ScalarType::I64);
285        let mut sink = Function::new(Signature::new("sink").with_return(i64_ty.clone()));
286        let mut sink_entry = Block::new("entry");
287        sink_entry.push(Instruction::Ret {
288            value: Some(Operand::Register(Register::Virtual(VirtualReg::gpr(0)))),
289        });
290        sink.add_block(sink_entry);
291
292        let mut bad = Function::new(
293            Signature::new("bad")
294                .with_return(i64_ty.clone())
295                .with_params(vec![Parameter::new(
296                    Register::Virtual(VirtualReg::gpr(0)),
297                    i64_ty.clone(),
298                )]),
299        );
300        let mut bad_entry = Block::new("entry");
301        bad_entry.push(Instruction::TailCall {
302            name: "sink".to_string(),
303            args: vec![Operand::Register(Register::Virtual(VirtualReg::gpr(0)))],
304        });
305        bad.add_block(bad_entry);
306
307        let mut module = Module::new("m");
308        module.add_function(sink);
309        module.add_function(bad);
310
311        let errs = module.validate().expect_err("arity mismatch");
312        assert!(
313            errs.iter()
314                .any(|e| e.contains("tailcall") && e.contains("sink")),
315            "{errs:?}"
316        );
317    }
318}
lamina_mir/module.rs

lamina_mir/
module.rs
