angr 9.2.139__py3-none-manylinux2014_x86_64.whl → 9.2.141__py3-none-manylinux2014_x86_64.whl

angr/analyses/reaching_definitions/function_handler.py

@@ -16,7 +16,6 @@ from angr.calling_conventions import SimCC
 from angr.sim_type import SimTypeFunction
 from angr.knowledge_plugins.key_definitions.definition import Definition
 from angr.knowledge_plugins.functions import Function
-from angr.analyses.reaching_definitions.dep_graph import FunctionCallRelationships
 from angr.code_location import CodeLocation, ExternalCodeLocation
 from angr.knowledge_plugins.key_definitions.constants import ObservationPointType
 from angr import SIM_LIBRARIES, SIM_TYPE_COLLECTIONS
@@ -246,7 +245,7 @@ class FunctionCallDataUnwrapped(FunctionCallData):
     @staticmethod
     @wraps
     def decorate(
-        f: Callable[[FunctionHandler, ReachingDefinitionsState, FunctionCallDataUnwrapped], None]
+        f: Callable[[FunctionHandler, ReachingDefinitionsState, FunctionCallDataUnwrapped], None],
     ) -> Callable[[FunctionHandler, ReachingDefinitionsState, FunctionCallData], None]:
         """
         Decorate a function handler method with this to make it take a FunctionCallDataUnwrapped instead of a
@@ -263,6 +262,20 @@ def _mk_wrapper(func, iself):
     return lambda *args, **kwargs: func(iself, *args, **kwargs)
 
 
+@dataclass
+class FunctionCallRelationships:
+    """
+    Produced by the function handler, provides associated callsite info and function input/output definitions.
+    """
+
+    callsite: CodeLocation
+    target: int | None
+    args_defns: list[set[Definition]]
+    other_input_defns: set[Definition]
+    ret_defns: set[Definition]
+    other_output_defns: set[Definition]
+
+
 # pylint: disable=unused-argument, no-self-use
 class FunctionHandler:
     """

angr/analyses/reaching_definitions/rd_state.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Any, TYPE_CHECKING, cast, overload
+from typing import Any, TYPE_CHECKING, overload
 from collections.abc import Iterable, Iterator
 import logging
 from typing_extensions import Self
@@ -7,7 +7,6 @@ from typing_extensions import Self
 import archinfo
 import claripy
 
-from angr.misc.ux import deprecated
 from angr.knowledge_plugins.key_definitions.environment import Environment
 from angr.knowledge_plugins.key_definitions.tag import Tag
 from angr.knowledge_plugins.key_definitions.heap_address import HeapAddress
@@ -541,41 +540,6 @@ class ReachingDefinitionsState:
         self.all_definitions = set()
         self.live_definitions.reset_uses()
 
