diff --git a/infernal.py b/infernal.py index 666066c..604c6f8 100644 --- a/infernal.py +++ b/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 diff --git a/tests/test_fib.py b/tests/test_fib.py index e9981d9..bfaf14e 100644 --- a/tests/test_fib.py +++ b/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() - for emu_ in emu: - pass - emu.text += emu.drawNames(0, emu.max_stack_size) + painter.drawState(emu) + for line_nr in emu: + painter.drawState(emu, line_nr) + 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}")