infernal-interpreter/infernal_interpreter/TikzPainter.py

import re

REG_STATUS_TO_COLOR = {
    'insert': 'green',
    'change': 'yellow',
    'remove': 'red',
    'none': 'black',
    'jump': 'orange',
}

DEFAULT_OPTIONS = """register_node/.style={rectangle, draw=black!30,
                    fill=black!5, very thick, minimum size=0.5,
                    minimum width=20mm}, text_node/.style={}"""


class TikzPainter:
    def __init__(
        self,
        registers=None,
        max_stack_size=8,
        stack_draw_mode='up',
        options=None,
    ):
        self.registers = registers if registers else REGISTERS
        self.text = []
        self.pos_x = 0
        self.max_stack_size = max_stack_size
        self.stack_draw_mode = stack_draw_mode
        self.options = options if options else DEFAULT_OPTIONS

    def addText(self, str, *args):
        text = re.sub(r'[^\\]%', '\\%', str.format(*args))
        self.text.append(text)

    def getRegColor(self, reg_state):
        return REG_STATUS_TO_COLOR[reg_state]

    def drawStackUpward(self, pos, emu, x):
        base_pointer = emu.getVal('%rbp')

        pos -= 1.5
        base_sp = min(emu.last_sp, emu.getVal('%rsp'))
        pos += 0.5 * max(0, emu.getVal('%rsp') - emu.last_sp)
        assert isinstance(emu.getVal('%rsp'), int)
        for index in range(
            base_sp,
            3 + min(emu.max_stack_size, emu.getVal('%rsp') + self.max_stack_size),
        ):
            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],
            )
            if base_pointer == index:
                self.addText(
                    '\t\\draw ({0},{2}) -- ({1},{2});\n',
                    x - 1.1,
                    x + 1.1,
                    pos - 0.25,
                )
                base_pointer = emu.stack[base_pointer]
            pos -= 0.5
        return pos, x

    def drawStackDownward(self, pos, emu, x):
        pos -= 1
        for index in range(
            emu.getVal('%rsp') - 8,
            max(emu.last_sp, emu.getVal('%rsp')),
        ):
            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
        return pos, x

    def saveState(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
        if self.stack_draw_mode == 'up':
            pos, x = self.drawStackUpward(pos, emu, x)
        else:
            pos, x = self.drawStackDownward(pos, emu, x)

        # 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 = 0
        pos = 0
        for reg_name in self.registers:
            self.addText(
                '\t\\node[text_node]() at ({}, {}){{{}}};\n',
                x,
                pos,
                reg_name[1:],
            )
            pos -= 0.5

        # Draw stack
        if self.stack_draw_mode == 'up':
            self.drawStackNamesUpward(pos, emu, x)
        else:
            self.drawStackNamesDownward(pos, emu, x)

    def drawStackNamesUpward(self, pos, emu, x):
        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 to_string(self, emu):
        PRE = (
            """ \\documentclass{standalone}
                  \\usepackage{tikz}
                  \\begin{document}
                  \\begin{tikzpicture}["""
            + self.options
            + ']'
        )
        POST = """\\end{tikzpicture}
                  \\end{document}"""
        self.drawNames(emu)
        return PRE + ''.join(self.text) + POST

    def __str__(self):
        return 'TikzPainter[]'