modelverse/kernel/modelverse_jit/cfg_to_tree.py

"""Lowers CFG-IR to tree-IR."""

from collections import defaultdict
import modelverse_jit.cfg_ir as cfg_ir
import modelverse_jit.tree_ir as tree_ir
import modelverse_jit.runtime as jit_runtime
import modelverse_jit.bytecode_to_tree as bytecode_to_tree

# The CFG deconstruction code here is based on the relooper algorithm
# as detailed in https://github.com/kripken/Relooper/blob/master/paper.pdf

class FlowGraphComponent(object):
    """Represents a control-flow graph component."""
    def __init__(self, entry_blocks, blocks, reachable):
        self.entry_blocks = entry_blocks
        self.blocks = blocks
        self.reachable = reachable

    def get_entry_reachable_blocks(self):
        """Gets the set of all blocks that are reachable from the entry points."""
        return set.union(*[self.get_reachable_blocks(block) for block in self.entry_blocks])

    def get_reachable_blocks(self, block):
        """Gets all blocks in this component that are reachable from the given block."""
        return set.intersection(self.reachable[block], self.blocks)

    def get_directly_reachable_blocks(self, block):
        """Gets all blocks in this component that are reachable from the given block by taking
           exactly one branch."""
        return set.intersection(cfg_ir.get_directly_reachable_blocks(block), self.blocks)

    def can_reach(self, source_block, target_block):
        """Tests if the first block can reach the second."""
        return target_block in self.reachable[source_block] and target_block in self.blocks

    def sub_component(self, new_entry_points):
        """Creates a sub-component of this component, with the given new entry points."""
        result = FlowGraphComponent(new_entry_points, self.blocks, self.reachable)
        result.blocks = result.get_entry_reachable_blocks()
        return result

def create_graph_component(entry_point):
    """Creates a flow graph component from the given entry point."""
    reachable = cfg_ir.get_all_reachable_blocks(entry_point)
    all_blocks = set(reachable[entry_point])
    all_blocks.add(entry_point)
    return FlowGraphComponent([entry_point], all_blocks, reachable)

class SimpleBlock(object):
    """A 'simple' block in the relooper algorithm."""
    def __init__(self, body, next_block):
        self.body = body
        self.next_block = next_block

    def lower(self, state):
        """Lowers this 'simple' block to a tree."""
        return tree_ir.create_block(
            state.lower_block(self.body),
            self.next_block.lower(state))

class EmptyBlock(object):
    """An empty relooper block."""
    def lower(self, _):
        """Lowers this empty block to a tree."""
        return tree_ir.EmptyInstruction()

class LoopBlock(object):
    """A 'loop' block in the relooper algorithm."""
    def __init__(self, inner, next_block):
        self.inner = inner
        self.next_block = next_block

    def lower(self, state):
        """Lowers this 'loop' block to a tree."""
        return tree_ir.create_block(
            tree_ir.LoopInstruction(self.inner.lower(state)),
            self.next_block.lower(state))

class MultipleBlock(object):
    """A 'multiple' block in the relooper algorithm that does _not_ require a loop."""
    def __init__(self, handled_blocks, next_block):
        self.handled_blocks = handled_blocks
        self.next_block = next_block

    def lower_handled_blocks(self, state):
        """Lowers the handled blocks of this 'multiple' block to a tree."""
        result = tree_ir.EmptyInstruction()
        for entry, block in self.handled_blocks:
            result = tree_ir.SelectInstruction(
                tree_ir.BinaryInstruction(
                    tree_ir.LoadLocalInstruction(state.label_variable),
                    '==',
                    tree_ir.LiteralInstruction(entry.index)),
                block.lower(state),
                result)
        return result

    def lower(self, state):
        """Lowers this 'multiple' block to a tree."""
        return tree_ir.create_block(
            self.lower_handled_blocks(state),
            self.next_block.lower(state))

class MultipleLoopBlock(MultipleBlock):
    """A 'multiple' block in the relooper algorithm."""
    def __init__(self, handled_blocks, next_block):
        MultipleBlock.__init__(self, handled_blocks, next_block)

    def lower(self, state):
        """Lowers this 'multiple' block to a tree."""
        return tree_ir.create_block(
            tree_ir.LoopInstruction(self.lower_handled_blocks(state)),
            self.next_block.lower(state))

