# -*- coding: UTF-8 -*-

"""
Control flow for the AST backend.

Adapted from Cython/Compiler/FlowControl.py
"""

import os
import ast
import sys
import copy
import inspect
import logging
import itertools
import subprocess

from meta import asttools

import numba
from numba import (error, transforms, closure, visitors, symtab, nodes,
                   typesystem)
from numba.utils import dump

from numba import *

debug = False
#debug = True
debug_cfg = False
#debug_cfg = True

logger = logging.getLogger(__name__)
if debug:
    logger.setLevel(logging.DEBUG)

if debug_cfg:
    dot_output_graph = os.path.expanduser("~/cfg.dot")
else:
    dot_output_graph = False

class ControlBlock(nodes.LowLevelBasicBlockNode):
    """
    Control flow graph node. Sequence of assignments and name references.
    This is simultaneously an AST node.

       children  set of children nodes
       parents   set of parent nodes
       positions set of position markers

       stats     list of block statements
       gen       dict of assignments generated by this block
       bound     set  of entries that are definitely bounded in this block

       Example:

        a = 1
        b = a + c # 'c' is already bounded or exception here

        stats = [Assignment(a), NameReference(a), NameReference(c),
                     Assignment(b)]
        gen = {Entry(a): Assignment(a), Entry(b): Assignment(b)}
        bound = set([Entry(a), Entry(c)])
    """

    _fields = ['phi_nodes', 'body']

    def __init__(self, id, label='empty', have_code=True,
                 is_expr=False, is_exit=False, pos=None,
                 is_fabricated=False):
        if pos:
            label = "%s_%s" % (label, error.format_pos(pos).rstrip(": "))
        super(ControlBlock, self).__init__(body=[], label=label)

        self.id = id

        self.children = set()
        self.parents = set()
        self.positions = set()

        self.stats = []
        self.gen = {}
        self.bound = set()

        # Same as i_input/i_output but for reaching defs with sets
        self.input = set()
        self.output = set()

        self.i_input = 0
        self.i_output = 0
        self.i_gen = 0
        self.i_kill = 0
        self.i_state = 0

        self.is_expr = is_expr
        self.is_exit = is_exit
        self.have_code = have_code

        # TODO: Make these bits
        # Set of blocks that dominate this block
        self.dominators = set()
        # Set of blocks where our dominance stops
        self.dominance_frontier = set()
        # SSA Φ locations. Maps Variables to a list of (basic_block, definition)
        # There can be only one reaching definition, since each variable is
        # assigned only once
        self.phis = {}
        self.phi_nodes = []

        # Promotions at the end of the block to have a consistent promoted
        # Φ type at one of our children.
        self.promotions = {} # (renamed_var_name, dst_type) -> promotion_node

        # LLVM entry and exit blocks. The entry block is the block before the
        # body is evaluated, the exit block the block after the body is
        # evaluated.
        self.exit_block = None
        self.phi_block = None
        self.exit_block = None
        self.promotions = set()

        self.symtab = None
        self.is_fabricated = is_fabricated
        # If set to True, branch from the previous basic block to this basic
        # block
        self.branch_here = False

    def empty(self):
        return (not self.stats and not self.positions and not self.phis)

    def detach(self):
        """Detach block from parents and children."""
        for child in self.children:
            child.parents.remove(self)
        for parent in self.parents:
            parent.children.remove(self)
        self.parents.clear()
        self.children.clear()

    def add_child(self, block):
        self.children.add(block)
        block.parents.add(self)

    def reparent(self, new_block):
        """
        Re-parent all children to the new block
        """
        for child in self.children:
            child.parents.remove(self)
            new_block.add_child(child)

    def delete(self, flow):
        """
        Delete a block from the cfg.
        """
        for parent in self.parents:
            parent.children.remove(self)
        for child in self.children:
            child.parents.remove(self)

        flow.blocks.remove(self)

    def __repr__(self):
        return 'Block(%d)' % self.id

    def __getattr__(self, attr):
        if attr in ('variable', 'type'):
            return getattr(self.body[0], attr)
        raise AttributeError

    def __getattr__(self, attr):
        if attr in ('variable', 'type', 'ctx'):
            return getattr(self.body[0], attr)
        raise AttributeError

    def __setattr__(self, attr, value):
        if attr in ('variable', 'type'):
            setattr(self.body[0], attr, value)
        else:
            super(ControlBlock, self).__setattr__(attr, value)


class ExitBlock(ControlBlock):
    """Non-empty exit point block."""

    def empty(self):
        return False


class AssignmentList:
    def __init__(self):
        self.stats = []


