bcomp 0.1.0

//! The emitter for turning the abstract syntax tree into Cranelift IR.

use cranelift::prelude::*;
use cranelift_codegen::{Context, isa, settings};
use cranelift_frontend::FunctionBuilderContext;
use cranelift_module::{DataDescription, Linkage, Module, default_libcall_names};
use cranelift_object::{ObjectBuilder, ObjectModule};
use target_lexicon::HOST;

use crate::parser::{
    ast::{PrintItem, Program, Statement},
    expression::{BinaryOperator, Expression, LiteralExpression},
};

fn main_signature(isa: &dyn isa::TargetIsa) -> Signature {
    // The `CallConv` defines how primitives in parameters and return values are handled.
    // Mainly which registers are used and when stack spills are used.
    //
    // In general, it's best to use `CallConv::Fast`.
    //
    // However, since the function we define is invoked from our targeted OS, we need to use
    // the calling convention the OS expects.
    let call_conv = isa.default_call_conv();

    Signature {
        call_conv,
        params: vec![],
        // Since we're linking to libc, we can return the exit code from main.
        returns: vec![AbiParam::new(types::I32)],
    }
}

/// The emitter state.
pub struct Emitter {
    module: ObjectModule,
    data_count: usize,
}

impl Emitter {
    /// Constructs a new instance of the emitter.
    ///
    /// This will error if any of the following are not met:
    /// - Unidentified ISA
    ///  - We used the `target_lexicon` crate for ISA identification, so if its not supported by that it will fail.
    /// - Invalid flags. If the user has passed flags for features not supported by an architecture it will fail.
    pub fn new() -> anyhow::Result<Self> {
        let shared_flags = settings::Flags::new(settings::builder());
        let isa_builder = isa::lookup(HOST)?;
        let isa = isa_builder.finish(shared_flags)?;
        let builder = ObjectBuilder::new(isa, "bcomp", default_libcall_names())?;

        Ok(Self {
            module: ObjectModule::new(builder),
            data_count: 0,
        })
    }

    /// Emits a single program to object code.
    ///
    /// This will consume the Emitter in order to produce the associated code. This means that a single emitter can only ever be used for one "program" (or file).
    pub fn emit_program(mut self, program: &Program) -> anyhow::Result<Vec<u8>> {
        let main_func_id = {
            let sig = main_signature(self.module.isa());

            // Add this function to our Module.
            self.module
                .declare_function("main", Linkage::Export, &sig)?
        };

        // much of the code here was sources from the examples here: https://github.com/simvux/cranelift-examples/blob/master/examples/

        // These contain the context needed for generating code for a function.
        //
        // It's a lot more efficient to construct them once, and then re-use them for all functions.
        let mut context = Context::new();
        let mut function_context = FunctionBuilderContext::new();

        let mut builder = FunctionBuilder::new(&mut context.func, &mut function_context);
        // always starts the basic program with a main function that calls each line sequentially.
        builder.func.signature = main_signature(self.module.isa());

        // Create the functions entry block.
        let block0 = builder.create_block();
        builder.switch_to_block(block0);

        // When we know that there are no more blocks to be written which may jump to this block, we want to seal
        // it. This improves the quality of code generation.
        builder.seal_block(block0);

        for line in &program.lines {
            for statement in &line.statements {
                self.emit_statement(&mut builder, statement)?;
            }
        }

        // returns 0 as a good return value for our program
        // we should always return 0 at the end, unless some other return value is emmited previously
        let zero = builder.ins().iconst(types::I32, 0);
        builder.ins().return_(&[zero]);

        // validates and finalizes the function generation
        if let Err(err) = codegen::verify_function(builder.func, self.module.isa()) {
            panic!("verifier error: {err}");
        }
        builder.finalize();
        self.module.define_function(main_func_id, &mut context)?;

        // Finalize the module to generate our `Product`.
        //
        // If we have additional information such as unwind information or DWARF debug information,
        // they can be added to `Product`. For this example, we'll skip such optional additions.
        let product = self.module.finish();

        // Generate the object file.
        let bytes = product.emit()?;
        Ok(bytes)
    }

    fn emit_statement(
        &mut self,
        builder: &mut FunctionBuilder,
        stmt: &Statement,
    ) -> anyhow::Result<()> {
        match stmt {
            // emits return logic if we hit something like this
            Statement::End | Statement::Return => {
                let zero = builder.ins().iconst(types::I32, 0);
                builder.ins().return_(&[zero]);
                Ok(())
            }
            Statement::Print { items } => {
                self.emit_print_statement(builder, items)?;
                Ok(())
            }
            Statement::Rem(_) => Ok(()),
            Statement::Expression(expr) => {
                self.emit_expression_inner(builder, expr)?;
                Ok(())
            }
            _ => todo!(),
        }
    }

