from rpython.rlib import jit

# TODO: Find heavily executed lambdas that do not participate in a loop in the
# callgraph.

NOT_RECURSIVE = 0b00
LOOP_HEADER = 0b01
LOOP_PARTICIPANT = 0b11

def join_states(s1, s2):
    return s1 | s2

class Namer(object): #pragma: no cover

    def __init__(self):
        self.names   = {}
        self.counter = 0

    def nameof(self, val):
        name = self.names.get(val, None)
        if name is not None:
            return name
        name = "A_" + str(self.counter)
        self.counter += 1
        self.names[val] = name
        return name

class CallGraph(object):
    def __init__(self):
        self.calls     = {}
        self.recursive = {}

    def register_call(self, lam, calling_app, cont, env):
        if jit.we_are_jitted():
            return
        if not calling_app:
            return
        calling_lam = calling_app.surrounding_lambda
        if not calling_lam:
            return
        subdct = self.calls.get(calling_lam, None)
        if subdct is None:
            self.calls[calling_lam] = subdct = {}
            lam_in_subdct = False
        else:
            lam_in_subdct = lam in subdct
        cont_ast = cont.get_next_executed_ast()
        config = env.pycketconfig()
        is_recursive = NOT_RECURSIVE
        if not lam_in_subdct:
            subdct[lam] = None
            is_recursive = self.is_recursive(calling_lam, lam)
            if is_recursive == LOOP_HEADER:
                calling_lam.enable_jitting()
        # It is possible to have multiple consuming continuations for a given
        # function body. This will attempt to mark them all.
        same_lambda = cont_ast and cont_ast.surrounding_lambda is calling_lam
        if same_lambda:
            # did not call is_recursive yet
            if lam_in_subdct:
                is_recursive = self.is_recursive(calling_lam, lam)
                if is_recursive != NOT_RECURSIVE:
                    cont_ast.set_should_enter()

    def is_recursive(self, lam, starting_from=None):
        # quatratic in theory, hopefully not very bad in practice
        status = self.recursive.get(lam, NOT_RECURSIVE)
        if status == LOOP_HEADER:
            return LOOP_HEADER

        if starting_from is None:
            starting_from = lam
        reachable = self.calls.get(starting_from, None)
        if reachable is None:
            return False
        init = Path(starting_from, None)
        todo = [(key, init) for key in reachable]
        visited = {}
        while todo:
            current, path = todo.pop()
            if current is lam:
                # all the lambdas in the path are recursive too
                while path:
                    status = self.recursive.get(path.node, NOT_RECURSIVE)
                    self.recursive[path.node] = join_states(status, LOOP_PARTICIPANT)
                    path = path.prev
                self.recursive[lam] = LOOP_HEADER
                return LOOP_HEADER
            status = self.recursive.get(current, NOT_RECURSIVE)
            if current in visited or status == LOOP_HEADER:
                continue
            reachable = self.calls.get(current, None)
            if reachable:
                path = Path(current, path)
                for key in reachable:
                    todo.append((key, path))
            visited[current] = None
        return False

    def write_dot_file(self, output): #pragma: no cover
        counter = 0
        output.write("digraph callgraph {\n")
        names = Namer()
        for node in self.calls.iterkeys():
            if node.body[0].should_enter:
                name = names.nameof(node)
                output.write(name)
                output.write(" [fillcolor=red,style=filled];\n")
        for src, subdct in self.calls.iteritems():
            for dst in subdct:
                srcname = names.nameof(src)
                dstname = names.nameof(dst)
                output.write(srcname)
                output.write(" -> ")
                output.write(dstname)
                if dst.body[0].should_enter:
                    output.write(" [color=blue]")
                output.write(";\n")
        output.write("}\n")

class Path(object):

    def __init__(self, node, prev):
        self.node = node
        self.prev = prev

    def __repr__(self):
        return "%s -> %s" % (self.node, self.prev)


@@ -2216,17 +2216,16 @@
 _HEADER%0A
-%0A
         
@@ -2387,21 +2387,29 @@
  return 
-False
+NOT_RECURSIVE
 %0A       
@@ -2683,21 +2683,27 @@
         
-while
+for node in
  path:%0A 
@@ -2749,21 +2749,16 @@
 ive.get(
-path.
 node, NO
@@ -2805,21 +2805,16 @@
 cursive%5B
-path.
 node%5D = 
@@ -2855,45 +2855,8 @@
 NT)%0A
-                    path = path.prev%0A
     
@@ -3019,51 +3019,51 @@
  if 
-current in visited or status == LOOP_HEADER
+status == LOOP_HEADER or current in visited
 :%0A  
@@ -3155,32 +3155,44 @@
     if reachable
+ is not None
 :%0A              
@@ -3354,21 +3354,29 @@
  return 
-False
+NOT_RECURSIVE
 %0A%0A    de
@@ -3738,24 +3738,310 @@
 illed%5D;%5Cn%22)%0A
+            else:%0A                status = self.recursive.get(node, NOT_RECURSIVE)%0A                if status == LOOP_PARTICIPANT:%0A                    name = names.nameof(node)%0A                    output.write(name)%0A                    output.write(%22 %5Bfillcolor=green,style=filled%5D;%5Cn%22)%0A
         for 
@@ -4492,16 +4492,50 @@
 ject):%0A%0A
+    __slots__ = ('node', 'prev')%0A%0A
     def 
@@ -4587,16 +4587,16 @@
  = node%0A
-
         
@@ -4613,16 +4613,130 @@
 = prev%0A%0A
+    def __iter__(self):%0A        while self is not None:%0A            yield self.node%0A            self = self.prev%0A%0A
     def 

Fix for return values and nicer printing of callgraph for graphviz