VM generator

vm_gen.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function, division
from contextlib import contextmanager
import sys
if sys.version_info.major < 3:
	input = raw_input

class CodeGenerator(object):
	def __init__(self, tab="    "):
		self.tab = tab
		self.level = 0
		self.code = ""
	def add(self, line):
		self.code += ("{0}{1}\n".format(self.tab*self.level, line))
	@contextmanager
	def indent(self):
		self.level += 1
		yield
		self.level -= 1
	@contextmanager
	def DEF(self, name, parameters = ()):
		self.add("\ndef {0}({1}):".format(name, ",".join(parameters)))
		self.level += 1
		yield
		self.level -= 1
	@contextmanager
	def WHILE(self, condition):
		self.add("while {0}:".format(condition))
		self.level += 1
		yield
		self.level -= 1
	@contextmanager
	def FOR(self, item, list):
		self.add("for {0} in {1}:".format(item, list))
		self.level += 1
		yield
		self.level -= 1
	@contextmanager
	def IF(self, condition):
		self.add("if {0}:".format(condition))
		self.level += 1
		yield
		self.level -= 1
	@contextmanager
	def ELSE(self):
		self.add("else:")
		self.level += 1
		yield
		self.level -= 1
	@contextmanager
	def DICT(self, name):
		self.add("{0} = {{".format(name))
		self.level += 1
		yield
		self.add("}")
		self.level -= 1

def generate_vm_function(c, line):
	print(line.split('\t'))
	bytecode, name, stack_effect = line.split('\t')
	name, _, parameters = name.partition(' ')
	parameters = parameters.split(' ')
	stack_effect = stack_effect.strip("(").rstrip(")")
	effect_start, effect_end = stack_effect.split("--")
	data_start, _, return_start = effect_start.partition("R:")
	data_end, _, return_end = effect_end.partition("R:")
	data_start, data_end = data_start.split(), data_end.split()
	return_start, return_end = return_start.split(), return_end.split()
	#And now, to the code-generation-mobile!
	#c.add("def vm_{0}():".format(name))
	with c.DEF("vm_%s"%name, ()):
		for var in reversed(data_start):
			c.add("{0} = pop()".format(var))
		for r_var in reversed(return_start):
			c.add("{0} = r_pop()".format(r_var))
		for var in data_end:
			if var in ("B", "W", "DW", "QW"):
				#c.add("print('pushed %s - %s')"%(var, name))
				c.add("inc_pc({0})".format({"B":1, "W":2, "DW":4, "QW":8}[var]))
			try:
				c.add("push({0})".format(int(var)))
			except ValueError:
				c.add("push({0})".format(var))
		for r_var in return_end:
			c.add("r_push({0})".format(r_var))
	#print(data_start, "R:", return_start, "--", data_end, "R:", return_end)
	return bytecode, name

def generate_pre(c):
	c.add("DATA, RETURN = [], []")
	c.add("PC = 0")
	c.add("B, W, DW, QW = 0, 0, 0, 0")
	with c.DEF("push", ("value", )):
		c.add("global DATA")
		c.add("DATA.append(value)")
	with c.DEF("pop"):
		c.add("global DATA")
		c.add("return DATA.pop()")
	with c.DEF("r_push", ("value", )):
		c.add("global RETURN")
		c.add("DATA.append(value)")
	with c.DEF("r_pop"):
		c.add("global RETURN")
		c.add("return RETURN.pop()")
	with c.DEF("inc_pc", ("x", )):
		c.add("global PC")
		c.add("PC += x")
	with c.DEF("vm_nop"):
		c.add("pass")
	with c.DEF("execute"):
		c.add("global PC, B, W, DW, QW")
		#c.add("PC = start")
		with c.WHILE("PC < len(M)"):
			c.add("instr = ord(M[PC])")
			c.add("if PC+1 < len(M): B = ord(M[PC+1])")
			c.add("if PC+2 < len(M): W = B<<8 | ord(M[PC+2])")
			#c.add("print('%s %s %s [%s] %02x %04x %04x %04x %02x'%(PC, PC+1, PC+2, len(M), B, W, B<<8, B<<8 | (ord(M[PC+2]) if PC+2<len(M) else 0xFF), (ord(M[PC+2]) if PC+2<len(M) else 0xFF)))")
			#c.add("print(instr, ord(instr), OPCODES[instr])")
			c.add("OPCODES.get(instr, vm_nop)()")
			c.add("PC += 1")

