import sys

from stringbuffer import StringBuffer

import llvmcpy.llvm as llvm

def isSystemCall(name):
    """returns a tuple consisting of the boolean and the cleaned up name.
    """
    if name and ord(name[0]) == 1:
        return (True, name[1:])
    else:
        return (False, name)

#I think there ia a LLVM C-API call for this. Maybe also the system call thingy.
def isIntrinsic(name):
    return name.startswith('llvm.')



class Node(object):

    def __init__(self, name, nid, prototype=None):
        self.name = name
        self.nid = nid
        self.prototype = prototype
        self.attributes = {}

    def set_attribute(self, key, val):
        self.attributes[key] = val

    def get_attribute(self, key, default=None):
        self.attributes.get(key, default)

    def toDotString(self, attributes = False, fill = None):
        cr = r'\n'
        sb = StringBuffer()
        has_attribute = False
        highlight = 'yellow'
        sb.append(self.nid).append(' [label="').append(self.name)
        if attributes:
            if isinstance(attributes, list):
                for key in self.attributes:
                    if key in attributes:
                        has_attribute = True
                        sb.append(cr).append(key).append(' = ').append(self.attributes[key])
            else:
                for key in self.attributes:
                    #has_attribute = True
                    sb.append(cr).append(key).append(' = ').append(self.attributes[key])
        sb.append('"')
        if fill:
            sb.append(', color=').append(fill).append(',style=filled')
        elif has_attribute:
            sb.append(', color=').append(highlight).append(',style=filled')

        sb.append('];')
        return str(sb)




