from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import claripy
import cle
from capstone import CS_GRP_CALL, CS_GRP_IRET, CS_GRP_JUMP, CS_GRP_RET

from angr import BP_BEFORE, BP_AFTER, sim_options
from angr.errors import AngrTracerError, SimIRSBNoDecodeError
from .base import ExplorationTechnique

if TYPE_CHECKING:
    from angr.sim_state import SimState


l = logging.getLogger(name=__name__)


class TracingMode:
    """
    :ivar Strict:       Strict mode, the default mode, where an exception is raised immediately if tracer's path
                        deviates from the provided trace.
    :ivar Permissive:   Permissive mode, where tracer attempts to force the path back to the provided trace when a
                        deviation happens. This does not always work, especially when the cause of deviation is related
                        to input that will later be used in exploit generation. But, it might work magically sometimes.
    :ivar CatchDesync:  CatchDesync mode, catch desync because of sim_procedures. It might be a sign of something
                        interesting.
    """

    Strict = "strict"
    Permissive = "permissive"
    CatchDesync = "catch_desync"


class TracerDesyncError(AngrTracerError):
    """
    An error class to report tracing Tracing desyncronization error
    """

    def __init__(self, msg, deviating_addr=None, deviating_trace_idx=None):
        super().__init__(msg)
        self.deviating_addr = deviating_addr
        self.deviating_trace_idx = deviating_trace_idx


class RepHook:
    """
    Hook rep movs/stos to speed up constraint solving
    TODO: This should be made an exploration technique later
    """

    def __init__(self, mnemonic):
        self.mnemonic = mnemonic

    @staticmethod
    def _inline_call(state, procedure, *arguments, **kwargs):
        e_args = [claripy.BVV(a, state.arch.bits) if isinstance(a, int) else a for a in arguments]
        p = procedure(project=state.project, **kwargs)
        return p.execute(state, None, arguments=e_args)

    def run(self, state):
        from angr import SIM_PROCEDURES  # pylint: disable=import-outside-toplevel

        dst = state.regs.edi if state.arch.name == "X86" else state.regs.rdi

        if self.mnemonic.startswith("stos"):
            # store a string
            if self.mnemonic == "stosb":
                val = state.regs.al
                multiplier = 1
            elif self.mnemonic == "stosw":
                val = state.regs.ax
                multiplier = 2
            elif self.mnemonic == "stosd":
                val = state.regs.eax
                multiplier = 4
            elif self.mnemonic == "stosq":
                val = state.regs.rax
                multiplier = 8
            else:
                raise NotImplementedError(f"Unsupported mnemonic {self.mnemonic}")

            size = (state.regs.ecx if state.arch.name == "X86" else state.regs.rcx) * multiplier

            memset = SIM_PROCEDURES["libc"]["memset"]
            memset().execute(state, arguments=[dst, val, size])

            if state.arch.name == "X86":
                state.regs.edi += size
                state.regs.ecx = 0
            else:
                state.regs.rdi += size
                state.regs.rcx = 0

        elif self.mnemonic.startswith("movs"):
            src = state.regs.esi if state.arch.name == "X86" else state.regs.rsi

            # copy a string
            if self.mnemonic == "movsb":
                multiplier = 1
            elif self.mnemonic == "movsw":
                multiplier = 2
            elif self.mnemonic == "movsd":
                multiplier = 4
            elif self.mnemonic == "movsq":
                multiplier = 8
            else:
                raise NotImplementedError(f"Unsupported mnemonic {self.mnemonic}")

            size = (state.regs.ecx if state.arch.name == "X86" else state.regs.rcx) * multiplier

            memcpy = SIM_PROCEDURES["libc"]["memcpy"]
            memcpy().execute(state, arguments=[dst, src, size])

            if state.arch.name == "X86":
                state.regs.edi += size
                state.regs.esi -= size
                state.regs.ecx = 0
            else:
                state.regs.rdi += size
                state.regs.rsi -= size
                state.regs.rcx = 0

        else:
            raise NotImplementedError(f"Unsupported mnemonic {self.mnemonic}")


