Brainfuck compiler with LLVM

brainfuck.cpp

#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Verifier.h"
#include "llvm/MC/TargetRegistry.h"
#include "llvm/Support/Host.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Target/TargetOptions.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include <string>
#include <vector>

using namespace llvm;

// quicksort in brainfuck
// source: https://codegolf.stackexchange.com/a/4702/40811
const std::string bf_code = ">>>>>>>>,[>,]<[[>>>+<<<-]>[<+>-]<+<]>[<<<<<<<<+>>>>>>>>-]<<<<<<<<[[>>+>+>>+<<<<<-]>>[<<+>>-]<[>+>>+>>+<<<<<-]>[<+>-]>>>>[-<->]+<[>->+<<-[>>-<<[-]]]>[<+>-]>[<<+>>-]<+<[->-<<[-]<[-]<<[-]<[[>+<-]<]>>[>]<+>>>>]>[-<<+[-[>+<-]<-[>+<-]>>>>>>>>[<<<<<<<<+>>>>>>>>-]<<<<<<]<<[>>+<<-]>[>[>+>>+<<<-]>[<+>-]>>>>>>[<+<+>>-]<[>+<-]<<<[>+>[<-]<[<]>>[<<+>[-]+>-]>-<<-]>>[-]+<<<[->>+<<]>>[->-<<<<<[>+<-]<[>+<-]>>>>>>>>[<<<<<<<<+>>>>>>>>-]<<]>[[-]<<<<<<[>>+>>>>>+<<<<<<<-]>>[<<+>>-]>>>>>[-[>>[<<<+>>>-]<[>+<-]<-[>+<-]>]<<[[>>+<<-]<]]>]<<<<<<-]>[>>>>>>+<<<<<<-]<<[[>>>>>>>+<<<<<<<-]>[<+>-]<+<]<[[>>>>>>>>+<<<<<<<<-]>>[<+>-]<+<<]>+>[<-<<[>+<-]<[<]>[[<+>-]>]>>>[<<<<+>>>>-]<<[<+>-]>>]<[-<<+>>]>>>]<<<<<<]>>>>>>>>>>>[.>]";

