angr 9.2.165__cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

angr/analyses/s_propagator.py

@@ -0,0 +1,542 @@
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Mapping
+from collections import defaultdict
+
+import networkx
+
+from angr.ailment.block import Block
+from angr.ailment.expression import (
+    Const,
+    VirtualVariable,
+    VirtualVariableCategory,
+    StackBaseOffset,
+    Load,
+    Convert,
+    Expression,
+)
+from angr.ailment.statement import Assignment, Store, Return, Jump, ConditionalJump
+
+from angr.knowledge_plugins.functions import Function
+from angr.code_location import CodeLocation, ExternalCodeLocation
+from angr.analyses import Analysis, register_analysis
+from angr.utils.ssa import (
+    get_vvar_uselocs,
+    get_vvar_deflocs,
+    has_ite_expr,
+    has_ite_stmt,
+    has_tmp_expr,
+    is_phi_assignment,
+    is_const_assignment,
+    is_const_and_vvar_assignment,
+    is_const_vvar_load_assignment,
+    is_const_vvar_load_dirty_assignment,
+    is_const_vvar_tmp_assignment,
+    is_vvar_propagatable,
+    get_tmp_uselocs,
+    get_tmp_deflocs,
+    phi_assignment_get_src,
+    has_store_stmt_in_between_stmts,
+)
+
+
+class SPropagatorModel:
+    """
+    The SPropagator model that stores replacements for virtual variables.
+    """
+
+    def __init__(self):
+        self.replacements: Mapping[CodeLocation, Mapping[Expression, Expression]] = {}
+        # store vvars that are definitely dead (but usually not removed by default because they are stack variables)
+        self.dead_vvar_ids: set[int] = set()
+
+
+class SPropagatorAnalysis(Analysis):
+    """
+    Constant and expression propagation that only supports SSA AIL graphs.
+    """
+
+    def __init__(  # pylint: disable=too-many-positional-arguments
+        self,
+        subject: Block | Function,
+        func_graph: networkx.DiGraph | None = None,
+        only_consts: bool = True,
+        stack_pointer_tracker=None,
+        func_args: set[VirtualVariable] | None = None,
+        func_addr: int | None = None,
+    ):
+        if isinstance(subject, Block):
+            self.block = subject
+            self.func = None
+            self.mode = "block"
+        elif isinstance(subject, Function):
+            self.block = None
+            self.func = subject
+            self.mode = "function"
+        else:
+            raise TypeError(f"Unsupported subject type {type(subject)}")
+
+        self.func_graph = func_graph
+        self.func_addr = func_addr
+        self.func_args = func_args
+        self.only_consts = only_consts
+        self._sp_tracker = stack_pointer_tracker
+
+        bp_as_gpr = False
+        the_func = None
+        if self.func is not None:
+            the_func = self.func
+        else:
+            if self.func_addr is not None:
+                with contextlib.suppress(KeyError):
+                    the_func = self.kb.functions.get_by_addr(self.func_addr)
+        if the_func is not None:
+            bp_as_gpr = the_func.info.get("bp_as_gpr", False)
+        self._bp_as_gpr = bp_as_gpr
+
+        # output
+        self.model = SPropagatorModel()
+
+        self._analyze()
+
+    @property
+    def replacements(self):
+        return self.model.replacements
+
+    @property
+    def dead_vvar_ids(self):
+        return self.model.dead_vvar_ids
+
+    def _analyze(self):
+        blocks: dict[tuple[int, int | None], Block]
+        match self.mode:
+            case "block":
+                assert self.block is not None
+                blocks = {(self.block.addr, self.block.idx): self.block}
+            case "function":
+                assert self.func_graph is not None
+                blocks = {(block.addr, block.idx): block for block in self.func_graph}
+            case _:
+                raise NotImplementedError
+
+        # find all vvar definitions
+        vvar_deflocs = get_vvar_deflocs(blocks.values())
+        # find all vvar uses
+        vvar_uselocs = get_vvar_uselocs(blocks.values())
+
+        # update vvar_deflocs using function arguments
+        if self.func_args:
+            for func_arg in self.func_args:
+                vvar_deflocs[func_arg.varid] = func_arg, ExternalCodeLocation()
+
+        # find all ret sites and indirect jump sites
+        retsites: set[tuple[int, int | None, int]] = set()
+        jumpsites: set[tuple[int, int | None, int]] = set()
+        for bb in blocks.values():
+            if bb.statements:
+                if isinstance(bb.statements[-1], Return):
+                    retsites.add((bb.addr, bb.idx, len(bb.statements) - 1))
+                elif isinstance(bb.statements[-1], Jump):
+                    jumpsites.add((bb.addr, bb.idx, len(bb.statements) - 1))
+
+        replacements = defaultdict(dict)
+
+        # find constant and other propagatable assignments
+        vvarid_to_vvar = {}
+        const_vvars: dict[int, Const] = {}
+        for vvar_id, (vvar, defloc) in vvar_deflocs.items():
+            if not vvar.was_reg and not vvar.was_parameter:
+                continue
+
+            vvarid_to_vvar[vvar_id] = vvar
+            if isinstance(defloc, ExternalCodeLocation):
+                continue
+
+            assert defloc.block_addr is not None
+            assert defloc.stmt_idx is not None
+
+            block = blocks[(defloc.block_addr, defloc.block_idx)]
+            stmt = block.statements[defloc.stmt_idx]
+            r, v = is_const_assignment(stmt)
+            if r:
+                # replace wherever it's used
+                assert v is not None
+                const_vvars[vvar_id] = v
+                for vvar_at_use, useloc in vvar_uselocs[vvar_id]:
+                    replacements[useloc][vvar_at_use] = v
+                continue
+
+            v = phi_assignment_get_src(stmt)
+            if v is not None:
+                src_varids = {vvar.varid if vvar is not None else None for _, vvar in v.src_and_vvars}
+                if None not in src_varids and all(varid in const_vvars for varid in src_varids):
+                    all_int_src_varids: set[int] = {varid for varid in src_varids if varid is not None}
+                    src_values = {
+                        (
+                            (const_vvars[varid].value, const_vvars[varid].bits)
+                            if isinstance(const_vvars[varid], Const)
+                            else const_vvars[varid]
+                        )
+                        for varid in all_int_src_varids
+                    }
+                    if len(src_values) == 1:
+                        # replace it!
+                        const_value = const_vvars[next(iter(all_int_src_varids))]
+                        const_vvars[vvar.varid] = const_value
+                        for vvar_at_use, useloc in vvar_uselocs[vvar.varid]:
+                            replacements[useloc][vvar_at_use] = const_value
+
+        # function mode only
+        if self.mode == "function":
+            assert self.func_graph is not None
+
+            for vvar_id, (vvar, defloc) in vvar_deflocs.items():
+                if vvar_id not in vvar_uselocs:
+                    continue
+                if vvar_id in const_vvars:
+                    continue
+                if isinstance(defloc, ExternalCodeLocation):
+                    continue
+
+                assert defloc.block_addr is not None
+                assert defloc.stmt_idx is not None
+
+                vvar_uselocs_set = set(vvar_uselocs[vvar_id])  # deduplicate
+
+                block = blocks[(defloc.block_addr, defloc.block_idx)]
+                stmt = block.statements[defloc.stmt_idx]
+                if (
+                    (vvar.was_reg or vvar.was_parameter)
+                    and len(vvar_uselocs_set) <= 2
+                    and isinstance(stmt, Assignment)
+                    and isinstance(stmt.src, Load)
+                ):
+                    # do we want to propagate this Load expression if it's used for less than twice?
+                    # it's often seen in the following pattern, where propagation will be beneficial:
+                    #    v0 = Load(...)
+                    #    if (!v0) {
+                    #       v1 = v0 + 1;
+                    #    }
+                    can_replace = True
+                    for _, vvar_useloc in vvar_uselocs_set:
+                        if has_store_stmt_in_between_stmts(self.func_graph, blocks, defloc, vvar_useloc):
+                            can_replace = False
+
+                    if can_replace:
+                        # we can propagate this load because there is no store between its def and use
+                        for vvar_used, vvar_useloc in vvar_uselocs_set:
+                            replacements[vvar_useloc][vvar_used] = stmt.src
+                        continue
+
+                if (
+                    (vvar.was_reg or vvar.was_stack)
+                    and len(vvar_uselocs_set) == 2
+                    and isinstance(stmt, Assignment)
+                    and not is_phi_assignment(stmt)
+                ):
+                    # a special case: in a typical switch-case construct, a variable may be used once for comparison
+                    # for the default case and then used again for constructing the jump target. we can propagate this
+                    # variable for such cases.
+                    uselocs = {loc for _, loc in vvar_uselocs_set}
+                    if self.is_vvar_used_for_addr_loading_switch_case(uselocs, blocks) and not has_tmp_expr(stmt.src):
+                        for vvar_used, vvar_useloc in vvar_uselocs_set:
+                            replacements[vvar_useloc][vvar_used] = stmt.src
+                        # mark the vvar as dead and should be removed
+                        self.model.dead_vvar_ids.add(vvar.varid)
+                        continue
+
+                if is_vvar_propagatable(vvar, stmt):
+                    if len(vvar_uselocs_set) == 1:
+                        vvar_used, vvar_useloc = next(iter(vvar_uselocs_set))
+                        if (
+                            is_const_vvar_load_assignment(stmt)
+                            and not has_store_stmt_in_between_stmts(self.func_graph, blocks, defloc, vvar_useloc)
+                            and not has_tmp_expr(stmt.src)
+                        ):
+                            # we can propagate this load because there is no store between its def and use
+                            replacements[vvar_useloc][vvar_used] = stmt.src
+                            continue
+
+                        if is_const_and_vvar_assignment(stmt) and not has_tmp_expr(stmt.src):
+                            # if the useloc is a phi assignment statement, ensure that stmt.src is the same as the phi
+                            # variable
+                            assert vvar_useloc.block_addr is not None
+                            assert vvar_useloc.stmt_idx is not None
+                            useloc_stmt = blocks[(vvar_useloc.block_addr, vvar_useloc.block_idx)].statements[
+                                vvar_useloc.stmt_idx
+                            ]
+                            if is_phi_assignment(useloc_stmt):
+                                if (
+                                    isinstance(stmt.src, VirtualVariable)
+                                    and stmt.src.oident == useloc_stmt.dst.oident
+                                    and stmt.src.category == useloc_stmt.dst.category
+                                ):
+                                    replacements[vvar_useloc][vvar_used] = stmt.src
+                            else:
+                                replacements[vvar_useloc][vvar_used] = stmt.src
+                            continue
+
+                    else:
+                        non_exitsite_uselocs = [
+                            loc
+                            for _, loc in vvar_uselocs_set
+                            if (loc.block_addr, loc.block_idx, loc.stmt_idx) not in (retsites | jumpsites)
+                        ]
+                        if is_const_and_vvar_assignment(stmt):
+                            if len(non_exitsite_uselocs) == 1:
+                                # this vvar is used once if we exclude its uses at ret sites or jump sites. we can
+                                # propagate it
+                                for vvar_used, vvar_useloc in vvar_uselocs_set:
+                                    replacements[vvar_useloc][vvar_used] = stmt.src
+                                continue
+
+                            if (
+                                len(set(non_exitsite_uselocs)) == 1
+                                and not has_ite_expr(stmt.src)
+                                and not has_tmp_expr(stmt.src)
+                            ):
+                                useloc = non_exitsite_uselocs[0]
+                                assert useloc.block_addr is not None
+                                assert useloc.stmt_idx is not None
+                                useloc_stmt = blocks[(useloc.block_addr, useloc.block_idx)].statements[useloc.stmt_idx]
+                                if stmt.src.depth <= 3 and not has_ite_stmt(useloc_stmt):
+                                    # remove duplicate use locs (e.g., if the variable is used multiple times by the
+                                    # same statement) - but ensure stmt is simple enough
+                                    for vvar_used, vvar_useloc in vvar_uselocs_set:
+                                        replacements[vvar_useloc][vvar_used] = stmt.src
+                                    continue
+
+                # special logic for global variables: if it's used once or multiple times, and the variable is never
+                # updated before it's used, we will propagate the load
+                if (vvar.was_reg or vvar.was_parameter) and isinstance(stmt, Assignment) and not has_tmp_expr(stmt.src):
+                    stmt_src = stmt.src
+                    # unpack conversions
+                    while isinstance(stmt_src, Convert):
+                        stmt_src = stmt_src.operand
+                    if (
+                        isinstance(stmt_src, Load)
+                        and isinstance(stmt_src.addr, Const)
+                        and isinstance(stmt_src.addr.value, int)
+                    ):
+                        gv_updated = False
+                        for _vvar_used, vvar_useloc in vvar_uselocs_set:
+                            gv_updated |= self.is_global_variable_updated(
+                                self.func_graph,
+                                blocks,
+                                vvar.varid,
+                                stmt_src.addr.value,
+                                stmt_src.size,
+                                defloc,
+                                vvar_useloc,
+                            )
+                        if not gv_updated:
+                            for vvar_used, vvar_useloc in vvar_uselocs_set:
+                                replacements[vvar_useloc][vvar_used] = stmt.src
+                            continue
+
+        for vvar_id, uselocs in vvar_uselocs.items():
+            vvar = next(iter(uselocs))[0] if vvar_id not in vvarid_to_vvar else vvarid_to_vvar[vvar_id]
+            vvar_uselocs_set = set(uselocs)  # deduplicate
+
+            if self._sp_tracker is not None and vvar.category == VirtualVariableCategory.REGISTER:
+                if vvar.oident == self.project.arch.sp_offset:
+                    sp_bits = (
+                        (self.project.arch.registers["sp"][1] * self.project.arch.byte_width)
+                        if "sp" in self.project.arch.registers
+                        else None
+                    )
+                    for vvar_at_use, useloc in vvar_uselocs_set:
+                        sb_offset = self._sp_tracker.offset_before(useloc.ins_addr, self.project.arch.sp_offset)
+                        if sb_offset is not None:
+                            v = StackBaseOffset(None, self.project.arch.bits, sb_offset)
+                            if sp_bits is not None and vvar.bits < sp_bits:
+                                # truncation needed
+                                v = Convert(None, sp_bits, vvar.bits, False, v)
+                            replacements[useloc][vvar_at_use] = v
+                    continue
+                if not self._bp_as_gpr and vvar.oident == self.project.arch.bp_offset:
+                    bp_bits = (
+                        (self.project.arch.registers["bp"][1] * self.project.arch.byte_width)
+                        if "bp" in self.project.arch.registers
+                        else None
+                    )
+                    for vvar_at_use, useloc in vvar_uselocs_set:
+                        sb_offset = self._sp_tracker.offset_before(useloc.ins_addr, self.project.arch.bp_offset)
+                        if sb_offset is not None:
+                            v = StackBaseOffset(None, self.project.arch.bits, sb_offset)
+                            if bp_bits is not None and vvar.bits < bp_bits:
+                                # truncation needed
+                                v = Convert(None, bp_bits, vvar.bits, False, v)
+                            replacements[useloc][vvar_at_use] = v
+                    continue
+
+        # find all tmp definitions
+        tmp_deflocs = get_tmp_deflocs(blocks.values())
+        # find all tmp uses
+        tmp_uselocs = get_tmp_uselocs(blocks.values())
+
+        for block_loc, tmp_and_uses in tmp_uselocs.items():
+            for tmp_atom, tmp_uses in tmp_and_uses.items():
+                # take a look at the definition and propagate the definition if supported
+                assert block_loc.block_addr is not None
+
+                block = blocks[(block_loc.block_addr, block_loc.block_idx)]
+                tmp_def_stmtidx = tmp_deflocs[block_loc][tmp_atom]
+
+                stmt = block.statements[tmp_def_stmtidx]
+                if isinstance(stmt, Assignment):
+                    r, v = is_const_assignment(stmt)
+                    if r:
+                        # we can propagate it!
+                        for tmp_used, tmp_use_stmtidx in tmp_uses:
+                            replacements[
+                                CodeLocation(block_loc.block_addr, tmp_use_stmtidx, block_idx=block_loc.block_idx)
+                            ][tmp_used] = stmt.src
+                        continue
+
+                    r = is_const_vvar_tmp_assignment(stmt)
+                    if r:
+                        # we can propagate it!
+                        if isinstance(stmt.src, VirtualVariable):
+                            v = const_vvars.get(stmt.src.varid, stmt.src)
+                        else:
+                            v = stmt.src
+
+                        for tmp_used, tmp_use_stmtidx in tmp_uses:
+                            replacements[
+                                CodeLocation(block_loc.block_addr, tmp_use_stmtidx, block_idx=block_loc.block_idx)
+                            ][tmp_used] = v
+                        continue
+
+                    if len(tmp_uses) <= 2 and is_const_vvar_load_dirty_assignment(stmt):
+                        for tmp_used, tmp_use_stmtidx in tmp_uses:
+                            same_inst = (
+                                block.statements[tmp_def_stmtidx].ins_addr == block.statements[tmp_use_stmtidx].ins_addr
+                            )
+                            has_store = any(
+                                isinstance(stmt_, Store)
+                                for stmt_ in block.statements[tmp_def_stmtidx + 1 : tmp_use_stmtidx]
+                            )
+                            if same_inst or not has_store:
+                                # we can propagate this load because either we do not consider memory aliasing problem
+                                # within the same instruction (blocks must be originally lifted with
+                                # CROSS_INSN_OPT=False), or there is no store between its def and use.
+                                replacements[
+                                    CodeLocation(block_loc.block_addr, tmp_use_stmtidx, block_idx=block_loc.block_idx)
+                                ][tmp_used] = stmt.src
+
+        self.model.replacements = replacements
+
+    @staticmethod
+    def is_global_variable_updated(
+        func_graph, block_dict, varid: int, gv_addr: int, gv_size: int, defloc: CodeLocation, useloc: CodeLocation
+    ) -> bool:
+        defblock = block_dict[(defloc.block_addr, defloc.block_idx)]
+        useblock = block_dict[(useloc.block_addr, useloc.block_idx)]
+
+        # traverse a graph slice from the def block to the use block and check if the global variable is updated
+        seen = {defblock}
+        queue = [defblock]
+        while queue:
+            block = queue.pop(0)
+
+            start_stmt_idx = defloc.stmt_idx if block is defblock else 0  # inclusive
+            end_stmt_idx = useloc.stmt_idx if block is useblock else len(block.statements)  # exclusive
+            assert start_stmt_idx is not None
+            assert end_stmt_idx is not None
+
+            for idx in range(start_stmt_idx, end_stmt_idx):
+                stmt = block.statements[idx]
+                if isinstance(stmt, Store) and isinstance(stmt.addr, Const):
+                    store_addr = stmt.addr.value
+                    store_size = stmt.size
+                    if gv_addr <= store_addr < gv_addr + gv_size or store_addr <= gv_addr < store_addr + store_size:
+                        return True
+
+            if block is useblock:
+                continue
+
+            for succ in func_graph.successors(block):
+                if succ not in seen:
+                    abort_path = False
+                    for stmt in succ.statements:
+                        if is_phi_assignment(stmt) and any(
+                            vvar.varid == varid for _, vvar in stmt.src.src_and_vvars if vvar is not None
+                        ):
+                            # the virtual variable is no longer live after this point
+                            abort_path = True
+                            break
+                    if abort_path:
+                        continue
+
+                    seen.add(succ)
+                    queue.append(succ)
+
+        return False
+
+    @staticmethod
+    def is_vvar_used_for_addr_loading_switch_case(uselocs: set[CodeLocation], blocks) -> bool:
+        """
+        Check if a virtual variable is used for loading an address in a switch-case construct.
+
+        :param uselocs: The use locations of the virtual variable.
+        :param blocks:  All blocks of the current function.
+        :return:        True if the virtual variable is used for loading an address in a switch-case construct, False
+                        otherwise.
+        """
+
+        if len(uselocs) != 2:
+            return False
+
+        useloc_0, useloc_1 = list(uselocs)
+        block_0 = blocks[(useloc_0.block_addr, useloc_0.block_idx)]
+        stmt_0 = block_0.statements[useloc_0.stmt_idx]
+        block_1 = blocks[(useloc_1.block_addr, useloc_1.block_idx)]
+        stmt_1 = block_1.statements[useloc_1.stmt_idx]
+
+        if isinstance(stmt_0, Jump):
+            stmt_0, stmt_1 = stmt_1, stmt_0
+            block_0, block_1 = block_1, block_0
+        if not isinstance(stmt_0, ConditionalJump) or not isinstance(stmt_1, Jump):
+            return False
+
+        # check if stmt_0 jumps to block_1
+        if not isinstance(stmt_0.true_target, Const) or not isinstance(stmt_0.false_target, Const):
+            return False
+        stmt_0_targets = {
+            (stmt_0.true_target.value, stmt_0.true_target_idx),
+            (stmt_0.false_target.value, stmt_0.false_target_idx),
+        }
+        return (block_1.addr, block_1.idx) in stmt_0_targets
+
+    @staticmethod
+    def vvar_dep_graph(blocks, vvar_def_locs, vvar_use_locs) -> networkx.DiGraph:
+        g = networkx.DiGraph()
+
+        for var_id in vvar_def_locs:
+            # where is it used?
+            for _, use_loc in vvar_use_locs[var_id]:
+                if isinstance(use_loc, ExternalCodeLocation):
+                    g.add_edge(var_id, "ExternalCodeLocation")
+                    continue
+                assert use_loc.block_addr is not None
+                assert use_loc.stmt_idx is not None
+                block = blocks[(use_loc.block_addr, use_loc.block_idx)]
+                stmt = block.statements[use_loc.stmt_idx]
+                if isinstance(stmt, Assignment):
+                    if isinstance(stmt.dst, VirtualVariable):
+                        g.add_edge(var_id, stmt.dst.varid)
+                    else:
+                        g.add_edge(var_id, f"Assignment@{stmt.ins_addr:#x}")
+                elif isinstance(stmt, Store):
+                    # store to memory
+                    g.add_edge(var_id, f"Store@{stmt.ins_addr:#x}")
+                else:
+                    # other statements
+                    g.add_edge(var_id, f"{stmt.__class__.__name__}@{stmt.ins_addr:#x}")
+
+        return g
+
+
+register_analysis(SPropagatorAnalysis, "SPropagator")