class ContinueFlow(cfg_ir.FlowInstruction):
    """Represents a control flow instruction which continues to the next loop iteration."""
    def __init__(self, loop):
        cfg_ir.FlowInstruction.__init__(self)
        self.loop = loop

    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 'continue'

class BreakFlow(cfg_ir.FlowInstruction):
    """Represents a control flow instruction which breaks out of a loop."""
    def __init__(self, loop):
        cfg_ir.FlowInstruction.__init__(self)
        self.loop = loop

    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 'break'

def solipsize(graph_component, loop, entry_blocks, next_blocks):
    """Replaces branches to entry blocks and next blocks by jumps to 'continue' and 'break'
       blocks, respectively."""
    all_blocks = set()
    reachable_from_inner = set()
    for block in graph_component.blocks:
        if block in next_blocks:
            continue

        all_blocks.add(block)
        for branch in block.flow.branches():
            if branch.block in entry_blocks:
                # Create a 'continue' block, and point to that.
                continue_block = cfg_ir.BasicBlock(block.counter)
                for param in branch.block.parameters:
                    continue_block.append_parameter(param)

                continue_block.flow = ContinueFlow(loop)
                branch.block = continue_block
                all_blocks.add(continue_block)
            elif branch.block in next_blocks:
                # Create a 'break' block, and point to that.
                break_block = cfg_ir.BasicBlock(block.counter)
                for param in branch.block.parameters:
                    break_block.append_parameter(param)

                break_block.flow = BreakFlow(loop)
                branch.block = break_block
                all_blocks.add(break_block)
                reachable_from_inner.add(branch.block)

    return reachable_from_inner, FlowGraphComponent(
        graph_component.entry_blocks,
        all_blocks,
        {
            block : cfg_ir.get_reachable_blocks(block)
            for block in all_blocks
        })

def to_relooper_loop(graph_component):
    """Converts the given graph component to a relooper 'loop'."""
    entry_blocks = graph_component.entry_blocks
    result = LoopBlock(None, None)
    inner_blocks = []
    next_blocks = []
    for block in graph_component.blocks:
        if any([graph_component.can_reach(block, ep) for ep in entry_blocks]):
            inner_blocks.append(block)
        else:
            next_blocks.append(block)

    reachable_from_inner, inner_component = solipsize(
        graph_component, result, entry_blocks, next_blocks)
    result.inner = reloop(inner_component)

    next_component = FlowGraphComponent(
        reachable_from_inner, next_blocks, graph_component.reachable)
    result.next_block = reloop(next_component)

    return result

def to_relooper_multiple_or_loop(graph_component):
    """Converts the given graph component to a relooper 'multiple' or 'loop'."""
    # From the Emscripten paper:
    #
    # If we have more than one entry, try to create a Multiple block:
    # For each entry, find all the labels it reaches that
    # cannot be reached by any other entry. If at least one entry
    # has such labels, return a Multiple block, whose Handled
    # blocks are blocks for those labels (and whose entries are
    # those labels), and whose Next block is all the rest. Entries
    # for the next block are entries that did not become part of
    # the Handled blocks, and also labels that can be reached
    # from the Handled blocks.
    entry_blocks = graph_component.entry_blocks
    if len(entry_blocks) <= 1:
        return to_relooper_loop(graph_component)

    entry_reachables = {ep : graph_component.get_reachable_blocks(ep) for ep in entry_blocks}
    exclusive_entries = {}
    for entry in entry_blocks:
        exclusive_blocks = set(entry_reachables[entry])
        for other_entry in entry_blocks:
            if other_entry != entry:
                exclusive_blocks.difference_update(entry_reachables[other_entry])
        if len(exclusive_blocks) > 0:
            exclusive_entries[entry] = exclusive_blocks

    if len(exclusive_entries) == 0:
        return to_relooper_loop(graph_component)

    next_entries = set(graph_component.entry_blocks)
    for block_set in exclusive_entries.values():
        for elem in block_set:
            directly_reachable = graph_component.get_directly_reachable_blocks(elem)
            directly_reachable.remove(elem)
            next_entries.update(directly_reachable)

    next_entries.difference_update(exclusive_entries.keys())

    result = MultipleLoopBlock({}, None)
    for entry, exclusive_blocks in exclusive_entries.items():
        other_blocks = set(graph_component.blocks)
        other_blocks.difference_update(exclusive_blocks)
        result.handled_blocks[entry] = reloop(
            solipsize(graph_component, result, set([entry]), other_blocks))

    result.next_block = reloop(graph_component.sub_component(next_entries))

    return result

