# pylint:disable=too-many-boolean-expressions
from __future__ import annotations
from typing import Any, NamedTuple, TYPE_CHECKING
import copy
import logging
import enum
from collections import defaultdict, namedtuple
from collections.abc import Iterable
from dataclasses import dataclass

import networkx
import capstone

from angr import ailment
from angr.ailment.block_walker import AILBlockViewer
from angr.ailment.expression import VirtualVariable
from angr.errors import AngrDecompilationError
from angr.knowledge_base import KnowledgeBase
from angr.knowledge_plugins.functions import Function
from angr.knowledge_plugins.cfg.memory_data import MemoryDataSort
from angr.knowledge_plugins.key_definitions import atoms
from angr.codenode import BlockNode
from angr.utils import timethis
from angr.utils.ssa import is_phi_assignment
from angr.utils.graph import GraphUtils
from angr.utils.types import dereference_simtype_by_lib
from angr.calling_conventions import SimRegArg, SimStackArg, SimFunctionArgument, SimCCUsercall
from angr.sim_type import (
    SimType,
    SimTypeChar,
    SimTypeInt,
    SimTypeLongLong,
    SimTypeShort,
    SimTypeFunction,
    SimTypeBottom,
    SimTypeFloat,
    SimTypePointer,
    SimStruct,
    SimTypeArray,
    SimCppClass,
)
from angr.analyses.stack_pointer_tracker import Register, OffsetVal
from angr.sim_variable import SimVariable, SimStackVariable, SimRegisterVariable, SimMemoryVariable
from angr.procedures.stubs.UnresolvableCallTarget import UnresolvableCallTarget
from angr.procedures.stubs.UnresolvableJumpTarget import UnresolvableJumpTarget
from angr.analyses import Analysis, register_analysis
from angr.analyses.cfg.cfg_base import CFGBase
from angr.analyses.reaching_definitions import ReachingDefinitionsAnalysis
from angr.analyses.typehoon import Typehoon
from .ail_simplifier import AILSimplifier
from .ssailification.ssailification import Ssailification
from .stack_item import StackItem, StackItemType
from .return_maker import ReturnMaker
from .ailgraph_walker import AILGraphWalker, RemoveNodeNotice
from .optimization_passes import (
    OptimizationPassStage,
    StackCanarySimplifier,
    TagSlicer,
    DUPLICATING_OPTS,
    CONDENSING_OPTS,
)
from .semantic_naming import SemanticNamingOrchestrator

if TYPE_CHECKING:
    from angr.knowledge_plugins.cfg import CFGModel
    from .notes import DecompilationNote
    from .decompilation_cache import DecompilationCache
    from .peephole_optimizations import PeepholeOptimizationStmtBase, PeepholeOptimizationExprBase

l = logging.getLogger(name=__name__)


BlockCache = namedtuple("BlockCache", ("rd", "prop"))


class ClinicMode(enum.Enum):
    """
    Analysis mode for Clinic.
    """

    DECOMPILE = 1
    COLLECT_DATA_REFS = 2


@dataclass
class DataRefDesc:
    """
    The fields of this class is compatible with items inside IRSB.data_refs.
    """

    data_addr: int
    data_size: int
    block_addr: int
    stmt_idx: int
    ins_addr: int
    data_type_str: str


class ClinicStage(enum.IntEnum):
    """
    Different stages of treating an ailment.
    """

    INITIALIZATION = 0
    AIL_GRAPH_CONVERSION = 1
    MAKE_RETURN_SITES = 2
    MAKE_ARGUMENT_LIST = 3
    PRE_SSA_LEVEL0_FIXUPS = 4
    SSA_LEVEL0_TRANSFORMATION = 5
    CONSTANT_PROPAGATION = 6
    TRACK_STACK_POINTERS = 7
    PRE_SSA_LEVEL1_SIMPLIFICATIONS = 8
    SSA_LEVEL1_TRANSFORMATION = 9
    POST_SSA_LEVEL1_SIMPLIFICATIONS = 10
    MAKE_CALLSITES = 11
    POST_CALLSITES = 12
    RECOVER_VARIABLES = 13
    SEMANTIC_VARIABLE_NAMING = 14
    COLLECT_EXTERNS = 15