class ControlFlow(object):
    """
    Control-flow graph.

       entry_point ControlBlock entry point for this graph
       exit_point  ControlBlock normal exit point
       block       ControlBlock current block
       blocks      set    children nodes
       entries     set    tracked entries
       loops       list   stack for loop descriptors
       exceptions  list   stack for exception descriptors

    """

    def __init__(self, source_descr):
        self.source_descr = source_descr

        self.blocks = []
        self.entries = set()
        self.loops = []
        self.exceptions = []

        self.entry_point = ControlBlock(-1, label='entry')
        self.exit_point = ExitBlock(0, label='exit')
        self.block = self.entry_point

    def newblock(self, parent=None, **kwargs):
        """
        Create floating block linked to `parent` if given.
        Does NOT set the current block to the new block.
        """
        block = ControlBlock(len(self.blocks), **kwargs)
        self.blocks.append(block)
        if parent:
            parent.add_child(block)

        return block

    def nextblock(self, parent=None, **kwargs):
        """
        Create child block linked to current or `parent` if given.
        Sets the current block to the new block.
        """
        block = self.newblock(parent, **kwargs)
        if not parent and self.block:
            self.block.add_child(block)

        self.block = block
        return block

    def exit_block(self, parent=None, **kwargs):
        """
        Create a floating exit block. This can later be added to self.blocks.
        This is useful to ensure topological order.
        """
        block = self.newblock(parent, have_code=False, is_exit=True, **kwargs)
        self.blocks.pop()
        return block

    def add_exit(self, exit_block):
        "Add an exit block after visiting the body"
        exit_block.id = len(self.blocks)
        self.blocks.append(exit_block)

    def is_tracked(self, entry):
        return entry.renameable

    def mark_position(self, node):
        """Mark position, will be used to draw graph nodes."""
        if self.block:
            src_descr = self.source_descr
            pos = (src_descr,) + getpos(node)
            self.block.positions.add(pos)

    def mark_assignment(self, lhs, rhs, entry, assignment, warn_unused=True):
        if self.block:
            if not self.is_tracked(entry):
                return
            assignment = NameAssignment(lhs, rhs, entry, assignment,
                                        warn_unused=warn_unused)
            self.block.stats.append(assignment)
            self.block.gen[entry] = assignment
            self.entries.add(entry)
            return assignment

    def mark_argument(self, lhs, rhs, entry):
        if self.block and self.is_tracked(entry):
            assignment = Argument(lhs, rhs, entry)
            self.block.stats.append(assignment)
            self.block.gen[entry] = assignment
            self.entries.add(entry)

    def mark_deletion(self, node, entry):
        if self.block and self.is_tracked(entry):
            assignment = NameDeletion(node, entry)
            self.block.stats.append(assignment)
            self.block.gen[entry] = Uninitialized
            self.entries.add(entry)

    def mark_reference(self, node, entry):
        if self.block and self.is_tracked(entry):
            self.block.stats.append(NameReference(node, entry))
            # Local variable is definitely bound after this reference
            if not allow_null(node):
                self.block.bound.add(entry)
            self.entries.add(entry)

    def normalize(self):
        """Delete unreachable and orphan blocks."""
        blocks = set(self.blocks)
        queue = set([self.entry_point])
        visited = set()
        while queue:
            root = queue.pop()
            visited.add(root)
            for child in root.children:
                if child not in visited:
                    queue.add(child)
        unreachable = blocks - visited
        for block in unreachable:
            block.detach()
        visited.remove(self.entry_point)
        for block in visited:
            if block.empty():
                for parent in block.parents: # Re-parent
                    for child in block.children:
                        parent.add_child(child)
                block.detach()
                unreachable.add(block)
        blocks -= unreachable
        self.blocks = [block for block in self.blocks if block in blocks]

    def initialize(self):
        """Set initial state, map assignments to bits."""
        self.assmts = {}

        offset = 0
        for entry in self.entries:
            assmts = AssignmentList()
            assmts.bit = 1 << offset
            assmts.mask = assmts.bit
            self.assmts[entry] = assmts
            offset += 1

        for block in self.blocks:
            block.stats = block.phis.values() + block.stats
            for stat in block.stats:
                if isinstance(stat, (PhiNode, NameAssignment)):
                    stat.bit = 1 << offset
                    assmts = self.assmts[stat.entry]
                    assmts.stats.append(stat)
                    assmts.mask |= stat.bit
                    offset += 1

        for block in self.blocks:
            for entry, stat in block.gen.items():
                assmts = self.assmts[entry]
                if stat is Uninitialized:
                    block.i_gen |= assmts.bit
                else:
                    block.i_gen |= stat.bit
                block.i_kill |= assmts.mask
            block.i_output = block.i_gen
            for entry in block.bound:
                block.i_kill |= self.assmts[entry].bit

        for assmts in self.assmts.itervalues():
            self.entry_point.i_gen |= assmts.bit
        self.entry_point.i_output = self.entry_point.i_gen

    def map_one(self, istate, entry):
        "Map the bitstate of a variable to the definitions it represents"
        ret = set()
        assmts = self.assmts[entry]
        if istate & assmts.bit:
            ret.add(Uninitialized)
        for assmt in assmts.stats:
            if istate & assmt.bit:
                ret.add(assmt)
        return ret

    def reaching_definitions(self):
        """Per-block reaching definitions analysis."""
        dirty = True
        while dirty:
            dirty = False
            for block in self.blocks:
                i_input = 0
                for parent in block.parents:
                    i_input |= parent.i_output
                i_output = (i_input & ~block.i_kill) | block.i_gen
                if i_output != block.i_output:
                    dirty = True
                block.i_input = i_input
                block.i_output = i_output

    def initialize_sets(self):
        """
        Set initial state, run after SSA. There is only ever one live
        definition of a variable in a block, so we can simply track input
        and output definitions as the Variable/Entry they came as.
        """
        for block in self.blocks:
            # Insert phi nodes from SSA stage into the assignments of the block
            for phi in block.phis:
                block.gen.setdefault(phi, []).insert(0, phi)

            # Update the kill set with the variables that are assigned to in
            # the block
            block.kill = set(block.gen)
            block.output = set(block.gen)
            #for entry in block.bound:
            #    block.i_kill |= self.assmts[entry].bit

        for assmts in self.assmts.itervalues():
            self.entry_point.i_gen |= assmts.bit
        self.entry_point.i_output = self.entry_point.i_gen

    def compute_dominators(self):
        """
        Compute the dominators for the CFG, i.e. for each basic block the
        set of basic blocks that dominate that block. This mean from the
        entry block to that block must go through the blocks in the dominator
        set.

        dominators(x) = {x} ∪ (∩ dominators(y) for y ∈ preds(x))
        """
        blocks = set(self.blocks)
        for block in self.blocks:
            block.dominators = blocks

        changed = True
        while changed:
            changed = False
            for block in self.blocks:
                parent_dominators = [parent.dominators for parent in block.parents]
                new_doms = set.intersection(block.dominators, *parent_dominators)
                new_doms.add(block)

                if new_doms != block.dominators:
                    block.dominators = new_doms
                    changed = True

    def immediate_dominator(self, x):
        """
        The dominator of x that is dominated by all other dominators of x.
        This is the block that has the largest dominator set.
        """
        candidates = x.dominators - set([x])
        if not candidates:
            return None

        result = max(candidates, key=lambda b: len(b.dominators))
        ndoms = len(result.dominators)
        assert len([b for b in candidates if len(b.dominators) == ndoms]) == 1
        return result

    def compute_dominance_frontier(self):
        """
        Compute the dominance frontier for all blocks. This indicates for
        each block where dominance stops in the CFG. We use this as the place
        to insert Φ functions, since at the dominance frontier there are
        multiple control flow paths to the block, which means multiple
        variable definitions can reach there.
        """
        if debug:
            print "Dominator sets:"
            for block in self.blocks:
                print block.id, sorted(block.dominators, key=lambda b: b.id)

        blocks = []
        for block in self.blocks:
            if block.parents:
                block.idom = self.immediate_dominator(block)
                block.visited = False
                blocks.append(block)

        self.blocks = blocks

        def visit(block, result):
            block.visited = True
            for child in block.children:
                if not child.visited:
                    visit(child, result)
            result.append(block)

        #postorder = []
        #visit(self.blocks[0], postorder)
        postorder = self.blocks[::-1]

        # Compute dominance frontier
        for x in postorder:
            for y in x.children:
                if y.idom is not x:
                    # We are not an immediate dominator of our successor, add
                    # to frontier
                    x.dominance_frontier.add(y)

            for z in self.blocks:
                if z.idom is x:
                    for y in z.dominance_frontier:
                        if y.idom is not x:
                            x.dominance_frontier.add(y)

    def update_for_ssa(self, ast, symbol_table):
        """
        1) Compute phi nodes

            for each variable v
                1) insert empty phi nodes in dominance frontier of each block
                   that defines v
                2) this phi defines a new assignment in each block in which
                   it is inserted, so propagate (recursively)

        2) Reaching definitions

            Set block-local symbol table for each block.
            This is a rudimentary form of reaching definitions, but we can
            do it in a single pass because all assignments are known (since
            we inserted the phi functions, which also count as assignments).
            This means the output set is known up front for each block
            and never changes. After setting all output sets, we can compute
            the input sets in a single pass:

                1) compute output sets for each block
                2) compute input sets for each block

        3) Update phis with incoming variables. The incoming variables are
           last assignments of the predecessor blocks in the CFG.
        """
        # Print dominance frontier
        if debug:
            print "Dominance frontier:"
            for block in self.blocks:
                print 'DF(%d) = %s' % (block.id, block.dominance_frontier)

        #
        ### 1) Insert phi nodes in the right places
        #
        for name, variable in symbol_table.iteritems():
            if not variable.renameable:
                continue

            defining = []
            for b in self.blocks:
                if variable in b.gen:
                    defining.append(b)

            for defining_block in defining:
                for f in defining_block.dominance_frontier:
                    phi = f.phis.get(variable, None)
                    if phi is None:
                        phi = PhiNode(f, variable)
                        f.phis[variable] = phi
                        defining.append(f)

        #
        ### 2) Reaching definitions and variable renaming
        #

        # Set originating block for each variable (as if each variable were
        # initialized at the start of the function) and start renaming of
        # variables
        symbol_table.counters = dict.fromkeys(symbol_table, -1) # var_name -> counter
        self.blocks[0].symtab = symbol_table
        for var_name, var in symbol_table.items():
            if var.renameable:
                new_var = symbol_table.rename(var, self.blocks[0])
                new_var.set_uninitialized = True

        self.rename_assignments(self.blocks[0])

        for block in self.blocks[1:]:
            block.symtab = symtab.Symtab(parent=block.idom.symtab)
            for var, phi_node in block.phis.iteritems():
                phi_node.variable = block.symtab.rename(var, block)
                phi_node.variable.name_assignment = phi_node
                phi_node.variable.is_phi = True

            self.rename_assignments(block)

        #
        ### 3) Update the phis with all incoming entries
        #
        for block in self.blocks:
            # Insert phis in AST
            block.phi_nodes = block.phis.values()
            for variable, phi in block.phis.iteritems():
                for parent in block.parents:
                    incoming_var = parent.symtab.lookup_most_recent(variable.name)
                    phi.incoming.add(incoming_var)

                    # Update def-use chain
                    incoming_var.cf_references.append(phi)

    def rename_assignments(self, block):
        lastvars = dict(block.symtab)
        for stat in block.stats:
            if isinstance(stat, NameAssignment) and stat.assignment_node:
                # print "setting", stat.lhs, hex(id(stat.lhs))
                stat.lhs.variable = block.symtab.rename(stat.entry, block)
                stat.lhs.variable.name_assignment = stat
            elif isinstance(stat, NameReference):
                current_var = block.symtab.lookup_most_recent(stat.entry.name)
                stat.node.variable = current_var
                current_var.cf_references.append(stat.node)