def reloop(graph_component):
    """Applies the relooper algorithm to the given graph component."""
    entry_blocks = graph_component.entry_blocks
    if len(entry_blocks) == 0:
        return EmptyBlock()

    reachable_set = graph_component.get_entry_reachable_blocks()
    if len(entry_blocks) == 1 and entry_blocks[0] not in reachable_set:
        graph_component.blocks.remove(entry_blocks[0])
        return SimpleBlock(
            entry_blocks[0],
            reloop(
                FlowGraphComponent(
                    graph_component.get_directly_reachable_blocks(entry_blocks[0]),
                    graph_component.blocks,
                    graph_component.reachable)))
    elif all([block in reachable_set for block in entry_blocks]):
        return to_relooper_loop(graph_component)
    else:
        return to_relooper_multiple_or_loop(graph_component)

def reloop_trivial(graph_component):
    """Converts the given control-flow graph to a 'multiple' block that contains only 'simple'
       blocks."""
    return MultipleLoopBlock(
        [(block, SimpleBlock(block, EmptyBlock())) for block in graph_component.blocks],
        EmptyBlock())

def reloop_function_body(entry_point):
    """Reloops the control-flow graph defined by the given entry point."""
    return reloop_trivial(create_graph_component(entry_point))

def find_inlinable_definitions(entry_point):
    """Computes a set of definitions which are eligible for inlining, i.e., they
       are used only once and are consumed in the same basic block in which they
       are defined."""
    all_blocks = list(cfg_ir.get_all_blocks(entry_point))

    # Find all definition users for each definition.
    def_users = defaultdict(set)
    def_blocks = {}
    for block in all_blocks:
        for parameter_def in block.parameters:
            def_blocks[parameter_def] = block
        for definition in block.definitions + [block.flow]:
            def_blocks[definition] = block
            for dependency in definition.get_all_dependencies():
                if not isinstance(dependency, cfg_ir.Branch):
                    def_users[dependency].add(definition)

    # Find all definitions which are eligible for inlining.
    eligible_defs = set()
    for definition, users in def_users.items():
        if len(users) == 1:
            single_user = next(iter(users))
            if def_blocks[single_user] == def_blocks[definition]:
                eligible_defs.add(definition)

    return eligible_defs