def generate_post(c):
	with c.DEF("check_eq_lists", ("a", "b")):
		with c.IF("len(a) != len(b)"):
			c.add("return False")
		c.add("return all(map(lambda i, j: i==j, a, b))")
	# with c.IF("TEST"):
		# c.add("print(TEST)")
		# c.add("execute(TEST)")
		# c.add("print('DATA\t%s'%DATA)")
		# c.add("print('RETURN\t%s'%RETURN)")
		# c.add("print('PC\t%s'%PC)")
		# c.add("print('B\t%02X\tW\t%04X'%(B, W))")
		# with c.IF("not check_eq_lists(TEST_RESULT[0], DATA) or not check_eq_lists(TEST_RESULT[1], RETURN)"):
			# c.add("print('\\nTest failed.')")
			# c.add("print('Expected:')")
			# c.add("print('\\tDATA\t%s'%TEST_RESULT[0])")
			# c.add("print('\\tRETURN\t%s'%TEST_RESULT[1])")
	with c.IF("TESTS"):
		with c.FOR("name, d_s, r_s, test, data, ret", "TESTS"):
			c.add("M = test")
			c.add("DATA, RETURN, PC = d_s, r_s, 0")
			c.add("execute()")
			with c.IF("VERBOSE"):
				c.add("print(test)")
				c.add("print('DATA\t%s'%DATA)")
				c.add("print('RETURN\t%s'%RETURN)")
				c.add("print('PC\t%s'%PC)")
				c.add("print('B\t%02X\tW\t%04X'%(B, W))")
			with c.IF("not check_eq_lists(data, DATA) or not check_eq_lists(ret, RETURN)"):
				c.add("print('Test failed (%s).'%name)")
				with c.IF("VERBOSE"):
					c.add("print('Expected:')")
					c.add("print('\\tDATA\t%s'%data)")
					c.add("print('\\tRETURN\t%s\\n'%ret)")
			with c.ELSE():
				c.add("print('Test OK (%s).'%name)")

def generate_vm(c, data):
	generate_pre(c)
	functions = []
	for line in data.split("\n"):
		if len(line) == 0: continue
		bytecode, name = generate_vm_function(c, line)
		functions.append((bytecode, name))
	with c.DICT("OPCODES"):
		for bytecode, name in functions:
			c.add("{0}: vm_{1},".format(int(bytecode, 16), name))
	generate_post(c)
	return c.code

def test_vm(code, tests):
	vm = compile(code, "<string>", 'exec')
	exec(vm, {"TESTS": tests, "VERBOSE": 0})

data = """
10\tpush_0\t( -- 0 )
11\tpush_1\t( -- 1 )
18\tpush_8\t( -- 8 )
1C\tpushb B\t( -- B )
1D\tpushw W\t( -- W )
20\tdrop\t( a -- )
21\t2drop\t( a b -- )
22\tdup\t( a -- a a )
23\t2dup\t( a b -- a b a b )
24\tswap\t( a b -- b a )
25\t2swap\t( a b c d -- c d a b )
26\trot\t( a b c -- c a b )
27\t2rot\t( a b c d e f -- e f a b c d )
28\tnip\t( a b -- b )
29\t2nip\t( a b c d -- c d )
2A\ttuck\t( a b -- b a b )
2B\t2tuck\t( a b c d -- c d a b c d )
2C\tover\t( a b -- a b a )
2D\t2over\t( a b c d -- a b c d a b )
30\tadd\t( a b -- a+b )
"""

if __name__ == "__main__":
	#data = "20\tdup\t( a -- a a )"
	c = CodeGenerator()
	vm = generate_vm(c, data)
	print(vm)
	#("", [], [], "", [], []),
	test_vm(vm, 
	[
		("Push_0", [11], [], "\x10", [11, 0], []),
		("Push_1", [11], [], "\x11", [11, 1], []),
		("Push_8", [11], [], "\x18", [11, 8], []),
		("Pushb 0x23", [11], [], "\x1C\x23", [11, 0x23], []),
		("Pushw 0x1234", [11], [], "\x1D\x12\x34", [11, 0x1234], []),
		("Drop", [11, 1], [], "\x20", [11], []),
		("Drop_2", [11, 1, 2], [], "\x21", [11], []),
		("Dup", [11, 1], [], "\x22", [11, 1, 1], []),
		("Dup_2", [11, 1, 2], [], "\x23", [11, 1, 2, 1, 2], []),
		("Swap", [11, 1, 2], [], "\x24", [11, 2, 1], []),
		("Swap_2", [11, 1, 2, 3, 4], [], "\x25", [11, 3, 4, 1, 2], []),
		("Rot", [11, 1, 2, 3], [], "\x26", [11, 3, 1, 2], []),
		("Rot_2", [11, 1, 2, 3, 4, 5, 6], [], "\x27", [11, 5, 6, 1, 2, 3, 4], []),
		
		("?", [], [], "\x00\x10\x18\x2A\x30\x30", [16], []),
	])