class StatementDescr(object):
    is_assignment = False

class LoopDescr(object):
    def __init__(self, next_block, loop_block):
        self.next_block = next_block
        self.loop_block = loop_block
        self.exceptions = []


class ExceptionDescr(object):
    """Exception handling helper.

    entry_point   ControlBlock Exception handling entry point
    finally_enter ControlBlock Normal finally clause entry point
    finally_exit  ControlBlock Normal finally clause exit point
    """

    def __init__(self, entry_point, finally_enter=None, finally_exit=None):
        self.entry_point = entry_point
        self.finally_enter = finally_enter
        self.finally_exit = finally_exit


class NameAssignment(object):

    is_assignment = True

    def __init__(self, lhs, rhs, entry, assignment_node, warn_unused=True):
        if not hasattr(lhs, 'cf_state'):
            lhs.cf_state = set()
        if not hasattr(lhs, 'cf_is_null'):
            lhs.cf_is_null = False

        self.lhs = lhs
        self.rhs = rhs
        self.assignment_node = assignment_node

        self.entry = entry
        self.pos = getpos(lhs)
        self.refs = set()
        self.is_arg = False
        self.is_deletion = False

        # NOTE: this is imperfect, since it means warnings are disabled for
        # *all* definitions in the function...
        self.entry.warn_unused = warn_unused

    def __repr__(self):
        return '%s(entry=%r)' % (self.__class__.__name__, self.entry)

    def infer_type(self, scope):
        return self.rhs.infer_type(scope)

    def type_dependencies(self, scope):
        return self.rhs.type_dependencies(scope)

