Moved drawing functionality to a new class "TikzPainter".

infernal.py

@@ -3,23 +3,92 @@ import re
 import sys
 
 REGISTERS=["%rip","%rbp","%rsp","","%rax","%rdi","%r8","%r9"]
-REG_STATUS_TO_COLOR = {"insert":"green","change":"yellow","remove":"red"}
+REG_STATUS_TO_COLOR = {
+	"insert": "green",
+	"change": "yellow",
+	"remove": "red",
+	"none":   "black"
+}
+
+class TikzPainter:
+
+	def __init__ (self, registers=None, max_stack_size=8):
+		self.registers = registers if registers else REGISTERS
+		self.text = []
+		self.pos_x = 0
+		self.max_stack_size = max_stack_size
+
+	def addText (self, str, *args):
+		self.text.append(str.format(*args))
+
+	def getRegColor (self, reg_state):
+		return REG_STATUS_TO_COLOR[reg_state]
+
+	def drawState (self, emu, line_nr=None):
+		x = (self.pos_x+1)*2.5-0.5
+		self.pos_x += 1
+
+		# Draw register cells
+		pos = 0.5
+		for reg_name in self.registers:
+			pos -= 0.5
+			if reg_name =="":
+				continue
+			reg_state = emu.regState(reg_name)
+			reg_dect = "register_node, fill="+self.getRegColor(reg_state)+"!10"
+			self.addText("\t\\node[{}]() at ({}, {}){{{}}};\n",
+					reg_dect, x, pos, emu.getVal(reg_name))
+
+		# Draw stack
+		pos -= 1.5
+		base_sp = max(emu.last_sp, emu.getVal('%rsp'))
+		if emu.last_sp - 1 == emu.getVal('%rsp'):
+			pos += 0.5
+		for index in range(base_sp, max(-1, base_sp-self.max_stack_size),-1):
+			reg_state = emu.regState("m"+str(index))
+			reg_dect = "register_node, fill="+self.getRegColor(reg_state)+"!10"
+			self.addText("\t\\node[{}]() at ({}, {}){{{}}};\n",
+					reg_dect, x, pos, emu.stack[index])
+			pos -= 0.5
+
+		# Draw line signature
+		if line_nr == None:
+			return
+		pos = 2
+		signature = emu.getLineSignature(line_nr)
+		for token in signature:
+			self.addText("\t\\node[text_node]() at ({}, {}){{{}}};\n", x-1.25, pos, token)
+			pos -= 0.5
+
+	def drawNames (self, emu):
+		x = -1
+		pos = 0
+		for reg_name in self.registers:
+			self.addText("\t\\node[text_node]() at ({}, {}){{{}}};\n", x, pos, reg_name[1:])
+			pos -= 0.5
+
+		self.addText("\t\\node[text_node]() at ({}, {}){{TOS+1}};\n", x, pos-0.5)
+		self.addText("\t\\node[text_node]() at ({}, {}){{TOS}};\n", x, pos-1.0)
+		pos -= 1
+		for index in range(1,self.max_stack_size):
+			pos -= 0.5
+			self.addText("\t\\node[text_node]() at ({}, {}){{TOS-{}}};\n", x, pos, index)
+
+	def __str__ (self):
+		return "".join(self.text)
+
 
 class Emulator:
 
 	def __init__ (self, source_text):
 		self.registers = {}
 		for reg_name in REGISTERS:
-			self.registers[reg_name] = 0
+			self.registers[reg_name] = "Junk"
 		self.stack = {}
 		self.code = []
-		self.codep = []
-		self.linesExec = []
 		self.labels = {}
 		self.changes = {}
 		self.status = {}
-		self.text = ""
-		self.x_coord = 0
 		self.last_sp = 0
 		self.max_stack_size = 0
 		index = 0
@@ -65,82 +134,22 @@ class Emulator:
 				tokens = tokens[:i]
 				break
 		self.code.append(tokens)
-		self.codep.append(line_text)
 		return index + 1
 
-	def getRegColor (self, reg_name):
+	def regState (self, reg_name):
 		if reg_name in self.changes:
-			return REG_STATUS_TO_COLOR[self.changes[reg_name]]
+			return self.changes[reg_name]
 		else:
-			return "black"
+			return "none"
 
 	def getLineSignature (self, line_nr):
 		return [ "\\"+token if token[0]=="$" or token[0]=="%" else token for token in self.code[line_nr] ]
 
