<commit_before>from collections import defaultdict
import logging

import networkx

import simuvex

from .pathprioritizer import PathPrioritizer
from .errors import AngrAnnotatedCFGError, AngrExitError
from .analyses.cfg_node import CFGNode

l = logging.getLogger("angr.annocfg")

class AnnotatedCFG(object):
    """
    AnnotatedCFG is a control flow graph with statement whitelists and exit whitelists to describe a slice of the
    program.
    """
    def __init__(self, project, cfg=None, detect_loops=False):
        """
        Constructor.

        :param project: The angr Project instance
        :param cfg: Control flow graph. Only used when path prioritizer is used.
        :param detect_loops: Only used when path prioritizer is used.
        """
        self._project = project

        self._cfg = None
        self._target = None

        self._run_statement_whitelist = defaultdict(list)
        self._exit_taken = defaultdict(list)
        self._addr_to_run = {}
        self._addr_to_last_stmt_id = {}
        self._loops = []
        self._path_merge_points = [ ]
        self._path_prioritizer = None


        if cfg is not None:
            self._cfg = cfg

            #if target_irsb_addr is not None:
            #    self._path_prioritizer = PathPrioritizer(self._cfg, self._target)

        if self._cfg is not None:
            for run in self._cfg.nodes():
                self._addr_to_run[self.get_addr(run)] = run

    #
    # Public methods
    #

    def from_digraph(self, digraph):
        """
        Initialize this AnnotatedCFG object with a networkx.DiGraph consisting of the following
        form of nodes:

        Tuples like (block address, statement ID)

        Those nodes are connected by edges indicating the execution flow.

        :param digraph: A networkx.DiGraph object
        """

        for n1 in digraph.nodes_iter():
            addr1, stmt_idx1 = n1
            self.add_statements_to_whitelist(addr1, (stmt_idx1,))

            successors = digraph[n1]
            for n2 in successors:
                addr2, stmt_idx2 = n2

                if addr1 != addr2:
                    # There is a control flow transition from block `addr1` to block `addr2`
                    self.add_exit_to_whitelist(addr1, addr2)

                self.add_statements_to_whitelist(addr2, (stmt_idx2,))

    def get_addr(self, run):
        if isinstance(run, CFGNode):
            return run.addr
        elif type(run) in (int, long):
            return run
        else:
            raise AngrAnnotatedCFGError("Unknown type '%s' of the 'run' argument" % type(run))

    def add_block_to_whitelist(self, block):
        addr = self.get_addr(block)
        self._run_statement_whitelist[addr] = True

    def add_statements_to_whitelist(self, block, stmt_ids):
        addr = self.get_addr(block)
        if type(stmt_ids) is bool:
            if type(self._run_statement_whitelist[addr]) is list and self._run_statement_whitelist[addr]:
                raise Exception("WTF")
            self._run_statement_whitelist[addr] = stmt_ids
        elif -1 in stmt_ids:
            self._run_statement_whitelist[addr] = True
        else:
            self._run_statement_whitelist[addr].extend(stmt_ids)
            self._run_statement_whitelist[addr] = \
                sorted(list(set(self._run_statement_whitelist[addr])))

    def add_exit_to_whitelist(self, run_from, run_to):
        addr_from = self.get_addr(run_from)
        addr_to = self.get_addr(run_to)
        self._exit_taken[addr_from].append(addr_to)

    def set_last_statement(self, block_addr, stmt_id):
        self._addr_to_last_stmt_id[block_addr] = stmt_id

    def add_loop(self, loop_tuple):
        """
        A loop tuple contains a series of IRSB addresses that form a loop. Ideally
        it always starts with the first IRSB that we meet during the execution.
        """
        self._loops.append(loop_tuple)

    def set_path_merge_points(self, points):
        self._path_merge_points = points.copy()

    def should_take_exit(self, addr_from, addr_to):
        if addr_from in self._exit_taken:
            return addr_to in self._exit_taken[addr_from]

        return False

    def should_execute_statement(self, addr, stmt_id):
        if self._run_statement_whitelist is None:
            return True
        elif addr in self._run_statement_whitelist:
            r = self._run_statement_whitelist[addr]
            if isinstance(r, bool):
                return r
            else:
                return stmt_id in self._run_statement_whitelist[addr]
        return False

    def get_run(self, addr):
        if addr in self._addr_to_run:
            return self._addr_to_run[addr]
        return None

    def get_whitelisted_statements(self, addr):
        """
        :returns: True if all statements are whitelisted
        """
        if addr in self._run_statement_whitelist:
            if self._run_statement_whitelist[addr] is True:
                return None # This is the default value used in SimuVEX to say
                            # we execute all statements in this basic block. A
                            # little weird...

            else:
                return self._run_statement_whitelist[addr]

        else:
            return []

    def get_last_statement_index(self, addr):
        if addr in self._addr_to_last_stmt_id:
            return self._addr_to_last_stmt_id[addr]
        elif addr in self._run_statement_whitelist:
            return max(self._run_statement_whitelist[addr])
        return None

    def get_loops(self):
        return self._loops<commit_msg>AnnotatedCFG: last_statement_index() returns None if the block has exits.<commit_after>from collections import defaultdict
import logging

import networkx

import simuvex