int main() {
	// ------------------------ Init LLVM, create empty module -------------------------------------

	LLVMContext llvmContext;
	Module module("brainfuck", llvmContext);
	IRBuilder<> builder(llvmContext);

	// ------------------------ Prepare target information -----------------------------------------

	InitializeAllTargetInfos();
	InitializeAllTargets();
	InitializeAllTargetMCs();
	InitializeAllAsmParsers();
	InitializeAllAsmPrinters();

	std::string error;
	auto targetTriple = sys::getDefaultTargetTriple(); // for me: x86_64-unknown-linux-gnu
	auto target = TargetRegistry::lookupTarget(targetTriple, error);
	auto targetMachine = target->createTargetMachine(targetTriple, "generic", "", TargetOptions{}, Optional<Reloc::Model>{});

	module.setTargetTriple(targetTriple);
	module.setDataLayout(targetMachine->createDataLayout());

	// ------------------------ Create function bf_quicksort  --------------------------------------

	// function type: void bf_quicksort(int32_t*, uint8_t*, uint8_t*)
	std::vector<Type*> argTypes = {
		builder.getPtrTy(), // (note: pointer types are opaque ("type-erased"))
		builder.getPtrTy(),
		builder.getPtrTy(),
	};

	FunctionType *functionType = FunctionType::get(builder.getVoidTy(), argTypes, false);
	Function *function = Function::Create(functionType, Function::ExternalLinkage, "bf_quicksort", module);

	// ------------------------ Code generation: copy arguments to "mutable local memory" ----------

	builder.SetInsertPoint(BasicBlock::Create(llvmContext, "entry", function));

	auto ptr_alloca = builder.CreateAlloca(builder.getPtrTy());
	auto input_alloca = builder.CreateAlloca(builder.getPtrTy());
	auto output_alloca = builder.CreateAlloca(builder.getPtrTy());
	builder.CreateStore(&function->args().begin()[0], ptr_alloca);
	builder.CreateStore(&function->args().begin()[1], input_alloca);
	builder.CreateStore(&function->args().begin()[2], output_alloca);

	// ------------------------ Code generation: main loop -----------------------------------------

	// need to remember previous blocks to handle nested loops
	std::vector<BasicBlock*> blocks;

	for (char c : bf_code) {
		switch (c) {
			case '+': { // *dataptr = *dataptr + 1
				auto ptr = builder.CreateLoad(builder.getPtrTy(), ptr_alloca);
				auto value = builder.CreateLoad(builder.getInt32Ty(), ptr);
				auto added = builder.CreateAdd(value, builder.getInt32(1));
				builder.CreateStore(added, ptr);
				break;
			}
			case '-': { // *dataptr = *dataptr - 1
				auto ptr = builder.CreateLoad(builder.getPtrTy(), ptr_alloca);
				auto value = builder.CreateLoad(builder.getInt32Ty(), ptr);
				auto added = builder.CreateSub(value, builder.getInt32(1));
				builder.CreateStore(added, ptr);
				break;
			}
			case '>': { // dataptr = dataptr + 1
				// GEP ("getelementptr") is how you do pointer arithmetic without converting pointers to integers
				auto ptr = builder.CreateLoad(builder.getPtrTy(), ptr_alloca);
				auto newptr = builder.CreateConstInBoundsGEP1_32(builder.getInt32Ty(), ptr, 1);
				builder.CreateStore(newptr, ptr_alloca);
				break;
			}
			case '<': { // dataptr = dataptr - 1
				auto ptr = builder.CreateLoad(builder.getPtrTy(), ptr_alloca);
				auto newptr = builder.CreateConstInBoundsGEP1_32(builder.getInt32Ty(), ptr, -1);
				builder.CreateStore(newptr, ptr_alloca);
				break;
			}
			case ',': { // *dataptr = *inputptr; inputptr = inputptr + 1;
				auto inputptr = builder.CreateLoad(builder.getPtrTy(), input_alloca);
				auto ptr = builder.CreateLoad(builder.getPtrTy(), ptr_alloca);
				auto value = builder.CreateLoad(builder.getInt8Ty(), inputptr);
				auto value_32 = builder.CreateSExt(value, builder.getInt32Ty());
				builder.CreateStore(value_32, ptr);
				auto gep_value = builder.CreateConstInBoundsGEP1_32(builder.getInt8Ty(), inputptr, 1);
				builder.CreateStore(gep_value, input_alloca);
				break;
			}
			case '.': { // *outputptr = *dataptr; outputptr = outputptr + 1;
				auto outputptr = builder.CreateLoad(builder.getPtrTy(), output_alloca);
				auto ptr = builder.CreateLoad(builder.getPtrTy(), ptr_alloca);
				auto value = builder.CreateLoad(builder.getInt8Ty(), ptr);
				auto value_8 = builder.CreateTrunc(value, builder.getInt8Ty());
				builder.CreateStore(value_8, outputptr);
				auto gep_value = builder.CreateConstInBoundsGEP1_32(builder.getInt8Ty(), outputptr, 1);
				builder.CreateStore(gep_value, output_alloca);
				break;
			}
			case '[': { // if (*dataptr != 0) { jump inside loop body } else { jump to block after ] }
				BasicBlock *loopBB = BasicBlock::Create(llvmContext, "loop", function);
				BasicBlock *afterBB = BasicBlock::Create(llvmContext, "afterloop", function);
				blocks.push_back(loopBB);
				blocks.push_back(afterBB);

				auto ptr = builder.CreateLoad(builder.getPtrTy(), ptr_alloca);
				auto value = builder.CreateLoad(builder.getInt32Ty(), ptr);
				auto cmp = builder.CreateICmpNE(value, builder.getInt32(0));
				builder.CreateCondBr(cmp, loopBB, afterBB);

				builder.SetInsertPoint(loopBB);
				break;
			}
			case ']': { // if (*dataptr != 0) { jump back to start of loop body } else { leave loop }
				BasicBlock *afterBB = blocks.back(); blocks.pop_back();
				BasicBlock *loopBB = blocks.back(); blocks.pop_back();

				auto ptr = builder.CreateLoad(builder.getPtrTy(), ptr_alloca);
				auto value = builder.CreateLoad(builder.getInt32Ty(), ptr);
				auto cmp = builder.CreateICmpNE(value, builder.getInt32(0));
				builder.CreateCondBr(cmp, loopBB, afterBB);

				builder.SetInsertPoint(afterBB);
				break;
			}
		}
	}
	builder.CreateRetVoid();
	verifyFunction(*function);

	// ------------------------ Run `-O2`-equivalent optimization passes ---------------------------

	legacy::PassManager optPM;
	PassManagerBuilder().populateModulePassManager(optPM);
	optPM.run(module);
	// print resulting optimized IR
	//module.print(errs(), nullptr);

	// ------------------------ Generate object file -----------------------------------------------

	std::error_code ec;
	raw_fd_ostream output("brainfuck.o", ec);

	legacy::PassManager outputPM;
	targetMachine->addPassesToEmitFile(outputPM, output, nullptr, CGFT_ObjectFile);

	outputPM.run(module);
	output.flush();
}

/*
usage of generated object file: `g++ main.cpp brainfuck.o`
example main.cpp:

#include <iostream>
#include <vector>

extern "C" { void bf_quicksort(int32_t* data, uint8_t *input, uint8_t* output); }

int main(){
	std::vector<int> data(100000, 0);
	auto dataptr = &data[50000]; // leave room to the left

	std::string output, input = "edbca";
	output.resize(1000);

	bf_quicksort(dataptr, (uint8_t*)&input[0], (uint8_t*)&output[0]);

	std::cout << output << "\n";
}
*/