infernal-interpreter/infernal_interpreter/Emulator.py

import re

from . import Junk
from .opcodes import OPCODES

REGISTERS = [
    '%rax',
    '%rbx',
    '%rcx',
    '%rdx',
    '%rsp',
    '%rbp',
    '%rsi',
    '%rdi',
    '%r8',
    '%r9',
    '%r10',
    '%r11',
    '%r12',
    '%r13',
    '%r14',
    '%r15',
]


class CodeParseException(BaseException):
    def __init__(self, line_nr, string):
        self.line_nr = line_nr
        self.string = string


class InternalExecutionException(BaseException):
    pass


class Emulator:
    def __init__(self, source_text, max_stack_size=1000):
        self.source = source_text
        self.registers = {}
        for reg_name in REGISTERS:
            self.registers[reg_name] = Junk.Junk()
        self.stack = {}
        for i in range(max_stack_size):
            self.stack[i] = Junk.Junk()
        self.code = []
        self.labels = {}
        self.changes = {}
        self.status = {'i': True}
        self.last_sp = 0
        self.max_stack_size = 0
        index = 0
        for line in iter(source_text.splitlines()):
            index = self.processSourceLine(line, index)

    def setStack(self, *stack_list, **kwargs):
        """Sets various stack elements, starting from 0 and going up.
        Automatically sets rsp. This can be disabled by passing set_rsp=False.
        """
        i = len(self.stack)
        for element in stack_list:
            i -= 1
            self.stack[i] = element
        self.setRegs(rsp=i)
        self.last_sp = i - 1

    def setRegs(self, **reg_dict):
        for reg_name, reg_val in reg_dict.items():
            assert reg_name[0] != '%'
            self.registers['%' + reg_name] = reg_val

    def getVal(self, val_text):
        if val_text[0] == '$':
            return int(val_text[1:])
        elif val_text[0] == '%':
            return self.registers[val_text]
        else:
            raise ValueError(f'{val_text} is not an usable value name')

    def compareVal(self, valT1, valT2):
        val1 = self.getVal(valT2)
        val2 = self.getVal(valT1)
        if isinstance(val1, Junk.Junk) or isinstance(val2, Junk.Junk):
            self.status['i'] = (True,)
            return
        self.status['g'] = val1 > val2
        self.status['l'] = val1 < val2
        self.status['e'] = val1 == val2
        self.status['i'] = False

    def changedRegisters(self, *args):
        for reg, val in args:
            self.changes[reg] = val

    def processSourceLine(self, line_text, index):
        tokens = re.findall(r'[^\s,]+', line_text)
        if len(tokens) <= 0 or tokens[0][0] == '#':
            return index
        if tokens[0][-1] == ':':
            self.labels[tokens[0][:-1]] = index
            tokens = tokens[1:]
        for i in range(len(tokens)):
            if tokens[i][0] == '#':
                tokens = tokens[:i]
                break
        self.code.append(tokens)
        return index + 1

    def regState(self, reg_name):
        if reg_name in self.changes:
            return self.changes[reg_name]
        else:
            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 pushToStack(self, new_element):
        self.registers['%rsp'] -= 1
        self.stack[self.registers['%rsp']] = new_element
        self.changedRegisters(
            ('%rsp', 'change'), ('m' + str(self.registers['%rsp']), 'insert'),
        )
        self.max_stack_size = max(self.registers['%rsp'], self.max_stack_size)

    def popFromStack(self):
        temp = self.stack[self.registers['%rsp']]
        self.changedRegisters(
            ('%rsp', 'change'), ('m' + str(self.registers['%rsp']), 'remove'),
        )
        self.registers['%rsp'] += 1
        return temp

    def jump(self, label, cond_op='or', **conditions):
        if len(conditions) > 0 and self.status['i']:
            raise Junk.JunkComparisonException(*self.status['i'])
        and_, or_ = True, False
        for cnd_name, cnd_val in conditions.items():
            if self.status[cnd_name] == cnd_val:
                or_ = True
            else:
                and_ = False
        if or_ if cond_op == 'or' else and_:
            self.registers['%rip'] = self.labels[label]
            return True
        else:
            return False

    def iterate(self):
        old_rip = self.registers['%rip']
        self.last_sp = self.registers['%rsp']
        if isinstance(old_rip, Junk.Junk) or old_rip >= len(self.code):
            return None
        if not isinstance(self.registers['%rip'], int):
            raise InternalExecutionException(
                'Register %rip should be integer, but was '
                + str(self.registers['%rip']),
            )
        instruct = self.code[self.registers['%rip']]
        opcode = instruct[0]
        self.changes = {}

        OPCODES[opcode](self, *instruct[1:])

        if self.registers['%rip'] == old_rip:
            self.registers['%rip'] += 1
        else:
            self.changedRegisters(('%rip', 'jump'))

        return old_rip

    def __iter__(self):
        return self

    def __next__(self):
        output = self.iterate()
        if output is None:
            raise StopIteration()
        else:
            return output

    def getUsedRegisters(self):
        return [
            reg_name
            for reg_name, reg_val in self.registers.items()
            if not isinstance(reg_val, Junk.Junk)
        ]