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

lamina_mir/
module.rs