class CallGraph(object):

    def __init__(self, gname):
        self.id_counter = 0
        self.gname = gname
        # the maps node_ids to nodes
        self.nid_to_node = {}
        # maps the name to nodes
        self.name_to_node = {}
        # the set of edges
        self.edges = set()
        # maps nodes to the set of edges that they participate in.
        self.nids_to_edges = {}
        # set of externals
        self.declarations = set()
        # what to do here?
        self.calls = 0
        self.indirect_calls = 0
        self.bitcasts = 0

    @staticmethod
    def fromModule(gname, module, skip_system_calls = False):
        callgraph = CallGraph(gname)


        def resolve_callee(call_inst):
            """Tries to resolve the function called in the call instruction, returns None if it can't."""
            callee = None
            called = instruction.get_called()
            fname = called.get_name()
            if fname:
                callee = fname
            elif called.is_a_load_inst():
                #print("No callee name: LOAD")
                callgraph.indirect_calls += 1
            elif called.is_a_bit_cast_inst():
                #print("No callee name: BITCAST")
                #print("called: ", called.print_value_to_string())
                callgraph.bitcasts += 1
            elif called.is_a_constant():
                opcode = called.get_const_opcode()
                assert(opcode == llvm.Opcode['BitCast'])
                nops = called.get_num_operands()
                assert(nops == 1)
                callee = called.get_operand(0).get_name()
            else:
                print("No callee name: unhandled case")
                asssert(False)
                #print("Called: ", called)
                #print("Called: ", called.print_value_to_string())
                #print("Instr: ", instruction.print_value_to_string())
                pass

            return callee



        for function in module.iter_functions():

            name = function.get_name()

            (isSysCall, name) = isSystemCall(name)

            node = callgraph.addNode(name, None, function.type_of().print_type_to_string())
            node_id = node.nid

            if function.is_declaration():
                callgraph.declarations.add(node.nid)
                continue

            basic_blocks = 0
            instructions = 0


            callees = set()
            for bb in function.iter_basic_blocks():
                basic_blocks += 1
                for instruction in bb.iter_instructions():
                    instructions += 1
                    if instruction.is_a_call_inst():
                        callgraph.calls += 1
                        #print(instruction.print_value_to_string())
                        #print(instruction.get_called().get_name())
                        callee = resolve_callee(instruction)
                        (isSysCall, callee) = isSystemCall(callee)
                        #ignore instrinsics and ...
                        if callee and not isIntrinsic(callee):
                            if not (skip_system_calls and isSysCall):
                                callees.add(callee)

            for callee in callees:
                callgraph.addEdge(node_id, callee)

            node.set_attribute('basic_blocks', basic_blocks)
            node.set_attribute('instructions', instructions)

        return callgraph


    def __str__(self):
        return self.graphInfo()


    def graphInfo(self, nodes = None):
        nids = set(self.nid_to_node.keys()) if nodes is None else self.toNidSet(nodes)
        eset = set()
        for e in self.edges:
            if (e[0] in nids) and (e[1] in nids):
                eset.add(e)
        statistics = (self.gname, len(nids), len(eset), self.calls, self.indirect_calls + self.bitcasts)
        return "CallGraph {0} with {1} nodes and {2} edges.\t({3} calls, {4} unresolved calls)\n".format(*statistics)


    def addNode(self, name, id=None, prototype=None):
        """creates and adds the node to graph (if necessary), and returns it.

        if you happen to know its id you can set that too (useful for making subgraphs).
        """
        node = None
        if name not in self.name_to_node:
            nid = self.id_counter if id is None else id
            node = Node(name, nid, prototype)
            self.name_to_node[name] = node
            self.nids_to_edges[nid] = set()
            self.nid_to_node[nid] = node
            self.id_counter += 1
        else:
            node = self.name_to_node[name]
        return node

    def getNodes(self, nids=None):

        if nids is None:
            return set(self.name_to_node.keys())
        else:
            retval = set()
            for nid in nids:
                retval.add(self.nid_to_node[nid].name)
            return retval

    def getNIDs(self):
        return set(self.nid_to_node.keys())

    def addEdge(self, source, target):
        """adds the edge from the (id of) source to the (id of) target.

        if the source or target is an integer it is treated as the id,
        if it is a string it is presumed to be the name of a node.
        """
        src_id = source if isinstance(source, int) else self.addNode(source).nid
        tgt_id = target if isinstance(target, int) else self.addNode(target).nid
        edge = (src_id, tgt_id)
        self.edges.add(edge)

        self.nids_to_edges[src_id].add(edge)
        self.nids_to_edges[tgt_id].add(edge)


        return edge


    def toDotString(self, nodes = None, attributes = False, highlights = None, fill = None):
        """produces the dot for either the entire graph, or just restricted to the nodes passed in.
        """

        highighted_nids = set() if highlights is None else self.toNidSet(highlights)

        nids = set(self.nid_to_node.keys()) if nodes is None else self.toNidSet(nodes)

        sb = StringBuffer()
        sb.append('digraph ').append(self.gname).append(' {\n')
        #sb.append('\tnode [shape=box];\n')
        for name in self.name_to_node:
            node = self.name_to_node[name]
            node_id = node.nid
            if node_id not in nids:
                continue
            if fill and (node_id in highighted_nids):
                sb.append('\t').append(node.toDotString(attributes, fill)).append('\n')
            else:
                sb.append('\t').append(node.toDotString(attributes)).append('\n')

        for edge in self.edges:
            src_id = edge[0]
            tgt_id = edge[1]
            if src_id not in nids:
                continue
            if tgt_id not in nids:
                continue
            sb.append('\t').append(src_id).append(' -> ').append(tgt_id).append(';\n')

        sb.append('}\n')
        return str(sb)

    def annotate(self, annotation, mapping):
        for name in mapping.keys():
            node = self.name_to_node.get(name, None)
            if node is not None:
                val = mapping[name]
                #print('node: {0} annotation: {1} value: {2}'.format(name, annotation, val))
                node.attributes[annotation] = val
            else:
                sys.stderr.write("skipping {0}\n".format(name))



    def remove_isolated(self, name):
        subgraph = CallGraph(self.name)
        has_an_edge = set()
        for edge in self.edges:
            has_an_edge.add(edge[0])
            has_an_edge.add(edge[1])

        for node, node_id in self.name_to_node.iteritems():
            if node_id in has_an_edge:
                subgraph.addNode(node, node_id)

        for edge in self.edges:
            subgraph.addEdge(*edge)

        return subgraph


    def restrict(self, name, nodes):
        """returns a subgraph with given name whose nodes must belong to the given set.

        """
        subgraph = CallGraph(name)

        node_ids = set()

        for node in nodes:
            if node in self.name_to_node:
                node_id = self.name_to_node[node]
                node_ids.add(subgraph.addNode(node, node_id).nid)

        for edge in self.edges:
            if (edge[0] in node_ids) and (edge[1] in node_ids):
                subgraph.addEdge(edge[0], edge[1])

        return subgraph



    def toNidSet(self, nodes):
        """converts an iterable of nodes (name, node, or nid) to a set of node ids."""
        answer = set()
        for node in nodes:
            if isinstance(node, Node):
                answer.add(node.nid)
            elif isinstance(node, int):
                answer.add(node)
            elif node in self.name_to_node:
                answer.add(self.name_to_node[node].nid)
