Brainfuck native compiler for i386, x86-64

bfcompiler.rb

#!/usr/bin/ruby
require 'tempfile'

CELL_SIZE = 30_000

ASM_STUB_X86_64 =<<EOF
.comm bf_cell, #{CELL_SIZE}
.global write_byte
.global sys_exit
.global _start

.text
write_byte:
	mov $1, %rax    # write()
	mov $1, %rdi    # fd = stdout = 1
	mov $1, %rdx    # sizeof() = 1
	syscall         # write(1, %rsi, 1);
ret

read_byte:
	mov $2, %rax    # read()
	mov $0, %rdi    # fs = stdin = 0
	mov $1, %rdx    # sizeof() = 1
	syscall         # read(1, %rsi, 1);
ret

# exit with status in %rsi
sys_exit:
	mov $60, %rax   # exit()
	mov (%rsi), %rdi
	syscall         # exit(%rsi)
ret

_start:    
	mov $bf_cell, %rsi
EOF

ASM_END_STUB =<<EOF
	call sys_exit
ret
EOF

class IrBuilder
    CELLOP = '+-'
    PTROP  = '<>'
    attr_reader :ir
    attr_reader :label

    def initialize
        @ir = []
        @label = []
        @label_cnt = 0
    end

    def build_read(times=1)
        @ir << { read: times }
    end

    def build_write(times=1)
        @ir << { write: times }
    end

    def build_operation(op, where=1)
        if @ir.last&.has_key?(op)
            @ir.last[op] += where
        else
            @ir << { op => where }
        end
    end

    def build_cell_change(what=1)
        build_operation(:cell_change, what)
    end

    def build_pointer_change(where)
        build_operation(:ptr_change, where)        
    end

    def build_label
        unless @label.empty?
            @label.push(@ir.size-1) 
        else
            @label.push(0)
        end
        build_operation(:label, @label_cnt)
        @label_cnt += 1
    end

    def build_jump
        puts "Don't know where to jump!" if @label.empty?
        @ir << { jump: @label.pop }
    end

    def assemble(prettify = true)
        return ''
    end

    def assemble_full(prettify = true)
        return ''
    end
end

class X86IrBuilder < IrBuilder
    ERROR_DEF_SUBCLASS = "method %s, called with (%s) has to be implemented in the subclass"

    def initialize
        super
        @ir_buffer = []
    end

    def compile
        @ir.each_with_index do |instr|
            if instr.has_key? :read
                @ir_buffer += [:call_read_byte] * instr[:read]
            elsif instr.has_key? :write
                @ir_buffer += [:call_write_byte] * instr[:write]
            else
                compile_operation(instr)
            end
        end
    end

    def compile_cell_change(n)
        return if n == 0
        #@ir_buffer << :mov_rsi_dl
        compile_math_op(:decb_deref_ptr, :incb_deref_ptr, :subb_deref_ptr, :addb_deref_ptr, n)
        #@ir_buffer << :mov_dl_rsi
    end

    def compile_ptr_change(n)
        return if n == 0
        compile_math_op(:dec_ptr, :inc_ptr, :sub_ptr, :add_ptr, n)
    end

    def compile_jump(where)
        @ir_buffer += [:cmpb_0_deref_ptr, [:jne, where]]
    end

    def compile_operation(instr)
        if instr.has_key? :label
            @ir_buffer << instr.flatten
        elsif instr.has_key? :ptr_change
            compile_ptr_change(instr[:ptr_change])
        elsif instr.has_key? :cell_change
            compile_cell_change(instr[:cell_change])
        elsif instr.has_key? :jump
            compile_jump(instr[:jump])
        else
            puts "Unknown instruction '#{instr}'"
        end
    end

    def assemble(prettify = true)
        asm = ""
        @ir_buffer.each do |instr|
            asm_line = assemble_instruction(instr)
            asm += asm_line unless asm_line.nil?
        end
        return asm
    end

    def assemble_full(prettify)
        asm = get_asm_stub
        asm += assemble(prettify)
        asm += ASM_END_STUB
        return asm
    end

    def assemble_one_arg(mnemonic, arg, prettify)
        asm = "#{mnemonic} #{arg}"
        return "\t#{asm}\n" if prettify
        return asm
    end

    def assemble_two_arg(mnemonic, arg0, arg1, prettify)
        asm = "#{mnemonic} #{arg0}, #{arg1}"
        return "\t#{asm}\n" if prettify
        return asm
    end

    def assemble_instruction(instr, prettify=true)
        if instr.is_a? Array
            (op, arg) = instr
            return ".L#{arg}:#{"\n" if prettify}" if op == :label
            istr = op.to_s.split(/_/, 2)
            mnemonic = istr[0]
            return case op
            when :add_ptr, :sub_ptr
                assemble_two_arg(mnemonic, "$#{arg}", get_ptr_name, prettify)
            when :addb_deref_ptr, :subb_deref_ptr
                assemble_two_arg(mnemonic, "$#{arg}", get_deref_ptr_name, prettify)
            when :jne
                assemble_one_arg(mnemonic, ".L#{arg}", prettify)
            else
                puts "Unknown instruction #{mnemonic}!"
                nil
            end
        else
            istr = instr.to_s.split(/_/, 2)
            mnemonic = istr[0]
            return case instr
            when :call_read_byte, :call_write_byte
                assemble_one_arg(mnemonic, istr[1], prettify)
            when :cmpb_0_deref_ptr
                assemble_two_arg(mnemonic, "$0", get_deref_ptr_name, prettify)
            when :inc_ptr, :dec_ptr
                assemble_one_arg(mnemonic, get_ptr_name, prettify)
            when :incb_deref_ptr, :decb_deref_ptr
                assemble_one_arg(mnemonic, get_deref_ptr_name, prettify)
            else
                puts "Unknown instruction #{mnemonic}!"
                nil
            end
        end
    end

    def dump
        print @ir_buffer
    end


    protected
        def throw_no_impl(*args)
            throw NotImplementedError.new(sprintf(ERROR_DEF_SUBCLASS, __method__, args))
        end

        alias_method :get_ptr_name, :throw_no_impl
        alias_method :get_deref_ptr_name, :throw_no_impl
        alias_method :get_asm_stub, :throw_no_impl
        alias_method :bit_width, :throw_no_impl

        def compile_math_op(dec_op, inc_op, sub_op, add_op, n)
            @ir_buffer << 
            (if n == -1
                dec_op
            elsif n == 1
                inc_op
            elsif n < -1
                [sub_op, n.abs]
            elsif n > 1
                [add_op, n]
            end)
        end