class Argument(NameAssignment):
    def __init__(self, lhs, rhs, entry):
        NameAssignment.__init__(self, lhs, rhs, entry)
        self.is_arg = True


class PhiNode(nodes.Node):

    def __init__(self, block, variable):
        self.block = block
        # Unrenamed variable. This will be replaced by the renamed version
        self.variable = variable
        self.type = None
        # self.incoming_blocks = []
        self.incoming = set()
        self.phis = set()

        self.assignment_node = self

    @property
    def entry(self):
        return self.variable

    def add_incoming_block(self, block):
        self.incoming_blocks.append(block)

    def add(self, block, assmnt):
        if assmnt is not self:
            self.phis.add((block, assmnt))

    def __repr__(self):
        lhs = self.variable.name
        if self.variable.renamed_name:
            lhs = self.variable.unmangled_name
        incoming = ", ".join("var(%s, %s)" % (var_in.unmangled_name, var_in.type)
                                 for var_in in self.incoming)
        if self.variable.type:
            type = str(self.variable.type)
        else:
            type = ""
        return "%s %s = phi(%s)" % (type, lhs, incoming)

    def find_incoming(self):
        for parent_block in self.block.parents:
            name = self.variable.name
            incoming_var = parent_block.symtab.lookup_most_recent(name)
            yield parent_block, incoming_var


class NameDeletion(NameAssignment):
    def __init__(self, lhs, entry):
        NameAssignment.__init__(self, lhs, lhs, entry)
        self.is_deletion = True

    def infer_type(self, scope):
        inferred_type = self.rhs.infer_type(scope)
        if (not inferred_type.is_pyobject and
            inferred_type.can_coerce_to_pyobject(scope)):
            return py_object_type
        return inferred_type


class Uninitialized(object):
    pass

def getpos(node):
    if isinstance(node, NameAssignment):
        return node.pos
    return node.lineno, node.col_offset

class NameReference(object):
    def __init__(self, node, entry):
        if not hasattr(node, 'cf_state'):
            node.cf_state = set()
        self.node = node
        self.entry = entry
        self.pos = getpos(node)

    def __repr__(self):
        return '%s(entry=%r)' % (self.__class__.__name__, self.entry)


class ControlFlowState(list):
    # Keeps track of Node's entry assignments
    #
    # cf_is_null        [boolean] It is uninitialized
    # cf_maybe_null     [boolean] May be uninitialized
    # is_single         [boolean] Has only one assignment at this point

    cf_maybe_null = False
    cf_is_null = False
    is_single = False

    def __init__(self, state):
        if Uninitialized in state:
            state.discard(Uninitialized)
            self.cf_maybe_null = True
            if not state:
                self.cf_is_null = True
        else:
            if len(state) == 1:
                self.is_single = True
        super(ControlFlowState, self).__init__(state)

    def one(self):
        return self[0]


class GVContext(object):
    """Graphviz subgraph object."""

    def __init__(self):
        self.blockids = {}
        self.nextid = 0
        self.children = []
        self.sources = {}

    def add(self, child):
        self.children.append(child)

    def nodeid(self, block):
        if block not in self.blockids:
            self.blockids[block] = 'block%d' % self.nextid
            self.nextid += 1
        return self.blockids[block]

    def extract_sources(self, block):
        if not block.positions:
            return ''
        start = min(block.positions)
        stop = max(block.positions)
        srcdescr = start[0]
        if not srcdescr in self.sources:
            self.sources[srcdescr] = list(srcdescr.get_lines())
        lines = self.sources[srcdescr]

        src_descr, begin_line, begin_col = start
        src_descr, end_line, end_col = stop
        lines = lines[begin_line - 1:end_line]
        if not lines:
            return ''
        #lines[0] = lines[0][begin_col:]
        #lines[-1] = lines[-1][:end_col]
        return '\\n'.join([line.strip() for line in lines if line.strip()])

    def render(self, fp, name, annotate_defs=False):
        """Render graphviz dot graph"""
        fp.write('digraph %s {\n' % name)
        fp.write(' node [shape=box];\n')
        for child in self.children:
            child.render(fp, self, annotate_defs)
        fp.write('}\n')

    def escape(self, text):
        return text.replace('"', '\\"').replace('\n', '\\n')


class GV(object):
    """
    Graphviz DOT renderer.
    """

    def __init__(self, name, flow):
        self.name = name
        self.flow = flow

    def format_phis(self, block):
        result = "\\l".join(str(phi) for var, phi in block.phis.iteritems())
        return result

    def render(self, fp, ctx, annotate_defs=False):
        fp.write(' subgraph %s {\n' % self.name)
        for block in self.flow.blocks:
            if block.have_code:
                code = ctx.extract_sources(block)
                if annotate_defs:
                    for stat in block.stats:
                        if isinstance(stat, NameAssignment):
                            code += '\n %s [definition]' % stat.entry.name
                        elif isinstance(stat, NameReference):
                            if stat.entry:
                                code += '\n %s [reference]' % stat.entry.name
            else:
                code = ""

            if block.have_code and block.label == 'empty':
                label = ''
            else:
                label = '%s: ' % block.label

            phis = self.format_phis(block)
            label = '%d\\l%s%s\\n%s' % (block.id, label, phis, code)

            pid = ctx.nodeid(block)
            fp.write('  %s [label="%s"];\n' % (pid, ctx.escape(label)))
        for block in self.flow.blocks:
            pid = ctx.nodeid(block)
            for child in block.children:
                fp.write('  %s -> %s;\n' % (pid, ctx.nodeid(child)))
        fp.write(' }\n')

class SourceDescr(object):
    def __init__(self, func, ast):
        self.func = func
        self.ast = ast

    def get_lines(self):
        if self.func:
            source = inspect.getsource(self.func)
        else:
            try:
                source = asttools.dump_python_source(self.ast)
            except Exception:
                source = ""

        source = "\n" * (self.ast.lineno - 2) + source
        return source.splitlines()