#            else:
#                print("Skipping {0}".format(node))
        return answer


    def close_up(self, nodes):
        # first convert the input to nids (if necessary).
        answer = self.toNidSet(nodes)
        new = True
        while new:
            before = len(answer)
            for nid in answer.copy():
                candidates = self.nids_to_edges[nid]
                for e in candidates:
                    if e[1] == nid:
                        answer.add(e[0])
            after = len(answer)
            new = after > before
        return answer


    def close_down(self, nodes):
        # first convert the input to nids (if necessary).
        answer = self.toNidSet(nodes)
        new = True
        while new:
            before = len(answer)
            for nid in answer.copy():
                candidates = self.nids_to_edges[nid]
                for e in candidates:
                    if e[0] == nid:
                        answer.add(e[1])
            after = len(answer)
            new = after > before
        return answer


    def dump_prototypes(self, path, nids = None):

        nids = set(self.nid_to_node.keys()) if nids is None else nids

        try:
            with open(path, 'w') as fp:
                for nid in nids:
                    node = self.nid_to_node[nid]
                    if node.prototype is not None:
                        fp.write(node.name)
                        fp.write(" :  ")
                        fp.write(node.prototype)
                        fp.write("\n")
        except Exception as e:
            print(e)
            sys.stderr.write('dump_prototypes("{0}") threw {1}\n'.format(path, str(e)))

import sys

from stringbuffer import StringBuffer

import llvmcpy.llvm as llvm

def isSystemCall(name):
    """returns a tuple consisting of the boolean and the cleaned up name.
    """
    if name and ord(name[0]) == 1:
        return (True, name[1:])
    else:
        return (False, name)

#I think there ia a LLVM C-API call for this. Maybe also the system call thingy.
def isIntrinsic(name):
    return name.startswith('llvm.')



class Node(object):

    def __init__(self, name, nid, prototype=None):
        self.name = name
        self.nid = nid
        self.prototype = prototype
        self.attributes = {}

    def set_attribute(self, key, val):
        self.attributes[key] = val

    def get_attribute(self, key, default=None):
        self.attributes.get(key, default)

    def toDotString(self, attributes = False, fill = None):
        cr = r'\n'
        sb = StringBuffer()
        has_attribute = False
        highlight = 'yellow'
        sb.append(self.nid).append(' [label="').append(self.name)
        if attributes:
            if isinstance(attributes, list):
                for key in self.attributes:
                    if key in attributes:
                        has_attribute = True
                        sb.append(cr).append(key).append(' = ').append(self.attributes[key])
            else:
                for key in self.attributes:
                    #has_attribute = True
                    sb.append(cr).append(key).append(' = ').append(self.attributes[key])
        sb.append('"')
        if fill:
            sb.append(', color=').append(fill).append(',style=filled')
        elif has_attribute:
            sb.append(', color=').append(highlight).append(',style=filled')

        sb.append('];')
        return str(sb)




