toy bf-llvm program

bf-llvm.cc

#include <vector>
#include <cstdint>
#include <iostream>

#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/Module.h>
#include <llvm/ExecutionEngine/GenericValue.h>
#include <llvm/ExecutionEngine/Interpreter.h>
#include <llvm/ExecutionEngine/JIT.h>
#include <llvm/Support/TargetSelect.h>

int main()
{
    // Stack size
    const uint64_t ssize = 50000;

    llvm::InitializeNativeTarget();
    llvm::LLVMContext& context = llvm::getGlobalContext();
    llvm::BasicBlock* instr_block;
    llvm::Module module("BrainFuck", context);
    llvm::Value* cursor;
    llvm::Value* cells;

    // Create a stack for our looping conditions
    std::vector<llvm::BasicBlock*> stack;

    // Wrapper to quickly construct ints
    auto make_int = [&context](size_t length)
    {
        return llvm::Type::getIntNTy(context, length);
    };

    // Wrapper to quickly construct pointers
    auto make_ptr_int = [&context](size_t length)
    {
        return llvm::Type::getIntNPtrTy(context, length);
    };

    // Wrapper to quickly construct voids
    auto make_void = [&context]()
    {
        return llvm::Type::getVoidTy(context);
    };

    // Externally link a new function for the llvm
    auto make_func = [&module](const llvm::Twine& name,
            llvm::Type* return_type, std::vector<llvm::Type*> arg_types)
    {
        return llvm::Function::Create(
            llvm::FunctionType::get(return_type, arg_types, false),
            llvm::Function::ExternalLinkage, name, &module);
    };

    // A block is a container of instructions, so we construct a builder
    // whenever we want a context for our instructions
    auto b = [&]()
    {
        return std::unique_ptr<llvm::IRBuilder<>>(
            new llvm::IRBuilder<>(instr_block));
    };

    // Retrieve the cursor value for this call
    auto v = [&]()
    {
        auto builder = b();
        std::vector<llvm::Value*> args = {
            llvm::ConstantInt::get(make_int(32), 0),
            builder->CreateLoad(cursor, "cursor")
        };
        return builder->CreateGEP(cells, args, "value_ptr");
    };

    // Setup functions that we need for our BrainFuck interpreter
    // This is our wrapper function for a BrainFuck program
    auto llvm_putchar = make_func("putchar", make_int(32), {make_int(32)});
    auto llvm_scoped = make_func("scoped", make_void(), {});
    auto llvm_bzero = make_func("bzero",
        make_void(), {make_ptr_int(8), make_int(32)});

    // Initializer
    // Set the current instruction block to point to the main function
    {
        instr_block = llvm::BasicBlock::Create(context, "entry", llvm_scoped);
        auto entry_builder = b();

        // Now create the moving pointer and storage
        cursor = entry_builder->CreateAlloca(make_int(64));
        entry_builder->CreateStore(
            llvm::ConstantInt::get(make_int(64), 0), cursor);

        // Now create the make-shift cells
        cells = entry_builder->CreateAlloca(
            llvm::ArrayType::get(make_int(8), ssize));

        // Now create the caller arguments
        std::vector<llvm::Value*> bzero_args = {
            v(), llvm::ConstantInt::get(make_int(32), ssize)
        };

        entry_builder->CreateCall(llvm_bzero, bzero_args);

    }

    // Add
    // Increment the load by size
    auto instr_incr = [&](size_t size)
    {
        auto builder = b();
        auto value_ptr = v();
        auto value = builder->CreateLoad(value_ptr, "value");
        builder->CreateStore(
            builder->CreateAdd(
                value, llvm::ConstantInt::get(value->getType(), size)),
            value_ptr);
    };

    // Sub
    // Decrement the load by size
    auto instr_decr = [&](size_t size)
    {
        auto builder = b();
        auto value_ptr = v();
        auto value = builder->CreateLoad(value_ptr, "value");
        builder->CreateStore(
            builder->CreateSub(
                value, llvm::ConstantInt::get(value->getType(), size)),
            value_ptr);
    };

    // We now create a stack / looping condition
    auto instr_open = [&]()
    {
        // Setup a while() / looping condition
        auto loop = llvm::BasicBlock::Create(context, "loop", llvm_scoped);
        auto builder = b();

        // Create an unconditional instruction for the loop
        builder->CreateBr(loop);
        instr_block = loop;

        // Retrieve new builder and value
        builder = b();
        auto value_ptr = v();

        // Build our conditional
        auto lhs = builder->CreateLoad(value_ptr, "value");
        auto rhs = llvm::ConstantInt::get(make_int(8), 0);
        auto conditional = builder->CreateICmpNE(lhs, rhs, "conditional");

        // Create the conditional instruction
        auto if_true = llvm::BasicBlock::Create(
            context, "if_true", llvm_scoped);
        auto if_false = llvm::BasicBlock::Create(
            context, "if_false", llvm_scoped);
        builder->CreateCondBr(conditional, if_true, if_false);

        // Push our current instructions on the stack and set our current
        // instruction to the if_true block
        stack.push_back(if_false);
        stack.push_back(loop);
        instr_block = if_true;
    };

    // Close the looping condition
    auto instr_close = [&]()
    {
        auto loop = stack.back(); stack.pop_back();
        auto if_false = stack.back(); stack.pop_back();
        b()->CreateBr(loop);
        instr_block = if_false;
    };

    // Left
    auto instr_left = [&](size_t size)
    {
        auto builder = b();
        auto value_ptr = cursor;
        auto value = builder->CreateLoad(value_ptr);
        builder->CreateStore(
            builder->CreateSub(
                value, llvm::ConstantInt::get(value->getType(), size)),
            value_ptr);
    };

    // right
    auto instr_right = [&](size_t size)
    {
        auto builder = b();
        auto value_ptr = cursor;
        auto value = builder->CreateLoad(value_ptr);
        builder->CreateStore(
            builder->CreateAdd(
                value, llvm::ConstantInt::get(value->getType(), size)),
            value_ptr);
    };


    // Putchar
    auto instr_putchar = [&]()
    {
        auto builder = b();
        auto value_ptr = v();
        auto value = builder->CreateLoad(value_ptr, "value");
        builder->CreateCall(
            llvm_putchar, builder->CreateSExt(value, make_int(32)));
    };

    // Read the program
    bool reset_wait_for = false;
    size_t wait_for_size = 0;
    char instr, wait_for = 0;
    while (std::cin >> instr) {

        // Optimize incr/decr instructions
        if (wait_for == '+' && instr == '-') {
            wait_for_size--;
        } else if (wait_for == '-' && instr == '+') {
            wait_for_size++;
        } else if (wait_for == '+' && instr != '+') {
            instr_incr(wait_for_size);
            reset_wait_for = true;
        } else if (wait_for == '-' && instr != '-') {
            instr_decr(wait_for_size);
            reset_wait_for = true;
        }

        // Optimize left/right instructions
        else if (wait_for == '<' && instr == '>') {
            wait_for_size--;
        } else if (wait_for == '<' && instr == '>') {
            wait_for_size--;
        } else if (wait_for == '<' && instr != '<') {
            instr_left(wait_for_size);
            reset_wait_for = true;
        } else if (wait_for == '>' && instr != '>') {
            instr_right(wait_for_size);
            reset_wait_for = true;
        }

        // Reset conditions
        if (reset_wait_for) {
            wait_for = 0;
            wait_for_size = 0;
            reset_wait_for = false;
        }

        switch (instr) {
            case '+': wait_for = '+'; wait_for_size++; break;
            case '-': wait_for = '-'; wait_for_size++; break;
            case '[': instr_open(); break;
            case ']': instr_close(); break;
            case '<': wait_for = '<'; wait_for_size++; break;
            case '>': wait_for = '>'; wait_for_size++; break;
            case '.': instr_putchar(); break;
            default: continue;
        };
    };

    // Force a void return for the llvm hook
    {
        auto builder = b();
        builder->CreateRetVoid();
    }

    // Finalize the module and call the program
    auto engine = llvm::EngineBuilder(&module).create();
    auto entry = (intptr_t)(engine->getPointerToFunction(llvm_scoped));

    // Run the compiled code
    ((void(*)())(entry))();

    return 0;
}