-	def draw (self, x, line_nr=None):
-		x = str(x)
-		text = ""
-		pos = 0
-		for reg_name in REGISTERS:
-			reg_dect = "register_node, fill=" + self.getRegColor(reg_name)+"!10"
-			if reg_name != "":
-				text += "\t\\node["+reg_dect+"]("+reg_name[1:]+") at ("+x+", "+str(pos)+"){"+str(self.registers[reg_name])+"};\n"
-			pos -= 0.5
-
-		pos -= 1
-		base_sp = max(self.last_sp,self.registers['%rsp'])
-		if self.last_sp - 1 == self.registers['%rsp']:
-			pos += 0.5
-		for index in range(base_sp, -1,-1):
-			mem_name = "m"+str(index)
-			reg_dect = "register_node, fill=" + self.getRegColor(mem_name)+"!10"
-			text += "\t\\node["+reg_dect+"]("+mem_name+") at ("+x+", "+str(pos)+"){"+str(self.stack[index])+"};\n"
-			pos -= 0.5
-
-		if line_nr != None:
-			pos = 2
-			signature = self.getLineSignature(line_nr)
-			try:
-				for i in range(3):
-					text += "\t\\node[text_node]() at ("+str(float(x)-1.25)+", "+str(pos)+"){"+signature[i]+"};\n"
-					pos -= 0.5
-			except:
-				pass
-		#for line_nr in self.linesExec:
-		#	text += "\t\\node[text_node]() at ("+x+", "+str(pos)+"){"+self.codep[line_nr]+"};\n"
-		self.linesExec = []
-
-		return text
-
-	def drawRegs (self, line_nr=None):
-		self.text += self.draw((self.x_coord+1)*2.5-0.5, line_nr)
-		self.text += "\n"
-		self.x_coord += 1
-
-	def drawNames (self, x, stack_size=8):
-		x = str(x)
-		text = ""
-		pos = 0
-		for reg_name in REGISTERS:
-			text += "\t\\node[text_node]("+reg_name[1:]+") at ("+x+", "+str(pos)+"){"+reg_name[1:]+"};\n"
-			pos -= 0.5
-
-		text += "\t\\node[text_node]() at ("+x+", "+str(pos-0.5)+"){TOS+1};\n"
-		text += "\t\\node[text_node]() at ("+x+", "+str(pos-1)+"){TOS};\n"
-		pos -= 1.5
-		for index in range(1,stack_size+1):
-			mem_name = "m"+str(index)
-			text += "\t\\node[text_node]() at ("+x+", "+str(pos)+"){TOS-"+str(index)+"};\n"
-			pos -= 0.5
-
-		return text
-
 	def iterate (self):
 		old_rip = self.registers['%rip']
 		self.last_sp = self.registers['%rsp']
-		#self.linesExec.append(old_rip)
 		if self.registers['%rip'] >= len(self.code):
-			return False
+			return None
 		instruct = self.code[self.registers['%rip']]
 		opcode = instruct[0]
 		self.changes = {}
@@ -192,13 +201,14 @@ class Emulator:
 		else:
 			self.changedRegisters(('%rip',"change"))
 
-		self.drawRegs(old_rip)
+		return old_rip
 
 	def __iter__(self):
 		return self
 
 	def next(self):
-		if self.iterate() == False:
+		output = self.iterate()
+		if output is None:
 			raise StopIteration()
 		else:
-			return self
+			return output

tests/test_fib.py

@@ -61,18 +61,19 @@ if __name__ == "__main__":
 
 	# Setup
 	emu = infernal.Emulator(fib_prog)
+	painter = infernal.TikzPainter()
 	emu.setStack("junk...","calling eip")
 	emu.setRegs( rip = 0, rbp = 'old bp', rsp = 1, rdi = a )
 
-	emu.drawRegs()
+	painter.drawState(emu)
-	for emu_ in emu:
+	for line_nr in emu:
-		pass
+		painter.drawState(emu, line_nr)
-	emu.text += emu.drawNames(0, emu.max_stack_size)
+	painter.drawNames(emu)
 
 	printf("fib({}) = {}", a, emu.registers["%rax"])
 
 	with open("tikz.tex","w") as f:
 		f.write("\\documentclass{standalone}\n\n\usepackage{tikz}\\begin{document}\n")
 		f.write("\\begin{tikzpicture}["+OPTIONS+"]\n")
-		f.write(emu.text)
+		f.write(str(painter))
 		f.write("\\end{tikzpicture}\n\\end{document}")