class Tracer(ExplorationTechnique):
    """
    An exploration technique that follows an angr path with a concrete input.
    The tracing result is the state at the last address of the trace, which can be found in the
    'traced' stash.

    If the given concrete input makes the program crash, you should provide crash_addr, and the
    crashing state will be found in the 'crashed' stash.

    :param trace:               The basic block trace.
    :param resiliency:          Should we continue to step forward even if qemu and angr disagree?
    :param keep_predecessors:   Number of states before the final state we should log.
    :param crash_addr:          If the trace resulted in a crash, provide the crashing instruction
                                pointer here, and the 'crashed' stash will be populated with the
                                crashing state.
    :param syscall_data:        Data related to various syscalls recorded by tracer for replaying
    :param copy_states:         Whether COPY_STATES should be enabled for the tracing state. It is
                                off by default because most tracing workloads benefit greatly from
                                not performing copying. You want to enable it if you want to see
                                the missed states. It will be re-added for the last 2% of the trace
                                in order to set the predecessors list correctly. If you turn this
                                on you may want to enable the LAZY_SOLVES option.
    :param mode:                Tracing mode.
    :param aslr:                Whether there are aslr slides. if not, tracer uses trace address
                                as state address.
    :param follow_unsat:        Whether unsatisfiable states should be treated as potential
                                successors or not.

    :ivar predecessors:         A list of states in the history before the final state.
    """

    def __init__(
        self,
        trace=None,
        resiliency=False,
        keep_predecessors=1,
        crash_addr=None,
        syscall_data=None,
        copy_states=False,
        fast_forward_to_entry=True,
        mode=TracingMode.Strict,
        aslr=True,
        follow_unsat=False,
    ):
        super().__init__()
        self._trace = trace
        self._resiliency = resiliency
        self._crash_addr = crash_addr
        self._syscall_data = syscall_data
        self._copy_states = copy_states
        self._mode = mode
        self._aslr = aslr
        self._follow_unsat = follow_unsat
        self._fast_forward_to_entry = fast_forward_to_entry

        self._aslr_slides: dict[cle.Backend, int] = {}
        self._current_slide = None

        self._fd_bytes = None

        # keep track of the last basic block we hit
        self.predecessors: list[SimState] = [None] * keep_predecessors
        self.last_state = None

        # whether we should follow the trace
        self._no_follow = self._trace is None

        # Keep track of count of termination point
        self._last_block_total_count = self._trace.count(self._trace[-1])
        self._last_block_seen_count = 0

        # sanity check: copy_states must be enabled in Permissive mode since we may need to backtrack from a previous
        # state.
        if self._mode == TracingMode.Permissive and not self._copy_states:
            raise ValueError('"copy_states" must be True when tracing in permissive mode.')

    def _locate_entry_point(self, angr_addr):
        # ...via heuristics
        indices = set()
        threshold = 0x40000
        while not indices and threshold > 0x2000:
            for idx, addr in enumerate(self._trace):
                if ((addr - angr_addr) & 0xFFF) == 0 and (idx == 0 or abs(self._trace[idx - 1] - addr) > threshold):
                    indices.add(idx)

            indices = {i for i in indices if self._filter_idx(angr_addr, i)}
            threshold //= 2

        return indices

    def _identify_aslr_slides(self):
        """
        libraries can be mapped differently in the original run(in the trace) and in angr
        this function identifies the difference(called aslr slides) of each library to help angr translate
        original address and address in angr back and forth
        """
        if self._aslr:
            # if we don't know whether there is any slide, we need to identify the slides via heuristics
            for obj in self.project.loader.all_objects:
                # do not analyze pseudo-objects
                if obj.binary_basename.startswith("cle##"):
                    continue

                # heuristic 1: non-PIC  objects are loaded without aslr slides
                if not obj.pic:
                    self._aslr_slides[obj] = 0
                    continue

                # heuristic 2: library objects with custom_base_addr are loaded at the correct locations
                if obj._custom_base_addr:
                    l.info("%s is assumed to be loaded at the address matching the one in the trace", obj)
                    self._aslr_slides[obj] = 0
                    continue

                # heuristic 3: entry point of an object should appear in the trace
                possibilities = None
                for entry in obj.initializers + ([obj.entry] if obj.is_main_bin else []):
                    indices = self._locate_entry_point(entry)
                    slides = {self._trace[idx] - entry for idx in indices}
                    if possibilities is None:
                        possibilities = slides
                    else:
                        possibilities.intersection_update(slides)

                if possibilities is None:
                    continue

                if len(possibilities) == 0:
                    raise AngrTracerError(
                        f"Trace does not seem to contain object initializers for {obj}. "
                        "Do you want to have a Tracer(aslr=False)?"
                    )
                if len(possibilities) == 1:
                    self._aslr_slides[obj] = next(iter(possibilities))
                else:
                    raise AngrTracerError(
                        f"Trace seems ambiguous with respect to what the ASLR slides are for {obj}. "
                        "This is surmountable, please open an issue."
                    )
        else:
            # if we know there is no slides, just trust the address in the loader
            for obj in self.project.loader.all_objects:
                # do not analyze pseudo-objects
                if obj.binary_basename.startswith("cle##"):
                    continue
                self._aslr_slides[obj] = 0
            self._current_slide = 0

    def _filter_idx(self, angr_addr, idx):
        slide = self._trace[idx] - angr_addr
        block = self.project.factory.block(angr_addr)
        legal_next = block.vex.constant_jump_targets
        if legal_next:
            return any(a + slide == self._trace[idx + 1] for a in legal_next)
        # the intuition is that if the first block of an initializer does an indirect jump,
        # it's probably a call out to another binary (notably __libc_start_main)
        # this is an awful fucking heuristic but it's as good as we've got
        return abs(self._trace[idx] - self._trace[idx + 1]) > 0x1000

    def set_fd_data(self, fd_data: dict[int, bytes]):
        """
        Set concrete bytes of various fds read by the program
        """

        self._fd_bytes = fd_data

    def setup(self, simgr):
        simgr.populate("missed", [])
        simgr.populate("traced", [])
        simgr.populate("crashed", [])
        simgr.populate("desync", [])

        if len(simgr.active) != 1:
            raise AngrTracerError("Tracer is being invoked on a SimulationManager without exactly one active state")

        self._identify_aslr_slides()

        if self._fast_forward_to_entry:
            idx = self._trace.index(self._translate_state_addr(self.project.entry))
            # step to entry point
            while simgr.one_active.addr != self.project.entry:
                simgr.step(extra_stop_points={self.project.entry})
                if len(simgr.active) == 0:
                    raise AngrTracerError("Could not step to the first address of the trace - simgr is empty")
                if len(simgr.active) > 1:
                    raise AngrTracerError(
                        "Could not step to the first address of the trace - state split. "
                        "Do you want to have a Tracer(fast_forward_to_entry=False)?"
                    )
                simgr.drop(stash="unsat")
        else:
            idx = 0

        # initialize the state info
        simgr.one_active.globals["trace_idx"] = idx
        simgr.one_active.globals["sync_idx"] = None
        simgr.one_active.globals["sync_timer"] = 0
        simgr.one_active.globals["is_desync"] = False

        # disable state copying!
        if not self._copy_states:
            # insulate our caller from this nonsense by making a single copy at the beginning
            simgr.active[0] = simgr.active[0].copy()
            simgr.active[0].options.remove(sim_options.COPY_STATES)

    def complete(self, simgr):
        return bool(simgr.traced)

    def filter(self, simgr, state, **kwargs):
        # check completion
        if state.globals["trace_idx"] >= len(self._trace) - 1:
            # if the the state is a desync state and the user wants to keep it,
            # then do what the user wants
            if self._mode == TracingMode.CatchDesync and self.project.is_hooked(state.addr):
                return "desync"
            # do crash windup if necessary
            if self._crash_addr is not None:
                self.last_state, crash_state = self.crash_windup(state, self._crash_addr)
                simgr.populate("crashed", [crash_state])
                self.predecessors.append(state)
                self.predecessors.pop(0)

            return "traced"

        return simgr.filter(state, **kwargs)

    def step(self, simgr, stash="active", **kwargs):
        simgr.drop(stash="missed")
        return simgr.step(stash=stash, syscall_data=self._syscall_data, fd_bytes=self._fd_bytes, **kwargs)

    def step_state(self, simgr, state, **kwargs):
        if state.history.jumpkind == "Ijk_Exit":
            return {"traced": [state]}

        # maintain the predecessors list
        self.predecessors.append(state)
        self.predecessors.pop(0)

        if state.globals["trace_idx"] > len(self._trace) * 0.98:
            state.options.add(sim_options.COPY_STATES)
            state.options.add(sim_options.LAZY_SOLVES)

        # optimization:
        # look forward, is it a rep stos/movs instruction?
        # if so, we add a temporary hook to speed up constraint solving
        if not self.project.is_hooked(state.addr):
            block = self.project.factory.block(state.addr)

            if len(block.capstone.insns) == 1 and (
                block.capstone.insns[0].mnemonic.startswith("rep m")
                or block.capstone.insns[0].mnemonic.startswith("rep s")
            ):
                insn = block.capstone.insns[0]
                self.project.hook(state.addr, RepHook(insn.mnemonic.split(" ")[1]).run, length=insn.size)

        # perform the step. ask qemu to stop at the termination point.
        # if termination point occurs multiple times in trace, pass details to SimEngineUnicorn's native interface so
        # that it can stop at last block
        if self._last_block_total_count > 1:
            stops = set(kwargs.pop("extra_stop_points", ()))
            last_block_details = {
                "addr": self._trace[-1],
                "tot_count": self._last_block_total_count,
                "curr_count": self._last_block_seen_count,
            }
        else:
            stops = set(kwargs.pop("extra_stop_points", ())) | {self._trace[-1]}
            last_block_details = None

        succs_dict = simgr.step_state(state, extra_stop_points=stops, last_block_details=last_block_details, **kwargs)
        if None not in succs_dict and simgr.errored:
            raise simgr.errored[-1].error
        sat_succs = succs_dict[None]  # satisfiable states
        succs = sat_succs + succs_dict["unsat"]  # both satisfiable and unsatisfiable states

        if not self._follow_unsat:
            # Only satisfiable states need to be checked for correct successor
            if len(sat_succs) == 1:
                try:
                    self._update_state_tracking(sat_succs[0])
                except TracerDesyncError as ex:
                    if self._mode == TracingMode.Permissive:
                        succs_dict = self._force_resync(simgr, state, ex.deviating_trace_idx, ex.deviating_addr, kwargs)
                    else:
                        raise
            elif len(sat_succs) == 0:
                raise Exception("No satisfiable state is available!")
            else:
                succ = self._pick_correct_successor(sat_succs)
                succs_dict[None] = [succ]
                succs_dict["missed"] = [s for s in sat_succs if s is not succ]
        else:
            # Check all states for correct successor
            if len(succs) == 1:
                self._update_state_tracking(succs[0])
            elif len(succs) == 0:
                raise Exception("All states disappeared!")
            else:
                succ = self._pick_correct_successor(succs)
                succs_dict[None] = [succ]
                succs_dict["missed"] = [s for s in succs if s is not succ]
        assert len(succs_dict[None]) == 1

        # if there is a catchable desync, we should return the last sync state
        if succs_dict[None][0].globals["is_desync"]:
            simgr.active[0].globals["trace_idx"] = len(self._trace)
            succs_dict[None][0] = state
        return succs_dict

    def _force_resync(self, simgr, state, deviating_trace_idx, deviating_addr, kwargs):
        """
        When a deviation happens, force the tracer to take the branch specified in the trace by manually setting the
        PC to the one in the trace. This method is only used in Permissive tracing mode.

        :param simgr:               The simulation manager instance.
        :param state:               The program state before the current step.
        :param deviating_trace_idx: The index of address in the trace where a desync happens.
        :param deviating_addr:      The address that tracer takes when the desync happens. Should be different from the
                                    one in the trace.
        :param kwargs:              Other keyword arguments that will be passed to step_state().
        :return:                    A new successor dict.
        :rtype:                     dict
        """

        # if unicorn engine is enabled, disable it. forced execution requires single-stepping in angr.
        unicorn_option_removed = False
        if sim_options.UNICORN in state.options:
            state.options.remove(sim_options.UNICORN)
            unicorn_option_removed = True

        # single step until right before the deviating state
        trace_idx = state.globals["trace_idx"]
        while trace_idx != deviating_trace_idx - 1:
            succs_dict = simgr.step_state(state, **kwargs)
            succs = succs_dict[None]
            assert len(succs) == 1
            self._update_state_tracking(succs[0])
            state = succs[0]
            trace_idx += 1

        # step the state further and then manually set the PC
        succs_dict = simgr.step_state(state, **kwargs)
        succs = succs_dict[None]
        if len(succs) != 1 or succs[0].addr != deviating_addr:
            raise TracerDesyncError("Address mismatch during single-stepping.")
        succ = succs[0]
        expected_addr = self._trace[deviating_trace_idx]
        current_obj = self.project.loader.find_object_containing(state.addr)
        assert current_obj is not None
        translated_addr = self._translate_trace_addr(expected_addr, current_obj)
        l.info(
            "Attempt to fix a deviation: Forcing execution from %#x to %#x (instead of %#x).",
            state.addr,
            succ.addr,
            translated_addr,
        )
        succ._ip = translated_addr

        succ.globals["trace_idx"] = trace_idx + 1
        succs_dict = {None: [succ]}

        if unicorn_option_removed:
            succ.options.add(sim_options.UNICORN)

        return succs_dict

    def _pick_correct_successor(self, succs):
        # there's been a branch of some sort. Try to identify which state stayed on the trace.
        assert len(succs) > 0
        idx = succs[0].globals["trace_idx"]

        res = []
        last_description = succs[0].history.descriptions[-1]
        if "Unicorn" in last_description:
            # Multiple new states were created in SimEngineUnicorn. State which has non-zero recent block count is a
            # valid successor since only correct successor is sync'd with native state
            for succ in succs:
                if succ.history.recent_block_count > 0:
                    res.append(succ)
        else:
            for succ in succs:
                try:
                    if self._compare_addr(self._trace[idx + 1], succ.addr):
                        res.append(succ)
                except AngrTracerError:
                    pass

        if not res:
            raise Exception("No states followed the trace?")

        if len(res) > 1:
            raise Exception("The state split but several successors have the same (correct) address?")

        self._update_state_tracking(res[0])
        return res[0]

    def _update_state_tracking(self, state: SimState):
        idx = state.globals["trace_idx"]
        sync = state.globals["sync_idx"]
        timer = state.globals["sync_timer"]

        self._last_block_seen_count += state.history.recent_bbl_addrs.count(self._trace[-1])

        if state.history.recent_block_count > 1:
            # multiple blocks were executed this step. they should follow the trace *perfectly*
            # or else something is up
            # "something else" so far only includes concrete transmits, or...
            # TODO: https://github.com/unicorn-engine/unicorn/issues/874
            # ^ this means we will see desyncs of the form unicorn suddenly skips a bunch of qemu blocks
            assert state.history.recent_block_count == len(state.history.recent_bbl_addrs)

            for addr_idx, addr in enumerate(state.history.recent_bbl_addrs):
                if addr in [
                    state.unicorn.cgc_transmit_addr,
                    state.unicorn.cgc_receive_addr,
                    state.unicorn.cgc_random_addr,
                ]:
                    continue

                if sync is not None and sync != "entry":
                    if self._compare_addr(self._trace[sync], addr):
                        # Found the address in trace. Start normal trace checks from next address
                        idx = sync + 1
                        state.globals["sync_idx"] = None
                        sync = None

                    continue

                if self._compare_addr(self._trace[idx], addr) or self._check_qemu_unicorn_large_block_split(
                    state, idx, addr_idx
                ):
                    idx += 1
                else:
                    is_contained, increment = self._check_qemu_block_in_unicorn_block(state, idx, addr_idx)
                    if is_contained:
                        idx += increment
                        # Big block is now skipped in qemu trace. Perform compare at correct index again.
                        if self._compare_addr(self._trace[idx], addr):
                            idx += 1
                            continue

                    raise TracerDesyncError(
                        "Oops! angr did not follow the trace", deviating_addr=addr, deviating_trace_idx=idx
                    )

            idx -= 1  # use normal code to do the last synchronization

        if sync == "entry":
            trace_addr = self._translate_state_addr(state.addr)
            # this address should only ever appear once in the trace. we verified this during setup.
            idx = self._trace.index(trace_addr)
            state.globals["trace_idx"] = idx
            state.globals["sync_idx"] = None
        elif sync is not None:
            timer -= 1
            if self._compare_addr(self._trace[sync], state.addr):
                state.globals["trace_idx"] = sync
                state.globals["sync_idx"] = None
                state.globals["sync_timer"] = 0
            elif timer > 0:
                state.globals["sync_timer"] = timer
            else:
                raise Exception(
                    f"Trace failed to synchronize! We expected it to hit {self._trace[sync]:#x} (trace addr), "
                    "but it failed to do this within a timeout"
                )

        elif state.history.jumpkind.startswith("Ijk_Exit"):
            # termination! will be handled by filter
            pass
        elif self.project.is_hooked(state.addr) and not self.project.loader.extern_object.contains_addr(state.addr):
            # handle simprocedures
            self._sync_return(state, idx)
        elif self._compare_addr(self._trace[idx + 1], state.addr):
            # normal case
            state.globals["trace_idx"] = idx + 1
        elif self.project.loader._extern_object is not None and self.project.loader.extern_object.contains_addr(
            state.addr
        ):
            # externs
            proc = self.project.hooked_by(state.addr)
            if proc is None:
                raise Exception("Extremely bad news: we're executing an unhooked address in the externs space")
            if proc.display_name == "LinuxLoader":
                state.globals["sync_idx"] = "entry"
            elif proc.is_continuation:
                orig_addr = self.project.loader.find_symbol(proc.display_name).rebased_addr
                obj = self.project.loader.find_object_containing(orig_addr)
                orig_trace_addr = self._translate_state_addr(orig_addr, obj)
                if 0 <= self._trace[idx + 1] - orig_trace_addr <= 0x10000:
                    # this is fine. we do nothing and then next round
                    # it'll get handled by the is_hooked(state.history.addr) case
                    pass
                else:
                    # this may also be triggered as a consequence of the unicorn issue linked above
                    raise Exception("BUG: State is returning to a continuation that isn't its own???")
            elif state.addr == getattr(self.project.simos, "vsyscall_addr", None):
                if not self._sync_callsite(state, idx, state.history.addr):
                    raise AngrTracerError("Could not synchronize following vsyscall")
            elif self.project.hooked_by(state.addr).display_name.startswith("IFuncResolver"):
                if not self._sync_return(state, idx):
                    raise AngrTracerError("Could not synchronize at ifunc return address")
            else:
                # see above
                pass
        elif state.history.jumpkind.startswith("Ijk_Sys"):
            # syscalls
            state.globals["sync_idx"] = idx + 1
            state.globals["sync_timer"] = 1
        elif self.project.is_hooked(state.history.addr):
            # simprocedures - is this safe..?
            self._fast_forward(state)
        elif state.addr == self._trace[-1]:
            # we may have prematurely stopped because of setting stop points. try to resync.
            state.globals["sync_idx"] = idx + 1
            state.globals["sync_timer"] = 1
        elif (
            self.project.is_hooked(state.addr)
            and self.project.loader.find_symbol(self.project.hooked_by(state.addr).display_name) is not None
            and self.project.loader.find_symbol(self.project.hooked_by(state.addr).display_name).subtype.value[0] == 10
        ):  # STT_GNU_IFUNC #pylint:disable=line-too-long
            if not self._sync_return(state, idx):
                raise AngrTracerError("Could not synchronize at ifunc return address")
        elif self._analyze_misfollow(state, idx):
            # misfollow analysis will set a sync point somewhere if it succeeds
            pass
        else:
            raise TracerDesyncError(
                "Oops! angr did not follow the trace", deviating_addr=state.addr, deviating_trace_idx=idx + 1
            )

        if state.globals["sync_idx"] is not None:
            l.debug(
                "Trace: %s-%s/%s synchronizing %s",
                state.globals["trace_idx"],
                state.globals["sync_idx"],
                len(self._trace),
                state.globals["sync_timer"],
            )
        else:
            l.debug("Trace: %s/%s", state.globals["trace_idx"], len(self._trace))

    def _translate_state_addr(self, state_addr, obj=None):
        if obj is None:
            obj = self.project.loader.find_object_containing(state_addr)
        if obj not in self._aslr_slides:
            raise Exception("Internal error: cannot translate address")
        return state_addr + self._aslr_slides[obj]

    def _translate_trace_addr(self, trace_addr, obj=None):
        if obj is None:
            for obj, slide in self._aslr_slides.items():  # pylint: disable=redefined-argument-from-local
                if obj.contains_addr(trace_addr - slide):
                    break
            else:
                raise Exception("Can't figure out which object this address belongs to")
        if obj not in self._aslr_slides:
            raise Exception("Internal error: object is untranslated")
        return trace_addr - self._aslr_slides[obj]

    def _compare_addr(self, trace_addr, state_addr):
        if self._current_slide is not None and trace_addr == state_addr + self._current_slide:
            return True

        current_bin = self.project.loader.find_object_containing(state_addr)
        if current_bin is self.project.loader._extern_object or current_bin is self.project.loader._kernel_object:
            return False
        if current_bin in self._aslr_slides:
            self._current_slide = self._aslr_slides[current_bin]
            return trace_addr == state_addr + self._current_slide
        if ((trace_addr - state_addr) & 0xFFF) == 0:
            self._aslr_slides[current_bin] = self._current_slide = trace_addr - state_addr
            return True
        # error handling
        if current_bin:
            raise AngrTracerError(
                f"Trace desynced on jumping into {current_bin.provides}. "
                "Did you load the right version of this library?"
            )
        raise AngrTracerError(f"Trace desynced on jumping into {state_addr:#x}, where no library is mapped!")

    def _check_qemu_block_in_unicorn_block(self, state: SimState, trace_curr_idx, state_desync_block_idx):
        """
        Check if desync occurred because unicorn block was split into multiple blocks in qemu tracer. If yes, find the
        correct increment for trace index
        """

        # We first find the block address where the trace and state's history match
        for trace_match_idx in range(trace_curr_idx - 1, -1, -1):
            if self._trace[trace_match_idx] == state.history.recent_bbl_addrs[state_desync_block_idx - 1]:
                break
        else:
            # Failed to find matching block address. qemu block is probably not contained in a previous block.
            return (False, -1)

        control_flow_insn_types = [CS_GRP_CALL, CS_GRP_IRET, CS_GRP_JUMP, CS_GRP_RET]
        big_block_start = self._trace[trace_match_idx]
        big_block_end = None
        curr_block_addr = big_block_start
        while True:
            curr_block = state.project.factory.block(self._translate_trace_addr(curr_block_addr))
            curr_block_last_insn = curr_block.capstone.insns[-1]
