modelverse/kernel/modelverse_jit/cfg_ir.py

"""Defines control flow graph IR data structures."""

# Let's just agree to disagree on map vs list comprehensions, pylint.
# pylint: disable=I0011,W0141

class SharedCounter(object):
    """Defines a shared counter."""
    def __init__(self):
        self.index = 0

    def next_value(self):
        """Gets the next value for this counter."""
        result = self.index
        self.index += 1
        return result

class BasicBlock(object):
    """Represents a basic block."""
    def __init__(self, counter):
        self.parameters = []
        self.definitions = []
        self.counter = counter
        self.index = counter.next_value()
        self.flow = UnreachableFlow()

    def append_parameter(self, parameter):
        """Appends a parameter to this basic block."""
        assert isinstance(parameter, BlockParameter)
        result = Definition(self.counter.next_value(), parameter)
        self.parameters.append(result)
        return result

    def prepend_definition(self, value):
        """Defines the given value in this basic block."""
        assert isinstance(value, Value)
        result = Definition(self.counter.next_value(), value)
        self.definitions.insert(0, result)
        return result

    def append_definition(self, value):
        """Defines the given value in this basic block."""
        assert isinstance(value, Value)
        result = Definition(self.counter.next_value(), value)
        self.definitions.append(result)
        return result

    def remove_definition(self, definition):
        """Removes the given definition from this basic block."""
        return self.definitions.remove(definition)

    def __str__(self):
        prefix = '!%d(%s):' % (self.index, ', '.join(map(str, self.parameters)))
        return '\n'.join(
            [prefix] +
            [' ' * 4 + str(item) for item in self.definitions + [self.flow]])

class Definition(object):
    """Maps a value to a variable."""
    def __init__(self, index, value):
        self.index = index
        self.value = value
        assert isinstance(value, Value)

    def redefine(self, new_value):
        """Tweaks this definition to take on the given new value."""
        self.value = new_value
        assert isinstance(new_value, Value)

    def ref_str(self):
        """Gets a string that represents a reference to this definition."""
        return '$%d' % self.index

    def __str__(self):
        return '$%d = %s' % (self.index, str(self.value))

class Instruction(object):
    """Represents an instruction."""
    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        raise NotImplementedError()

    def get_all_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends,
           along with any dependencies of dependencies."""
        results = list(self.get_dependencies())
        for item in results:
            results.extend(item.get_all_dependencies())

        return results

class Branch(Instruction):
    """Represents a branch from one basic block to another."""
    def __init__(self, block, arguments):
        self.block = block
        assert isinstance(block, BasicBlock)
        self.arguments = arguments
        assert all([isinstance(arg, Definition) for arg in arguments])

    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return self.arguments

    def __str__(self):
        return '!%d(%s)' % (self.block.index, ', '.join(map(str, self.arguments)))

class FlowInstruction(Instruction):
    """Represents a control flow instruction which terminates a basic block."""
    def branches(self):
        """Gets a list of basic blocks targeted by this flow instruction."""
        raise NotImplementedError()

class JumpFlow(FlowInstruction):
    """Represents a control flow instruction which jumps directly to a basic block."""
    def __init__(self, branch):
        FlowInstruction.__init__(self)
        self.branch = branch
        assert isinstance(branch, Branch)

    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return self.branches()

    def branches(self):
        """Gets a list of basic blocks targeted by this flow instruction."""
        return [self.branch]

    def __str__(self):
        return 'jump %s' % self.branch

class SelectFlow(FlowInstruction):
    """Represents a control flow instruction which jumps to one of two basic blocks depending
       on whether a condition is truthy or not."""
    def __init__(self, condition, if_branch, else_branch):
        FlowInstruction.__init__(self)
        self.condition = condition
        assert isinstance(condition, Definition)
        self.if_branch = if_branch
        assert isinstance(if_branch, Branch)
        self.else_branch = else_branch
        assert isinstance(else_branch, Branch)

    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return [self.condition] + self.branches()

    def branches(self):
        """Gets a list of basic blocks targeted by this flow instruction."""
        return [self.if_branch, self.else_branch]

    def __str__(self):
        return 'select %s, %s, %s' % (self.condition.ref_str(), self.if_branch, self.else_branch)

class ReturnFlow(FlowInstruction):
    """Represents a control flow instruction which terminates the execution of the current
       function and returns a value."""
    def __init__(self, value):
        FlowInstruction.__init__(self)
        self.value = value
        assert isinstance(value, Value)

    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return [self.value]

    def branches(self):
        """Gets a list of basic blocks targeted by this flow instruction."""
        return []

    def __str__(self):
        return 'return %s' % self.value.ref_str()

class UnreachableFlow(FlowInstruction):
    """Represents a control flow instruction which is unreachable."""
    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return []

    def branches(self):
        """Gets a list of basic blocks targeted by this flow instruction."""
        return []

    def __str__(self):
        return 'unreachable'

class Value(Instruction):
    """A value: an instruction that produces some result."""
    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        raise NotImplementedError()

    def has_side_effects(self):
        """Tells if this instruction has side-effects."""
        return False

class BlockParameter(Value):
    """A basic block parameter."""
    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return []

    def __str__(self):
        return 'block-parameter'

class FunctionParameter(Value):
    """A function parameter."""
    def __init__(self, name):
        Value.__init__(self)
        self.name = name

    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return []

    def __str__(self):
        return 'func-parameter %s' % self.name

class Literal(Value):
    """A literal value."""
    def __init__(self, literal):
        Value.__init__(self)
        self.literal = literal

    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return []

    def __str__(self):
        return 'literal %r' % self.literal

class JitFunctionCall(Value):
    """A value that is the result of a function call."""
    def __init__(self, target, argument_list):
        Value.__init__(self)
        assert isinstance(target, Definition)
        self.target = target
        assert all([isinstance(val, Definition) for val in argument_list])
        self.argument_list = argument_list

    def has_side_effects(self):
        """Tells if this instruction has side-effects."""
        return True

    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return [self.target] + [val for _, val in self.argument_list]

    def __str__(self):
        return 'call %s(%s)' % (
            self.target.ref_str(),
            ', '.join(['%s=%s' % (key, val.ref_str()) for key, val in self.argument_list]))