class DefinitionScheduler(object):
    """Schedules definitions within a basic block."""
    # We now have the opportunity to schedule definitions in a somewhat intelligent way.
    # We ought to make sure that we don't reorder operations in an observable way, though.
    # Specifically:
    #
    #     We are allowed to reorder operations which have no side-effects,
    #     as long as they do not depend on each other. So this is fine:
    #
    #         $a = load $x
    #         $b = load $y
    #         $c = call-indirect 'jit' f($b, $a)
    #
    #     ==>
    #
    #         $c = call-indirect 'jit' f(load $y, load $x)
    #
    #     But we are not allowed to reorder operations which do have side-effects.
    #     This is _not_ okay:
    #
    #         $a = load $x
    #         $b = call-indirect 'jit' g()
    #         $c = store $x, $a
    #
    #     ==>
    #
    #         $b = call-indirect 'jit' g()
    #         $c = store $x, load $x
    #
    # We can combine the concepts of depending on the results of other instructions
    # and depending on the potential side-effects of instructions by extending the notion
    # of what a dependency is: every instruction depends on its explicit dependencies,
    # and has an implicit dependency on all effectful instructions that precede it.
    # Additionally, every instruction with side-effects also depends on every instruction
    # that precedes it, regardless of whether those instructions have side-effects or not.
    #
    # The advantage of this way of looking at things is that we can now construct a dependency
    # graph. For example, this is the dependency graph for the second example.
    #
    #                                $c = store $x, $a
    #                               /                 \
    #                              /                   \
    #                             v                     v
    #                    $a = load $x <------------- $b = call-indirect 'jit' g()
    #
    # Any topological sort of the dependency graph is a valid instruction scheduling, but we
    # specifically want to grow a tree that uses few temporaries.
    #
    # We can accomplish this by growing the tree in reverse: we pick a definition on which no other
    # definition depends, and remove it from the dependency graph. We then iterate over the
    # definition's dependencies in reverse. If a dependency is encountered that is eligible for
    # inlining, i.e., it is used only once and is defined in the basic block where it is used,
    # then we really want to schedule it inline -- doing so allows us to avoid defining a temporary.
    # Here's the algorithm that we'll use to schedule dependencies.
    #
    # while (dependency has not been scheduled) and (another definition depends on the dependency):
    #     pick a definition that depends on the dependency and schedule it
    #     store the scheduled definition's value in a temporary
    #
    # if (dependency has not been scheduled) and (dependency is eligible for inlining):
    #     schedule the dependency inline
    # else if (dependency has not been scheduled):
    #     schedule the dependency, and store it in a temporary
    # else:
    #     load the (already scheduled) dependency's value from a temporary
    #
    # Note that it is possible for another definition to also depend on a dependency, despite having
    # the dependency used only once. This merely means that at least one dependency is based on an
    # instruction's side-effects.
    #
    # Okay, so that's the theory. In this class, we will use the following data structure to
    # improve performance:
    #
    #     * We will implement our graph as a map of definitions to a
    #       (outgoing edges, incoming edges) tuple. This will allow us to quickly test if a
    #       definition is a dependency of some other definition, and will also allow us to
    #       erase definitions from the graph efficiently.
    #
    #     * We will not create implicit (side-effect--based) dependencies beyond the last
    #       definition with side-effects. Said definition already depends on every definition
    #       that precedes it, so those edges are simply redundant.
    #
    def __init__(self, lowering_state, definitions, flow):
        self.lowering_state = lowering_state
        self.dependency_graph = {}
        self.construct_dependency_graph(definitions + [flow])

    def add_dependency(self, definition, dependency):
        """Makes the given definition dependent on the given dependency in the
           dependency graph."""
        def_outgoing_edges, _ = self.dependency_graph[definition]
        _, dep_incoming_edges = self.dependency_graph[dependency]
        def_outgoing_edges.add(dependency)
        dep_incoming_edges.add(definition)

    def construct_dependency_graph(self, definitions):
        """Construct the dependency graph for the given set of definitions."""
        def_set = set(definitions)

        # Initialize the dependency graph for every definition in the basic block.
        for definition in definitions:
            self.dependency_graph[definition] = (set(), set())

        # Construct the dependency graph.
        last_def_batch = []
        last_effectful_def = None
        for definition in definitions:
            # Explicit dependencies.
            for dep in definition.get_all_dependencies():
                if dep in def_set:
                    self.add_dependency(definition, dep)

            # Implicit dependency on the previous definition with side-effects.
            if last_effectful_def is not None:
                self.add_dependency(definition, last_effectful_def)

            # Implicit dependency of definitions with side-effects on all definitions
            # that precede them. Implementation detail: we don't actually make
            # definitions dependent on _all_ definitions that precede them. It suffices
            # to insert a dependency on every definition between the current definition
            # and the previous definition with side-effects.
            if definition.has_side_effects():
                for implicit_dependency in last_def_batch:
                    self.add_dependency(definition, implicit_dependency)
                last_def_batch = []
                last_effectful_def = definition
            else:
                last_def_batch.append(definition)

    def remove_definition(self, definition):
        """Removes the given definition from the graph."""
        assert not self.has_dependents(definition)

        # Remove the definition from the graph.
        outgoing_edges, _ = self.dependency_graph[definition]
        del self.dependency_graph[definition]

        # Remove the definition's outgoing edges from the graph.
        for dependency in outgoing_edges:
            _, incoming_edges = self.dependency_graph[dependency]
            incoming_edges.remove(definition)

    def has_scheduled(self, definition):
        """Tests if the given definition has already been scheduled."""
        return definition not in self.dependency_graph

    def get_schedulable_definition(self):
        """Finds a schedulable definition. Returns None if the dependency graph is empty."""
        if len(self.dependency_graph) == 0:
            return None

        for node in self.dependency_graph.keys():
            if not self.has_dependents(node):
                return node

        raise ValueError(
            'Could not find an instruction to schedule. Dependency graph: %s' %
            self.dependency_graph)

    def has_dependents(self, definition):
        """Tests if there is at least one definition in the dependency graph that
           depends on the given definition."""
        return len(self.get_dependents(definition)) > 0

    def get_dependents(self, definition):
        """Gets the set of definitions that depend on the given definition."""
        _, incoming_edges = self.dependency_graph[definition]
        return incoming_edges

    def use(self, definition):
        """Creates a tree that produces the given definition's value."""
        def_value = cfg_ir.get_def_value(definition)
        if isinstance(def_value, LoweringState.inline_value_types):
            return self.lowering_state.lower_value(def_value)
        elif isinstance(def_value, cfg_ir.AllocateRootNode):
            return tree_ir.LoadLocalInstruction(
                self.lowering_state.get_root_node_name(def_value))
        elif not self.has_scheduled(definition):
            return self.schedule(definition)
        elif definition.has_value():
            return self.lowering_state.create_definition_load(definition)
        else:
            return tree_ir.EmptyInstruction()

    def __schedule_definition(self, definition, store_and_forget):
        """Schedules the given definition, excluding definitions that are dependent on it.
           A list of trees is returned."""
        assert not self.has_scheduled(definition)
        statements = []

        # print('Scheduling %s; store_and_forget: %s' % (definition, store_and_forget))

        self.remove_definition(definition)
        if isinstance(definition, cfg_ir.FlowInstruction):
            statements.append(self.lowering_state.lower_value(definition))
        else:
            # We're sure that we're dealing with a bona fide cfg_ir.Definition now.
            definition_value = definition.value
            if cfg_ir.is_value_def(definition_value, LoweringState.inline_value_types):
                pass
            elif (not store_and_forget
                  and definition in self.lowering_state.inlinable_definitions):
                statements.append(self.lowering_state.lower_value(definition_value))
            else:
                lowered_value = self.lowering_state.lower_value(definition_value)
                if definition.has_value():
                    def_load = self.lowering_state.create_definition_load(definition)
                    statements.append(
                        tree_ir.IgnoreInstruction(def_load.create_store(lowered_value)))
                    if not store_and_forget:
                        statements.append(def_load)
                else:
                    statements.append(lowered_value)

        # print('Done scheduling %s; result: %s' % (definition, tree_ir.create_block(*statements)))

        return statements

    def __schedule_top_of_stack(self, stack, statements, store_and_forget):
        """Tries to schedule the given stack's top-of-stack element."""
        definition = stack[-1]
        if self.has_scheduled(definition):
            # Definition has already been scheduled. Load it if it is the only
            # element on the stack. Otherwise, do nothing.
            stack.pop()
            if (not store_and_forget
                    and definition.has_value()
                    and not cfg_ir.is_value_def(definition, LoweringState.inline_value_types)):
                statements.append(self.lowering_state.create_definition_load(definition))
            return statements

        dependents = self.get_dependents(definition)
        if len(dependents) > 0:
            dependent = next(iter(dependents))
            stack.append(dependent)
            return statements
        else:
            # Definition has not been scheduled, and all of its dependent definitions
            # have been scheduled.
            stack.pop()
            return self.__schedule_definition(definition, store_and_forget) + statements

    def schedule(self, definition, store_and_forget=False):
        """Schedules the given definition. The resulting tree is returned."""
        assert not self.has_scheduled(definition)

        statements = []
        schedule_stack = [definition]

        while len(schedule_stack) > 0:
            statements = self.__schedule_top_of_stack(
                schedule_stack, statements, store_and_forget or len(schedule_stack) > 1)

        return tree_ir.create_block(*statements)