class MessageCollection:
    """Collect error/warnings messages first then sort"""

    def __init__(self):
        self.messages = []

    def error(self, node, message):
        self.messages.append((getpos(node), node, True, message))

    def warning(self, node, message):
        self.messages.append((getpos(node), node, False, message))

    def report(self):
        self.messages.sort()
        errors = []
        for pos, node, is_error, message in self.messages:
            if is_error:
                errors.append((node, message))
            warning(node, message)

        if errors:
            raise error.NumbaError(*errors[0])

def warning(node, message):
    # printing allows us to test the code
    print "Warning %s%s" % (error.format_pos(node), message)
    # logger.warning("Warning %s: %s", error.format_postup(getpos(node)), message)

def warn_unreachable(node):
    if hasattr(node, 'lineno'):
        print "Warning, unreachable code at %s" % error.format_pos(node).rstrip(': ')

def allow_null(node):
    return False

def check_definitions(flow, compiler_directives):
    flow.initialize()
    flow.reaching_definitions()

    # Track down state
    assignments = set()
    # Node to entry map
    references = {}
    assmt_nodes = set()

    for block in flow.blocks:
        i_state = block.i_input
        for stat in block.stats:
            i_assmts = flow.assmts[stat.entry]
            state = flow.map_one(i_state, stat.entry)
            if isinstance(stat, NameAssignment):
                stat.lhs.cf_state.update(state)
                assmt_nodes.add(stat.lhs)
                i_state = i_state & ~i_assmts.mask
                if stat.is_deletion:
                    i_state |= i_assmts.bit
                else:
                    i_state |= stat.bit
                assignments.add(stat)
                # if stat.rhs is not fake_rhs_expr:
                stat.entry.cf_assignments.append(stat)
            elif isinstance(stat, NameReference):
                references[stat.node] = stat.entry
                stat.entry.cf_references.append(stat)
                stat.node.cf_state.update(state)
                if not allow_null(stat.node):
                    i_state &= ~i_assmts.bit
                state.discard(Uninitialized)
                for assmt in state:
                    assmt.refs.add(stat)

    # Check variable usage
    warn_maybe_uninitialized = compiler_directives['warn.maybe_uninitialized']
    warn_unused_result = compiler_directives['warn.unused_result']
    warn_unused = compiler_directives['warn.unused']
    warn_unused_arg = compiler_directives['warn.unused_arg']

    messages = MessageCollection()

    # assignment hints
    for node in assmt_nodes:
        if Uninitialized in node.cf_state:
            node.cf_maybe_null = True
            if len(node.cf_state) == 1:
                # node.cf_state == set([Uninitialized])
                node.cf_is_null = True
            else:
                node.cf_is_null = False
        else:
            node.cf_is_null = False
            node.cf_maybe_null = False

    # Find uninitialized references and cf-hints
    for node, entry in references.iteritems():
        if Uninitialized in node.cf_state:
            node.cf_maybe_null = True
            from_closure = False # entry.from_closure
            if not from_closure and len(node.cf_state) == 1:
                node.cf_is_null = True
            if allow_null(node) or from_closure: # or entry.is_pyclass_attr:
                pass # Can be uninitialized here
            elif node.cf_is_null:
                is_object = True #entry.type.is_pyobject
                is_unspecified = False #entry.type.is_unspecified
                error_on_uninitialized = False #entry.error_on_uninitialized
                if entry.renameable and (is_object or is_unspecified or
                                         error_on_uninitialized):
                    messages.error(
                        node,
                        "local variable '%s' referenced before assignment"
                        % entry.name)
                else:
                    messages.warning(
                        node,
                        "local variable '%s' referenced before assignment"
                        % entry.name)
            elif warn_maybe_uninitialized:
                messages.warning(
                    node,
                    "local variable '%s' might be referenced before assignment"
                    % entry.name)
        else:
            node.cf_is_null = False
            node.cf_maybe_null = False

    # Unused result
    for assmt in assignments:
        if not assmt.refs: # and not assmt.entry.is_pyclass_attr
                           # and not assmt.entry.in_closure):
            if assmt.entry.cf_references and warn_unused_result:
                if assmt.is_arg:
                    messages.warning(assmt, "Unused argument value '%s'" %
                                     assmt.entry.name)
                else:
                    messages.warning(assmt, "Unused result in '%s'" %
                                     assmt.entry.name)
            assmt.lhs.cf_used = False

    # Unused entries
    for entry in flow.entries:
        if (not entry.cf_references and not entry.is_cellvar and
                entry.renameable): # and not entry.is_pyclass_attr
            if entry.is_arg:
                if warn_unused_arg:
                    messages.warning(entry, "Unused argument '%s'" %
                                     entry.name)
            else:
                if warn_unused and entry.warn_unused and flow.is_tracked(entry):
                    messages.warning(entry, "Unused variable '%s'" %
                                     entry.name)
            entry.cf_used = False

    messages.report()

    for node in assmt_nodes:
        node.cf_state = None #ControlFlowState(node.cf_state)
    for node in references:
        node.cf_state = None #ControlFlowState(node.cf_state)


class FuncDefExprNode(nodes.Node):
    """
    Wraps an inner function node until the closure code kicks in.
    """

    _fields = ['func_def']

