llvm_scratch/core/module/
module.rs

1use fmt::Formatter;
2use std::collections::{BTreeMap, BTreeSet};
3use std::fmt;
4
5use id_arena::Arena;
6
7use crate::core::{
8    function::{Function, FunctionId, ReturnType},
9    intrinsics::Intrinsic,
10    target_datalayout::TargetDataLayout,
11    target_triple::TargetTriple,
12    global_variable::*,
13};
14
15/// An representation of LLVMModule.
16#[allow(dead_code)]
17pub struct Module {
18    target_triple: Option<TargetTriple>,
19
20    target_datalayout: Option<TargetDataLayout>,
21
22    function_arena: Arena<Function>,
23    functions: BTreeSet<FunctionId>,
24
25    intrinsics: BTreeSet<Intrinsic>,
26
27    global_vars: BTreeMap<String, GlobalVariable>,
28}
29
30impl fmt::Display for Module {
31    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
32        if let Some(layout) = &self.target_datalayout {
33            writeln!(f, "target datalayout = \"{}\"", layout)?;
34        }
35        if let Some(target_triple) = &self.target_triple {
36            writeln!(f, "target triple = \"{}\"", target_triple)?;
37        }
38
39        for (var_name, var) in self.global_vars.iter() {
40            writeln!(f, "@{} = {}", var_name, var)?;
41        }
42
43        for func_id in self.functions.iter() {
44            let func = self.function_arena.get(*func_id).unwrap();
45
46            writeln!(f, "{}", func)?;
47        }
48
49        for i in self.intrinsics.iter() {
50            writeln!(f, "declare {}", i)?;
51        }
52
53        Ok(())
54    }
55}
56
57impl Module {
58    pub fn set_target_triple(&mut self, triple: TargetTriple) {
59        self.target_triple = Some(triple);
60    }
61
62    pub fn set_target_datalayout(&mut self, layout: TargetDataLayout) {
63        self.target_datalayout = Some(layout);
64    }
65
66    pub fn get_function_ref_as_mut(&mut self, func_id: FunctionId) -> &mut Function {
67        self.function_arena.get_mut(func_id).unwrap()
68    }
69
70    pub fn get_function_ref(&self, func_id: FunctionId) -> &Function {
71        self.function_arena.get(func_id).unwrap()
72    }
73
74    pub fn new_function(&mut self, func_name: &str, ret_type: ReturnType) -> FunctionId {
75        let func_id = self
76            .function_arena
77            .alloc(Function::new(func_name, ret_type));
78
79        assert!(self.functions.insert(func_id));
80
81        func_id
82    }
83
84    pub fn add_intrinsic(&mut self, i: Intrinsic) {
85        self.intrinsics.insert(i);
86    }
87
88    pub fn add_global_var(&mut self, name: String, var: GlobalVariable) {
89        self.global_vars.insert(name, var);
90    }
91}
92
93impl Default for Module {
94    fn default() -> Self {
95        Self {
96            target_triple: None,
97            target_datalayout: None,
98
99            function_arena: Arena::new(),
100            functions: BTreeSet::new(),
101
102            intrinsics: BTreeSet::new(),
103            global_vars: BTreeMap::new(),
104        }
105    }
106}