angr/analyses/s_reaching_definitions/__init__.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from .s_rda_model import SRDAModel
+from .s_rda_view import SRDAView
+from .s_reaching_definitions import SReachingDefinitionsAnalysis
+
+
+__all__ = (
+    "SRDAModel",
+    "SRDAView",
+    "SReachingDefinitionsAnalysis",
+)

angr/analyses/s_reaching_definitions/s_rda_model.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Generator
+from typing import Any, Literal, overload
+
+from angr.ailment.expression import VirtualVariable, Tmp
+
+from angr.knowledge_plugins.key_definitions import atoms, Definition
+from angr.code_location import CodeLocation
+
+
+class SRDAModel:
+    """
+    The model for SRDA.
+    """
+
+    def __init__(self, func_graph, func_args, arch):
+        self.func_graph = func_graph
+        self.func_args = func_args
+        self.arch = arch
+        self.varid_to_vvar: dict[int, VirtualVariable] = {}
+        self.all_vvar_definitions: dict[int, CodeLocation] = {}
+        self.all_vvar_uses: dict[int, list[tuple[VirtualVariable | None, CodeLocation]]] = defaultdict(list)
+        self.all_tmp_definitions: dict[CodeLocation, dict[atoms.Tmp, int]] = defaultdict(dict)
+        self.all_tmp_uses: dict[CodeLocation, dict[atoms.Tmp, set[tuple[Tmp, int]]]] = defaultdict(dict)
+        self.phi_vvar_ids: set[int] = set()
+        self.phivarid_to_varids_with_unknown: dict[int, set[int | None]] = {}
+        self.phivarid_to_varids: dict[int, set[int]] = {}
+        self.vvar_uses_by_loc: dict[CodeLocation, list[int]] = {}
+
+    def add_vvar_use(self, vvar_id: int, expr: VirtualVariable | None, loc: CodeLocation) -> None:
+        self.all_vvar_uses[vvar_id].append((expr, loc))
+        if loc not in self.vvar_uses_by_loc:
+            self.vvar_uses_by_loc[loc] = []
+        self.vvar_uses_by_loc[loc].append(vvar_id)
+
+    @property
+    def all_definitions(self) -> Generator[Definition]:
+        for vvar_id, defloc in self.all_vvar_definitions.items():
+            vvar = self.varid_to_vvar[vvar_id]
+            yield Definition(atoms.VirtualVariable(vvar_id, vvar.size, vvar.category, vvar.oident), defloc)
+
+    def is_phi_vvar_id(self, idx: int) -> bool:
+        return idx in self.phi_vvar_ids
+
+    def get_all_definitions(self, block_loc: CodeLocation) -> set[Definition]:
+        s = set()
+        for vvar_id, codeloc in self.all_vvar_definitions.items():
+            vvar = self.varid_to_vvar[vvar_id]
+            if codeloc.block_addr == block_loc.block_addr and codeloc.block_idx == block_loc.block_idx:
+                s.add(Definition(atoms.VirtualVariable(vvar_id, vvar.size, vvar.category, vvar.oident), codeloc))
+        return s | self.get_all_tmp_definitions(block_loc)
+
+    def get_all_tmp_definitions(self, block_loc: CodeLocation) -> set[Definition]:
+        s = set()
+        for tmp_atom, stmt_idx in self.all_tmp_definitions[block_loc].items():
+            s.add(Definition(tmp_atom, CodeLocation(block_loc.block_addr, stmt_idx, block_idx=block_loc.block_idx)))
+        return s
+
+    @overload
+    def get_uses_by_location(self, loc: CodeLocation, exprs: Literal[True]) -> set[tuple[Definition, Any | None]]: ...
+
+    @overload
+    def get_uses_by_location(self, loc: CodeLocation, exprs: Literal[False] = ...) -> set[Definition]: ...
+
+    def get_uses_by_location(
+        self, loc: CodeLocation, exprs: bool = False
+    ) -> set[Definition] | set[tuple[Definition, Any | None]]:
+        """
+        Retrieve all definitions that are used at a given location.
+
+        :param loc:     The code location.
+        :return:        A set of definitions that are used at the given location.
+        """
+        if exprs:
+            def_with_exprs: set[tuple[Definition, Any]] = set()
+            if loc not in self.vvar_uses_by_loc:
+                return def_with_exprs
+            for vvar_id in self.vvar_uses_by_loc[loc]:
+                vvar = self.varid_to_vvar[vvar_id]
+                def_with_exprs.add(
+                    (
+                        Definition(
+                            atoms.VirtualVariable(vvar_id, vvar.size, vvar.category, vvar.oident),
+                            self.all_vvar_definitions[vvar_id],
+                        ),
+                        vvar,
+                    )
+                )
+            return def_with_exprs
+
+        defs: set[Definition] = set()
+        if loc not in self.vvar_uses_by_loc:
+            return defs
+        for vvar_id in self.vvar_uses_by_loc[loc]:
+            vvar = self.varid_to_vvar[vvar_id]
+            defs.add(
+                Definition(
+                    atoms.VirtualVariable(vvar_id, vvar.size, vvar.category, vvar.oident),
+                    self.all_vvar_definitions[vvar_id],
+                )
+            )
+        return defs
+
+    def get_vvar_uses(self, obj: VirtualVariable | atoms.VirtualVariable) -> set[CodeLocation]:
+        if obj.varid in self.all_vvar_uses:
+            return {loc for _, loc in self.all_vvar_uses[obj.varid]}
+        return set()
+
+    def get_vvar_uses_with_expr(
+        self, obj: VirtualVariable | atoms.VirtualVariable
+    ) -> set[tuple[VirtualVariable | None, CodeLocation]]:
+        if obj.varid in self.all_vvar_uses:
+            return set(self.all_vvar_uses[obj.varid])
+        return set()
+
+    def get_tmp_uses(self, obj: atoms.Tmp, block_loc: CodeLocation) -> set[CodeLocation]:
+        if block_loc not in self.all_tmp_uses:
+            return set()
+        if obj not in self.all_tmp_uses[block_loc]:
+            return set()
+        s = set()
+        for _, stmt_idx in self.all_tmp_uses[block_loc][obj]:
+            s.add(CodeLocation(block_loc.block_addr, stmt_idx, block_idx=block_loc.block_idx))
+        return s
+
+    def get_uses_by_def(self, def_: Definition) -> set[CodeLocation]:
+        if isinstance(def_.atom, atoms.Tmp):
+            return self.get_tmp_uses(
+                def_.atom,
+                CodeLocation(def_.codeloc.block_addr, def_.codeloc.stmt_idx, block_idx=def_.codeloc.block_idx),
+            )
+        if isinstance(def_.atom, atoms.VirtualVariable):
+            return self.get_vvar_uses(def_.atom)
+        return set()