class ControlFlowAnalysis(visitors.NumbaTransformer):
    """
    Control flow analysis pass that builds the CFG and injects the blocks
    into the AST (but not all blocks are injected).

    The CFG must be build in topological DFS order, e.g. the 'if' condition
    block must precede the clauses and the clauses must precede the exit.
    """

    graphviz = False
    gv_ctx = None
    source_descr = None

    function_level = 0

    def __init__(self, context, func, ast, allow_rebind_args, **kwargs):
        super(ControlFlowAnalysis, self).__init__(context, func, ast, **kwargs)
        self.visitchildren = self.generic_visit
        self.current_directives = kwargs.get('directives', None) or {}
        self.current_directives['warn'] = kwargs.get('warn', True)
        self.set_default_directives()
        self.symtab = self.initialize_symtab(allow_rebind_args)

        self.graphviz = self.current_directives['control_flow.dot_output']
        if self.graphviz:
            self.gv_ctx = GVContext()
            self.source_descr = SourceDescr(func, ast)

        # Stack of control flow blocks
        self.stack = []
        self.flow = ControlFlow(self.source_descr)

    def set_default_directives(self):
        "Set some defaults for warnings"
        warn = self.current_directives['warn']
        self.current_directives.setdefault('warn.maybe_uninitialized', warn)
        self.current_directives.setdefault('warn.unused_result', False)
        self.current_directives.setdefault('warn.unused', warn)
        self.current_directives.setdefault('warn.unused_arg', warn)
        self.current_directives.setdefault('control_flow.dot_output', dot_output_graph)
        self.current_directives.setdefault('control_flow.dot_annotate_defs', False)

    def initialize_symtab(self, allow_rebind_args):
        """
        Populate the symbol table with variables and set their renaming status.

        Variables appearing in locals, or arguments typed through the 'jit'
        decorator are not renameable.
        """
        symbols = symtab.Symtab(self.symtab)
        for var_name in self.local_names:
            variable = symtab.Variable(None, name=var_name, is_local=True)

            # Set cellvar status. Free variables are not assignments, and
            # are caught in the type inferencer
            variable.is_cellvar = var_name in self.cellvars
            # variable.is_freevar = var_name in self.freevars

            variable.renameable = (
                var_name not in self.locals and not
                (variable.is_cellvar or variable.is_freevar) and
                (var_name not in self.argnames or allow_rebind_args))

            symbols[var_name] = variable

        return symbols

    def visit(self, node):
        if hasattr(node, 'lineno'):
            self.mark_position(node)
        if not self.flow.block:
            # Unreachable code
            # NOTE: removing this here means there is no validation of the
            # unreachable code!
            warn_unreachable(node)
            return None

        return super(ControlFlowAnalysis, self).visit(node)

    def handle_inner_function(self, node):
        "Create assignment code for inner functions and mark the assignment"
        lhs = ast.Name(node.name, ast.Store())
        ast.copy_location(lhs, node)

        rhs = FuncDefExprNode(func_def=node)
        ast.copy_location(rhs, node)

        fields = rhs._fields
        rhs._fields = []
        assmnt = ast.Assign(targets=[lhs], value=rhs)
        result = self.visit(assmnt)
        rhs._fields = fields

        return result

    def visit_FunctionDef(self, node):
        #for arg in node.args:
        #    if arg.default:
        #        self.visitchildren(arg)
        if self.function_level:
            return self.handle_inner_function(node)

        self.function_level += 1

        self.visitlist(node.decorator_list)
        self.stack.append(self.flow)

        # Collect all entries
        for var_name, var in self.symtab.iteritems():
            if var_name not in self.locals:
                self.flow.entries.add(var)

        self.flow.nextblock(label='entry')
        self.mark_position(node)

        # Function body block
        node.body_block = self.flow.nextblock()
        for arg in node.args.args:
            self.visit_Name(arg)

        self.visitlist(node.body)
        self.function_level -= 1

        # Exit point
        self.flow.add_exit(self.flow.exit_point)
        if self.flow.block:
            self.flow.block.add_child(self.flow.exit_point)

        # Cleanup graph
        # self.flow.normalize()
        check_definitions(self.flow, self.current_directives)

        # self.render_gv(node)

        self.flow.compute_dominators()
        self.flow.compute_dominance_frontier()
        self.flow.update_for_ssa(self.ast, self.symtab)

        return node

    def render_gv(self, node):
        render_gv(node, self.gv_ctx, self.flow, self.current_directives)

    def mark_assignment(self, lhs, rhs=None, assignment=None, warn_unused=True):
        assert self.flow.block

        if self.flow.exceptions:
            exc_descr = self.flow.exceptions[-1]
            self.flow.block.add_child(exc_descr.entry_point)
            self.flow.nextblock()

        if not rhs:
            rhs = None

        lhs = self.visit(lhs)
        name_assignment = None
        if isinstance(lhs, ast.Name):
            name_assignment = self.flow.mark_assignment(
                    lhs, rhs, self.symtab[lhs.name], assignment,
                    warn_unused=warn_unused)

        if self.flow.exceptions:
            exc_descr = self.flow.exceptions[-1]
            self.flow.block.add_child(exc_descr.entry_point)
            self.flow.nextblock()

        return lhs, name_assignment

    def mark_position(self, node):
        """Mark position if DOT output is enabled."""
        if self.current_directives['control_flow.dot_output']:
            self.flow.mark_position(node)

    def visit_Assign(self, node):
        node.value = self.visit(node.value)
        if len(node.targets) == 1 and isinstance(node.targets[0],
                                                 (ast.Tuple, ast.List)):
            node.targets = node.targets[0].elts

        for i, target in enumerate(node.targets):
            # target = self.visit(target)

            maybe_unused_node = isinstance(target, nodes.MaybeUnusedNode)
            if maybe_unused_node:
                target = target.name_node

            lhs, name_assignment = self.mark_assignment(target, node.value,
                                                        assignment=node,
                                                        warn_unused=not maybe_unused_node)
            node.targets[i] = lhs

        return node

    def visit_AugAssign(self, node):
        """
        Inplace assignment.

        Resolve a += b to a = a + b. Set 'inplace_op' attribute of the
        Assign node so later stages may recognize inplace assignment.

        Do this now, so that we can correctly mark the RHS reference.
        """
        target = node.target

        rhs_target = copy.deepcopy(target)
        rhs_target.ctx = ast.Load()
        ast.fix_missing_locations(rhs_target)

        bin_op = ast.BinOp(rhs_target, node.op, node.value)
        assignment = ast.Assign([target], bin_op)
        assignment.inplace_op = node.op
        return self.visit(assignment)

    def visit_Name(self, old_node):
        node = nodes.Name(old_node.id, old_node.ctx)
        ast.copy_location(node, old_node)

        # Set some defaults
        node.cf_maybe_null = True
        node.cf_is_null = False
        node.allow_null = False

        node.name = node.id
        if isinstance(node.ctx, ast.Param):
            var = self.symtab[node.name]
            var.is_arg = True
            self.flow.mark_assignment(node, None, var, assignment=None)
        elif isinstance(node.ctx, ast.Load):
            var = self.symtab.lookup(node.name)
            if var:
                # Local variable
                self.flow.mark_reference(node, var)

        # Set position of assignment of this definition
        if isinstance(node.ctx, (ast.Param, ast.Store)):
            var = self.symtab[node.name]
            if var.lineno == -1:
                var.lineno = getattr(node, "lineno", 0)
                var.col_offset = getattr(node, "col_offset", 0)

        return node

    def visit_MaybeUnusedNode(self, node):
        self.symtab[node.name_node.id].warn_unused = False
        return self.visit(node.name_node)

    def visit_Suite(self, node):
        if self.flow.block:
            for stat in node.body:
                self.visit(stat)
                if not self.flow.block:
                    stat.is_terminator = True
                    break

        return node

    def visit_ImportFrom(self, node):
        for name, target in node.names:
            if name != "*":
                self.mark_assignment(target, assignment=node)

        self.visitchildren(node)
        return node

    def exit_block(self, exit_block, node):
        node.exit_block = exit_block
        self.flow.add_exit(exit_block)
        if exit_block.parents:
            self.flow.block = exit_block
        else:
            self.flow.block = None

        return node

    def visit_If(self, node):
        exit_block = self.flow.exit_block(label='exit_if', pos=node)

        # Condition
        cond_block = self.flow.nextblock(self.flow.block, label='if_cond',
                                         is_expr=True, pos=node.test)
        node.test = self.visit(node.test)

        # Body
        if_block = self.flow.nextblock(label='if_body', pos=node.body[0])
        self.visitlist(node.body)
        if self.flow.block:
            self.flow.block.add_child(exit_block)

        # Else clause
        if node.orelse:
            else_block = self.flow.nextblock(cond_block,
                                             label='else_body',
                                             pos=node.orelse[0])
            self.visitlist(node.orelse)
            if self.flow.block:
                self.flow.block.add_child(exit_block)
        else:
            cond_block.add_child(exit_block)
            else_block = None

        node = nodes.build_if(cond_block=cond_block, test=node.test,
                              if_block=if_block, body=node.body,
                              else_block=else_block, orelse=node.orelse,
                              exit_block=exit_block)
        return self.exit_block(exit_block, node)

    def _visit_loop_body(self, node, if_block=None, is_for=None):
        """
        Visit body of while and for loops and handle 'else' clause
        """
        loop_name = "for" if is_for else "while"
        if if_block:
           node.if_block = if_block
        else:
            node.if_block = self.flow.nextblock(label="%s_body" % loop_name,
                                                pos=node.body[0])
        self.visitlist(node.body)
        self.flow.loops.pop()

        if is_for:
            if self.flow.block:
                self.flow.block.add_child(node.incr_block)
            # Ensure topological dominator order
            self.flow.blocks.pop(node.incr_block.id)
            self.flow.blocks.append(node.incr_block)
            self.flow.block = node.incr_block

        if self.flow.block:
            # Add back-edge
            self.flow.block.add_child(node.cond_block)

        # Else clause
        if node.orelse:
            node.else_block = self.flow.nextblock(
                        parent=node.cond_block,
                        label="else_clause_%s" % loop_name,
                        pos=node.orelse[0])
            self.visitlist(node.orelse)
            if self.flow.block:
                self.flow.block.add_child(node.exit_block)
        else:
            node.cond_block.add_child(node.exit_block)

        self.exit_block(node.exit_block, node)

    def visit_While(self, node):
        node.cond_block = self.flow.nextblock(label='while_condition',
                                              pos=node.test)
        node.exit_block = self.flow.exit_block(label='exit_while', pos=node)

        # Condition block
        self.flow.loops.append(LoopDescr(node.exit_block, node.cond_block))
        node.test = self.visit(node.test)

        self._visit_loop_body(node)
        return nodes.build_while(**vars(node))

    def visit_For(self, node):
        # Evaluate iterator in previous block
        node.iter = self.visit(node.iter)

        # Start condition block
        node.cond_block = self.flow.nextblock(label='for_condition',
                                              pos=node.iter)
        node.exit_block = self.flow.exit_block(label='exit_for', pos=node)

        # Increment temporary variable, continue should branch here
        node.incr_block = self.flow.newblock(label="for_increment", pos=node)
        node.incr_block.branch_here = True

        self.flow.loops.append(LoopDescr(node.exit_block, node.incr_block))

        # Target assignment
        if_block = self.flow.nextblock(label='loop_body', pos=node.body[0])
        #node.target_block = self.flow.nextblock(label='for_target',
        #                                        pos=node.target)
        node.target, name_assignment = self.mark_assignment(
                    node.target, assignment=None, warn_unused=False)
        self._visit_loop_body(node, if_block=if_block, is_for=True)
        node = nodes.For(**vars(node))
        if name_assignment:
            name_assignment.assignment_node = node
        return node

    def visit_ListComp(self, node):
        """
        Rewrite list comprehensions to the equivalent for loops.

        AST syntax:

            ListComp(expr elt, comprehension* generators)
            comprehension = (expr target, expr iter, expr* ifs)

            'ifs' represent a chain of ANDs
        """
        assert len(node.generators) > 0

        # Create innermost body, i.e. list.append(expr)
        # TODO: size hint for PyList_New
        list_create = ast.List(elts=[], ctx=ast.Load())
        list_create.type = object_ # typesystem.ListType()
        list_create = nodes.CloneableNode(list_create)
        list_value = nodes.CloneNode(list_create)
        list_append = ast.Attribute(list_value, "append", ast.Load())
        append_call = ast.Call(func=list_append, args=[node.elt],
                               keywords=[], starargs=None, kwargs=None)

        # Build up the loops from inwards to outwards
        body = append_call
        for comprehension in reversed(node.generators):
            # Hanlde the 'if' clause
            ifs = comprehension.ifs
            if len(ifs) > 1:
                make_boolop = lambda op1, op2: ast.BoolOp(op=ast.And(),
                                                          values=[op1, op2])
                if_test = reduce(make_boolop, ifs)
            elif len(ifs) == 1:
                if_test, = ifs
            else:
                if_test = None

            if if_test is not None:
                body = ast.If(test=if_test, body=[body], orelse=[])

            # Wrap list.append() call or inner loops
            body = ast.For(target=comprehension.target,
                           iter=comprehension.iter, body=[body], orelse=[])

        expr = nodes.ExpressionNode(stmts=[list_create, body], expr=list_value)
        return self.visit(expr)

    def visit_GeneratorExp(self, node):
        raise error.NumbaError(
                node, "Generator comprehensions are not yet supported")

    def visit_SetComp(self, node):
        raise error.NumbaError(
                node, "Set comprehensions are not yet supported")

    def visit_DictComp(self, node):
        raise error.NumbaError(
                node, "Dict comprehensions are not yet supported")

    def visit_With(self, node):
        node.context_expr = self.visit(node.context_expr)
        if node.optional_vars:
            # TODO: Mark these as assignments!
            node.optional_vars = self.visit(node.optional_vars)

        self.visitlist(node.body)
        return node

    def visit_Raise(self, node):
        self.visitchildren(node)
        if self.flow.exceptions:
            self.flow.block.add_child(self.flow.exceptions[-1].entry_point)

        self.flow.block = None
        return node

    def visit_Return(self, node):
        self.visitchildren(node)

        for exception in self.flow.exceptions[::-1]:
            if exception.finally_enter:
                self.flow.block.add_child(exception.finally_enter)
                if exception.finally_exit:
                    exception.finally_exit.add_child(self.flow.exit_point)
                break
        else:
            if self.flow.block:
                self.flow.block.add_child(self.flow.exit_point)

        self.flow.block = None
        return node

    def visit_Break(self, node):
        if not self.flow.loops:
            #error(node.pos, "break statement not inside loop")
            return node

        loop = self.flow.loops[-1]
        for exception in loop.exceptions[::-1]:
            if exception.finally_enter:
                self.flow.block.add_child(exception.finally_enter)
                if exception.finally_exit:
                    exception.finally_exit.add_child(loop.next_block)
                break
        else:
            self.flow.block.add_child(loop.next_block)

        #self.flow.nextblock(parent=loop.next_block)
        self.flow.block = None
        return node

    def visit_Continue(self, node):
        if not self.flow.loops:
            #error(node.pos, "continue statement not inside loop")
            return node

        loop = self.flow.loops[-1]
        for exception in loop.exceptions[::-1]:
            if exception.finally_enter:
                self.flow.block.add_child(exception.finally_enter)
                if exception.finally_exit:
                    exception.finally_exit.add_child(loop.loop_block)
                break
        else:
            self.flow.block.add_child(loop.loop_block)

        self.flow.block = None
        return node

    def visit_Print(self, node):
        self.generic_visit(node)
        return node

