path: root/patio-to-cranelift/src/codegen/context.rs

use std::{collections::HashMap, path::Path};

use cranelift::prelude::Configurable;
use cranelift_module::Module;

use crate::errors::CompileError;

pub(crate) struct CompilerContext {
    codegen_context: cranelift_codegen::Context,
    func_name_to_func_id: HashMap<String, cranelift_module::FuncId>,
    module: cranelift_object::ObjectModule,
    pub(crate) target_triple: target_lexicon::Triple,
}

impl CompilerContext {
    pub(crate) fn for_object(
        source_path: &Path,
        target_triple: target_lexicon::Triple,
    ) -> Result<Self, CompileError> {
        let isa = Self::isa(target_triple.clone())?;

        let object_builder = cranelift_object::ObjectBuilder::new(
            isa,
            source_path.as_os_str().as_encoded_bytes(),
            cranelift_module::default_libcall_names(),
        )
        .map_err(CompileError::ObjectBuilderNewFailed)?;

        let module = cranelift_object::ObjectModule::new(object_builder);

        Ok(Self::new(module, target_triple))
    }

    fn isa(
        target: target_lexicon::Triple,
    ) -> Result<std::sync::Arc<dyn cranelift::prelude::isa::TargetIsa>, CompileError> {
        let mut flags_builder = cranelift::prelude::settings::builder();
        flags_builder
            .enable("is_pic")
            .map_err(CompileError::EnablePicFailed)?;
        cranelift::prelude::isa::lookup(target)
            .map_err(CompileError::LookupIsa)?
            .finish(cranelift::prelude::settings::Flags::new(flags_builder))
            .map_err(CompileError::BuildIsa)
    }

    fn new(module: cranelift_object::ObjectModule, target_triple: target_lexicon::Triple) -> Self {
        CompilerContext {
            codegen_context: module.make_context(),
            func_name_to_func_id: Default::default(),
            module,
            target_triple,
        }
    }

    pub(crate) fn into_object(self) -> cranelift_object::ObjectProduct {
        self.module.finish()
    }

    pub(crate) fn declare_function(
        &mut self,
        name: &str,
        linkage: cranelift_module::Linkage,
        signature: &cranelift::prelude::Signature,
    ) {
        let func_id = self
            .module
            .declare_function(name, linkage, signature)
            .unwrap();
        self.func_name_to_func_id.insert(name.to_string(), func_id);
    }

    pub(crate) fn enter_function(
        &mut self,
        name: &str,
        signature: cranelift::prelude::Signature,
        generate_function_body: impl Fn(&mut CompilerFunctionContext),
    ) {
        // TODO: clear context at start or end of function?
        self.codegen_context.clear();

        let func_id = self.func_name_to_func_id[name];

        // TODO: reuse function builder context
        let mut func_builder_context = cranelift::prelude::FunctionBuilderContext::new();
        let mut function_builder = cranelift::prelude::FunctionBuilder::new(
            &mut self.codegen_context.func,
            &mut func_builder_context,
        );
        function_builder.func.signature = signature;

        let entry_block = function_builder.create_block();
        function_builder.append_block_params_for_function_params(entry_block);
        function_builder.switch_to_block(entry_block);
        function_builder.seal_block(entry_block);

        let mut context = CompilerFunctionContext {
            func_name_to_func_id: &self.func_name_to_func_id,
            function_builder,
            module: &mut self.module,
            target_triple: self.target_triple.clone(),
            local_name_to_value: Default::default(),
            local_name_to_type: Default::default(),
            var_name_to_variable: Default::default(),
            loop_blocks: Default::default(),
        };
        generate_function_body(&mut context);

        // TODO: seal blocks individually for efficiency
        context.function_builder.seal_all_blocks();

        self.module
            .define_function(func_id, &mut self.codegen_context)
            .unwrap();
    }

    pub(crate) fn declare_data(
        &mut self,
        name: &str,
        linkage: cranelift_module::Linkage,
        mutability: Mutability,
    ) {
        self.module
            .declare_data(name, linkage, mutability.is_mutable(), false)
            .unwrap();
    }

    pub(crate) fn define_data(&mut self, name: &str, bytes: Box<[u8]>) {
        // TODO: deduplicate finding data_id from name
        let data_id = match self.module.declarations().get_name(name) {
            Some(cranelift_module::FuncOrDataId::Data(data_id)) => data_id,
            Some(cranelift_module::FuncOrDataId::Func(_)) | None => {
                todo!("no data with name: {name}")
            }
        };
        self.module
            .define_data(
                data_id,
                &cranelift_module::DataDescription {
                    // TODO: need to handle relocations?
                    init: cranelift_module::Init::Bytes { contents: bytes },
                    function_decls: Default::default(),
                    data_decls: Default::default(),
                    function_relocs: Default::default(),
                    data_relocs: Default::default(),
                    custom_segment_section: None,
                    align: None,
                    used: false,
                },
            )
            .unwrap();
    }
}

