use fmt::Formatter;
use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use id_arena::Arena;
use crate::core::{
function::{Function, FunctionId, ReturnType},
intrinsics::Intrinsic,
target_datalayout::TargetDataLayout,
target_triple::TargetTriple,
global_variable::*,
};
#[allow(dead_code)]
pub struct Module {
target_triple: Option<TargetTriple>,
target_datalayout: Option<TargetDataLayout>,
function_arena: Arena<Function>,
functions: BTreeSet<FunctionId>,
intrinsics: BTreeSet<Intrinsic>,
global_vars: BTreeMap<String, GlobalVariable>,
}
impl fmt::Display for Module {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
if let Some(layout) = &self.target_datalayout {
writeln!(f, "target datalayout = \"{}\"", layout)?;
}
if let Some(target_triple) = &self.target_triple {
writeln!(f, "target triple = \"{}\"", target_triple)?;
}
for (var_name, var) in self.global_vars.iter() {
writeln!(f, "@{} = {}", var_name, var)?;
}
for func_id in self.functions.iter() {
let func = self.function_arena.get(*func_id).unwrap();
writeln!(f, "{}", func)?;
}
for i in self.intrinsics.iter() {
writeln!(f, "declare {}", i)?;
}
Ok(())
}
}
impl Module {
pub fn set_target_triple(&mut self, triple: TargetTriple) {
self.target_triple = Some(triple);
}
pub fn set_target_datalayout(&mut self, layout: TargetDataLayout) {
self.target_datalayout = Some(layout);
}
pub fn get_function_ref_as_mut(&mut self, func_id: FunctionId) -> &mut Function {
self.function_arena.get_mut(func_id).unwrap()
}
pub fn get_function_ref(&self, func_id: FunctionId) -> &Function {
self.function_arena.get(func_id).unwrap()
}
pub fn new_function(&mut self, func_name: &str, ret_type: ReturnType) -> FunctionId {
let func_id = self
.function_arena
.alloc(Function::new(func_name, ret_type));
assert!(self.functions.insert(func_id));
func_id
}
pub fn add_intrinsic(&mut self, i: Intrinsic) {
self.intrinsics.insert(i);
}
pub fn add_global_var(&mut self, name: String, var: GlobalVariable) {
self.global_vars.insert(name, var);
}
}
impl Default for Module {
fn default() -> Self {
Self {
target_triple: None,
target_datalayout: None,
function_arena: Arena::new(),
functions: BTreeSet::new(),
intrinsics: BTreeSet::new(),
global_vars: BTreeMap::new(),
}
}
}