angr 9.2.159__cp310-abi3-macosx_11_0_arm64.whl → 9.2.161__cp310-abi3-macosx_11_0_arm64.whl

angr/analyses/variable_recovery/engine_base.py

@@ -525,6 +525,10 @@ class SimEngineVRBase(
     def _store_to_stack(
         self, stack_offset, data: RichR[claripy.ast.BV | claripy.ast.FP], size, offset=0, atom=None, endness=None
     ):
+        """
+        Store data to a stack location. We limit the size of the data to store to 256 bytes for performance reasons.
+        """
+
         if atom is None:
             existing_vars = self.state.variable_manager[self.func_addr].find_variables_by_stmt(
                 self.block.addr, self.stmt_idx, "memory"
@@ -550,7 +554,11 @@ class SimEngineVRBase(
             variable, variable_offset = next(iter(existing_vars))
 
         if isinstance(stack_offset, int):
-            expr = self.state.annotate_with_variables(data.data, [(variable_offset, variable)])
+            expr = data.data
+            if isinstance(expr, claripy.ast.Bits) and expr.size() > 1024:
+                # we don't write more than 256 bytes to the stack at a time for performance reasons
+                expr = expr[expr.size() - 1 : expr.size() - 1024]
+            expr = self.state.annotate_with_variables(expr, [(variable_offset, variable)])
             stack_addr = self.state.stack_addr_from_offset(stack_offset)
             self.state.stack_region.store(stack_addr, expr, endness=endness)
 

angr/analyses/variable_recovery/variable_recovery_base.py

@@ -5,6 +5,8 @@ from collections.abc import Generator, Iterable
 import logging
 from collections import defaultdict
 
+import networkx
+
 import archinfo
 import claripy
 from claripy.annotation import Annotation
@@ -86,8 +88,18 @@ class VariableRecoveryBase(Analysis):
     The base class for VariableRecovery and VariableRecoveryFast.
     """
 
-    def __init__(self, func, max_iterations, store_live_variables: bool, vvar_to_vvar: dict[int, int] | None = None):
+    def __init__(
+        self,
+        func,
+        max_iterations,
+        store_live_variables: bool,
+        vvar_to_vvar: dict[int, int] | None = None,
+        func_graph: networkx.DiGraph | None = None,
+        entry_node_addr: int | tuple[int, int | None] | None = None,
+    ):
         self.function = func
+        self.func_graph = func_graph
+        self.entry_node_addr = entry_node_addr
         self.variable_manager = self.kb.variables
 
         self._max_iterations = max_iterations
@@ -120,7 +132,23 @@ class VariableRecoveryBase(Analysis):
 
     def initialize_dominance_frontiers(self):
         # Computer the dominance frontier for each node in the graph
-        df = self.project.analyses.DominanceFrontier(self.function)
+        func_entry = None
+        if self.func_graph is not None:
+            entry_node_addr = self.entry_node_addr if self.entry_node_addr is not None else self.function.addr
+            assert entry_node_addr is not None
+            if isinstance(entry_node_addr, int):
+                func_entry = next(iter(node for node in self.func_graph if node.addr == entry_node_addr))
+            elif isinstance(entry_node_addr, tuple):
+                func_entry = next(
+                    iter(
+                        node
+                        for node in self.func_graph
+                        if node.addr == entry_node_addr[0] and node.idx == entry_node_addr[1]
+                    )
+                )
+            else:
+                raise TypeError(f"Unsupported entry node address type: {type(entry_node_addr)}")
+        df = self.project.analyses.DominanceFrontier(self.function, func_graph=self.func_graph, entry=func_entry)
         self._dominance_frontiers = defaultdict(set)
         for b0, domfront in df.frontiers.items():
             for d in domfront:

angr/analyses/variable_recovery/variable_recovery_fast.py

@@ -237,6 +237,7 @@ class VariableRecoveryFast(ForwardAnalysis, VariableRecoveryBase):  # pylint:dis
         self,
         func: Function | str | int,
         func_graph: networkx.DiGraph | None = None,
+        entry_node_addr: int | tuple[int, int | None] | None = None,
         max_iterations: int = 2,
         low_priority=False,
         track_sp=True,
@@ -259,10 +260,18 @@ class VariableRecoveryFast(ForwardAnalysis, VariableRecoveryBase):  # pylint:dis
         function_graph_visitor = visitors.FunctionGraphVisitor(func, graph=func_graph)
 
         # Make sure the function is not empty
-        if not func.block_addrs_set or func.startpoint is None:
+        if (not func.block_addrs_set or func.startpoint is None) and not func_graph:
             raise AngrVariableRecoveryError(f"Function {func!r} is empty.")
 
-        VariableRecoveryBase.__init__(self, func, max_iterations, store_live_variables, vvar_to_vvar=vvar_to_vvar)
+        VariableRecoveryBase.__init__(
+            self,
+            func,
+            max_iterations,
+            store_live_variables,
+            vvar_to_vvar=vvar_to_vvar,
+            func_graph=func_graph_with_calls,
+            entry_node_addr=entry_node_addr,
+        )
         ForwardAnalysis.__init__(
             self, order_jobs=True, allow_merging=True, allow_widening=False, graph_visitor=function_graph_visitor
         )

angr/emulator.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import logging
+from enum import Enum
+
+from angr.engines.concrete import ConcreteEngine, HeavyConcreteState
+from angr.errors import AngrError
+
+
+log = logging.getLogger(name=__name__)
+
+
+class EmulatorException(AngrError):
+    """Base class for exceptions raised by the Emulator."""
+
+
+class EngineException(EmulatorException):
+    """Exception raised when the emulator encounters an unhandlable error in the engine."""
+
+
+class StateDivergedException(EmulatorException):
+    """Exception raised when an engine returns multiple successors."""
+
+
+class EmulatorStopReason(Enum):
+    """
+    Enum representing the reason for stopping the emulator.
+    """
+
+    INSTRUCTION_LIMIT = "instruction_limit"
+    BREAKPOINT = "breakpoint"
+    NO_SUCCESSORS = "no_successors"
+    MEMORY_ERROR = "memory_error"
+    FAILURE = "failure"
+    EXIT = "exit"
+
+
+class Emulator:
+    """
+    Emulator is a utility that adapts an angr `ConcreteEngine` to a more
+    user-friendly interface for concrete execution. It only supports concrete
+    execution and requires a ConcreteEngine.
+
+    Saftey: This class is not thread-safe. It should only be used in a
+    single-threaded context. It can be safely shared between multiple threads,
+    provided that only one thread is using it at a time.
+    """
+
+    _engine: ConcreteEngine
+    _state: HeavyConcreteState
+
+    def __init__(self, engine: ConcreteEngine, init_state: HeavyConcreteState):
+        """
+        :param engine: The `ConcreteEngine` to use for emulation.
+        :param init_state: The initial state to use for emulation.
+        """
+        self._engine = engine
+        self._state = init_state
+
+    @property
+    def state(self) -> HeavyConcreteState:
+        """
+        The current state of the emulator.
+        """
+        return self._state
+
+    @property
+    def breakpoints(self) -> set[int]:
+        """
+        The set of currently set breakpoints.
+        """
+        return self._engine.get_breakpoints()
+
+    def add_breakpoint(self, addr: int) -> None:
+        """
+        Add a breakpoint at the given address.
+
+        :param addr: The address to set the breakpoint at.
+        """
+        self._engine.add_breakpoint(addr)
+
+    def remove_breakpoint(self, addr: int) -> None:
+        """
+        Remove a breakpoint at the given address, if present.
+
+        :param addr: The address to remove the breakpoint from.
+        """
+        self._engine.remove_breakpoint(addr)
+
+    def run(self, num_inst: int | None = None) -> EmulatorStopReason:
+        """
+        Execute the emulator.
+        """
+        completed_engine_execs = 0
+        num_inst_executed: int = 0
+        while self._state.history.jumpkind != "Ijk_Exit":
+            # Check if there is a breakpoint at the current address
+            if completed_engine_execs > 0 and self._state.addr in self._engine.get_breakpoints():
+                return EmulatorStopReason.BREAKPOINT
+
+            # Check if we've already executed the requested number of instructions
+            if num_inst is not None and num_inst_executed >= num_inst:
+                return EmulatorStopReason.INSTRUCTION_LIMIT
+
+            # Calculate remaining instructions for this engine execution
+            remaining_inst: int | None = None
+            if num_inst is not None:
+                remaining_inst = num_inst - num_inst_executed
+
+            # Run the engine to get successors
+            try:
+                successors = self._engine.process(self._state, num_inst=remaining_inst)
+            except EngineException as e:
+                raise EngineException(f"Engine encountered an error: {e}") from e
+
+            # Handle cases with an unexpected number of successors
+            if len(successors.successors) == 0:
+                return EmulatorStopReason.NO_SUCCESSORS
+            if len(successors.successors) > 1:
+                log.warning("Concrete engine returned multiple successors")
+
+            # Set the state before raising further exceptions
+            self._state = successors.successors[0]
+
+            # Track the number of instructions executed using the state's history
+            if self._state.history.recent_instruction_count > 0:
+                num_inst_executed += self._state.history.recent_instruction_count
+
+            if successors.successors[0].history.jumpkind == "Ijk_SigSEGV":
+                return EmulatorStopReason.MEMORY_ERROR
+
+            completed_engine_execs += 1
+
+        return EmulatorStopReason.EXIT
+
+
+__all__ = (
+    "Emulator",
+    "EmulatorException",
+    "EmulatorStopReason",
+    "EngineException",
+    "StateDivergedException",
+)

angr/engines/concrete.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import logging
+import typing
+from abc import abstractmethod, ABCMeta
+from typing_extensions import override
+
+import claripy
+from angr.engines.successors import SimSuccessors, SuccessorsEngine
+from angr.sim_state import SimState
+
+log = logging.getLogger(__name__)
+
+
+HeavyConcreteState = SimState[int, int]
+
+
+class ConcreteEngine(SuccessorsEngine, metaclass=ABCMeta):
+    """
+    ConcreteEngine extends SuccessorsEngine and adds APIs for managing breakpoints.
+    """
+
+    @abstractmethod
+    def get_breakpoints(self) -> set[int]:
+        """Return the set of currently set breakpoints."""
+
+    @abstractmethod
+    def add_breakpoint(self, addr: int) -> None:
+        """Add a breakpoint at the given address."""
+
+    @abstractmethod
+    def remove_breakpoint(self, addr: int) -> None:
+        """Remove a breakpoint at the given address, if present."""
+
+    @abstractmethod
+    def process_concrete(self, state: HeavyConcreteState, num_inst: int | None = None) -> HeavyConcreteState:
+        """
+        Process the concrete state and return a HeavyState object.
+
+        :param state: The concrete state to process.
+        :return: A HeavyState object representing the processed state.
+        """
+
+    @override
+    def process_successors(
+        self, successors: SimSuccessors, *, num_inst: int | None = None, **kwargs: dict[str, typing.Any]
+    ):
+        if len(kwargs) > 0:
+            log.warning("ConcreteEngine.process_successors received unknown kwargs: %s", kwargs)
+
+        # TODO: Properly error here when the state is not a HeavyConcreteState
+        # Alternatively, we could make SimSuccessors generic over the state type too
+        concrete_state = typing.cast(HeavyConcreteState, self.state)
+
+        concrete_successor = self.process_concrete(concrete_state, num_inst=num_inst)
+        successors.add_successor(
+            concrete_successor,
+            concrete_successor.ip,
+            claripy.true(),
+            concrete_successor.history.jumpkind,
+            add_guard=False,
+        )
+        successors.processed = True
+
+
+__all__ = ["ConcreteEngine", "HeavyConcreteState"]

angr/engines/icicle.py

@@ -5,17 +5,16 @@ from __future__ import annotations
 import logging
 import os
 from dataclasses import dataclass
+from typing_extensions import override
 
-import claripy
 import pypcode
 from archinfo import Arch, Endness
 
+from angr.engines.concrete import ConcreteEngine, HeavyConcreteState
 from angr.engines.failure import SimEngineFailure
 from angr.engines.hook import HooksMixin
-from angr.engines.successors import SuccessorsEngine
 from angr.engines.syscall import SimEngineSyscall
 from angr.rustylib.icicle import Icicle, VmExit, ExceptionCode
-from angr.sim_state import SimState
 
 log = logging.getLogger(__name__)
 
@@ -30,12 +29,13 @@ class IcicleStateTranslationData:
     to an angr state.
     """
 
-    base_state: SimState
+    base_state: HeavyConcreteState
     registers: set[str]
     writable_pages: set[int]
+    initial_cpu_icount: int
 
 
-class IcicleEngine(SuccessorsEngine):
+class IcicleEngine(ConcreteEngine):
     """
     An angr engine that uses Icicle to execute concrete states. The purpose of
     this implementation is to provide a high-performance concrete execution
@@ -55,6 +55,16 @@ class IcicleEngine(SuccessorsEngine):
     intends to provide a more complete set of features, such as hooks and syscalls.
     """
 
+    breakpoints: set[int]
+
+    def __init__(self, *args, **kwargs):
+        """
+        Initialize the IcicleEngine. This sets up the breakpoints set and
+        initializes the parent class.
+        """
+        super().__init__(*args, **kwargs)
+        self.breakpoints = set()
+
     @staticmethod
     def __make_icicle_arch(arch: Arch) -> str | None:
         """
@@ -81,7 +91,7 @@ class IcicleEngine(SuccessorsEngine):
         return IcicleEngine.__is_arm(icicle_arch) and addr & 1 == 1
 
     @staticmethod
-    def __get_pages(state: SimState) -> set[int]:
+    def __get_pages(state: HeavyConcreteState) -> set[int]:
         """
         Unfortunately, the memory model doesn't have a way to get all pages.
         Instead, we can get all of the backers from the loader, then all of the
@@ -104,7 +114,7 @@ class IcicleEngine(SuccessorsEngine):
         return pages
 
     @staticmethod
-    def __convert_angr_state_to_icicle(state: SimState) -> tuple[Icicle, IcicleStateTranslationData]:
+    def __convert_angr_state_to_icicle(state: HeavyConcreteState) -> tuple[Icicle, IcicleStateTranslationData]:
         icicle_arch = IcicleEngine.__make_icicle_arch(state.arch)
         if icicle_arch is None:
             raise ValueError("Unsupported architecture")
@@ -161,12 +171,15 @@ class IcicleEngine(SuccessorsEngine):
             base_state=state,
             registers=copied_registers,
             writable_pages=writable_pages,
+            initial_cpu_icount=emu.cpu_icount,
         )
 
         return (emu, translation_data)
 
     @staticmethod
-    def __convert_icicle_state_to_angr(emu: Icicle, translation_data: IcicleStateTranslationData) -> SimState:
+    def __convert_icicle_state_to_angr(
+        emu: Icicle, translation_data: IcicleStateTranslationData, status: VmExit
+    ) -> HeavyConcreteState:
         state = translation_data.base_state.copy()
 
         # 1. Copy the register values
@@ -181,20 +194,8 @@ class IcicleEngine(SuccessorsEngine):
             addr = page_num * state.memory.page_size
             state.memory.store(addr, emu.mem_read(addr, state.memory.page_size))
 
-        return state
-
-    def process_successors(self, successors, *, num_inst=0, **kwargs):
-        if len(kwargs) > 0:
-            log.warning("IcicleEngine.process_successors received unknown kwargs: %s", kwargs)
-
-        emu, translation_data = self.__convert_angr_state_to_icicle(self.state)
-
-        if num_inst > 0:
-            emu.icount_limit = num_inst
-
-        status = emu.run()  # pylint: ignore=assignment-from-no-return (pylint bug)
+        # 3. Set history.jumpkind
         exc = emu.exception_code
-
         if status == VmExit.UnhandledException:
             if exc in (
                 ExceptionCode.ReadUnmapped,
@@ -203,22 +204,57 @@ class IcicleEngine(SuccessorsEngine):
                 ExceptionCode.WritePerm,
                 ExceptionCode.ExecViolation,
             ):
-                jumpkind = "Ijk_SigSEGV"
+                state.history.jumpkind = "Ijk_SigSEGV"
             elif exc == ExceptionCode.Syscall:
-                jumpkind = "Ijk_Syscall"
+                state.history.jumpkind = "Ijk_Syscall"
             elif exc == ExceptionCode.Halt:
-                jumpkind = "Ijk_Exit"
+                state.history.jumpkind = "Ijk_Exit"
             elif exc == ExceptionCode.InvalidInstruction:
-                jumpkind = "Ijk_NoDecode"
+                state.history.jumpkind = "Ijk_NoDecode"
             else:
-                jumpkind = "Ijk_EmFail"
+                state.history.jumpkind = "Ijk_EmFail"
         else:
-            jumpkind = "Ijk_Boring"
+            state.history.jumpkind = "Ijk_Boring"
+
+        # 4. Set history.recent_instruction_count
+        state.history.recent_instruction_count = emu.cpu_icount - translation_data.initial_cpu_icount
+
+        return state
+
+    @override
+    def get_breakpoints(self) -> set[int]:
+        """Return the set of currently set breakpoints."""
+        return self.breakpoints
+
+    @override
+    def add_breakpoint(self, addr: int) -> None:
+        """Add a breakpoint at the given address."""
+        self.breakpoints.add(addr)
+
+    @override
+    def remove_breakpoint(self, addr: int) -> None:
+        """Remove a breakpoint at the given address, if present."""
+        self.breakpoints.discard(addr)
+
+    @override
+    def process_concrete(self, state: HeavyConcreteState, num_inst: int | None = None) -> HeavyConcreteState:
+        emu, translation_data = self.__convert_angr_state_to_icicle(state)
+
+        # Set breakpoints, skip the current PC. This assumes that if running
+        # with a breakpoint at the current PC, then the user has already done
+        # the necessary handling and is resuming execution.
+        for addr in self.breakpoints:
+            if emu.pc != addr:
+                emu.add_breakpoint(addr)
+
+        # Set the instruction count limit
+        if num_inst is not None and num_inst > 0:
+            emu.icount_limit = num_inst
 
-        successor_state = IcicleEngine.__convert_icicle_state_to_angr(emu, translation_data)
-        successors.add_successor(successor_state, successor_state.ip, claripy.true(), jumpkind, add_guard=False)
+        # Run it
+        status = emu.run()
 
-        successors.processed = True
+        return IcicleEngine.__convert_icicle_state_to_angr(emu, translation_data, status)
 
 
 class UberIcicleEngine(SimEngineFailure, SimEngineSyscall, HooksMixin, IcicleEngine):

angr/exploration_techniques/driller_core.py

@@ -94,7 +94,7 @@ class DrillerCore(ExplorationTechnique):
     @staticmethod
     def _has_false(state):
         # Check if the state is unsat even if we remove preconstraints.
-        if state.scratch.guard.identical(claripy.false()):
+        if claripy.is_false(state.scratch.guard):
             return True
 
-        return any(c.identical(claripy.false()) for c in state.solver.constraints)
+        return any(claripy.is_false(c) for c in state.solver.constraints)

angr/project.py

@@ -297,11 +297,18 @@ class Project:
 
         # Step 1: get the set of libraries we are allowed to use to resolve unresolved symbols
         missing_libs = []
+        missing_wincore_dlls = False
         for lib_name in self.loader.missing_dependencies:
             try:
                 missing_libs.extend(SIM_LIBRARIES[lib_name])
             except KeyError:
                 l.info("There are no simprocedures for missing library %s :(", lib_name)
+                if lib_name.startswith("api-ms-win-"):
+                    missing_wincore_dlls = True
+        if missing_wincore_dlls and "kernel32.dll" not in self.loader.missing_dependencies:
+            # some of the missing api-ms-win-*.dll libraries are actually provided by kernel32.dll
+            missing_libs.extend(SIM_LIBRARIES["kernel32.dll"])
+
         # additionally provide libraries we _have_ loaded as a fallback fallback
         # this helps in the case that e.g. CLE picked up a linux arm libc to satisfy an android arm binary
         for lib in self.loader.all_objects:

angr/sim_type.py

@@ -609,33 +609,41 @@ class SimTypeWideChar(SimTypeReg):
 
     _base_name = "char"
 
-    def __init__(self, signed=True, label=None):
+    def __init__(self, signed=True, label=None, endness: Endness = Endness.BE):
         """
         :param label: the type label.
         """
         SimTypeReg.__init__(self, 16, label=label)
         self.signed = signed
+        self.endness = endness
 
     def __repr__(self):
         return "wchar"
 
     def store(self, state, addr, value: StoreType):
-        self._size = state.arch.byte_width
         try:
             super().store(state, addr, value)
         except TypeError:
             if isinstance(value, bytes) and len(value) == 2:
-                value = claripy.BVV(value[0], state.arch.byte_width)
+                inner = (
+                    ((value[0] << state.arch.byte_width) | value[1])
+                    if self.endness == Endness.BE
+                    else ((value[1] << state.arch.byte_width) | value[0])
+                )
+                value = claripy.BVV(inner, state.arch.byte_width * 2)
                 super().store(state, addr, value)
             else:
                 raise
 
     def extract(self, state, addr, concrete=False) -> Any:
-        self._size = state.arch.byte_width
-
-        out = super().extract(state, addr, concrete)
+        out = state.memory.load(addr, 2)
         if concrete:
-            return bytes([out])
+            data = state.solver.eval(out, cast_to=bytes)
+            fmt_str = "utf-16be" if self.endness == Endness.BE else "utf-16le"
+            try:
+                return data.decode(fmt_str)
+            except UnicodeDecodeError:
+                return data
         return out
 
     def _init_str(self):
@@ -645,7 +653,7 @@ class SimTypeWideChar(SimTypeReg):
         )
 
     def copy(self):
-        return self.__class__(signed=self.signed, label=self.label)
+        return self.__class__(signed=self.signed, label=self.label, endness=self.endness)
 
 
 class SimTypeBool(SimTypeReg):

angr/state_plugins/unicorn_engine.py

@@ -394,14 +394,14 @@ class _VexArchInfo(ctypes.Structure):
 
 def _load_native():
     if sys.platform == "darwin":
-        libfile = "angr_native.dylib"
+        libfile = "unicornlib.dylib"
     elif sys.platform in {"win32", "cygwin"}:
-        libfile = "angr_native.dll"
+        libfile = "unicornlib.dll"
     else:
-        libfile = "angr_native.so"
+        libfile = "unicornlib.so"
 
     try:
-        angr_path = str(importlib.resources.files("angr") / "lib" / libfile)
+        angr_path = str(importlib.resources.files("angr") / libfile)
         h = ctypes.CDLL(angr_path)
 
         VexArch = ctypes.c_int

angr/utils/graph.py

@@ -534,11 +534,30 @@ class Dominators:
 
     def _pd_eval(self, v):
         assert self._ancestor is not None
+        assert self._semi is not None
         assert self._label is not None
 
         if self._ancestor[v.index] is None:
             return v
-        self._pd_compress(v)
+
+        # pd_compress without recursion
+        queue = []
+        current = v
+        ancestor = self._ancestor[current.index]
+        greater_ancestor = self._ancestor[ancestor.index]
+        while greater_ancestor is not None:
+            queue.append(current)
+            current, ancestor = ancestor, greater_ancestor
+            greater_ancestor = self._ancestor[ancestor.index]
+
+        for vv in reversed(queue):
+            if (
+                self._semi[self._label[self._ancestor[vv.index].index].index].index
+                < self._semi[self._label[vv.index].index].index
+            ):
+                self._label[vv.index] = self._label[self._ancestor[vv.index].index]
+            self._ancestor[vv.index] = self._ancestor[self._ancestor[vv.index].index]
+
         return self._label[v.index]
 
     def _pd_compress(self, v):

angr/utils/ssa/__init__.py

@@ -89,7 +89,7 @@ def get_reg_offset_base(reg_offset, arch, size=None, resilient=True):
 
 
 def get_vvar_deflocs(
-    blocks, phi_vvars: dict[int, set[int]] | None = None
+    blocks, phi_vvars: dict[int, set[int | None]] | None = None
 ) -> dict[int, tuple[VirtualVariable, CodeLocation]]:
     vvar_to_loc: dict[int, tuple[VirtualVariable, CodeLocation]] = {}
     for block in blocks:
@@ -100,7 +100,7 @@ def get_vvar_deflocs(
                 )
                 if phi_vvars is not None and isinstance(stmt.src, Phi):
                     phi_vvars[stmt.dst.varid] = {
-                        vvar_.varid for src, vvar_ in stmt.src.src_and_vvars if vvar_ is not None
+                        vvar_.varid if vvar_ is not None else None for src, vvar_ in stmt.src.src_and_vvars
                     }
             elif isinstance(stmt, Call):
                 if isinstance(stmt.ret_expr, VirtualVariable):
@@ -161,7 +161,7 @@ def get_tmp_uselocs(blocks) -> dict[CodeLocation, dict[atoms.Tmp, set[tuple[Tmp,
     return tmp_to_loc
 
 
-def is_const_assignment(stmt: Statement) -> tuple[bool, Const | None]:
+def is_const_assignment(stmt: Statement) -> tuple[bool, Const | StackBaseOffset | None]:
     if isinstance(stmt, Assignment) and isinstance(stmt.src, (Const, StackBaseOffset)):
         return True, stmt.src
     return False, None