class CallGraph(object):

    def __init__(self, gname):
        self.id_counter = 0
        self.gname = gname
        # the maps node_ids to nodes
        self.nid_to_node = {}
        # maps the name to nodes
        self.name_to_node = {}
        # the set of edges
        self.edges = set()
        # maps nodes to the set of edges that they participate in.
        self.nids_to_edges = {}
        # set of externals
        self.declarations = set()
        # what to do here?
        self.calls = 0
        self.indirect_calls = 0
        self.bitcasts = 0

    @staticmethod
    def fromModule(gname, module, skip_system_calls = False):
        callgraph = CallGraph(gname)


        def resolve_callee(call_inst):
            """Tries to resolve the function called in the call instruction, returns None if it can't."""
            callee = None
            called = instruction.get_called()
            fname = called.get_name()
            if fname:
                callee = fname
            elif called.is_a_load_inst():
                #print("No callee name: LOAD")
                callgraph.indirect_calls += 1
            elif called.is_a_bit_cast_inst():
                #print("No callee name: BITCAST")
                #print("called: ", called.print_value_to_string())
                callgraph.bitcasts += 1
            elif called.is_a_constant():
                opcode = called.get_const_opcode()
                assert(opcode == llvm.Opcode['BitCast'])
                nops = called.get_num_operands()
                assert(nops == 1)
                callee = called.get_operand(0).get_name()
            else:
                print("No callee name: unhandled case")
                asssert(False)
                #print("Called: ", called)
                #print("Called: ", called.print_value_to_string())
                #print("Instr: ", instruction.print_value_to_string())
                pass

            return callee



        for function in module.iter_functions():

            name = function.get_name()

            (isSysCall, name) = isSystemCall(name)

            node = callgraph.addNode(name, None, function.type_of().print_type_to_string())
            node_id = node.nid

            if function.is_declaration():
                callgraph.declarations.add(node.nid)
                continue

            basic_blocks = 0
            instructions = 0


            callees = set()
            for bb in function.iter_basic_blocks():
                basic_blocks += 1
                for instruction in bb.iter_instructions():
                    instructions += 1
                    if instruction.is_a_call_inst():
                        callgraph.calls += 1
                        #print(instruction.print_value_to_string())
                        #print(instruction.get_called().get_name())
                        callee = resolve_callee(instruction)
                        (isSysCall, callee) = isSystemCall(callee)
                        #ignore instrinsics and ...
                        if callee and not isIntrinsic(callee):
                            if not (skip_system_calls and isSysCall):
                                callees.add(callee)

            for callee in callees:
                callgraph.addEdge(node_id, callee)

            node.set_attribute('basic_blocks', basic_blocks)
            node.set_attribute('instructions', instructions)

        return callgraph


    def __str__(self):
        return self.graphInfo()


    def graphInfo(self, nodes = None):
        nids = set(self.nid_to_node.keys()) if nodes is None else self.toNidSet(nodes)
        eset = set()
        for e in self.edges:
            if (e[0] in nids) and (e[1] in nids):
                eset.add(e)
        statistics = (self.gname, len(nids), len(eset), self.calls, self.indirect_calls + self.bitcasts)
        return "CallGraph {0} with {1} nodes and {2} edges.\t({3} calls, {4} unresolved calls)\n".format(*statistics)


    def addNode(self, name, id=None, prototype=None):
        """creates and adds the node to graph (if necessary), and returns it.

        if you happen to know its id you can set that too (useful for making subgraphs).
        """
        node = None
        if name not in self.name_to_node:
            nid = self.id_counter if id is None else id
            node = Node(name, nid, prototype)
            self.name_to_node[name] = node
            self.nids_to_edges[nid] = set()
            self.nid_to_node[nid] = node
            self.id_counter += 1
        else:
            node = self.name_to_node[name]
        return node

    def getNodes(self, nids=None):

        if nids is None:
            return set(self.name_to_node.keys())
        else:
            retval = set()
            for nid in nids:
                retval.add(self.nid_to_node[nid].name)
            return retval

    def getNIDs(self):
        return set(self.nid_to_node.keys())

    def addEdge(self, source, target):
        """adds the edge from the (id of) source to the (id of) target.

        if the source or target is an integer it is treated as the id,
        if it is a string it is presumed to be the name of a node.
        """
        src_id = source if isinstance(source, int) else self.addNode(source).nid
        tgt_id = target if isinstance(target, int) else self.addNode(target).nid
        edge = (src_id, tgt_id)
        self.edges.add(edge)

        self.nids_to_edges[src_id].add(edge)
        self.nids_to_edges[tgt_id].add(edge)


        return edge


    def toDotString(self, nodes = None, attributes = False, highlights = None, fill = None):
        """produces the dot for either the entire graph, or just restricted to the nodes passed in.
        """

        highighted_nids = set() if highlights is None else self.toNidSet(highlights)

        nids = set(self.nid_to_node.keys()) if nodes is None else self.toNidSet(nodes)

        sb = StringBuffer()
        sb.append('digraph ').append(self.gname).append(' {\n')
        #sb.append('\tnode [shape=box];\n')
        for name in self.name_to_node:
            node = self.name_to_node[name]
            node_id = node.nid
            if node_id not in nids:
                continue
            if fill and (node_id in highighted_nids):
                sb.append('\t').append(node.toDotString(attributes, fill)).append('\n')
            else:
                sb.append('\t').append(node.toDotString(attributes)).append('\n')

        for edge in self.edges:
            src_id = edge[0]
            tgt_id = edge[1]
            if src_id not in nids:
                continue
            if tgt_id not in nids:
                continue
            sb.append('\t').append(src_id).append(' -> ').append(tgt_id).append(';\n')

        sb.append('}\n')
        return str(sb)

    def annotate(self, annotation, mapping):
        for name in mapping.keys():
            node = self.name_to_node.get(name, None)
            if node is not None:
                val = mapping[name]
                #print('node: {0} annotation: {1} value: {2}'.format(name, annotation, val))
                node.attributes[annotation] = val
            else:
                sys.stderr.write("skipping {0}\n".format(name))



    def remove_isolated(self, name):
        subgraph = CallGraph(self.name)
        has_an_edge = set()
        for edge in self.edges:
            has_an_edge.add(edge[0])
            has_an_edge.add(edge[1])

        for node, node_id in self.name_to_node.iteritems():
            if node_id in has_an_edge:
                subgraph.addNode(node, node_id)

        for edge in self.edges:
            subgraph.addEdge(*edge)

        return subgraph


    def restrict(self, name, nodes):
        """returns a subgraph with given name whose nodes must belong to the given set.

        """
        subgraph = CallGraph(name)

        node_ids = set()

        for node in nodes:
            if node in self.name_to_node:
                node_id = self.name_to_node[node]
                node_ids.add(subgraph.addNode(node, node_id).nid)

        for edge in self.edges:
            if (edge[0] in node_ids) and (edge[1] in node_ids):
                subgraph.addEdge(edge[0], edge[1])

        return subgraph



    def toNidSet(self, nodes):
        """converts an iterable of nodes (name, node, or nid) to a set of node ids."""
        answer = set()
        for node in nodes:
            if isinstance(node, Node):
                answer.add(node.nid)
            elif isinstance(node, int):
                answer.add(node)
            elif node in self.name_to_node:
                answer.add(self.name_to_node[node].nid)