class DeleteStatement(visitors.NumbaVisitor):
    """
    Delete a (compound) statement that contains basic blocks.
    The statement must be at the start of the entry block.

    idom: the immediate dominator of
    """

    def __init__(self, flow):
        self.flow = flow

    def visit_If(self, node):
        self.generic_visit(node)

        # Visit ControlBlocks
        self.visit(node.cond_block)
        self.visit(node.if_block)
        if node.orelse:
            self.visit(node.else_block)
        if node.exit_block:
            self.visit(node.exit_block)

    visit_While = visit_If

    def visit_For(self, node):
        self.generic_visit(node)

        # Visit ControlBlocks
        self.visit(node.cond_block)
        self.visit(node.if_block)
        self.visit(node.incr_block)
        if node.orelse:
            self.visit(node.else_block)
        if node.exit_block:
            self.visit(node.exit_block)


    def visit_ControlBlock(self, node):
        #print "deleting block", node
        for phi in node.phi_nodes:
            for incoming in phi.incoming:
                #print "deleting", incoming, phi
                incoming.cf_references.remove(phi)

        self.generic_visit(node)
        node.delete(self.flow)

    def visit_Name(self, node):
        references = node.variable.cf_references
        if isinstance(node.ctx, ast.Load) and node in references:
            references.remove(node)