class LoweringState(object):
    """Stores information related to the relooper->tree conversion."""
    def __init__(self, jit, inlinable_definitions):
        self.jit = jit
        self.label_variable = tree_ir.VariableName('__label')
        self.definition_loads = {}
        self.local_name_map = bytecode_to_tree.LocalNameMap()
        self.root_edge_names = {}
        self.inlinable_definitions = inlinable_definitions
        self.scheduler = None

    def __get_root_edge_name(self, root_node):
        """Gets the name of the given root edge's variable."""
        return self.get_root_node_name(root_node) + '_edge'

    def get_root_node_name(self, root_node):
        """Gets the name of the given root node's variable."""
        if isinstance(root_node, cfg_ir.Definition):
            return self.get_root_node_name(root_node.value)

        if root_node in self.root_edge_names:
            return self.root_edge_names[root_node]

        result = 'jit_locals%d' % len(self.root_edge_names)
        self.root_edge_names[root_node] = result
        return result

    def create_definition_load(self, definition):
        """Loads the given definition's variable."""
        if definition in self.definition_loads:
            return self.definition_loads[definition]

        result = tree_ir.LoadLocalInstruction(None)
        self.definition_loads[definition] = result
        return result

    def use_definition(self, definition):
        """Loads the given definition's value."""
        return self.scheduler.use(definition)

    def lower_block(self, block):
        """Lowers the given (relooped) block to a tree."""
        statements = []
        self.scheduler = DefinitionScheduler(self, block.definitions, block.flow)
        while 1:
            schedulable_def = self.scheduler.get_schedulable_definition()
            if schedulable_def is None:
                break

            statements.append(self.scheduler.schedule(schedulable_def))

        self.scheduler = None

        statements.reverse()
        return tree_ir.create_block(*statements)

    def lower_value(self, value):
        """Lowers the given instruction to a tree."""
        value_type = type(value)
        if value_type in LoweringState.instruction_lowerings:
            return LoweringState.instruction_lowerings[value_type](self, value)
        else:
            raise jit_runtime.JitCompilationFailedException(
                "Unknown CFG instruction: '%s'" % value)

    def lower_literal(self, value):
        """Lowers the given literal value."""
        return tree_ir.LiteralInstruction(value.literal)

    def lower_check_local_exists(self, value):
        """Lowers a 'check-value-exists' value."""
        return tree_ir.LocalExistsInstruction(
            self.local_name_map.get_local_name(value.variable.node_id))

    def lower_declare_local(self, value):
        """Lowers a 'declare-local' value."""
        local_name = self.local_name_map.get_local_name(value.variable.node_id)
        return tree_ir.create_block(
            tree_ir.create_new_local_node(
                local_name,
                self.use_definition(value.root_node)),
            tree_ir.LoadLocalInstruction(local_name))

    def lower_resolve_local(self, value):
        """Lowers a 'resolve-local' value."""
        return tree_ir.LoadLocalInstruction(
            self.local_name_map.get_local_name(value.variable.node_id))

    def lower_declare_global(self, value):
        """Lowers a 'declare-global' value."""
        #
        #     global_var, = yield [("CN", [])]
        #     _globals, = yield [("RD", [task_root, "globals"])]
        #     yield [("CD", [_globals, var_name, global_var])]
        #
        global_var = tree_ir.StoreLocalInstruction(None, tree_ir.CreateNodeInstruction())
        return tree_ir.create_block(
            global_var.create_store(
                tree_ir.CreateNodeInstruction()),
            tree_ir.CreateDictionaryEdgeInstruction(
                tree_ir.ReadDictionaryValueInstruction(
                    bytecode_to_tree.load_task_root(),
                    tree_ir.LiteralInstruction('globals')),
                tree_ir.LiteralInstruction(
                    value.variable.name),
                global_var.create_load()),
            global_var.create_load())

    def lower_resolve_global(self, value):
        """Lowers a 'resolve-global' value."""
        #
        #     _globals, = yield [("RD", [task_root, "globals"])]
        #     global_var, = yield [("RD", [_globals, var_name])]
        #
        return tree_ir.ReadDictionaryValueInstruction(
            tree_ir.ReadDictionaryValueInstruction(
                bytecode_to_tree.load_task_root(),
                tree_ir.LiteralInstruction('globals')),
            tree_ir.LiteralInstruction(value.variable.name))

    def lower_function_parameter(self, value):
        """Lowers a 'function-parameter' value."""
        return tree_ir.LoadLocalInstruction(value.name)

    def lower_alloc_root_node(self, value):
        """Lowers an 'alloc-root-node' value."""
        local_name = tree_ir.VariableName(self.get_root_node_name(value))
        return tree_ir.create_block(
            tree_ir.create_new_local_node(
                local_name,
                tree_ir.LoadIndexInstruction(
                    tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
                    tree_ir.LiteralInstruction('task_root')),
                self.__get_root_edge_name(value)),
            tree_ir.LoadLocalInstruction(local_name))

    def lower_free_root_node(self, value):
        """Lowers a 'free-root-node' value."""
        return tree_ir.DeleteEdgeInstruction(
            tree_ir.LoadLocalInstruction(self.__get_root_edge_name(value.root_node)))

    def lower_load_pointer(self, value):
        """Lowers a 'load' value."""
        return bytecode_to_tree.create_access(self.use_definition(value.pointer))

    def lower_store_pointer(self, value):
        """Lowers a 'store' value."""
        ptr, val = self.use_definition(value.pointer), self.use_definition(value.value)
        return bytecode_to_tree.create_assign(ptr, val)

    def lower_read(self, value):
        """Lowers a 'read' value."""
        return tree_ir.ReadValueInstruction(
            self.use_definition(value.node))

    def lower_create_node(self, value):
        """Lowers a 'create-node' value."""
        if value.value.has_value():
            return tree_ir.CreateNodeWithValueInstruction(
                self.use_definition(value.value))
        else:
            return tree_ir.CreateNodeInstruction()

    def lower_binary(self, value):
        """Lowers a 'binary' value."""
        lhs, rhs = self.use_definition(value.lhs), self.use_definition(value.rhs)
        return tree_ir.BinaryInstruction(lhs, value.operator, rhs)

    def lower_input(self, _):
        """Lowers an 'input' value."""
        return bytecode_to_tree.create_input(self.jit.use_input_function)

    def lower_output(self, value):
        """Lowers an 'output' value."""
        return bytecode_to_tree.create_output(self.use_definition(value.value))

    def lower_direct_call(self, value):
        """Lowers a direct function call."""
        calling_convention = value.calling_convention
        if calling_convention in LoweringState.call_lowerings:
            return LoweringState.call_lowerings[calling_convention](self, value)
        else:
            raise jit_runtime.JitCompilationFailedException(
                "Unknown calling convention: '%s' in instruction '%s'" %
                (calling_convention, value))

    def lower_indirect_call(self, value):
        """Lowers an indirect function call."""
        target = self.use_definition(value.target)
        args = [(name, self.use_definition(arg))
                for name, arg in value.argument_list]
        return bytecode_to_tree.create_indirect_call(target, args)

    def lower_simple_positional_call(self, value):
        """Lowers a direct call that uses the 'simple-positional' calling convention."""
        return tree_ir.CallInstruction(
            tree_ir.LoadGlobalInstruction(value.target_name),
            [self.use_definition(arg) for _, arg in value.argument_list])

    def lower_jit_call(self, value):
        """Lowers a direct call that uses the 'jit' calling convention."""
        arg_list = [(name, self.use_definition(arg))
                    for name, arg in value.argument_list]
        intrinsic = self.jit.get_intrinsic(value.target_name)
        if intrinsic is not None:
            return bytecode_to_tree.apply_intrinsic(intrinsic, arg_list)
        else:
            return tree_ir.create_jit_call(
                tree_ir.LoadGlobalInstruction(value.target_name),
                arg_list,
                tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME))

    def lower_macro_call(self, value):
        """Expands a macro call."""
        arg_list = [self.use_definition(arg) for _, arg in value.argument_list]
        if value.target_name in LoweringState.macro_lowerings:
            return LoweringState.macro_lowerings[value.target_name](*arg_list)
        else:
            raise jit_runtime.JitCompilationFailedException(
                "Unknown macro: '%s' in instruction '%s'" %
                (value.target_name, value))

    def lower_jump(self, flow):
        """Lowers the given 'jump' flow instruction to a tree."""
        return self.lower_branch(flow.branch)

    def lower_select(self, flow):
        """Lowers the given 'select' flow instruction to a tree."""
        # Schedule all branch arguments, so their definitions don't end up in the
        # 'if' statement.
        statements = []
        for branch in (flow.if_branch, flow.else_branch):
            for arg in branch.arguments:
                if not self.scheduler.has_scheduled(arg):
                    statements.append(self.scheduler.schedule(arg, True))

        statements.append(
            tree_ir.SelectInstruction(
                self.use_definition(flow.condition),
                self.lower_branch(flow.if_branch),
                self.lower_branch(flow.else_branch)))

        return tree_ir.create_block(*statements)

    def lower_return(self, flow):
        """Lowers the given 'return' flow instruction to a tree."""
        return tree_ir.ReturnInstruction(self.use_definition(flow.value))

    def lower_throw(self, flow):
        """Lowers the given 'throw' flow instruction to a tree."""
        return tree_ir.RaiseInstruction(self.use_definition(flow.exception))

    def lower_unreachable(self, _):
        """Lowers the given 'unreachable' flow instruction to a tree."""
        return tree_ir.IgnoreInstruction(
            tree_ir.CallInstruction(
                tree_ir.LoadGlobalInstruction(jit_runtime.UNREACHABLE_FUNCTION_NAME), []))

    def lower_break(self, _):
        """Lowers the given 'break' flow instruction to a tree."""
        return tree_ir.BreakInstruction()

    def lower_continue(self, _):
        """Lowers the given 'continue' flow instruction to a tree."""
        return tree_ir.ContinueInstruction()

    def lower_branch(self, branch):
        """Lowers the given (relooped) branch to a tree."""
        instructions = []
        for param, arg in zip(branch.block.parameters, branch.arguments):
            instructions.append(
                tree_ir.IgnoreInstruction(
                    self.create_definition_load(param).create_store(self.use_definition(arg))))

        instructions.append(
            tree_ir.IgnoreInstruction(
                tree_ir.StoreLocalInstruction(
                    self.label_variable,
                    tree_ir.LiteralInstruction(branch.block.index))))

        return tree_ir.create_block(*instructions)

    inline_value_types = (cfg_ir.Literal, cfg_ir.ResolveLocal, cfg_ir.FunctionParameter)

    instruction_lowerings = {
        cfg_ir.Literal : lower_literal,
        cfg_ir.CheckLocalExists : lower_check_local_exists,
        cfg_ir.DeclareLocal : lower_declare_local,
        cfg_ir.ResolveLocal : lower_resolve_local,
        cfg_ir.DeclareGlobal : lower_declare_global,
        cfg_ir.ResolveGlobal : lower_resolve_global,
        cfg_ir.FunctionParameter : lower_function_parameter,
        cfg_ir.AllocateRootNode : lower_alloc_root_node,
        cfg_ir.DeallocateRootNode : lower_free_root_node,
        cfg_ir.LoadPointer : lower_load_pointer,
        cfg_ir.StoreAtPointer : lower_store_pointer,
        cfg_ir.Read : lower_read,
        cfg_ir.CreateNode : lower_create_node,
        cfg_ir.Binary : lower_binary,
        cfg_ir.Input : lower_input,
        cfg_ir.Output : lower_output,
        cfg_ir.DirectFunctionCall : lower_direct_call,
        cfg_ir.IndirectFunctionCall : lower_indirect_call,
        cfg_ir.JumpFlow : lower_jump,
        cfg_ir.SelectFlow : lower_select,
        cfg_ir.ReturnFlow : lower_return,
        cfg_ir.ThrowFlow : lower_throw,
        cfg_ir.UnreachableFlow : lower_unreachable,
        BreakFlow : lower_break,
        ContinueFlow : lower_continue
    }

    call_lowerings = {
        cfg_ir.SIMPLE_POSITIONAL_CALLING_CONVENTION : lower_simple_positional_call,
        cfg_ir.JIT_CALLING_CONVENTION : lower_jit_call,
        cfg_ir.MACRO_POSITIONAL_CALLING_CONVENTION : lower_macro_call
    }

    macro_lowerings = {
        cfg_ir.PRINT_MACRO_NAME: tree_ir.PrintInstruction
    }

def lower_flow_graph(entry_point, jit):
    """Lowers the control-flow graph defined by the given entry point to tree IR."""
    cfg_lowerer = LoweringState(jit, find_inlinable_definitions(entry_point))
    lowered_body = reloop_function_body(entry_point).lower(cfg_lowerer)
    lowered_body = tree_ir.CompoundInstruction(
        cfg_lowerer.lower_jump(cfg_ir.create_jump(entry_point)), lowered_body)
    return tree_ir.map_and_simplify(lambda x: x, lowered_body)