from .pathprioritizer import PathPrioritizer
from .errors import AngrAnnotatedCFGError, AngrExitError
from .analyses.cfg_node import CFGNode

l = logging.getLogger("angr.annocfg")

class AnnotatedCFG(object):
    """
    AnnotatedCFG is a control flow graph with statement whitelists and exit whitelists to describe a slice of the
    program.
    """
    def __init__(self, project, cfg=None, detect_loops=False):
        """
        Constructor.

        :param project: The angr Project instance
        :param cfg: Control flow graph. Only used when path prioritizer is used.
        :param detect_loops: Only used when path prioritizer is used.
        """
        self._project = project

        self._cfg = None
        self._target = None

        self._run_statement_whitelist = defaultdict(list)
        self._exit_taken = defaultdict(list)
        self._addr_to_run = {}
        self._addr_to_last_stmt_id = {}
        self._loops = []
        self._path_merge_points = [ ]
        self._path_prioritizer = None


        if cfg is not None:
            self._cfg = cfg

            #if target_irsb_addr is not None:
            #    self._path_prioritizer = PathPrioritizer(self._cfg, self._target)

        if self._cfg is not None:
            for run in self._cfg.nodes():
                self._addr_to_run[self.get_addr(run)] = run

    #
    # Public methods
    #

    def from_digraph(self, digraph):
        """
        Initialize this AnnotatedCFG object with a networkx.DiGraph consisting of the following
        form of nodes:

        Tuples like (block address, statement ID)

        Those nodes are connected by edges indicating the execution flow.

        :param digraph: A networkx.DiGraph object
        """

        for n1 in digraph.nodes_iter():
            addr1, stmt_idx1 = n1
            self.add_statements_to_whitelist(addr1, (stmt_idx1,))

            successors = digraph[n1]
            for n2 in successors:
                addr2, stmt_idx2 = n2

                if addr1 != addr2:
                    # There is a control flow transition from block `addr1` to block `addr2`
                    self.add_exit_to_whitelist(addr1, addr2)

                self.add_statements_to_whitelist(addr2, (stmt_idx2,))

    def get_addr(self, run):
        if isinstance(run, CFGNode):
            return run.addr
        elif type(run) in (int, long):
            return run
        else:
            raise AngrAnnotatedCFGError("Unknown type '%s' of the 'run' argument" % type(run))

    def add_block_to_whitelist(self, block):
        addr = self.get_addr(block)
        self._run_statement_whitelist[addr] = True

    def add_statements_to_whitelist(self, block, stmt_ids):
        addr = self.get_addr(block)
        if type(stmt_ids) is bool:
            if type(self._run_statement_whitelist[addr]) is list and self._run_statement_whitelist[addr]:
                raise Exception("WTF")
            self._run_statement_whitelist[addr] = stmt_ids
        elif -1 in stmt_ids:
            self._run_statement_whitelist[addr] = True
        else:
            self._run_statement_whitelist[addr].extend(stmt_ids)
            self._run_statement_whitelist[addr] = \
                sorted(list(set(self._run_statement_whitelist[addr])))

    def add_exit_to_whitelist(self, run_from, run_to):
        addr_from = self.get_addr(run_from)
        addr_to = self.get_addr(run_to)
        self._exit_taken[addr_from].append(addr_to)

    def set_last_statement(self, block_addr, stmt_id):
        self._addr_to_last_stmt_id[block_addr] = stmt_id

    def add_loop(self, loop_tuple):
        """
        A loop tuple contains a series of IRSB addresses that form a loop. Ideally
        it always starts with the first IRSB that we meet during the execution.
        """
        self._loops.append(loop_tuple)

    def set_path_merge_points(self, points):
        self._path_merge_points = points.copy()

    def should_take_exit(self, addr_from, addr_to):
        if addr_from in self._exit_taken:
            return addr_to in self._exit_taken[addr_from]

        return False

    def should_execute_statement(self, addr, stmt_id):
        if self._run_statement_whitelist is None:
            return True
        elif addr in self._run_statement_whitelist:
            r = self._run_statement_whitelist[addr]
            if isinstance(r, bool):
                return r
            else:
                return stmt_id in self._run_statement_whitelist[addr]
        return False

    def get_run(self, addr):
        if addr in self._addr_to_run:
            return self._addr_to_run[addr]
        return None

    def get_whitelisted_statements(self, addr):
        """
        :returns: True if all statements are whitelisted
        """
        if addr in self._run_statement_whitelist:
            if self._run_statement_whitelist[addr] is True:
                return None # This is the default value used in SimuVEX to say
                            # we execute all statements in this basic block. A
                            # little weird...

            else:
                return self._run_statement_whitelist[addr]

        else:
            return []

    def get_last_statement_index(self, addr):
        if addr in self._exit_taken:
            return None
        if addr in self._addr_to_last_stmt_id:
            return self._addr_to_last_stmt_id[addr]
        elif addr in self._run_statement_whitelist:
            return max(self._run_statement_whitelist[addr])
        return None

    def get_loops(self):
        return self._loops
{'fraction_non_alphanumeric': 0.07218572054314498, 'fraction_numerical': 0.003504161191414805, 'mean_word_length': 2.2040415380297502, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6134053', 'n_tokens_mistral': 3427, 'n_tokens_neox': 3298, 'n_words': 875}