from typing import List, Optional


class CallSite:
    """
    Describes a call site on a CFG.
    """

    __slots__ = ('caller_func_addr', 'block_addr', 'callee_func_addr', )

    def __init__(self, caller_func_addr: int, block_addr:int, callee_func_addr: int):
        self.caller_func_addr = caller_func_addr
        self.callee_func_addr = callee_func_addr
        self.block_addr = block_addr

    def __repr__(self):
        return "<CallSite in function %#x, calling #%x at block %#x>" % (
            self.caller_func_addr,
            self.callee_func_addr,
            self.block_addr
        )


class CallTrace:
    """
    Describes a series of functions calls to get from one function (current_function_address()) to another function or
    a basic block (self.target).
    """

    __slots__ = ('callsites', 'target', )

    def __init__(self, target: int):
        self.target = target
        self.callsites: List[CallSite] = [ ]

    def __repr__(self):
        return "<Trace with %d callsites>" % len(self.callsites)

    def current_function_address(self) -> int:
        if not self.callsites:
            return self.target
        return self.callsites[-1].caller_func_addr

    def step_back(self, caller_func_addr: int, block_addr: Optional[int], callee_func_addr) -> 'CallTrace':
        # create a new CallSite object
        site = CallSite(caller_func_addr, block_addr, callee_func_addr)
        t = self.copy()
        t.callsites.append(site)
        return t

    def includes_function(self, func_addr: int) -> bool:
        if self.target == func_addr:
            return True
        if any(cs.caller_func_addr == func_addr for cs in self.callsites):
            return True
        return False

    def copy(self) -> 'CallTrace':
        t = CallTrace(self.target)
        t.callsites = self.callsites[::]
        return t