#            else:
#                print("Skipping {0}".format(node))
        return answer


    def close_up(self, nodes):
        # first convert the input to nids (if necessary).
        answer = self.toNidSet(nodes)
        new = True
        while new:
            before = len(answer)
            for nid in answer.copy():
                candidates = self.nids_to_edges[nid]
                for e in candidates:
                    if e[1] == nid:
                        answer.add(e[0])
            after = len(answer)
            new = after > before
        return answer


    def close_down(self, nodes):
        # first convert the input to nids (if necessary).
        answer = self.toNidSet(nodes)
        new = True
        while new:
            before = len(answer)
            for nid in answer.copy():
                candidates = self.nids_to_edges[nid]
                for e in candidates:
                    if e[0] == nid:
                        answer.add(e[1])
            after = len(answer)
            new = after > before
        return answer


    def dump_prototypes(self, path, nids = None):

        nids = set(self.nid_to_node.keys()) if nids is None else nids

        try:
            with open(path, 'w') as fp:
                for nid in nids:
                    node = self.nid_to_node[nid]
                    if node.prototype is not None:
                        fp.write(node.name)
                        fp.write(" :  ")
                        fp.write(node.prototype)
                        fp.write("\n")
        except Exception as e:
            print(e)
            sys.stderr.write('dump_prototypes("{0}") threw {1}\n'.format(path, str(e)))

returns a tuple consisting of the boolean and the cleaned up name. #I think there ia a LLVM C-API call for this. Maybe also the system call thingy. #has_attribute = True # the maps node_ids to nodes # maps the name to nodes # the set of edges # maps nodes to the set of edges that they participate in. # set of externals # what to do here? Tries to resolve the function called in the call instruction, returns None if it can't. #print("No callee name: LOAD") #print("No callee name: BITCAST") #print("called: ", called.print_value_to_string()) #print("Called: ", called) #print("Called: ", called.print_value_to_string()) #print("Instr: ", instruction.print_value_to_string()) #print(instruction.print_value_to_string()) #print(instruction.get_called().get_name()) #ignore instrinsics and ... creates and adds the node to graph (if necessary), and returns it.          if you happen to know its id you can set that too (useful for making subgraphs). adds the edge from the (id of) source to the (id of) target.          if the source or target is an integer it is treated as the id,         if it is a string it is presumed to be the name of a node. produces the dot for either the entire graph, or just restricted to the nodes passed in. #sb.append('\tnode [shape=box];\n') #print('node: {0} annotation: {1} value: {2}'.format(name, annotation, val)) returns a subgraph with given name whose nodes must belong to the given set. converts an iterable of nodes (name, node, or nid) to a set of node ids. #            else: #                print("Skipping {0}".format(node)) # first convert the input to nids (if necessary). # first convert the input to nids (if necessary).