end

class X86Ir32Builder < X86IrBuilder
    attr_reader :bit_width

    def initialize
        super
        @bit_width = 64
    end

    def get_ptr_name
        return '%ecx'
    end

    def get_deref_ptr_name
        return "(%ecx)"
    end

    def get_asm_stub
        asm = ASM_STUB_X86_64
        subst = [['%rax', '%eax'], ['%rdi', '%ebx'], ['%rdx', '%edx'],
            ['%rsi', '%ecx'], ['syscall', 'int $0x80']]
        subst.each do |sub|
            asm.gsub!(sub[0], sub[1])
        end
        return asm
    end
end

class X86Ir64Builder < X86IrBuilder
    attr_reader :bit_width
    def initialize
        super
        @bit_width = 64
    end

    def get_ptr_name
        return '%rsi'
    end

    def get_deref_ptr_name
        return "(%rsi)"
    end

    def get_asm_stub
        return ASM_STUB_X86_64
    end
end

class BfCompiler
    attr_accessor :program
    attr_accessor :builder
    attr_accessor :assembler

    def initialize(program)
       @program = program 
       @builder = nil
       @assembler = 'gcc'
    end

    def compile
        throw 'Need a builder!' if @builder.nil?
        @curr_op = nil
        @cnt_op  = 0
        @program.each_char do |ch|
            if @curr_op.nil?
                @curr_op = ch
                set_count
            else
                if @curr_op == ch
                    set_count
                else
                    compile_operation
                    @cnt_op  = 0
                    @curr_op = ch

                    set_count if pair_with?(ch)
                end
            end
        end
        compile_operation
        @builder.compile
    end

    def assemble(prettify=true)
        @builder.assemble(prettify)
    end

    def assemble_full(prettify=true)
        asm = "# Generated at '#{Time.now}' from this Brainfuck program:\n# "
        asm += @program.gsub(/\n/, "\n# ") + "\n\n"
        asm += @builder.assemble_full(prettify)
        puts asm
        return asm
    end

    def assemble_to_file(file, prettify=true)
        if file.is_a? String
            file = File.open(file, 'w')
        end
        file.puts(assemble_full(prettify))
        file.close
    end

    def assemble_to_tempfile(prettify=true)
        t = Tempfile.new ['bf', '.S']
        assemble_to_file(t, prettify)
        return t
    end

    def read_from_file(file)
        if file.is_a? String
            file = File.open(file, 'r')
        end
        @program = file.readlines.join
        file.close
    end

    def compile_file(input_file_name, output_file_name=nil, march=64)
        if output_file_name.nil?
            output_file_name = input_file_name.split('.')[0]
        end
        read_from_file(input_file_name)
        case march
        when 32
            @builder = X86Ir32Builder.new
        when 64
            @builder = X86Ir64Builder.new
        else
            throw "Unknown architecture #{march}"
        end

        compile
        t = assemble_to_tempfile
        `#{@assembler} #{t.path} -o #{output_file_name} -m#{march} -nostdlib`
        #t.unlink
    end

    private
        def pair_with?(op)
            return true if @curr_op == op
            return (CELLOP.include?(op) && CELLOP.include?(@curr_op)) ||
                (PTROP.include?(op) && PTROP.include?(@curr_op))
        end 

        def set_count
            case @curr_op
            when '+', '>', '.', ','
                @cnt_op += 1
            when '-', '<'
                @cnt_op -= 1
            else
                @cnt_op = 1
            end
        end

        def compile_operation
            return if @cnt_op == 0 || @curr_op.nil?

            case @curr_op
            when '+', '-'
                @builder.build_operation(:cell_change, @cnt_op)
            when '<', '>'
                @builder.build_operation(:ptr_change, @cnt_op)
            when '.'
                @builder.build_write(@cnt_op)
            when ','
                @builder.build_read(@cnt_op)
            when '['
                @builder.build_label
            when ']'
                @builder.build_jump
            end
        end
end

unless ARGV[0].nil?
    program = ARGV[0]
else
    print "Need a Brainfuck file!"
    exit 1
end

cc = BfCompiler.new(program)
cc.compile_file(program)