#----------------------------------------------------------------------------
# Graphviz Rendering
#----------------------------------------------------------------------------

def get_png_output_name(dot_output):
    prefix, ext = os.path.splitext(dot_output)
    i = 0
    while True:
        png_output = "%s%d.png" % (prefix, i)
        if not os.path.exists(png_output):
            break

        i += 1

    return png_output


def write_dotfile(current_directives, dot_output, gv_ctx):
    annotate_defs = current_directives['control_flow.dot_annotate_defs']
    fp = open(dot_output, 'wt')
    try:
        gv_ctx.render(fp, 'module', annotate_defs=annotate_defs)
    finally:
        fp.close()


def write_image(dot_output):
    png_output = get_png_output_name(dot_output)
    fp = open(png_output, 'wb')
    try:
        p = subprocess.Popen(['dot', '-Tpng', dot_output],
                             stdout=fp.fileno(),
                             stderr=subprocess.PIPE)
        p.wait()
    except EnvironmentError, e:
        logger.warn("Unable to write png: %s. Wrote %s" % (e, dot_output))
    else:
        logger.warn("Wrote %s" % png_output)
    finally:
        fp.close()


def render_gv(node, gv_ctx, flow, current_directives):
    gv_ctx.add(GV(node.name, flow))
    dot_output = current_directives['control_flow.dot_output']
    if dot_output:
        write_dotfile(current_directives, dot_output, gv_ctx)
        write_image(dot_output)
