modelverse/kernel/modelverse_jit/bytecode_ir.py

"""Provides data structures that represent parsed Modelverse bytecode graphs."""

class Instruction(object):
    """Represents a Modelverse bytecode instruction."""
    def __init__(self):
        self.next_instruction = None
        self.debug_information = None

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction, excluding the
           next instruction."""
        raise NotImplementedError()

    def get_reachable(self, filter_children=lambda _: True):
        """Gets the set of all instructions that are reachable from the given instruction, including
           this instruction. A function can be used to filter out certain instructions' children."""
        results = set()
        stack = [self]
        while len(stack) > 0:
            instr = stack.pop()
            results.add(instr)
            next_instr = instr.next_instruction
            if next_instr is not None and next_instr not in results:
                stack.append(next_instr)

            if filter_children(instr):
                for other in instr.get_directly_reachable():
                    if other not in results:
                        assert other is not None
                        stack.append(other)

        return results

class VariableNode(object):
    """Represents a variable node, which has an identifier and an optional name."""
    def __init__(self, node_id, name):
        self.node_id = node_id
        self.name = name

    def __str__(self):
        return 'var(%d, %s)' % (self.node_id, self.name)

    def __repr__(self):
        return 'VariableNode(%r, %r)' % (self.node_id, self.name)

class SelectInstruction(Instruction):
    """Represents an 'if/else' instruction."""
    def __init__(self, condition, if_clause, else_clause):
        Instruction.__init__(self)
        self.condition = condition
        self.if_clause = if_clause
        self.else_clause = else_clause

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        if self.else_clause is None:
            return (self.condition, self.if_clause)
        else:
            return (self.condition, self.if_clause, self.else_clause)

    constructor_parameters = (
        ('cond', Instruction),
        ('then', Instruction),
        ('else', Instruction))

    def __repr__(self):
        return '@%r: SelectInstruction(@%r, @%r, @%r)' % (
            id(self), id(self.condition), id(self.if_clause), id(self.else_clause))

class WhileInstruction(Instruction):
    """Represents a 'while' instruction."""
    def __init__(self, condition, body):
        Instruction.__init__(self)
        self.condition = condition
        self.body = body

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return (self.condition, self.body)

    constructor_parameters = (
        ('cond', Instruction),
        ('body', Instruction))

    def __repr__(self):
        return '@%r: WhileInstruction(@%r, @%r)' % (
            id(self), id(self.condition), id(self.body))

class BreakInstruction(Instruction):
    """Represents a 'break' instruction."""
    def __init__(self, loop):
        Instruction.__init__(self)
        self.loop = loop

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return (self.loop,)

    constructor_parameters = (('while', WhileInstruction),)

    def __repr__(self):
        return '@%r: BreakInstruction(@%r)' % (
            id(self), id(self.loop))

class ContinueInstruction(Instruction):
    """Represents a 'continue' instruction."""
    def __init__(self, loop):
        Instruction.__init__(self)
        self.loop = loop

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return (self.loop,)

    constructor_parameters = (('while', WhileInstruction),)

    def __repr__(self):
        return '@%r: ContinueInstruction(@%r)' % (
            id(self), id(self.loop))

class ReturnInstruction(Instruction):
    """Represents a 'return' instruction, which terminates the current function
       and optionally returns a value."""
    def __init__(self, value):
        Instruction.__init__(self)
        self.value = value

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        if self.value is None:
            return ()
        else:
            return (self.value,)

    constructor_parameters = (('value', Instruction),)

    def __repr__(self):
        return '@%r: ReturnInstruction(@%r)' % (
            id(self), id(self.value))

class CallInstruction(Instruction):
    """Represents a 'call' instruction, which calls a function with an argument
       list, encoded as a list of name-instruction tuples."""
    def __init__(self, target, argument_list):
        Instruction.__init__(self)
        self.target = target
        self.argument_list = argument_list

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return (self.target,) + tuple((value for _, value in self.argument_list))

    def __repr__(self):
        return '@%r: CallInstruction(@%r, [%s])' % (
            id(self), id(self.target),
            ', '.join(['%s=@%r' % (name, id(value)) for name, value in self.argument_list]))

