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

@@ -229,19 +229,30 @@
 ck_addr:
+ Optional%5B
 int
+%5D
 , callee
@@ -439,20 +439,22 @@
       re
-turn
+sult =
  %22%3CCallS
@@ -489,25 +489,304 @@
  #%25x
- at block %25#x%3E%22 %25
+%22 %25 (self.caller_func_addr,%0A                                                             self.callee_func_addr,)%0A        if self.block_addr is not None:%0A            result += %22at block %25#x%22 %25 self.block_addr%0A        result += %22%3E%22%0A        return result%0A%0A    def __eq__(self, other):%0A        return
  (%0A 
@@ -809,33 +809,62 @@
 caller_func_addr
-,
+ == other.caller_func_addr and
 %0A            sel
@@ -885,26 +885,75 @@
 addr
-,%0A            self
+ == other.callee_func_addr and%0A            self.block_addr == other
 .blo

Improve CallSite (#2259)