    /// Emits a print statement by calling the printf function in the c std lib
    ///
    /// TODO: Tweak this such that it uses a slightly better call with features like variable expansion and such.
    /// Its fine for right now for just dumping shit into stdout. I don't really like the way we call printf differently.
    /// I think we can simplify this.
    fn emit_print_statement(
        &mut self,
        builder: &mut FunctionBuilder,
        items: &[PrintItem],
    ) -> anyhow::Result<()> {
        if items.is_empty() {
            let value = self.emit_expression_inner(builder, &Expression::char('\n'))?;
            self.emit_printf_i32_call(builder, value)?;
            return Ok(());
        }

        for item in items {
            match item {
                PrintItem::String(value) => {
                    let value = self.emit_expression_inner(
                        builder,
                        &Expression::string(value.trim_matches('"').to_string()),
                    )?;
                    self.emit_printf_call(builder, value)?;
                }
                PrintItem::Char(value) => {
                    let value = Expression::char(*value);
                    let value = self.emit_expression_inner(builder, &value)?;
                    self.emit_printf_i32_call(builder, value)?;
                }
                PrintItem::Expression(value) => {
                    let value = self.emit_expression_inner(builder, value)?;
                    self.emit_printf_i32_call(builder, value)?;
                }
            }
        }

        Ok(())
    }

    /// Emits a printf call which knows that it's calling on a single string
    fn emit_printf_call(
        &mut self,
        builder: &mut FunctionBuilder,
        value: Value,
    ) -> anyhow::Result<()> {
        // establishes the signature for printf which we will be calling
        let ptr_type = self.module.target_config().pointer_type();
        let mut signature = Signature::new(self.module.isa().default_call_conv());
        signature.params.push(AbiParam::new(ptr_type));
        signature.returns.push(AbiParam::new(types::I32));
        let printf_id = self
            .module
            .declare_function("printf", Linkage::Import, &signature)?;
        let printf = self.module.declare_func_in_func(printf_id, builder.func);

        // sets the format param if we're passing in a number
        // TODO: Change this to support other types like float and such

        builder.ins().call(printf, &[value]);
        Ok(())
    }

    /// Emits a printf call which knows its trying to output an integer
    ///
    /// TODO: I really hate how this is done. I want there to be a single printf call function, and not have two that have repeated code.
    /// I also don't want to do too much work higher up when emitting code. I know I can't have it both ways, so for now we have two functions that
    /// are almost identical but do slightly different work.
    fn emit_printf_i32_call(
        &mut self,
        builder: &mut FunctionBuilder,
        value: Value,
    ) -> anyhow::Result<()> {
        // establishes the signature for printf which we will be calling
        let ptr_type = self.module.target_config().pointer_type();
        let mut signature = Signature::new(self.module.isa().default_call_conv());
        signature.params.push(AbiParam::new(ptr_type));
        signature.params.push(AbiParam::new(types::I32));

        signature.returns.push(AbiParam::new(types::I32));
        let printf_id = self
            .module
            .declare_function("printf", Linkage::Import, &signature)?;
        let printf = self.module.declare_func_in_func(printf_id, builder.func);

        // sets the format param if we're passing in a number
        // TODO: Change this to support other types like float and such
        let format = self.emit_string_pointer(builder, "%d\n")?;

        builder.ins().call(printf, &[format, value]);
        Ok(())
    }

    fn emit_string_pointer(
        &mut self,
        builder: &mut FunctionBuilder,
        value: &str,
    ) -> anyhow::Result<Value> {
        let name = format!("string_{}", self.data_count);
        self.data_count += 1;

        let data_id = self
            .module
            .declare_data(&name, Linkage::Local, false, false)?;
        let mut bytes = value.as_bytes().to_vec();
        bytes.push(0);

        let mut data = DataDescription::new();
        data.define(bytes.into_boxed_slice());
        self.module.define_data(data_id, &data)?;

        let data = self.module.declare_data_in_func(data_id, builder.func);
        let ptr_type = self.module.target_config().pointer_type();
        Ok(builder.ins().global_value(ptr_type, data))
    }

    fn emit_expression_inner(
        &mut self,
        builder: &mut FunctionBuilder,
        expr: &Expression,
    ) -> anyhow::Result<Value> {
        let value = match expr {
            Expression::Literal(literal) => self.emit_literal(builder, literal)?,
            Expression::Binary(binary) => {
                let left = self.emit_expression_inner(builder, binary.left())?;
                let right = self.emit_expression_inner(builder, binary.right())?;

                match binary.operator() {
                    BinaryOperator::Plus => builder.ins().iadd(left, right),
                    BinaryOperator::Minus => builder.ins().isub(left, right),
                    BinaryOperator::Multiply => builder.ins().imul(left, right),
                    BinaryOperator::Divide => builder.ins().sdiv(left, right),
                    BinaryOperator::Equal => builder.ins().icmp(IntCC::Equal, left, right),
                    BinaryOperator::Less => builder.ins().icmp(IntCC::SignedLessThan, left, right),
                    BinaryOperator::Greater => {
                        builder.ins().icmp(IntCC::SignedGreaterThan, left, right)
                    }
                    BinaryOperator::LessEqual => {
                        builder
                            .ins()
                            .icmp(IntCC::SignedLessThanOrEqual, left, right)
                    }
                    BinaryOperator::GreaterEqual => {
                        builder
                            .ins()
                            .icmp(IntCC::SignedGreaterThanOrEqual, left, right)
                    }
                    BinaryOperator::NotEqual => builder.ins().icmp(IntCC::NotEqual, left, right),
                }
            }
            Expression::Grouping(grouping) => {
                self.emit_expression_inner(builder, grouping.expression())?
            }
            Expression::Identifier(identifier) => {
                panic!(
                    "identifier expressions are not supported yet: {}",
                    identifier.name()
                )
            }
        };

        Ok(value)
    }

