glob:brainfuck_tests

brainfuck.py

    My awesome Brainfuck compiler. Generates `flat assembler`_ source code for either

    Linux (using ELF formatter and syscalls for I/O) or Windows

    (using PE console formatter and msvcrt's getchar/putchar for I/O).

.. _`flat assembler`: http://flatassembler.net

# Simple brainfuck-to-asm "compiler"

# by PiotrLegnica

# WTFPL licensed

import hashlib, os

def optimisable(fn):

    fn.optimisable = True

    return fn

class CodeGen(object):

    def __init__(self): self.loops = []

    def template(self, code): raise NotImplementedError

    def incr(self):     raise NotImplementedError

    def decr(self):     raise NotImplementedError

    def forward(self):  raise NotImplementedError

    def backward(self): raise NotImplementedError

    def outp(self):     raise NotImplementedError

    def inp(self):      raise NotImplementedError

    def loopb(self):    raise NotImplementedError

    def loope(self):    raise NotImplementedError

class X86CodeGen(CodeGen):

    @optimisable

    def incr(self, count = None):

        if count is not None:

            return ['add byte [ebx], {0:d}'.format(count)]

        return ['inc byte [ebx]']

    @optimisable

    def decr(self, count = None):

        if count is not None:

            return ['sub byte [ebx], {0:d}'.format(count)]

        return ['dec byte [ebx]']

    @optimisable

    def forward(self, count = None):

        if count is not None:

            return ['add ebx, {0:d}'.format(count)]

        return ['inc ebx']

    @optimisable

    def backward(self, count = None):

        if count is not None:

            return ['sub ebx, {0:d}'.format(count)]

        return ['dec ebx']

    def loopb(self):

        loopID = hashlib.sha1(os.urandom(16)).hexdigest()

        self.loops.append(loopID)

        return [

            'loop_{0}:'.format(loopID),

            'cmp byte [ebx], 0',

            'je end_loop_{0}'.format(loopID),

        ]

    def loope(self):

        assert len(self.loops) > 0, 'Unbalanced loops'

        loopID = self.loops.pop()

        return [

            'cmp byte [ebx], 0',

            'jne loop_{0}'.format(loopID),

            'end_loop_{0}:'.format(loopID),

        ]

class Win32X86CodeGen(X86CodeGen):

    'Generates Windows console program, using msvcrt routines for input and output'

    def template(self, code):

        tpl = [

            'format PE console', 'entry start',

            "include 'win32a.inc'",

            "section '.text' code readable executable",

            'start:', 'mov ebx, cell', 'brainfuck:'

        ]

        tpl.extend(code)

        tpl.extend([

            'endcat:', 'invoke ExitProcess, 0',

            "section '.data' data readable writable",

            'cell rb 65536',

            "section '.idata' import data readable writeable",

            "library msvcrt, 'msvcrt.dll', kernel32, 'kernel32.dll'",

            "import msvcrt, putch, 'putchar', getch, 'getchar'",

            "import kernel32, ExitProcess, 'ExitProcess'"

        ])

        return tpl

    def inp(self):

        return ['ccall [getch]', 'mov [ebx], eax']

    def outp(self):

        return ['ccall [putch], [ebx]']

class LinuxX86CodeGen(X86CodeGen):

    'Generates Linux program, using syscalls for input and output'

    def template(self, code):

        tpl = [

            'format elf executable', 'entry start',

            'stdin equ 0', 'stdout equ 1',

            'sys_exit equ 1', 'sys_read equ 3', 'sys_write equ 4',

            'segment readable executable',

            'start:', 'mov ebx, cell', 'brainfuck:'

        ]

        tpl.extend(code)

        tpl.extend([

            'endcat:', 'mov eax, sys_exit', 'xor ebx, ebx', 'int $80',

            'segment readable writable', 'cell rb 65536'

        ])

        return tpl

    def inp(self):

        return [

            'push ebx', 'mov ecx, ebx',

            'mov eax, sys_read', 'mov ebx, stdin',

            'mov edx, 1', 'int $80', 'pop ebx'

        ]

    def outp(self):

        return [

            'push ebx', 'mov ecx, ebx',

            'mov eax, sys_write', 'mov ebx, stdout',

            'mov edx, 1', 'int $80', 'pop ebx'

        ]

def compile(code, codegen):

    import itertools

    opcodes = []

    instrs = {

        '+': codegen.incr,

        '-': codegen.decr,

        '>': codegen.forward,

        '<': codegen.backward,

        ',': codegen.inp,

        '.': codegen.outp,

        '[': codegen.loopb,

        ']': codegen.loope,

    }

    for instr in code:

        if instr not in instrs: continue

        opcodes.append(instrs[instr])

    # lame code size optimiser

    generated = []

    for opcode, group in itertools.groupby(opcodes):

        if not getattr(opcode, 'optimisable', False):

            for opcode in group:

                generated.extend(opcode())

        else:

            group = list(group)

            generated.extend(opcode(len(group)))

    return codegen.template(generated)

if __name__ == '__main__':

    import sys

    if len(sys.argv) < 4:

        print 'brainfuck.py linux/win32 input.bf output.asm'

        sys.exit(1)

    codegen = {

        'linux': LinuxX86CodeGen, 'win32': Win32X86CodeGen

    }[sys.argv[1]]()

    import contextlib

    with contextlib.nested(open(sys.argv[2], 'rb'), open(sys.argv[3], 'wb')) as (input, output):

        code = input.read()

        asm  = compile(code, codegen)

        output.write('\n'.join(asm))