A miniature LLVM-based compiler

llvm_fun.rs

#![feature(plugin)]
#![plugin(peg_syntax_ext)]

extern crate llvm_sys as llvm;

use std::ptr;
use std::ffi::CString;
use std::collections::HashMap;

use llvm::prelude::*;

peg! parser(r#"
    use super::{AstNode, Statement};

    #[pub]
    program -> Vec<Statement>
        = statement ** "\n"

    statement -> Statement
        = "return" " " e:expression { Statement::Return(e) }
        / variable:identifier "=" e:expression { Statement::Assign(variable, e) }

    expression -> AstNode
        = sum

    sum -> AstNode
        = a:product "+" b:sum { AstNode::Add(Box::new(a), Box::new(b)) }
        / product

    product -> AstNode
        = a:num_or_id "*" b:product { AstNode::Mul(Box::new(a), Box::new(b)) }
        / num_or_id

    num_or_id -> AstNode
        = number
        / id:identifier { AstNode::Ref(id) }

    identifier -> String
        = [a-zA-Z]+ { match_str.to_owned() }

    number -> AstNode
        = [0-9]+ { AstNode::Literal(match_str.to_owned()) }  
"#);

pub enum AstNode {
    Literal(String),
    Add(Box<AstNode>, Box<AstNode>),
    Mul(Box<AstNode>, Box<AstNode>),
    Ref(String),
}

pub enum Statement {
    Return(AstNode),
    Assign(String, AstNode),
}

fn main() {
    let input = "a=4+6\nb=a+1\nreturn b+a";
    let parsed = parser::program(input).unwrap();

    codegen(parsed);
}

fn codegen(program: Vec<Statement>) {
    unsafe {
        let context = llvm::core::LLVMContextCreate();
        let module = llvm::core::LLVMModuleCreateWithName(b"nop\0".as_ptr() as *const _);
        let builder = llvm::core::LLVMCreateBuilderInContext(context);

        let int_type = llvm::core::LLVMInt64TypeInContext(context);

        let function_type = llvm::core::LLVMFunctionType(int_type, ptr::null_mut(), 0, 0);
        let function = llvm::core::LLVMAddFunction(module, b"main\0".as_ptr() as *const _,
                                                   function_type);

        let bb = llvm::core::LLVMAppendBasicBlockInContext(context, function,
                                                           b"entry\0".as_ptr() as *const _);
        llvm::core::LLVMPositionBuilderAtEnd(builder, bb);

        // let a = llvm::core::LLVMConstInt(int_type, value, 0);
        // let b = llvm::core::LLVMConstInt(int_type, 2, 0);

        // let name = CString::new("addtmp").unwrap();
        // let return_value = llvm::core::LLVMBuildAdd(builder, a, b, name.as_ptr());

        // let return_value = codegen_expr(context, builder, &ast);

        // let return_value = llvm::core::LLVMConstInt(int_type, 42, 0);
        // llvm::core::LLVMBuildRet(builder, return_value);

        let mut names = HashMap::new();

        for statement in program {
            codegen_statement(context, builder, &statement, &mut names);
        }

        let out_file = CString::new("out.ll").unwrap();
        llvm::core::LLVMPrintModuleToFile(module, out_file.as_ptr(), ptr::null_mut());

        llvm::core::LLVMDisposeBuilder(builder);
        llvm::core::LLVMDisposeModule(module);
        llvm::core::LLVMContextDispose(context);
    }
}

unsafe fn codegen_statement(context: LLVMContextRef, builder: LLVMBuilderRef, statement: &Statement, names: &mut HashMap<String, LLVMValueRef>) {
    match *statement {
        Statement::Return(ref expr) => {
            let return_value = codegen_expr(context, builder, expr, names);
            llvm::core::LLVMBuildRet(builder, return_value);
        },

        Statement::Assign(ref var_name, ref expr) => {
            let int_type = llvm::core::LLVMInt64TypeInContext(context);
            let var_ref = llvm::core::LLVMBuildAlloca(builder, int_type, CString::new(var_name.to_owned()).unwrap().as_ptr());
            let value = codegen_expr(context, builder, expr, names);

            llvm::core::LLVMBuildStore(builder, value, var_ref);

            names.insert(var_name.to_owned(), var_ref);
        }
    }
}

unsafe fn codegen_expr(context: LLVMContextRef, builder: LLVMBuilderRef, expr: &AstNode, names: &HashMap<String, LLVMValueRef>) -> LLVMValueRef {
    match *expr {
        AstNode::Literal(ref int_literal) => {
            let int_type = llvm::core::LLVMInt64TypeInContext(context);
            llvm::core::LLVMConstInt(int_type, int_literal.parse().unwrap(), 0)
        },

        AstNode::Add(ref lhs, ref rhs) => {
            let lhs = codegen_expr(context, builder, lhs, names);
            let rhs = codegen_expr(context, builder, rhs, names);

            llvm::core::LLVMBuildAdd(builder, lhs, rhs, b"addtmp\0".as_ptr() as *const _)
        },

        AstNode::Mul(ref lhs, ref rhs) => {
            let lhs = codegen_expr(context, builder, lhs, names);
            let rhs = codegen_expr(context, builder, rhs, names);

            llvm::core::LLVMBuildMul(builder, lhs, rhs, b"multmp\0".as_ptr() as *const _)
        },

        AstNode::Ref(ref var_name) => {
            llvm::core::LLVMBuildLoad(builder, *names.get(var_name).unwrap(), CString::new(var_name.to_owned()).unwrap().as_ptr())
        },
    }
}