class ConstantInstruction(Instruction):
    """Represents a 'constant' instruction, which produces a reference
       to a constant node."""
    def __init__(self, constant_id):
        Instruction.__init__(self)
        self.constant_id = constant_id
        assert self.constant_id is not None

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return ()

    constructor_parameters = (('node', int),)

    def __repr__(self):
        return '@%r: ConstantInstruction(%r)' % (id(self), self.constant_id)

class InputInstruction(Instruction):
    """Represents an 'input' instruction, which pops a node from the input
       queue."""
    def __init__(self):
        Instruction.__init__(self)

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return ()

    constructor_parameters = ()

    def __repr__(self):
        return '@%r: InputInstruction()' % id(self)

class OutputInstruction(Instruction):
    """Represents an 'output' instruction, which pushes a node onto the output
       queue."""
    def __init__(self, value):
        Instruction.__init__(self)
        self.value = value

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return (self.value,)

    constructor_parameters = (('value', Instruction),)

    def __repr__(self):
        return '@%r: OutputInstruction(@%r)' % (
            id(self), id(self.value))

class DeclareInstruction(Instruction):
    """Represents a 'declare' instruction, which declares a local variable."""
    def __init__(self, variable):
        Instruction.__init__(self)
        self.variable = variable

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return ()

    constructor_parameters = (('var', VariableNode),)

    def __repr__(self):
        return '@%r: DeclareInstruction(%r)' % (
            id(self), self.variable)

class GlobalInstruction(Instruction):
    """Represents a 'global' instruction, which declares a global variable."""
    def __init__(self, variable):
        Instruction.__init__(self)
        self.variable = variable

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return ()

    constructor_parameters = (('var', VariableNode),)

    def __repr__(self):
        return '@%r: GlobalInstruction(%r)' % (
            id(self), self.variable)

class ResolveInstruction(Instruction):
    """Represents a 'resolve' instruction, which resolves a variable node/name as
       either a local or global variable."""
    def __init__(self, variable):
        Instruction.__init__(self)
        self.variable = variable

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return ()

    constructor_parameters = (('var', VariableNode),)

    def __repr__(self):
        return '@%r: ResolveInstruction(%r)' % (
            id(self), self.variable)

class AccessInstruction(Instruction):
    """Represents an 'access' instruction, which loads the node pointed to by a
       pointer node."""
    def __init__(self, pointer):
        Instruction.__init__(self)
        self.pointer = pointer

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return (self.pointer,)

    constructor_parameters = (('var', Instruction),)

    def __repr__(self):
        return '@%r: AccessInstruction(@%r)' % (
            id(self), id(self.pointer))

class AssignInstruction(Instruction):
    """Represents an 'assign' instruction, which sets the node pointed to by a
       pointer node to the given node."""
    def __init__(self, pointer, value):
        Instruction.__init__(self)
        self.pointer = pointer
        self.value = value

    def get_directly_reachable(self):
        """Gets all instructions that are directly reachable from this instruction."""
        return (self.pointer, self.value)

    constructor_parameters = (
        ('var', Instruction),
        ('value', Instruction))

    def __repr__(self):
        return '@%r: AssignInstruction(@%r, @%r)' % (
            id(self), id(self.pointer), id(self.value))

INSTRUCTION_TYPE_MAPPING = {
    'if' : SelectInstruction,
    'while' : WhileInstruction,
    'return' : ReturnInstruction,
    'constant' : ConstantInstruction,
    'resolve' : ResolveInstruction,
    'declare' : DeclareInstruction,
    'global' : GlobalInstruction,
    'assign' : AssignInstruction,
    'access' : AccessInstruction,
    'output' : OutputInstruction,
    'input' : InputInstruction,
    'call' : CallInstruction,
    'break' : BreakInstruction,
    'continue' : ContinueInstruction
}
"""Maps instruction names to types."""