class Clinic(Analysis):
    """
    A Clinic deals with AILments.
    """

    _ail_manager: ailment.Manager

    def __init__(
        self,
        func: Function,
        remove_dead_memdefs=False,
        exception_edges=False,
        sp_tracker_track_memory=True,
        fold_callexprs_into_conditions=False,
        insert_labels=True,
        optimization_passes=None,
        cfg=None,
        peephole_optimizations: None | (
            Iterable[type[PeepholeOptimizationStmtBase] | type[PeepholeOptimizationExprBase]]
        ) = None,  # pylint:disable=line-too-long
        must_struct: set[str] | None = None,
        variable_kb: KnowledgeBase | None = None,
        reset_variable_names=False,
        rewrite_ites_to_diamonds=True,
        cache: DecompilationCache | None = None,
        mode: ClinicMode = ClinicMode.DECOMPILE,
        sp_shift: int = 0,
        inline_functions: set[Function] | None = None,
        inlined_counts: dict[int, int] | None = None,
        inlining_parents: set[int] | None = None,
        vvar_id_start: int = 0,
        optimization_scratch: dict[str, Any] | None = None,
        desired_variables: set[str] | None = None,
        force_loop_single_exit: bool = True,
        refine_loops_with_single_successor: bool = False,
        complete_successors: bool = False,
        max_type_constraints: int = 100_000,
        type_constraint_set_degradation_threshold: int = 150,
        ail_graph: networkx.DiGraph | None = None,
        arg_vvars: dict[int, tuple[ailment.Expr.VirtualVariable, SimVariable]] | None = None,
        start_stage: ClinicStage | None = ClinicStage.INITIALIZATION,
        end_stage: ClinicStage | None = None,
        skip_stages: tuple[ClinicStage, ...] = (),
        notes: dict[str, DecompilationNote] | None = None,
        static_vvars: dict | None = None,
        static_buffers: dict | None = None,
        semvar_naming: bool = True,
    ):
        if not func.normalized and mode == ClinicMode.DECOMPILE:
            raise ValueError("Decompilation must work on normalized function graphs.")

        self.function = func

        self.graph = None
        self.cc_graph: networkx.DiGraph | None = None
        self.unoptimized_graph: networkx.DiGraph | None = None
        self.arg_list = None
        self.arg_vvars: dict[int, tuple[ailment.Expr.VirtualVariable, SimVariable]] | None = None
        self.func_args = None
        self.func_ret_var = SimVariable(0, "__retvar", "__retvar")
        self.variable_kb = variable_kb
        self.externs: set[SimMemoryVariable] = set()
        self.data_refs: dict[int, list[DataRefDesc]] = {}  # data address to data reference description
        self.optimization_scratch = optimization_scratch if optimization_scratch is not None else {}

        self._func_graph: networkx.DiGraph | None = None
        self._init_ail_graph = ail_graph
        self._init_arg_vvars = arg_vvars
        self._start_stage = start_stage if start_stage is not None else ClinicStage.INITIALIZATION
        self._end_stage = end_stage if end_stage is not None else max(ClinicStage.__members__.values())
        self._skip_stages = skip_stages

        self._blocks_by_addr_and_size = {}
        self.entry_node_addr: tuple[int, int | None] = self.function.addr, None

        self._fold_callexprs_into_conditions = fold_callexprs_into_conditions
        self._insert_labels = insert_labels
        self._remove_dead_memdefs = remove_dead_memdefs
        self._exception_edges = exception_edges
        self._sp_tracker_track_memory = sp_tracker_track_memory
        self._cfg: CFGModel | None = cfg
        self.peephole_optimizations = peephole_optimizations
        self._must_struct = must_struct
        self._reset_variable_names = reset_variable_names
        self._rewrite_ites_to_diamonds = rewrite_ites_to_diamonds
        self.reaching_definitions: ReachingDefinitionsAnalysis | None = None
        self._cache = cache
        self._mode = mode
        self._max_type_constraints = max_type_constraints
        self._type_constraint_set_degradation_threshold = type_constraint_set_degradation_threshold
        self.vvar_id_start = vvar_id_start
        self.vvar_to_vvar: dict[int, int] | None = None
        # during SSA conversion, we create secondary stack variables because they overlap and are larger than the
        # actual stack variables. these secondary stack variables can be safely eliminated if not used by anything.
        self.secondary_stackvars: set[int] = set()

        self.notes = notes if notes is not None else {}
        self.static_vvars = static_vvars if static_vvars is not None else {}
        self.static_buffers = static_buffers if static_buffers is not None else {}
        self._semvar_naming = semvar_naming

        if not semvar_naming and ClinicStage.SEMANTIC_VARIABLE_NAMING not in self._skip_stages:
            self._skip_stages += (ClinicStage.SEMANTIC_VARIABLE_NAMING,)

        #
        # intermediate variables used during decompilation
        #

        self._ail_graph: networkx.DiGraph = None  # type:ignore
        self._spt = None
        # cached block-level reaching definition analysis results and propagator results
        self._block_simplification_cache: dict[ailment.Block, NamedTuple] | None = {}
        self._preserve_vvar_ids: set[int] = set()
        self._type_hints: list[tuple[atoms.VirtualVariable | atoms.MemoryLocation, str]] = []

        # inlining help
        self._sp_shift = sp_shift
        self._max_stack_depth = 0
        self._inline_functions = inline_functions if inline_functions else set()
        self._inlined_counts = {} if inlined_counts is None else inlined_counts
        self._inlining_parents = inlining_parents or ()
        self._desired_variables = desired_variables
        self._force_loop_single_exit = force_loop_single_exit
        self._refine_loops_with_single_successor = refine_loops_with_single_successor
        self._complete_successors = complete_successors

        self._register_save_areas_removed: bool = False
        self.edges_to_remove: list[tuple[tuple[int, int | None], tuple[int, int | None]]] = []
        self.copied_var_ids: set[int] = set()

        self._new_block_addrs: set[int] = set()

        # a reference to the Typehoon type inference engine; useful for debugging and loading stats post decompilation
        self.typehoon: Typehoon | None = None

        # sanity checks
        if not self.kb.functions:
            l.warning("No function is available in kb.functions. It will lead to a suboptimal conversion result.")

        if optimization_passes is not None:
            self._optimization_passes = optimization_passes
        else:
            self._optimization_passes = []

        self.stack_items: dict[int, StackItem] = {}
        if self.project.arch.call_pushes_ret:
            self.stack_items[0] = StackItem(0, self.project.arch.bytes, "ret_addr", StackItemType.RET_ADDR)

        # Set up the function graph according to configurations
        self._set_function_graph()

        if self._mode == ClinicMode.DECOMPILE:
            self._analyze_for_decompiling()
            if self._end_stage >= ClinicStage.MAKE_CALLSITES:
                self._constrain_callee_prototypes()
        elif self._mode == ClinicMode.COLLECT_DATA_REFS:
            self._analyze_for_data_refs()
        else:
            raise TypeError(f"Unsupported analysis mode {self._mode}")

    #
    # Public methods
    #

    def block(self, addr, size):
        """
        Get the converted block at the given specific address with the given size.

        :param int addr:
        :param int size:
        :return:
        """

        try:
            return self._blocks_by_addr_and_size[(addr, size)] if self._blocks_by_addr_and_size is not None else None
        except KeyError:
            return None

    def dbg_repr(self):
        """

        :return:
        """
        assert self.graph is not None

        s = ""

        for block in sorted(self.graph.nodes(), key=lambda x: x.addr):
            s += str(block) + "\n\n"

        return s

    #
    # Private methods
    #

    def _analyze_for_decompiling(self):
        # initialize the AIL conversion manager
        self._ail_manager = ailment.Manager(arch=self.project.arch)

        ail_graph = self._init_ail_graph if self._init_ail_graph is not None else self._decompilation_graph_recovery()
        if not ail_graph:
            return
        if self._start_stage <= ClinicStage.INITIALIZATION:
            ail_graph = self._decompilation_fixups(ail_graph)

        if self._inline_functions:
            self._max_stack_depth += self.calculate_stack_depth()
            ail_graph = self._inline_child_functions(ail_graph)

        ail_graph = self._decompilation_simplifications(ail_graph)

        if self._desired_variables:
            ail_graph = self._slice_variables(ail_graph)
        self.graph = ail_graph

    def _decompilation_graph_recovery(self):
        is_pcode_arch = ":" in self.project.arch.name

        # Remove alignment blocks
        self._update_progress(5.0, text="Removing alignment blocks")
        self._remove_alignment_blocks()

        # if the graph is empty, don't continue
        if not self._func_graph:
            return None

        # Make sure calling conventions of all functions that the current function calls have been recovered
        if not is_pcode_arch:
            self._update_progress(10.0, text="Recovering calling conventions")
            self._recover_calling_conventions()

        # Convert VEX blocks to AIL blocks and then simplify them

        self._update_progress(20.0, text="Converting VEX to AIL")
        self._convert_all()

        return self._make_ailgraph()

    def _decompilation_fixups(self, ail_graph):
        is_pcode_arch = ":" in self.project.arch.name

        self._remove_redundant_jump_blocks(ail_graph)
        # _fix_abnormal_switch_case_heads may re-lift from VEX blocks, so it should be placed as high up as possible
        self._fix_abnormal_switch_case_heads(ail_graph)
        if self._rewrite_ites_to_diamonds:
            self._rewrite_ite_expressions(ail_graph)
        self._remove_redundant_jump_blocks(ail_graph)
        if self._insert_labels:
            self._insert_block_labels(ail_graph)

        # Run simplification passes
        self._update_progress(22.0, text="Optimizing fresh ailment graph")
        ail_graph = self._run_simplification_passes(ail_graph, OptimizationPassStage.AFTER_AIL_GRAPH_CREATION)

        # Fix "fake" indirect jumps and calls
        self._update_progress(25.0, text="Analyzing simple indirect jumps")
        ail_graph = self._replace_single_target_indirect_transitions(ail_graph)

        # Fix tail calls
        self._update_progress(26.0, text="Analyzing tail calls")
        ail_graph = self._replace_tail_jumps_with_calls(ail_graph)

        # Fix special calls
        self._update_progress(28.0, text="Analyzing special calls")
        ail_graph = self._fix_special_call_calling_conventions(ail_graph)

        if is_pcode_arch:
            self._update_progress(29.0, text="Recovering calling conventions (AIL mode)")
            self._recover_calling_conventions(func_graph=ail_graph)

        return self._apply_callsite_prototype_and_calling_convention(ail_graph)

    def _slice_variables(self, ail_graph: networkx.DiGraph[ailment.Block]) -> networkx.DiGraph[ailment.Block]:
        assert self.variable_kb is not None and self._desired_variables is not None

        nodes_index = {(n.addr, n.idx): n for n in ail_graph.nodes()}

        vfm = self.variable_kb.variables.function_managers[self.function.addr]
        for v_name in self._desired_variables:
            v = next(iter(vv for vv in vfm._unified_variables if vv.name == v_name))
            for va in vfm.get_variable_accesses(v):
                assert va.location.block_addr is not None
                assert va.location.stmt_idx is not None
                nodes_index[(va.location.block_addr, va.location.block_idx)].statements[va.location.stmt_idx].tags[
                    "keep_in_slice"
                ] = True

        a = TagSlicer(
            self.function,
            graph=ail_graph,
            variable_kb=self.variable_kb,
        )
        if a.out_graph:
            # use the new graph
            ail_graph = a.out_graph
        return ail_graph

    def _inline_child_functions(self, ail_graph):
        for blk in ail_graph.nodes():
            for idx, stmt in enumerate(blk.statements):
                if isinstance(stmt, ailment.Stmt.Call) and isinstance(stmt.target, ailment.Expr.Const):
                    assert self.function._function_manager is not None
                    callee = self.function._function_manager.function(stmt.target.value)
                    if (
                        callee is None
                        or callee.addr == self.function.addr
                        or callee.addr in self._inlining_parents
                        or callee not in self._inline_functions
                        or callee.is_plt
                        or callee.is_simprocedure
                    ):
                        continue

                    ail_graph = self._inline_call(ail_graph, blk, idx, callee)
        return ail_graph

    @staticmethod
    def _inline_fix_block_phi_stmts(block: ailment.Block, new_block_idx: int) -> None:
        # update the source block ID of all phi variables
        for idx, stmt in enumerate(block.statements):
            if is_phi_assignment(stmt):
                new_src_and_vvars: list[tuple[tuple[int, int | None], VirtualVariable | None]] = [
                    ((src_block_addr, new_block_idx), vvar) for (src_block_addr, _), vvar in stmt.src.src_and_vvars
                ]
                new_src = ailment.Expr.Phi(stmt.src.idx, stmt.src.bits, new_src_and_vvars, **stmt.src.tags)
                new_stmt = ailment.Stmt.Assignment(stmt.idx, stmt.dst, new_src, **stmt.tags)
                block.statements[idx] = new_stmt

    def _inline_call(self, ail_graph: networkx.DiGraph, caller_block: ailment.Block, call_idx: int, callee: Function):
        callee_clinic = self.project.analyses.Clinic(
            callee,
            mode=ClinicMode.DECOMPILE,
            inline_functions=self._inline_functions,
            inlining_parents=(*self._inlining_parents, self.function.addr),
            inlined_counts=self._inlined_counts,
            optimization_passes=[StackCanarySimplifier],
            sp_shift=self._max_stack_depth,
            vvar_id_start=self.vvar_id_start,
            fail_fast=self._fail_fast,  # type: ignore
        )
        self.vvar_id_start = callee_clinic.vvar_id_start + 1
        self._max_stack_depth = callee_clinic._max_stack_depth
        callee_graph = callee_clinic.copy_graph()

        # uniquely mark all the blocks in case of duplicates (e.g., foo(); foo();)
        self._inlined_counts.setdefault(callee.addr, 0)
        block_idx = self._inlined_counts[callee.addr]
        for blk in callee_graph.nodes():
            blk.idx = block_idx
            self._inline_fix_block_phi_stmts(blk, block_idx)
        self._inlined_counts[callee.addr] += 1

        # figure out where the callee should start at and return to
        callee_start = next(n for n in callee_graph if n.addr == callee.addr)
        caller_successors = list(ail_graph.out_edges(caller_block, data=True))
        assert len(caller_successors) == 1
        caller_successor = caller_successors[0][1]
        ail_graph.remove_edge(caller_block, caller_successor)

        # update all callee return nodes with caller successor
        ail_graph = networkx.union(ail_graph, callee_graph)
        for blk in callee_graph.nodes():
            for idx, stmt in enumerate(list(blk.statements)):
                if isinstance(stmt, ailment.Stmt.Return):
                    # replace the return statement with an assignment to the return register
                    blk.statements.pop(idx)

                    if stmt.ret_exprs and self.project.arch.ret_offset is not None:
                        assign_to_retreg = ailment.Stmt.Assignment(
                            self._ail_manager.next_atom(),
                            ailment.Expr.Register(
                                self._ail_manager.next_atom(),
                                None,
                                self.project.arch.ret_offset,
                                self.project.arch.bits,
                            ),
                            stmt.ret_exprs[0],
                            **stmt.tags,
                        )
                        blk.statements.insert(idx, assign_to_retreg)
                        idx += 1
                    ail_graph.add_edge(blk, caller_successor)
                    break

        # update the call edge
        # first, remove the call statement. this is a type error but will be resolved later
        caller_block.statements[call_idx] = None  # type: ignore
        if (
            isinstance(caller_block.statements[call_idx - 2], ailment.Stmt.Store)
            and caller_block.statements[call_idx - 2].data.value == caller_successor.addr
        ):
            # don't push the return address
            caller_block.statements.pop(call_idx - 5)  # t6 = rsp<8>
            caller_block.statements.pop(call_idx - 5)  # t5 = (t6 - 0x8<64>)
            caller_block.statements.pop(call_idx - 5)  # rsp<8> = t5
            caller_block.statements.pop(
                call_idx - 5
            )  # STORE(addr=t5, data=0x40121b<64>, size=8, endness=Iend_LE, guard=None)
            caller_block.statements.pop(call_idx - 5)  # t7 = (t5 - 0x80<64>) <- wtf is this??
        elif (
            isinstance(caller_block.statements[call_idx - 1], ailment.Stmt.Store)
            and caller_block.statements[call_idx - 1].addr.base == "stack_base"
            and caller_block.statements[call_idx - 1].data.value == caller_successor.addr
        ):
            caller_block.statements.pop(call_idx - 1)  # s_10 =L 0x401225<64><8>

        # update caller_block to setup parameters
        if callee_clinic.arg_vvars:
            for arg_idx in sorted(callee_clinic.arg_vvars.keys()):
                param_vvar, reg_arg = callee_clinic.arg_vvars[arg_idx]
                if isinstance(reg_arg, SimRegisterVariable):
                    reg_offset = reg_arg.reg
                    stmt = ailment.Stmt.Assignment(
                        self._ail_manager.next_atom(),
                        param_vvar,
                        ailment.Expr.Register(
                            self._ail_manager.next_atom(),
                            None,
                            reg_offset,
                            reg_arg.bits,
                            ins_addr=caller_block.addr + caller_block.original_size,
                        ),
                        ins_addr=caller_block.addr + caller_block.original_size,
                    )
                    caller_block.statements.append(stmt)
                else:
                    raise NotImplementedError("Unsupported parameter type")

        caller_block.statements = [s for s in caller_block.statements if s is not None]

        ail_graph.add_edge(caller_block, callee_start)

        return ail_graph

    def calculate_stack_depth(self):
        # we need to reserve space for our own stack
        spt = self._track_stack_pointers()
        stack_offsets = spt.offsets_for(self.project.arch.sp_offset)
        if max(stack_offsets) > 2 ** (self.project.arch.bits - 1):
            # why is this unsigned...
            depth = min(s for s in stack_offsets if s > 2 ** (self.project.arch.bits - 1)) - 2**self.project.arch.bits
        else:
            depth = min(stack_offsets)

        if spt.inconsistent_for(self.project.arch.sp_offset):
            l.warning("Inconsistency found during stack pointer tracking. Stack depth may be incorrect.")
            depth -= 0x1000

        return depth

    def _decompilation_simplifications(self, ail_graph):
        self.arg_vvars = self._init_arg_vvars if self._init_arg_vvars is not None else {}
        self.func_args = {arg_vvar for arg_vvar, _ in self.arg_vvars.values()}
        self._ail_graph = ail_graph

        stages = {
            ClinicStage.MAKE_RETURN_SITES: self._stage_make_return_sites,
            ClinicStage.MAKE_ARGUMENT_LIST: self._stage_make_function_argument_list,
            ClinicStage.PRE_SSA_LEVEL0_FIXUPS: self._stage_pre_ssa_level0_fixups,
            ClinicStage.SSA_LEVEL0_TRANSFORMATION: self._stage_transform_to_ssa_level0,
            ClinicStage.CONSTANT_PROPAGATION: self._stage_constant_propagation,
            ClinicStage.TRACK_STACK_POINTERS: self._stage_track_stack_pointers,
            ClinicStage.PRE_SSA_LEVEL1_SIMPLIFICATIONS: self._stage_pre_ssa_level1_simplifications,
            ClinicStage.SSA_LEVEL1_TRANSFORMATION: self._stage_transform_to_ssa_level1,
            ClinicStage.POST_SSA_LEVEL1_SIMPLIFICATIONS: self._stage_post_ssa_level1_simplifications,
            ClinicStage.MAKE_CALLSITES: self._stage_make_function_callsites,
            ClinicStage.POST_CALLSITES: self._stage_post_callsite_simplifications,
            ClinicStage.RECOVER_VARIABLES: self._stage_recover_variables,
            ClinicStage.SEMANTIC_VARIABLE_NAMING: self._stage_semantic_variable_naming,
            ClinicStage.COLLECT_EXTERNS: self._stage_collect_externs,
        }

        for stage in sorted(stages):
            if stage < self._start_stage or stage > self._end_stage or stage in self._skip_stages:
                continue
            stages[stage]()

        # remove empty nodes from the graph
        self._ail_graph = self.remove_empty_nodes(self._ail_graph)
        # note that there are still edges to remove before we can structure this graph!

        self.cc_graph = self.copy_graph(self._ail_graph)
        return self._ail_graph

    def _stage_make_return_sites(self) -> None:
        self._update_progress(30.0, text="Making return sites")
        if self.function.prototype is None or not isinstance(self.function.prototype.returnty, SimTypeBottom):
            self._ail_graph = self._make_returns(self._ail_graph)
        self._ail_graph = self._run_simplification_passes(
            self._ail_graph, stage=OptimizationPassStage.BEFORE_SSA_LEVEL0_TRANSFORMATION
        )

    def _stage_make_function_argument_list(self) -> None:
        self._update_progress(33.0, text="Making argument list")
        self.arg_list = self._make_argument_list()
        self.arg_vvars = self._create_function_argument_vvars(self.arg_list)
        self.func_args = {arg_vvar for arg_vvar, _ in self.arg_vvars.values()}

    def _stage_pre_ssa_level0_fixups(self) -> None:
        # duplicate orphaned conditional jump blocks
        self._ail_graph = self._duplicate_orphaned_cond_jumps(self._ail_graph)
        # rewrite jmp_rax function calls
        self._ail_graph = self._rewrite_jump_rax_calls(self._ail_graph)

    def _stage_transform_to_ssa_level0(self) -> None:
        self._update_progress(35.0, text="Transforming to partial-SSA form (registers)")
        assert self.func_args is not None
        self._ail_graph = self._transform_to_ssa_level0(self._ail_graph, self.func_args)

    def _stage_constant_propagation(self) -> None:
        # full-function constant-only propagation
        self._update_progress(36.0, text="Constant propagation")
        self._simplify_function(
            self._ail_graph,
            remove_dead_memdefs=False,
            unify_variables=False,
            narrow_expressions=False,
            only_consts=True,
            fold_callexprs_into_conditions=self._fold_callexprs_into_conditions,
            max_iterations=1,
        )

    def _stage_track_stack_pointers(self) -> None:
        self._spt = self._track_stack_pointers()

    def _stage_transform_to_ssa_level1(self) -> None:
        self._update_progress(37.0, text="Transforming to partial-SSA form (stack variables)")
        # rewrite (qualified) stack variables into SSA form
        assert self.func_args is not None
        self._ail_graph = self._transform_to_ssa_level1(self._ail_graph, self.func_args)

        # Run simplification passes
        self._update_progress(49.0, text="Running simplifications 1.5")
        self._ail_graph = self._run_simplification_passes(
            self._ail_graph, stage=OptimizationPassStage.AFTER_SSA_LEVEL1_TRANSFORMATION
        )

        # register save area has been removed at this point - we should no longer use callee-saved registers in RDA
        self._register_save_areas_removed = True
        # clear the cached RDA result
        self.reaching_definitions = None

    def _stage_pre_ssa_level1_simplifications(self) -> None:
        # Simplify blocks
        # we never remove dead memory definitions before making callsites. otherwise stack arguments may go missing
        # before they are recognized as stack arguments.
        self._update_progress(38.0, text="Simplifying blocks 1")
        self._ail_graph = self._simplify_blocks(
            self._ail_graph,
            stack_pointer_tracker=self._spt,
            cache=self._block_simplification_cache,
            preserve_vvar_ids=self._preserve_vvar_ids,
            type_hints=self._type_hints,
        )
        self._rewrite_alloca(self._ail_graph)

        # Run simplification passes
        self._update_progress(40.0, text="Running simplifications 1")
        self._ail_graph = self._run_simplification_passes(
            self._ail_graph,
            stack_pointer_tracker=self._spt,
            stack_items=self.stack_items,
            stage=OptimizationPassStage.AFTER_SINGLE_BLOCK_SIMPLIFICATION,
        )

        # Simplify the entire function for the first time
        self._update_progress(45.0, text="Simplifying function 1")
        self._simplify_function(
            self._ail_graph,
            remove_dead_memdefs=False,
            unify_variables=False,
            narrow_expressions=False,
            fold_callexprs_into_conditions=self._fold_callexprs_into_conditions,
            arg_vvars=self.arg_vvars,
        )

        # Run simplification passes again. there might be more chances for peephole optimizations after function-level
        # simplification
        self._update_progress(48.0, text="Simplifying blocks 2")
        self._ail_graph = self._simplify_blocks(
            self._ail_graph,
            stack_pointer_tracker=self._spt,
            cache=self._block_simplification_cache,
            preserve_vvar_ids=self._preserve_vvar_ids,
            type_hints=self._type_hints,
        )

        # Run simplification passes
        self._update_progress(49.0, text="Running simplifications 2")
        self._ail_graph = self._run_simplification_passes(
            self._ail_graph, stage=OptimizationPassStage.BEFORE_SSA_LEVEL1_TRANSFORMATION
        )

    def _stage_post_ssa_level1_simplifications(self) -> None:
        # Rust-specific; only call this on Rust binaries when we can identify language and compiler
        self._ail_graph = self._rewrite_rust_probestack_call(self._ail_graph)
        # Windows-specific
        self._ail_graph = self._rewrite_windows_chkstk_call(self._ail_graph)

    def _stage_make_function_callsites(self) -> None:
        assert self.func_args is not None

        # Make call-sites
        self._update_progress(50.0, text="Making callsites")
        _, stackarg_offsets, removed_vvar_ids = self._make_callsites(
            self._ail_graph, self.func_args, stack_pointer_tracker=self._spt, preserve_vvar_ids=self._preserve_vvar_ids
        )

        # Run simplification passes
        self._update_progress(53.0, text="Running simplifications 2")
        self._ail_graph = self._run_simplification_passes(
            self._ail_graph, stage=OptimizationPassStage.AFTER_MAKING_CALLSITES
        )

        # Simplify the entire function for the second time
        self._update_progress(55.0, text="Simplifying function 2")
        self._simplify_function(
            self._ail_graph,
            remove_dead_memdefs=self._remove_dead_memdefs,
            stack_arg_offsets=stackarg_offsets,
            unify_variables=True,
            narrow_expressions=True,
            fold_callexprs_into_conditions=self._fold_callexprs_into_conditions,
            removed_vvar_ids=removed_vvar_ids,
            arg_vvars=self.arg_vvars,
            preserve_vvar_ids=self._preserve_vvar_ids,
        )

        # After global optimization, there might be more chances for peephole optimizations.
        # Simplify blocks for the second time
        self._update_progress(60.0, text="Simplifying blocks 3")
        self._ail_graph = self._simplify_blocks(
            self._ail_graph,
            stack_pointer_tracker=self._spt,
            cache=self._block_simplification_cache,
            preserve_vvar_ids=self._preserve_vvar_ids,
            type_hints=self._type_hints,
        )

        # Run simplification passes
        self._update_progress(65.0, text="Running simplifications 3")
