sprite-rs/src/jit/translate.rs

// Copyright (c) 2022 Marceline Cramer
// SPDX-License-Identifier: GPL-3.0-or-later

use crate::parse::ast;
use cranelift::prelude::*;

pub struct FunctionTranslator<'a> {
    pub int: types::Type,
    pub builder: FunctionBuilder<'a>,
}

impl<'a> FunctionTranslator<'a> {
    pub fn translate(mut self, fn_impl: &ast::FnImpl) {
        let signature = &fn_impl.def.signature;

        let return_info = if let Some(tail_expr) = &fn_impl.body.tail_expr {
            match signature.return_type {
                Some(t) => Some((tail_expr, t)),
                None => panic!("Function has tail expression but no return type"),
            }
        } else if let Some(_) = signature.return_type {
            panic!("Function has return type but no tail expression");
        } else {
            None
        };

        for _ in signature.args.iter() {
            self.builder
                .func
                .signature
                .params
                .push(AbiParam::new(self.int));
        }

        if let Some((_, _return_type)) = return_info {
            self.builder
                .func
                .signature
                .returns
                .push(AbiParam::new(self.int));
        }

        let entry_block = self.builder.create_block();
        self.builder
            .append_block_params_for_function_params(entry_block);
        self.builder.switch_to_block(entry_block);
        self.builder.seal_block(entry_block);

        if fn_impl.body.statements.len() > 0 {
            unimplemented!("Function body statements are unimplemented");
        }

        if let Some((tail_expr, _)) = return_info {
            let return_value = self.translate_expr(tail_expr);
            self.builder.ins().return_(&[return_value]);
        }

        println!("{}", self.builder.func.display());
        self.builder.finalize();
    }

    pub fn translate_expr(&mut self, expr: &ast::Expr) -> Value {
        use ast::Expr::*;
        match expr {
            Literal(ast::Literal::DecimalInteger(literal)) => {
                // TODO parse integers while building ast so that codegen doesn't have to
                let val: i64 = literal.parse().unwrap();
                self.builder.ins().iconst(self.int, val)
            }
            BinaryOp(op, terms) => {
                let lhs = self.translate_expr(&terms.0);
                let rhs = self.translate_expr(&terms.1);
                self.translate_binary_op(op, lhs, rhs)
            }
            // TODO the AST doesn't need this either
            Group(expr) => self.translate_expr(expr),
            _ => unimplemented!("Expression: {:#?}", expr),
        }
    }

    pub fn translate_binary_op(&mut self, op: &ast::BinaryOp, lhs: Value, rhs: Value) -> Value {
        use ast::BinaryOp::*;
        let ins = self.builder.ins();
        match op {
            Add => ins.iadd(lhs, rhs),
            Sub => ins.isub(lhs, rhs),
            Mul => ins.imul(lhs, rhs),
            Div => ins.udiv(lhs, rhs),
            _ => unimplemented!("Binary operation: {:#?}", op),
        }
    }
}