    fn emit_literal(
        &mut self,
        builder: &mut FunctionBuilder,
        literal: &LiteralExpression,
    ) -> anyhow::Result<Value> {
        let value = match literal {
            LiteralExpression::Integer(value) => builder.ins().iconst(types::I32, *value),
            LiteralExpression::Char(value) => builder
                .ins()
                .iconst(types::I32, i64::from(u32::from(*value))),
            LiteralExpression::String(value) => {
                self.emit_string_pointer(builder, &format!("{value}\n"))?
            }
            LiteralExpression::Float(_) => {
                panic!("only integer, char, and string literals are supported by the emitter")
            }
        };

        Ok(value)
    }
}

#[cfg(test)]
mod tests {
    use cranelift::prelude::*;
    use cranelift_module::Module;

    use crate::{
        emitter::{Emitter, main_signature},
        parser::{
            ast::{Line, PrintItem, Program, Statement},
            expression::{BinaryOperator, Expression},
        },
    };

    fn emit_expression_ir(expr: &Expression) -> anyhow::Result<String> {
        let mut context = cranelift_codegen::Context::new();
        let mut function_context = cranelift_frontend::FunctionBuilderContext::new();
        let mut builder = FunctionBuilder::new(&mut context.func, &mut function_context);
        builder.func.signature = main_signature(Emitter::new()?.module.isa());

        let block = builder.create_block();
        builder.switch_to_block(block);
        builder.seal_block(block);

        let mut emitter = Emitter::new()?;
        let value = emitter.emit_expression_inner(&mut builder, expr)?;
        builder.ins().return_(&[value]);
        builder.finalize();

        Ok(context.func.to_string())
    }

    #[test]
    fn emit_integer_expression() -> anyhow::Result<()> {
        let ir = emit_expression_ir(&Expression::integer(42))?;

        assert!(ir.contains("iconst.i32 42"));
        assert!(ir.contains("return v0"));

        Ok(())
    }

    #[test]
    fn emit_char_expression() -> anyhow::Result<()> {
        let ir = emit_expression_ir(&Expression::char('A'))?;

        assert!(ir.contains("iconst.i32 65"));

        Ok(())
    }

    #[test]
    fn emit_string_expression_program() -> anyhow::Result<()> {
        let program = Program {
            lines: vec![Line {
                number: None,
                statements: vec![Statement::Expression(Expression::string(
                    "\"HELLO\"".to_string(),
                ))],
            }],
        };

        let bytes = Emitter::new()?.emit_program(&program)?;

        assert!(!bytes.is_empty());

        Ok(())
    }

    #[test]
    fn emit_add_expression() -> anyhow::Result<()> {
        let ir = emit_expression_ir(&Expression::binary(
            Expression::integer(1),
            BinaryOperator::Plus,
            Expression::integer(2),
        ))?;

        assert!(ir.contains("iadd"));

        Ok(())
    }

    #[test]
    fn emit_multiply_expression() -> anyhow::Result<()> {
        let ir = emit_expression_ir(&Expression::binary(
            Expression::integer(3),
            BinaryOperator::Multiply,
            Expression::integer(4),
        ))?;

        assert!(ir.contains("imul"));

        Ok(())
    }

    #[test]
    fn emit_grouping_expression() -> anyhow::Result<()> {
        let ir = emit_expression_ir(&Expression::grouping(Expression::integer(7)))?;

        assert!(ir.contains("iconst.i32 7"));

        Ok(())
    }

    #[test]
    fn emit_print_string_program() -> anyhow::Result<()> {
        let program = Program {
            lines: vec![Line {
                number: None,
                statements: vec![Statement::Print {
                    items: vec![PrintItem::String("\"HELLO\"".to_string())],
                }],
            }],
        };

        let bytes = Emitter::new()?.emit_program(&program)?;

        assert!(!bytes.is_empty());

        Ok(())
    }

    #[test]
    fn emit_print_expression_program() -> anyhow::Result<()> {
        let program = Program {
            lines: vec![Line {
                number: None,
                statements: vec![Statement::Print {
                    items: vec![PrintItem::Expression(Expression::binary(
                        Expression::integer(1),
                        BinaryOperator::Plus,
                        Expression::integer(2),
                    ))],
                }],
            }],
        };

        let bytes = Emitter::new()?.emit_program(&program)?;

        assert!(!bytes.is_empty());

        Ok(())
    }
}