pub(crate) struct CompilerFunctionContext<'a> {
    func_name_to_func_id: &'a HashMap<String, cranelift_module::FuncId>,
    pub(crate) function_builder: cranelift::prelude::FunctionBuilder<'a>,
    module: &'a mut cranelift_object::ObjectModule,
    pub(crate) target_triple: target_lexicon::Triple,
    local_name_to_value: HashMap<String, cranelift_codegen::ir::Value>,
    local_name_to_type: HashMap<String, cranelift_codegen::ir::Type>,
    var_name_to_variable: HashMap<String, cranelift::prelude::Variable>,
    loop_blocks: Vec<Loop>,
}

impl<'a> CompilerFunctionContext<'a> {
    pub(crate) fn func_id(&mut self, func_name: &str) -> cranelift_codegen::ir::FuncRef {
        let func_id = self.func_name_to_func_id[func_name];
        self.module
            .declare_func_in_func(func_id, self.function_builder.func)
    }

    pub(crate) fn declare_local(&mut self, name: &str, type_: cranelift_codegen::ir::Type) {
        self.local_name_to_type.insert(name.to_string(), type_);
    }

    pub(crate) fn local_store(&mut self, name: &str, value: cranelift_codegen::ir::Value) {
        // TODO: error if already defined.
        self.local_name_to_value.insert(name.to_string(), value);
    }

    pub(crate) fn local_load(&self, name: &str) -> cranelift_codegen::ir::Value {
        self.local_name_to_value[name]
    }

    pub(crate) fn declare_var(&mut self, name: &str, type_: cranelift_codegen::ir::Type) {
        let variable = self.function_builder.declare_var(type_);
        self.var_name_to_variable.insert(name.to_string(), variable);
    }

    pub(crate) fn var(&mut self, name: &str) -> cranelift::prelude::Variable {
        self.var_name_to_variable[name]
    }

    pub(crate) fn local_type(&mut self, name: &str) -> cranelift_codegen::ir::Type {
        self.local_name_to_type[name]
    }

    pub(crate) fn data_ptr(&mut self, name: &str) -> cranelift_codegen::ir::GlobalValue {
        let data_id = match self.module.declarations().get_name(name) {
            Some(cranelift_module::FuncOrDataId::Data(data_id)) => data_id,
            Some(cranelift_module::FuncOrDataId::Func(_)) | None => {
                todo!("no data with name: {name}")
            }
        };
        self.module
            .declare_data_in_func(data_id, &mut self.function_builder.func)
    }

    pub(crate) fn pointer_type(&self) -> cranelift_codegen::ir::Type {
        cranelift_codegen::ir::Type::triple_pointer_type(&self.target_triple)
    }

    pub(crate) fn create_block(&mut self) -> cranelift_codegen::ir::Block {
        self.function_builder.create_block()
    }

    pub(crate) fn switch_to_block(&mut self, block: cranelift_codegen::ir::Block) {
        self.function_builder.switch_to_block(block);
    }

    pub(crate) fn seal_block(&mut self, block: cranelift_codegen::ir::Block) {
        self.function_builder.seal_block(block);
    }

    pub(crate) fn current_block(&self) -> cranelift_codegen::ir::Block {
        self.function_builder.current_block().unwrap()
    }

    pub(crate) fn enter_loop(
        &mut self,
        start_block: cranelift_codegen::ir::Block,
        end_block: cranelift_codegen::ir::Block,
    ) {
        self.loop_blocks.push(Loop {
            start_block,
            end_block,
        });
    }

    pub(crate) fn current_loop(&self) -> &Loop {
        self.loop_blocks.last().unwrap()
    }

    pub(crate) fn exit_loop(&mut self) {
        self.loop_blocks.pop();
    }
}

pub(crate) enum Mutability {
    Immutable,
    Mutable,
}

impl Mutability {
    fn is_mutable(&self) -> bool {
        match self {
            Mutability::Immutable => false,
            Mutability::Mutable => true,
        }
    }
}

pub(crate) struct Loop {
    pub(crate) start_block: cranelift_codegen::ir::Block,
    pub(crate) end_block: cranelift_codegen::ir::Block,
}