-    @deprecated("deref")
-    def pointer_to_atoms(self, pointer: MultiValues, size: int, endness: str) -> set[MemoryLocation]:
-        """
-        Given a MultiValues, return the set of atoms that loading or storing to the pointer with that value
-        could define or use.
-        """
-        result = set()
-        for vs in pointer.values():
-            for value in vs:
-                atom = self.pointer_to_atom(value, size, endness)
-                if atom is not None:
-                    result.add(atom)
-
-        return result
-
-    @deprecated("deref")
-    def pointer_to_atom(self, value: claripy.ast.BV, size: int, endness: str) -> MemoryLocation | None:
-        if self.is_top(value):
-            return None
-
-        stack_offset = self.get_stack_offset(value)
-        if stack_offset is not None:
-            addr = SpOffset(len(value), stack_offset)
-        else:
-            heap_offset = self.get_heap_offset(value)
-            if heap_offset is not None:
-                addr = HeapAddress(heap_offset)
-            elif value.op == "BVV":
-                addr = cast(int, value.args[0])
-            else:
-                # cannot resolve
-                return None
-
-        return MemoryLocation(addr, size, endness)
-
     @overload
     def deref(
         self,

angr/analyses/reaching_definitions/reaching_definitions.py

@@ -16,7 +16,6 @@ from angr.knowledge_plugins.functions import Function
 from angr.knowledge_plugins.key_definitions import ReachingDefinitionsModel, LiveDefinitions
 from angr.knowledge_plugins.key_definitions.constants import OP_BEFORE, OP_AFTER, ObservationPointType, ObservationPoint
 from angr.code_location import CodeLocation, ExternalCodeLocation
-from angr.misc.ux import deprecated
 from angr.analyses.forward_analysis.visitors.graph import NodeType
 from angr.analyses.analysis import Analysis
 from .engine_ail import SimEngineRDAIL
@@ -50,13 +49,13 @@ class ReachingDefinitionsAnalysis(
 
     def __init__(
         self,
-        subject: Subject | ailment.Block | Block | Function | str = None,
+        subject: Subject | ailment.Block | Block | Function | str,
         func_graph=None,
         max_iterations=30,
         track_tmps=False,
         track_consts=True,
         observation_points: Iterable[ObservationPoint] | None = None,
-        init_state: ReachingDefinitionsState = None,
+        init_state: ReachingDefinitionsState | None = None,
         init_context=None,
         state_initializer: RDAStateInitializer | None = None,
         cc=None,
@@ -242,10 +241,6 @@ class ReachingDefinitionsAnalysis(
     def visited_blocks(self):
         return self._visited_blocks
 
-    @deprecated(replacement="get_reaching_definitions_by_insn")
-    def get_reaching_definitions(self, ins_addr, op_type):
-        return self.get_reaching_definitions_by_insn(ins_addr, op_type)
-
     def get_reaching_definitions_by_insn(self, ins_addr, op_type):
         key = "insn", ins_addr, op_type
         if key not in self.observed_results:
@@ -280,29 +275,22 @@ class ReachingDefinitionsAnalysis(
         :param node_idx:    ID of the node. Used in AIL to differentiate blocks with the same address.
         """
 
-        key = None
-
+        key: ObservationPoint | None = None
         observe = False
 
         if self._observe_all:
             observe = True
-            key: ObservationPoint = (
-                ("node", node_addr, op_type) if node_idx is None else ("node", (node_addr, node_idx), op_type)
-            )
+            key = ("node", node_addr, op_type) if node_idx is None else ("node", (node_addr, node_idx), op_type)
         elif self._observation_points is not None:
-            key: ObservationPoint = (
-                ("node", node_addr, op_type) if node_idx is None else ("node", (node_addr, node_idx), op_type)
-            )
+            key = ("node", node_addr, op_type) if node_idx is None else ("node", (node_addr, node_idx), op_type)
             if key in self._observation_points:
                 observe = True
         elif self._observe_callback is not None:
             observe = self._observe_callback("node", addr=node_addr, state=state, op_type=op_type, node_idx=node_idx)
             if observe:
-                key: ObservationPoint = (
-                    ("node", node_addr, op_type) if node_idx is None else ("node", (node_addr, node_idx), op_type)
-                )
+                key = ("node", node_addr, op_type) if node_idx is None else ("node", (node_addr, node_idx), op_type)
 
-        if observe:
+        if observe and key:
             self.observed_results[key] = state.live_definitions
 
     def insn_observe(
@@ -321,14 +309,14 @@ class ReachingDefinitionsAnalysis(
         :param op_type:     Type of the observation point. Must be one of the following: OP_BEORE, OP_AFTER.
         """
 
-        key = None
+        key: ObservationPoint | None = None
         observe = False
 
         if self._observe_all:
             observe = True
-            key: ObservationPoint = ("insn", insn_addr, op_type)
+            key = ("insn", insn_addr, op_type)
         elif self._observation_points is not None:
-            key: ObservationPoint = ("insn", insn_addr, op_type)
+            key = ("insn", insn_addr, op_type)
             if key in self._observation_points:
                 observe = True
         elif self._observe_callback is not None:
@@ -336,9 +324,9 @@ class ReachingDefinitionsAnalysis(
                 "insn", addr=insn_addr, stmt=stmt, block=block, state=state, op_type=op_type
             )
             if observe:
-                key: ObservationPoint = ("insn", insn_addr, op_type)
+                key = ("insn", insn_addr, op_type)
 
-        if not observe:
+        if not (observe and key):
             return
 
         if isinstance(stmt, pyvex.stmt.IRStmt):
@@ -533,6 +521,7 @@ class ReachingDefinitionsAnalysis(
             ]
             if node.addr == self.subject.content.addr:
                 node_parents += [ExternalCodeLocation()]
+            assert block is not None
             self.model.at_new_block(
                 CodeLocation(block.addr, 0, block_idx=block.idx if isinstance(block, ailment.Block) else None),
                 node_parents,

angr/analyses/s_propagator.py

@@ -35,6 +35,7 @@ from angr.utils.ssa import (
     get_tmp_uselocs,
     get_tmp_deflocs,
     phi_assignment_get_src,
+    has_store_stmt_in_between_stmts,
 )
 
 
@@ -186,6 +187,8 @@ class SPropagatorAnalysis(Analysis):
 
         # function mode only
         if self.mode == "function":
+            assert self.func_graph is not None
+
             for vvar, defloc in vvar_deflocs.items():
                 if vvar.varid not in vvar_uselocs:
                     continue
@@ -213,7 +216,7 @@ class SPropagatorAnalysis(Analysis):
                     #    }
                     can_replace = True
                     for _, vvar_useloc in vvar_uselocs[vvar.varid]:
-                        if self.has_store_stmt_in_between(blocks, defloc, vvar_useloc):
+                        if has_store_stmt_in_between_stmts(self.func_graph, blocks, defloc, vvar_useloc):
                             can_replace = False
 
                     if can_replace:
@@ -241,8 +244,8 @@ class SPropagatorAnalysis(Analysis):
                 if vvar.was_reg or vvar.was_parameter:
                     if len(vvar_uselocs[vvar.varid]) == 1:
                         vvar_used, vvar_useloc = next(iter(vvar_uselocs[vvar.varid]))
-                        if is_const_vvar_load_assignment(stmt) and not self.has_store_stmt_in_between(
-                            blocks, defloc, vvar_useloc
+                        if is_const_vvar_load_assignment(stmt) and not has_store_stmt_in_between_stmts(
+                            self.func_graph, blocks, defloc, vvar_useloc
                         ):
                             # we can propagate this load because there is no store between its def and use
                             replacements[vvar_useloc][vvar_used] = stmt.src
@@ -462,44 +465,6 @@ class SPropagatorAnalysis(Analysis):
 
         return False
 
-    def has_store_stmt_in_between(
-        self, blocks: dict[tuple[int, int | None], Block], defloc: CodeLocation, useloc: CodeLocation
-    ) -> bool:
-        assert defloc.block_addr is not None
-        assert defloc.stmt_idx is not None
-        assert useloc.block_addr is not None
-        assert useloc.stmt_idx is not None
-        assert self.func_graph is not None
-
-        use_block = blocks[(useloc.block_addr, useloc.block_idx)]
-        def_block = blocks[(defloc.block_addr, defloc.block_idx)]
-
-        # traverse the graph, go from use_block until we reach def_block, and look for Store statements
-        seen = {use_block}
-        queue = [use_block]
-        while queue:
-            block = queue.pop(0)
-
-            starting_stmt_idx, ending_stmt_idx = 0, len(block.statements)
-            if block is def_block:
-                starting_stmt_idx = defloc.stmt_idx + 1
-            if block is use_block:
-                ending_stmt_idx = useloc.stmt_idx
-
-            for i in range(starting_stmt_idx, ending_stmt_idx):
-                if isinstance(block.statements[i], Store):
-                    return True
-
-            if block is def_block:
-                continue
-
-            for pred in self.func_graph.predecessors(block):
-                if pred not in seen:
-                    seen.add(pred)
-                    queue.append(pred)
-
-        return False
-
     @staticmethod
     def is_vvar_used_for_addr_loading_switch_case(uselocs: set[CodeLocation], blocks) -> bool:
         """

angr/analyses/s_reaching_definitions/s_rda_model.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from collections import defaultdict
 from collections.abc import Generator
-from typing import Any
+from typing import Any, Literal, overload
 
 from ailment.expression import VirtualVariable, Tmp
 
@@ -48,6 +48,12 @@ class SRDAModel:
             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]]:

angr/analyses/s_reaching_definitions/s_rda_view.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Callable
 from collections import defaultdict
 
+from ailment import Block
 from ailment.statement import Statement, Assignment, Call, Label
 from ailment.expression import VirtualVariable, VirtualVariableCategory, Expression
 
@@ -22,7 +24,7 @@ class RegVVarPredicate:
     Implements a predicate that is used in get_reg_vvar_by_stmt_idx and get_reg_vvar_by_insn.
     """
 
-    def __init__(self, reg_offset: int, vvars: set[VirtualVariable], arch):
+    def __init__(self, reg_offset: int, vvars: list[VirtualVariable], arch):
         self.reg_offset = reg_offset
         self.vvars = vvars
         self.arch = arch
@@ -47,7 +49,8 @@ class RegVVarPredicate:
             and stmt.dst.was_reg
             and stmt.dst.reg_offset == self.reg_offset
         ):
-            self.vvars.add(stmt.dst)
+            if stmt.dst not in self.vvars:
+                self.vvars.append(stmt.dst)
             return True
         if isinstance(stmt, Call):
             if (
@@ -55,7 +58,8 @@ class RegVVarPredicate:
                 and stmt.ret_expr.was_reg
                 and stmt.ret_expr.reg_offset == self.reg_offset
             ):
-                self.vvars.add(stmt.ret_expr)
+                if stmt.ret_expr not in self.vvars:
+                    self.vvars.append(stmt.ret_expr)
                 return True
             # is it clobbered maybe?
             clobbered_regs = self._get_call_clobbered_regs(stmt)
@@ -69,7 +73,7 @@ class StackVVarPredicate:
     Implements a predicate that is used in get_stack_vvar_by_stmt_idx and get_stack_vvar_by_insn.
     """
 
-    def __init__(self, stack_offset: int, size: int, vvars: set[VirtualVariable]):
+    def __init__(self, stack_offset: int, size: int, vvars: list[VirtualVariable]):
         self.stack_offset = stack_offset
         self.size = size
         self.vvars = vvars
@@ -82,7 +86,8 @@ class StackVVarPredicate:
             and stmt.dst.stack_offset <= self.stack_offset < stmt.dst.stack_offset + stmt.dst.size
             and stmt.dst.stack_offset <= self.stack_offset + self.size <= stmt.dst.stack_offset + stmt.dst.size
         ):
-            self.vvars.add(stmt.dst)
+            if stmt.dst not in self.vvars:
+                self.vvars.append(stmt.dst)
             return True
         return False
 
@@ -96,7 +101,13 @@ class SRDAView:
         self.model = model
 
     def _get_vvar_by_stmt(
-        self, block_addr: int, block_idx: int | None, stmt_idx: int, op_type: ObservationPointType, predicate
+        self,
+        block_addr: int,
+        block_idx: int | None,
+        stmt_idx: int,
+        op_type: ObservationPointType,
+        predicate: Callable,
+        consecutive: bool = False,
     ):
         # find the starting block
         for block in self.model.func_graph:
@@ -107,7 +118,10 @@ class SRDAView:
             return
 
         traversed = set()
-        queue = [(the_block, stmt_idx if op_type == ObservationPointType.OP_BEFORE else stmt_idx + 1)]
+        queue: list[tuple[Block, int | None]] = [
+            (the_block, stmt_idx if op_type == ObservationPointType.OP_BEFORE else stmt_idx + 1)
+        ]
+        predicate_returned_true = False
         while queue:
             block, start_stmt_idx = queue.pop(0)
             traversed.add(block)
@@ -115,7 +129,8 @@ class SRDAView:
             stmts = block.statements[:start_stmt_idx] if start_stmt_idx is not None else block.statements
 
             for stmt in reversed(stmts):
-                should_break = predicate(stmt)
+                r = predicate(stmt)
+                should_break = (predicate_returned_true and r is False) if consecutive else r
                 if should_break:
                     break
             else:
@@ -129,7 +144,7 @@ class SRDAView:
         self, reg_offset: int, block_addr: int, block_idx: int | None, stmt_idx: int, op_type: ObservationPointType
     ) -> VirtualVariable | None:
         reg_offset = get_reg_offset_base(reg_offset, self.model.arch)
-        vvars = set()
+        vvars = []
         predicater = RegVVarPredicate(reg_offset, vvars, self.model.arch)
         self._get_vvar_by_stmt(block_addr, block_idx, stmt_idx, op_type, predicater.predicate)
 
@@ -137,14 +152,14 @@ class SRDAView:
             # not found - check function arguments
             for func_arg in self.model.func_args:
                 if isinstance(func_arg, VirtualVariable):
-                    func_arg_category = func_arg.oident[0]
+                    func_arg_category = func_arg.parameter_category
                     if func_arg_category == VirtualVariableCategory.REGISTER:
-                        func_arg_regoff = func_arg.oident[1]
+                        func_arg_regoff = func_arg.parameter_reg_offset
                         if func_arg_regoff == reg_offset:
-                            vvars.add(func_arg)
+                            vvars.append(func_arg)
 
         assert len(vvars) <= 1
-        return next(iter(vvars), None)
+        return vvars[0] if vvars else None
 
     def get_stack_vvar_by_stmt(  # pylint: disable=too-many-positional-arguments
         self,
@@ -155,21 +170,24 @@ class SRDAView:
         stmt_idx: int,
         op_type: ObservationPointType,
     ) -> VirtualVariable | None:
-        vvars = set()
+        vvars = []
         predicater = StackVVarPredicate(stack_offset, size, vvars)
-        self._get_vvar_by_stmt(block_addr, block_idx, stmt_idx, op_type, predicater.predicate)
+        self._get_vvar_by_stmt(block_addr, block_idx, stmt_idx, op_type, predicater.predicate, consecutive=True)
 
         if not vvars:
             # not found - check function arguments
             for func_arg in self.model.func_args:
                 if isinstance(func_arg, VirtualVariable):
-                    func_arg_category = func_arg.oident[0]
+                    func_arg_category = func_arg.parameter_category
                     if func_arg_category == VirtualVariableCategory.STACK:
-                        func_arg_stackoff = func_arg.oident[1]
+                        func_arg_stackoff = func_arg.oident[1]  # type: ignore
                         if func_arg_stackoff == stack_offset and func_arg.size == size:
-                            vvars.add(func_arg)
-        assert len(vvars) <= 1
-        return next(iter(vvars), None)
+                            vvars.append(func_arg)
+        # there might be multiple vvars; we prioritize the one whose size fits the best
+        for v in vvars:
+            if v.stack_offset == stack_offset and v.size == size:
+                return v
+        return vvars[0] if vvars else None
 
     def _get_vvar_by_insn(self, addr: int, op_type: ObservationPointType, predicate, block_idx: int | None = None):
         # find the starting block
@@ -202,23 +220,23 @@ class SRDAView:
         self, reg_offset: int, addr: int, op_type: ObservationPointType, block_idx: int | None = None
     ) -> VirtualVariable | None:
         reg_offset = get_reg_offset_base(reg_offset, self.model.arch)
-        vvars = set()
+        vvars = []
         predicater = RegVVarPredicate(reg_offset, vvars, self.model.arch)
 
         self._get_vvar_by_insn(addr, op_type, predicater.predicate, block_idx=block_idx)
 
         assert len(vvars) <= 1
-        return next(iter(vvars), None)
+        return vvars[0] if vvars else None
 
     def get_stack_vvar_by_insn(  # pylint: disable=too-many-positional-arguments
         self, stack_offset: int, size: int, addr: int, op_type: ObservationPointType, block_idx: int | None = None
     ) -> VirtualVariable | None:
-        vvars = set()
+        vvars = []
         predicater = StackVVarPredicate(stack_offset, size, vvars)
         self._get_vvar_by_insn(addr, op_type, predicater.predicate, block_idx=block_idx)
 
         assert len(vvars) <= 1
-        return next(iter(vvars), None)
+        return vvars[0] if vvars else None
 
     def get_vvar_value(self, vvar: VirtualVariable) -> Expression | None:
         if vvar not in self.model.all_vvar_definitions:
@@ -227,7 +245,7 @@ class SRDAView:
 
         for block in self.model.func_graph:
             if block.addr == codeloc.block_addr and block.idx == codeloc.block_idx:
-                if codeloc.stmt_idx < len(block.statements):
+                if codeloc.stmt_idx is not None and codeloc.stmt_idx < len(block.statements):
                     stmt = block.statements[codeloc.stmt_idx]
                     if isinstance(stmt, Assignment) and stmt.dst.likes(vvar):
                         return stmt.src

angr/analyses/stack_pointer_tracker.py

@@ -22,6 +22,7 @@ try:
     from angr.engines import pcode
 except ImportError:
     pypcode = None
+    pcode = None
 
 if TYPE_CHECKING:
     from angr.block import Block
@@ -93,6 +94,11 @@ class Register:
             return self.offset == other.offset
         return False
 
+    def __add__(self, other) -> OffsetVal:
+        if type(other) is Constant:
+            return OffsetVal(self, other.val)
+        raise CouldNotResolveException
+
     def __repr__(self):
         return str(self.offset)
 
@@ -232,6 +238,7 @@ class StackPointerTrackerState:
     def give_up_on_memory_tracking(self):
         self.memory = {}
         self.is_tracking_memory = False
+        return self
 
     def store(self, addr, val):
         # strong update
@@ -370,7 +377,8 @@ class StackPointerTracker(Analysis, ForwardAnalysis):
         self._mem_merge_cache = {}
 
         if initial_reg_values:
-            self._reg_value_at_block_start[func.addr if func is not None else block.addr] = initial_reg_values
+            block_start_addr = func.addr if func is not None else block.addr  # type: ignore
+            self._reg_value_at_block_start[block_start_addr] = initial_reg_values
 
         _l.debug("Running on function %r", self._func)
         self._analyze()
@@ -461,9 +469,13 @@ class StackPointerTracker(Analysis, ForwardAnalysis):
         return any(self.inconsistent_for(r) for r in self.reg_offsets)
 
     def inconsistent_for(self, reg):
+        if self._func is None:
+            raise ValueError("inconsistent_for() is only supported in function mode")
         return any(self.offset_after_block(endpoint.addr, reg) is TOP for endpoint in self._func.endpoints)
 
     def offsets_for(self, reg):
+        if self._func is None:
+            raise ValueError("offsets_for() is only supported in function mode")
         return [
             o for block in self._func.blocks if (o := self.offset_after_block(block.addr, reg)) not in (TOP, BOTTOM)
         ]
@@ -481,7 +493,7 @@ class StackPointerTracker(Analysis, ForwardAnalysis):
     def _post_analysis(self):
         pass
 
-    def _get_register(self, offset):
+    def _get_register(self, offset) -> Register:
         name = self.project.arch.register_names[offset]
         size = self.project.arch.registers[name][1]
         return Register(offset, size * self.project.arch.byte_width)
@@ -557,7 +569,7 @@ class StackPointerTracker(Analysis, ForwardAnalysis):
         output_state = state.freeze()
         return None, output_state
 
-    def _process_vex_irsb(self, node, vex_block: pyvex.IRSB, state: StackPointerTrackerState) -> int:
+    def _process_vex_irsb(self, node, vex_block: pyvex.IRSB, state: StackPointerTrackerState) -> int | None:
         tmps = {}
         curr_stmt_start_addr = None
 
@@ -704,21 +716,16 @@ class StackPointerTracker(Analysis, ForwardAnalysis):
             if callees:
                 if len(callees) == 1:
                     callee = callees[0]
+                    track_rax = False
+                    if (
+                        (callee.info.get("is_rust_probestack", False) and self.project.arch.name == "AMD64")
+                        or (callee.info.get("is_alloca_probe", False) and self.project.arch.name == "AMD64")
+                        or callee.name == "__chkstk"
+                    ):
+                        # sp = sp - rax right after returning from the call
+                        track_rax = True
 
-                    if callee.info.get("is_rust_probestack", False) is True and self.project.arch.name == "AMD64":
-                        # special-case for rust_probestack: sp = sp - rax right after returning from the call, so we
-                        # need to keep track of rax
-                        for stmt in reversed(vex_block.statements):
-                            if (
-                                isinstance(stmt, pyvex.IRStmt.Put)
-                                and stmt.offset == self.project.arch.registers["rax"][0]
-                                and isinstance(stmt.data, pyvex.IRExpr.Const)
-                            ):
-                                state.put(stmt.offset, Constant(stmt.data.con.value), force=True)
-                                break
-                    elif callee.name == "__chkstk":
-                        # special-case for __chkstk: sp = sp - rax right after returning from the call, so we need to
-                        # keep track of rax
+                    if track_rax:
                         for stmt in reversed(vex_block.statements):
                             if (
                                 isinstance(stmt, pyvex.IRStmt.Put)
@@ -737,18 +744,20 @@ class StackPointerTracker(Analysis, ForwardAnalysis):
                     # found callee clean-up cases...
                     try:
                         v = state.get(self.project.arch.sp_offset)
+                        incremented = None
                         if v is BOTTOM:
                             incremented = BOTTOM
                         elif callee_cleanups[0].prototype is not None:
                             num_args = len(callee_cleanups[0].prototype.args)
                             incremented = v + Constant(self.project.arch.bytes * num_args)
-                        state.put(self.project.arch.sp_offset, incremented)
+                        if incremented is not None:
+                            state.put(self.project.arch.sp_offset, incremented)
                     except CouldNotResolveException:
                         pass
 
         return curr_stmt_start_addr
 
-    def _process_pcode_irsb(self, node, pcode_irsb: pcode.lifter.IRSB, state: StackPointerTrackerState) -> int:
+    def _process_pcode_irsb(self, node, pcode_irsb: pcode.lifter.IRSB, state: StackPointerTrackerState) -> int | None:
         unique = {}
         curr_stmt_start_addr = None
 
@@ -830,18 +839,20 @@ class StackPointerTracker(Analysis, ForwardAnalysis):
                     # found callee clean-up cases...
                     try:
                         v = state.get(self.project.arch.sp_offset)
+                        incremented = None
                         if v is BOTTOM:
                             incremented = BOTTOM
                         elif callee_cleanups[0].prototype is not None:
                             num_args = len(callee_cleanups[0].prototype.args)
                             incremented = v + Constant(self.project.arch.bytes * num_args)
-                        state.put(self.project.arch.sp_offset, incremented)
+                        if incremented is not None:
+                            state.put(self.project.arch.sp_offset, incremented)
                     except CouldNotResolveException:
                         pass
 
         return curr_stmt_start_addr
 
-    def _widen_states(self, *states):
+    def _widen_states(self, *states: FrozenStackPointerTrackerState):
         assert len(states) == 2
         merged, _ = self._merge_states(None, *states)
         if len(merged.memory) > 5:
@@ -849,13 +860,16 @@ class StackPointerTracker(Analysis, ForwardAnalysis):
             merged = merged.unfreeze().give_up_on_memory_tracking().freeze()
         return merged
 
-    def _merge_states(self, node, *states: StackPointerTrackerState):
+    def _merge_states(self, node, *states: FrozenStackPointerTrackerState):
         merged_state = states[0]
         for other in states[1:]:
             merged_state = merged_state.merge(other, node.addr, self._reg_merge_cache, self._mem_merge_cache)
         return merged_state, merged_state == states[0]
 
     def _find_callees(self, node) -> list[Function]:
+        if self._func is None:
+            raise ValueError("find_callees() is only supported in function mode")
+
         callees: list[Function] = []
         for _, dst, data in self._func.transition_graph.out_edges(node, data=True):
             if data.get("type") == "call" and isinstance(dst, Function):