Brainf*** interpreter

interp-tarpit.py

#!/usr/bin/env python

import sys
from collections import deque

SINGLE, SHIFT_LEFT_SEQ, SHIFT_RIGHT_SEQ, INC_SEQ, DEC_SEQ = range(5)

SHIFT_RIGHT   = ord('>') # increment the data pointer
SHIFT_LEFT    = ord('<') # decrement the data pointer
INC           = ord('+') # increment the current byte
DEC           = ord('-') # decrement the current byte
WRITE         = ord('.') # output the ascii char corresponding to the current byte
READ          = ord(',') # read one byte of input
JUMP_FORWARD  = ord('[') # jump forward to matching ] if current byte is 0
JUMP_BACKWARD = ord(']') # jump backward to matching [ if current byte is not 0

def short_to_bytes(val):
	return (val / 256, val % 256)
def bytes_to_short(hi, lo):
	return hi * 256 + lo

def compile(source):
	bytecode = bytearray()
	state = SINGLE
	brackets = []
	for c in source:
		c = ord(c)
		if c is SHIFT_RIGHT:
			if state is SHIFT_RIGHT_SEQ and bytecode[-1] < 255:
				bytecode[-1] += 1
			else:
				bytecode.extend((SHIFT_RIGHT, 1))
				state = SHIFT_RIGHT_SEQ
		elif c is SHIFT_LEFT:
			if state is SHIFT_LEFT_SEQ and bytecode[-1] < 255:
				bytecode[-1] += 1
			else:
				bytecode.extend((SHIFT_LEFT, 1))
				state = SHIFT_LEFT_SEQ
		elif c is INC:
			if state is INC_SEQ and bytecode[-1] < 255:
				bytecode[-1] += 1
			else:
				bytecode.extend((INC, 1))
				state = INC_SEQ
		elif c is DEC:
			if state is DEC_SEQ and bytecode[-1] < 255:
				bytecode[-1] += 1
			else:
				bytecode.extend((DEC, 1))
				state = DEC_SEQ
		elif c is READ or c is WRITE:
			bytecode.append(c)
			state = SINGLE
		elif c is JUMP_FORWARD:
			bytecode.extend((JUMP_FORWARD, 0, 0))
			brackets.append(len(bytecode) - 1) # -1 so that when the ip increments, it's correct
			state = SINGLE
		elif c is JUMP_BACKWARD:
			addr = brackets.pop()
			addrhi, addrlo = short_to_bytes(addr)
			bytecode.extend((JUMP_BACKWARD, addrhi, addrlo))
			here = short_to_bytes(len(bytecode) - 1) # -1 so that when the ip increments, it's correct
			bytecode[addr-1:addr+1] = here
			state = SINGLE
	return bytecode

def disassemble(bytecode):
	k = 0
	while k < len(bytecode):
		if bytecode[k] is SHIFT_RIGHT:
			k += 1
			print "%2d SHR   %d" % (k - 1, bytecode[k])
		elif bytecode[k] is SHIFT_LEFT:
			k += 1
			print "%2d SHL   %d" % (k - 1, bytecode[k])
		elif bytecode[k] is INC:
			k += 1
			print "%2d INC   %d" % (k - 1, bytecode[k])
		elif bytecode[k] is DEC:
			k += 1
			print "%2d DEC   %d" % (k - 1, bytecode[k])
		elif bytecode[k] is READ:
			print "%2d READ" % k
		elif bytecode[k] is WRITE:
			print "%2d WRITE" % k
		elif bytecode[k] is JUMP_FORWARD:
			k += 2
			print "%2d JMPZ  %d" % (k - 2, bytes_to_short(*bytecode[k-1:k+1]) - 2)
		elif bytecode[k] is JUMP_BACKWARD:
			k += 2
			print "%2d JMPNZ %d" % (k - 2, bytes_to_short(*bytecode[k-1:k+1]) - 2)
		else:
			print "%2d WTF" % k
		k += 1

def execute(bytecode):
	memory = bytearray(30000)
	ip = 0
	dp = 0
	buf = ""
	while True:
		if bytecode[ip] is SHIFT_RIGHT:
			ip += 1
			dp = (dp + bytecode[ip]) % len(memory)
		elif bytecode[ip] is SHIFT_LEFT:
			ip += 1
			dp = (dp - bytecode[ip]) % len(memory)
		elif bytecode[ip] is INC:
			ip += 1
			memory[dp] = (memory[dp] + bytecode[ip]) % 256
		elif bytecode[ip] is DEC:
			ip += 1
			memory[dp] = (memory[dp] - bytecode[ip]) % 256
		elif bytecode[ip] is READ:
			try:
				while len(buf) == 0:
					buf = deque(raw_input() + "\n")
				memory[dp] = ord(buf.popleft())
			except EOFError:
				memory[dp] = 0xff
		elif bytecode[ip] is WRITE:
			sys.stdout.write(chr(memory[dp]))
		elif bytecode[ip] is JUMP_FORWARD:
			ip += 2
			if memory[dp] == 0:
				ip = bytes_to_short(*bytecode[ip-1:ip+1])
		elif bytecode[ip] is JUMP_BACKWARD:
			ip += 2
			if memory[dp] != 0:
				ip = bytes_to_short(*bytecode[ip-1:ip+1])
		ip += 1
		if ip == len(bytecode):
			break
	return memory

def get_program(name):
	with open(name) as f:
		return f.read();

def dump_memory(name, memory):
	with open(name, "w") as f:
		f.write(memory)

if __name__ == "__main__":
	if len(sys.argv) == 1 or len(sys.argv) > 3:
		print "usage: %s <filename> [<dumpfile>]" % sys.argv[0]
	filename = sys.argv[1]
	source = get_program(filename)
	bytecode = compile(source)
	memory = execute(bytecode)
	if len(sys.argv) == 3:
		dump_memory(sys.argv[2], memory)