llvm_scratch/core/function/

function.rs

use fmt::Formatter;
use std::fmt;

use id_arena::{Arena, Id};

use crate::core::{
    basic_block::{BasicBlock, BasicBlockId, BasicBlockKind},
    function::param_attrs::*,
    instruction::Instruction,
    llvm_string::*,
    llvm_type::*,
};

pub type FunctionId = Id<Function>;

/// An representation of LLVMFunction.
/// LLVMFunction を表す構造体
pub struct Function {
    return_type: ReturnType,
    name: LLVMString,
    arg_list: ParameterSet,

    entry_block: BasicBlockId,
    bb_allocator: Arena<BasicBlock>,
}

impl fmt::Display for Function {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        writeln!(
            f,
            "define {} @{} {} {{",
            self.return_type, self.name, self.arg_list
        )?;

        let entry_bb = self.bb_allocator.get(self.entry_block).unwrap();

        writeln!(f, "{}", entry_bb)?;

        {
            entry_bb.print_successors(f, &self.bb_allocator)?;
        }

        writeln!(f, "}}")?;

        Ok(())
    }
}

impl Function {
    pub fn new(func_name: &str, ret_type: ReturnType) -> Self {
        let mut bb_alloc = Arena::new();
        let entry_block_id = bb_alloc.alloc(BasicBlock::new(
            LLVMString::from("entry"),
            None,
            BasicBlockKind::TERMINATED,
        ));
        Self {
            return_type: ret_type,
            name: LLVMString::from(func_name),
            arg_list: Default::default(),
            entry_block: entry_block_id,
            bb_allocator: bb_alloc,
        }
    }

    pub fn entry_bb_is_empty(&self) -> bool {
        let entry_bb = self.bb_allocator.get(self.entry_block).unwrap();
        entry_bb.insts_empty()
    }

    pub fn insert_inst(&mut self, bb: BasicBlockId, inst: Instruction) {
        if let Some(insert_bb) = self.bb_allocator.get_mut(bb) {
            insert_bb.new_inst(inst);
        }
    }

    pub fn returns_void(&self) -> bool {
        self.return_type.is_void()
    }

    pub fn get_name_ref(&self) -> &LLVMString {
        &self.name
    }

    pub fn get_entry_bb(&self) -> BasicBlockId {
        self.entry_block
    }

    pub fn args_empty(&self) -> bool {
        self.arg_list.is_empty()
    }

    pub fn new_argument(&mut self, arg: Parameter) {
        self.arg_list.add_arg(arg);
    }

    pub fn new_basic_block(
        &mut self,
        l: &str,
        pred: Option<BasicBlockId>,
        k: BasicBlockKind,
    ) -> BasicBlockId {
        self.bb_allocator
            .alloc(BasicBlock::new(LLVMString::from(l), pred, k))
    }
}

/// each function has a type signature.
#[derive(Eq, PartialEq, PartialOrd, Ord, Hash)]
pub struct FunctionType {
    return_type: ReturnType,
    arg_list: ParameterSet,
}

impl FunctionType {
    pub fn new(ret_type: ReturnType, args: ParameterSet) -> Self {
        Self {
            return_type: ret_type,
            arg_list: args,
        }
    }
}

impl fmt::Display for FunctionType {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{} {}", self.return_type, self.arg_list)
    }
}

/// each llvm function must have a return type.
/// LLVMFunctionに必ず存在する，返り値の型
#[derive(Eq, PartialEq, PartialOrd, Ord, Hash)]
pub struct ReturnType {
    return_type: LLVMType,
    /// a ReturnType may have a set of Attributes
    /// return_type は Attribute を持つ可能性がある
    attributes: Option<ParameterAttributes>,
}

impl ReturnType {
    pub fn new(ret_type: LLVMType, attrs: Option<ParameterAttributes>) -> Self {
        Self {
            return_type: ret_type,
            attributes: attrs,
        }
    }

    fn is_void(&self) -> bool {
        match &self.return_type.kind {
            LLVMTypeKind::VOID => true,
            _ => false,
        }
    }
}

impl fmt::Display for ReturnType {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match &self.attributes {
            Some(attrs) => write!(f, "{} {}", attrs, self.return_type),
            None => write!(f, "{}", self.return_type),
        }
    }
}

/// each llvm function must have a set of Parameter.
/// LLVMFunction が持つ引数リストの実装
#[derive(Eq, PartialEq, PartialOrd, Ord, Hash)]
pub struct ParameterSet {
    params: Vec<Parameter>,
}

impl Default for ParameterSet {
    fn default() -> Self {
        Self { params: Vec::new() }
    }
}

impl ParameterSet {
    pub fn new(args: Vec<Parameter>) -> Self {
        Self { params: args }
    }

    pub fn add_arg(&mut self, arg: Parameter) {
        self.params.push(arg);
    }

    fn is_empty(&self) -> bool {
        self.params.is_empty()
    }
}

impl fmt::Display for ParameterSet {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let fmt_string = self
            .params
            .iter()
            .map(|attr| attr.to_string())
            .collect::<Vec<String>>()
            .join(", ");

        write!(f, "({})", fmt_string)
    }
}

/// LLVMFunction が持つ各引数
#[derive(Eq, PartialEq, PartialOrd, Ord, Hash)]
pub struct Parameter {
    name: LLVMString,
    attributes: Option<ParameterAttributes>,
    param_type: LLVMType,
}

impl Parameter {
    pub fn new(
        param_name: LLVMString,
        attrs: Option<ParameterAttributes>,
        arg_type: LLVMType,
    ) -> Self {
        Self {
            name: param_name,
            attributes: attrs,
            param_type: arg_type,
        }
    }
}

impl fmt::Display for Parameter {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match &self.attributes {
            Some(attrs) => write!(f, "{} {} {}", self.param_type, attrs, self.name),
            None => write!(f, "{} {}", self.param_type, self.name),
        }
    }
}