# pylint:disable=wrong-import-position,broad-exception-caught,ungrouped-imports,import-outside-toplevel
from __future__ import annotations
import pathlib
import copy
from types import FunctionType
from typing import Any, TYPE_CHECKING
from collections.abc import Iterable
import logging

import networkx

import angr
from angr import ailment
from angr.ailment.block import Block
from angr.analyses.decompiler.counters.call_counter import AILBlockCallCounter
from angr.analyses.decompiler.peephole_optimizations.base import (
    PeepholeOptimizationExprBase,
    PeepholeOptimizationMultiStmtBase,
)
from angr.utils.ail import is_phi_assignment
from .seq_to_blocks import SequenceToBlocks

if TYPE_CHECKING:
    from angr.ailment import Address

_l = logging.getLogger(__name__)


def remove_last_statement(node):
    stmt = None

    if type(node) is CodeNode:
        stmt = remove_last_statement(node.node)
    elif type(node) is ailment.Block:
        stmt = node.statements[-1]
        node.statements = node.statements[:-1]
    elif type(node) is MultiNode or type(node) is SequenceNode:
        if node.nodes:
            stmt = remove_last_statement(node.nodes[-1])
            if BaseNode.test_empty_node(node.nodes[-1]):
                node.nodes = node.nodes[:-1]
    elif type(node) is ConditionNode:
        if node.true_node is None and node.false_node is not None:
            stmt = remove_last_statement(node.false_node)
        elif node.true_node is not None and node.false_node is None:
            stmt = remove_last_statement(node.true_node)
        else:
            raise NotImplementedError("More than one last statement exist")
    elif type(node) is LoopNode:
        stmt = remove_last_statement(node.sequence_node)
    else:
        raise NotImplementedError(type(node))

    return stmt


def remove_last_statements(node) -> bool:
    if type(node) is CodeNode:
        return remove_last_statements(node.node)
    if type(node) is ailment.Block:
        if not node.statements:
            return False
        node.statements = node.statements[:-1]
        return True
    if type(node) is MultiNode or type(node) is SequenceNode:
        if node.nodes:
            remove_last_statements(node.nodes[-1])
            if BaseNode.test_empty_node(node.nodes[-1]):
                node.nodes = node.nodes[:-1]
            return True
        return False
    if type(node) is ConditionNode:
        r = False
        if node.true_node is None and node.false_node is not None:
            r |= remove_last_statements(node.false_node)
        if node.true_node is not None and node.false_node is None:
            r |= remove_last_statements(node.true_node)
        return r
    if type(node) is LoopNode:
        return remove_last_statements(node.sequence_node)
    raise NotImplementedError(type(node))


def append_statement(node, stmt):
    if type(node) is CodeNode:
        append_statement(node.node, stmt)
        return
    if type(node) is ailment.Block:
        node.statements.append(stmt)
        return
    if type(node) is MultiNode:
        if node.nodes:
            append_statement(node.nodes[-1], stmt)
        else:
            raise NotImplementedError("MultiNode without nodes")
        return
    if type(node) is SequenceNode:
        if node.nodes:
            append_statement(node.nodes[-1], stmt)
        else:
            raise NotImplementedError("SequenceNode without nodes")
        return

    raise NotImplementedError(type(node))


def replace_last_statement(node, old_stmt, new_stmt):
    if type(node) is CodeNode:
        replace_last_statement(node.node, old_stmt, new_stmt)
        return
    if type(node) is ailment.Block:
        if node.statements[-1] is old_stmt:
            node.statements[-1] = new_stmt
        return
    if type(node) is MultiNode:
        if node.nodes:
            replace_last_statement(node.nodes[-1], old_stmt, new_stmt)
        return
    if type(node) is SequenceNode:
        if node.nodes:
            replace_last_statement(node.nodes[-1], old_stmt, new_stmt)
        return
    if type(node) is ConditionNode:
        if node.true_node is not None:
            replace_last_statement(node.true_node, old_stmt, new_stmt)
        if node.false_node is not None:
            replace_last_statement(node.false_node, old_stmt, new_stmt)
        return

    raise NotImplementedError(type(node))


def extract_jump_targets(stmt):
    """
    Extract concrete goto targets from a Jump or a ConditionalJump statement.

    :param stmt:    The statement to analyze.
    :return:        A list of known concrete jump targets.
    :rtype:         list
    """

    targets = []

    if isinstance(stmt, ailment.Stmt.Jump):
        if isinstance(stmt.target, ailment.Expr.Const):
            targets.append(stmt.target.value)
    elif isinstance(stmt, ailment.Stmt.ConditionalJump):
        if isinstance(stmt.true_target, ailment.Expr.Const):
            targets.append(stmt.true_target.value)
        if isinstance(stmt.false_target, ailment.Expr.Const):
            targets.append(stmt.false_target.value)

    return targets


def switch_extract_cmp_bounds(
    last_stmt: ailment.Stmt.ConditionalJump | ailment.Stmt.Statement,
) -> tuple[Any, int, int] | None:
    """
    Check the last statement of the switch-case header node, and extract lower+upper bounds for the comparison.

    :param last_stmt:   The last statement of the switch-case header node.
    :return:            A tuple of (comparison expression, lower bound, upper bound), or None
    """

    if not isinstance(last_stmt, ailment.Stmt.ConditionalJump):
        return None
    return switch_extract_cmp_bounds_from_condition(last_stmt.condition)


def switch_extract_cmp_bounds_from_condition(cond: ailment.Expr.Expression) -> tuple[Any, int, int] | None:
    # TODO: Add more operations
    if isinstance(cond, ailment.Expr.BinaryOp):
        op = cond.op
        op0, op1 = cond.operands
        if not isinstance(op1, ailment.Expr.Const):
            # swap them
            match op:
                case "CmpLE":
                    op = "CmpGE"
                case "CmpLT":
                    op = "CmpGT"
                case "CmpGE":
                    op = "CmpLE"
                case "CmpGT":
                    op = "CmpLT"
                case _:
                    # unsupported
                    return None
            op0, op1 = op1, op0

        if op in {"CmpLE", "CmpLT"}:
            if not (isinstance(op1, ailment.Expr.Const) and isinstance(op1.value, int)):
                return None
            cmp_ub = op1.value if op == "CmpLE" else op1.value - 1
            cmp_lb = 0
            cmp = op0
            if (
                isinstance(cmp, ailment.Expr.BinaryOp)
                and cmp.op == "Sub"
                and isinstance(cmp.operands[1], ailment.Expr.Const)
                and isinstance(cmp.operands[1].value, int)
            ):
                cmp_ub += cmp.operands[1].value
                cmp_lb += cmp.operands[1].value
                cmp = cmp.operands[0]
            return cmp, cmp_lb, cmp_ub

        if op in {"CmpGE", "CmpGT"}:
            # We got the negated condition here
            #  CmpGE -> CmpLT
            #  CmpGT -> CmpLE
            if not (isinstance(op1, ailment.Expr.Const) and isinstance(op1.value, int)):
                return None
            cmp_ub = op1.value if op == "CmpGT" else op1.value - 1
            cmp_lb = 0
            cmp = op0
            if (
                isinstance(cmp, ailment.Expr.BinaryOp)
                and cmp.op == "Sub"
                and isinstance(cmp.operands[1], ailment.Expr.Const)
                and isinstance(cmp.operands[1].value, int)
            ):
                cmp_ub += cmp.operands[1].value
                cmp_lb += cmp.operands[1].value
                cmp = cmp.operands[0]
            return cmp, cmp_lb, cmp_ub

    return None


def switch_extract_switch_expr_from_jump_target(target: ailment.Expr.Expression) -> ailment.Expr.Expression | None:
    """
    Extract the switch expression from the indirect jump target expression.

    :param target:  The target of the indirect jump statement.
    :return:        The extracted expression if successful, or None otherwise.
    """

    # e.g.: Jump (Conv(32->64, (Load(addr=((0x140000000<64> + (vvar_229{reg 80} * 0x4<64>)) + 0x2290<64>),
    #               size=4,
    #               endness=Iend_LE
    #             ) + 0x140000000<32>)))

    found_load = False
    while True:
        if isinstance(target, ailment.Expr.Convert):
            if target.from_bits < target.to_bits:
                target = target.operand
            else:
                return None
        elif isinstance(target, ailment.Expr.BinaryOp):
            if target.op == "Add":
                # it must be adding the target expr with a constant
                if isinstance(target.operands[0], ailment.Expr.Const):
                    target = target.operands[1]
                elif isinstance(target.operands[1], ailment.Expr.Const):
                    target = target.operands[0]
                else:
                    return None
            elif target.op == "Mul":
                # it must be multiplying the target expr with a constant
                if isinstance(target.operands[0], ailment.Expr.Const):
                    target = target.operands[1]
                elif isinstance(target.operands[1], ailment.Expr.Const):
                    target = target.operands[0]
                else:
                    return None
            elif target.op == "And":
                # it must be and-ing the target expr with a constant
                if (
                    isinstance(target.operands[1], ailment.Expr.VirtualVariable)
                    and isinstance(target.operands[0], ailment.Expr.Const)
                ) or (
                    isinstance(target.operands[0], ailment.Expr.VirtualVariable)
                    and isinstance(target.operands[1], ailment.Expr.Const)
                ):
                    break
                return None
            else:
                return None
        elif isinstance(target, ailment.Expr.Load):
            # we want the address!
            found_load = True
            target = target.addr
        elif isinstance(target, ailment.Expr.VirtualVariable):
            break
        else:
            return None
    return target if found_load else None


def switch_extract_bitwiseand_jumptable_info(last_stmt: ailment.Stmt.Jump) -> tuple[Any, int, int] | None:
    """
    Check the last statement of the switch-case header node (whose address is loaded from a jump table and computed
    using an index) and extract necessary information for rebuilding the switch-case construct.

    An example of the statement:

    Goto(Conv(32->s64, (
        Load(addr=(0x4530e4<64> + (Conv(32->64, (Conv(64->32, vvar_287{reg 32}) & 0x3<32>)) * 0x4<64>)),
             size=4, endness=Iend_LE) + 0x4530e4<32>))
    )

    Another example:

    Load(addr=(((vvar_9{reg 36} & 0x3<32>) * 0x4<32>) + 0x42cd28<32>), size=4, endness=Iend_LE)

    :param last_stmt:   The last statement of the switch-case header node.
    :return:            A tuple of (index expression, lower bound, upper bound), or None
    """

    if not isinstance(last_stmt, ailment.Stmt.Jump):
        return None

    # unpack the target expression
    target = last_stmt.target
    jump_addr_offset = None
    jumptable_load_addr = None
    while True:
        if isinstance(target, ailment.Expr.Convert) and (
            (target.from_bits == 32 and target.to_bits == 64) or (target.from_bits == 16 and target.to_bits == 32)
        ):
            target = target.operand
            continue
        if isinstance(target, ailment.Expr.BinaryOp) and target.op == "Add":
            if isinstance(target.operands[0], ailment.Expr.Const) and isinstance(target.operands[1], ailment.Expr.Load):
                jump_addr_offset = target.operands[0].value
                jumptable_load_addr = target.operands[1].addr
                break
            if isinstance(target.operands[1], ailment.Expr.Const) and isinstance(target.operands[0], ailment.Expr.Load):
                jump_addr_offset = target.operands[1].value
                jumptable_load_addr = target.operands[0].addr
                break
            return None
        if isinstance(target, ailment.Expr.Const):
            return None
        if isinstance(target, ailment.Expr.Load):
            jumptable_load_addr = target.addr
            jump_addr_offset = 0
            break
        break

    if jump_addr_offset is None or jumptable_load_addr is None:
        return None

    # parse jumptable_load_addr
    jumptable_offset = None
    jumptable_base_addr = None
    if isinstance(jumptable_load_addr, ailment.Expr.BinaryOp) and jumptable_load_addr.op == "Add":
        if isinstance(jumptable_load_addr.operands[0], ailment.Expr.Const):
            jumptable_base_addr = jumptable_load_addr.operands[0]
            jumptable_offset = jumptable_load_addr.operands[1]
        elif isinstance(jumptable_load_addr.operands[1], ailment.Expr.Const):
            jumptable_offset = jumptable_load_addr.operands[0]
            jumptable_base_addr = jumptable_load_addr.operands[1]

    if jumptable_offset is None or jumptable_base_addr is None:
        return None

    # parse jumptable_offset
    expr = jumptable_offset
    coeff = None
    index_expr = None
    lb = None
    ub: int | None = None
    while expr is not None:
        if isinstance(expr, ailment.Expr.BinaryOp):
            if expr.op == "Mul":
                if isinstance(expr.operands[1], ailment.Expr.Const):
                    coeff = expr.operands[1].value
                    expr = expr.operands[0]
                elif isinstance(expr.operands[0], ailment.Expr.Const):
                    coeff = expr.operands[0].value
                    expr = expr.operands[1]
                else:
                    return None
            elif expr.op == "And":
                masks = {0x1, 0x3, 0x7, 0xF, 0x1F, 0x3F, 0x7F, 0xFF, 0x1FF, 0x3FF}
                if isinstance(expr.operands[1], ailment.Expr.Const) and expr.operands[1].value in masks:
                    lb = 0
                    ub = expr.operands[1].value  # type:ignore
                    index_expr = expr
                    break
                if isinstance(expr.operands[0], ailment.Expr.Const) and expr.operands[1].value in masks:
                    lb = 0
                    ub = expr.operands[0].value  # type:ignore
                    index_expr = expr
                    break
                return None
            else:
                return None
        elif isinstance(expr, ailment.Expr.Convert):
            if expr.is_signed is False:
                expr = expr.operand
            else:
                return None
        else:
            break

    if coeff is not None and index_expr is not None and lb is not None and ub is not None:
        return index_expr, lb, ub
    return None


def get_ast_subexprs(claripy_ast):
    queue = [claripy_ast]
    while queue:
        ast = queue.pop(0)
        if ast.op == "And":
            queue += ast.args[1:]
            yield ast.args[0]
        else:
            yield ast


def insert_node(parent, insert_location: str, node, node_idx: int, label=None):
    if insert_location not in {"before", "after"}:
        raise ValueError('"insert_location" must be either "before" or "after"')

    if isinstance(parent, SequenceNode):
        if insert_location == "before":
            parent.nodes.insert(node_idx, node)
        else:  # if insert_location == "after":
            parent.nodes.insert(node_idx + 1, node)
    elif isinstance(parent, CodeNode):
        # Make a new sequence node
        if insert_location == "before":
            seq = SequenceNode(parent.addr, nodes=[node, parent.node])
        else:  # if insert_location == "after":
            seq = SequenceNode(parent.addr, nodes=[parent.node, node])
        parent.node = seq
    elif isinstance(parent, MultiNode):
        if insert_location == "before":
            parent.nodes.insert(node_idx, node)
        else:
            parent.nodes.insert(node_idx + 1, node)
    elif isinstance(parent, ConditionNode):
        if node_idx == 0:
            # true node
            if not isinstance(parent.true_node, SequenceNode):
                if parent.true_node is None:
                    parent.true_node = SequenceNode(parent.addr, nodes=[])
                else:
                    parent.true_node = SequenceNode(parent.true_node.addr, nodes=[parent.true_node])
            insert_node(parent.true_node, insert_location, node, 0)
        else:
            # false node
            if not isinstance(parent.false_node, SequenceNode):
                if parent.false_node is None:
                    parent.false_node = SequenceNode(parent.addr, nodes=[])
                else:
                    parent.false_node = SequenceNode(parent.false_node.addr, nodes=[parent.false_node])
            insert_node(parent.false_node, insert_location, node, 0)
    elif isinstance(parent, CascadingConditionNode):
        cond, child_node = parent.condition_and_nodes[node_idx]
        if not isinstance(child_node, SequenceNode):
            child_node = SequenceNode(child_node.addr, nodes=[child_node])
            parent.condition_and_nodes[node_idx] = (cond, child_node)
        insert_node(child_node, insert_location, node, 0)
    elif isinstance(parent, SwitchCaseNode):
        # note that this case will be hit only when the parent node is not a container, such as SequenceNode or
        # MultiNode. we always need to create a new SequenceNode and replace the original node in place.

        if label == "switch_expr":
            raise TypeError("You cannot insert a node after an expression.")
        if label == "case":
            # node_idx is the case number.
            if insert_location == "after":
                new_nodes = [parent.cases[node_idx], node]
            elif insert_location == "before":
                new_nodes = [node, parent.cases[node_idx]]
            else:
                raise TypeError(f'Unsupported insert_location value "{insert_location}".')
            seq = SequenceNode(new_nodes[0].addr, nodes=new_nodes)
            parent.cases[node_idx] = seq
        elif label == "default":
            if insert_location == "after":
                new_nodes = [parent.default_node, node]
            elif insert_location == "before":
                new_nodes = [node, parent.default_node]
            else:
                raise TypeError(f"Unsupported 'insert_location' value {insert_location!r}.")
            seq = SequenceNode(new_nodes[0].addr, nodes=new_nodes)
            parent.default_node = seq
        else:
            raise TypeError(
                f'Unsupported label value "{label}". Must be one of the following: switch_expr, case, ' f"default."
            )
    elif isinstance(parent, LoopNode):
        if label == "condition":
            raise ValueError("Cannot insert nodes into a condition expression.")
        if label == "body":
            if not isinstance(parent.sequence_node, SequenceNode):
                parent.sequence_node = SequenceNode(parent.sequence_node.addr, nodes=[parent.sequence_node])
            insert_node(parent.sequence_node, insert_location, node, node_idx)
        else:
            raise NotImplementedError(label)
    else:
        raise NotImplementedError(type(parent))


def _merge_ail_nodes(graph, node_a: ailment.Block, node_b: ailment.Block) -> ailment.Block:
    in_edges = list(graph.in_edges(node_a, data=True))
    out_edges = list(graph.out_edges(node_b, data=True))

    a_ogs = graph.nodes[node_a].get("original_nodes", [])
    b_ogs = graph.nodes[node_b].get("original_nodes", [])
    new_node = node_a.copy() if node_a.addr <= node_b.addr else node_b.copy()
    old_node = node_b if new_node == node_a else node_a
    # remove jumps in the middle of nodes when merging
    if new_node.statements and isinstance(new_node.statements[-1], ailment.Stmt.Jump):
        new_node.statements = new_node.statements[:-1]
    new_node.statements += old_node.statements
    new_node.original_size += old_node.original_size

    graph.remove_node(node_a)
    graph.remove_node(node_b)

    if new_node is not None:
        graph.add_node(new_node, original_nodes=a_ogs + b_ogs)
        for src, _, data in in_edges:
            if src is node_b:
                src = new_node
            graph.add_edge(src, new_node, **data)

        for _, dst, data in out_edges:
            if dst is node_a:
                dst = new_node
            graph.add_edge(new_node, dst, **data)

    return new_node


def to_ail_supergraph(transition_graph: networkx.DiGraph, allow_fake=False) -> networkx.DiGraph:
    """
    Takes an AIL graph and converts it into a AIL graph that treats calls and redundant jumps
    as parts of a bigger block instead of transitions. Calls to returning functions do not terminate basic blocks.

    Based on region_identifier super_graph

    :return: A converted super transition graph
    """
    # make a copy of the graph
    transition_graph = networkx.DiGraph(transition_graph)
    networkx.set_node_attributes(transition_graph, {node: [node] for node in transition_graph.nodes}, "original_nodes")

    while True:
        for src, dst, data in transition_graph.edges(data=True):
            type_ = data.get("type", None)

            if len(list(transition_graph.successors(src))) == 1 and len(list(transition_graph.predecessors(dst))) == 1:
                # calls in the middle of blocks OR boring jumps
                if (type_ == "fake_return") or (src.addr + src.original_size == dst.addr) or allow_fake:
                    _merge_ail_nodes(transition_graph, src, dst)
                    break

            # calls to functions with no return
            elif type_ == "call":
                transition_graph.remove_node(dst)
                break
        else:
            break

    return transition_graph


def is_empty_node(node) -> bool:
    if isinstance(node, ailment.Block):
        return not node.statements
    if isinstance(node, MultiNode):
        return all(is_empty_node(n) for n in node.nodes)
    if isinstance(node, SequenceNode):
        return all(is_empty_node(n) for n in node.nodes)
    return False


def is_empty_or_label_only_node(node) -> bool:
    if isinstance(node, ailment.Block):
        return not has_nonlabel_statements(node)
    if isinstance(node, MultiNode):
        return all(is_empty_or_label_only_node(n) for n in node.nodes)
    if isinstance(node, SequenceNode):
        return all(is_empty_or_label_only_node(n) for n in node.nodes)
    return False


def has_nonlabel_statements(block: ailment.Block) -> bool:
    return bool(block.statements and any(not isinstance(stmt, ailment.Stmt.Label) for stmt in block.statements))


def has_nonlabel_nonphi_statements(block: ailment.Block) -> bool:
    return bool(
        block.statements
        and any(not (isinstance(stmt, ailment.Stmt.Label) or is_phi_assignment(stmt)) for stmt in block.statements)
    )


def first_nonlabel_statement(block: ailment.Block | MultiNode) -> ailment.Stmt.Statement | None:
    if isinstance(block, MultiNode):
        for n in block.nodes:
            stmt = first_nonlabel_statement(n)
            if stmt is not None:
                return stmt
        return None

    for stmt in block.statements:
        if not isinstance(stmt, ailment.Stmt.Label):
            return stmt
    return None


def first_nonlabel_statement_id(block: ailment.Block) -> int | None:
    for idx, stmt in enumerate(block.statements):
        if not isinstance(stmt, ailment.Stmt.Label):
            return idx
    return len(block.statements)


def first_nonlabel_nonphi_statement(block: ailment.Block | MultiNode) -> ailment.Stmt.Statement | None:
    if isinstance(block, MultiNode):
        for n in block.nodes:
            stmt = first_nonlabel_nonphi_statement(n)
            if stmt is not None:
                return stmt
        return None

    for stmt in block.statements:
        if not (isinstance(stmt, ailment.Stmt.Label) or is_phi_assignment(stmt)):
            return stmt
    return None


def last_nonlabel_statement(block: ailment.Block) -> ailment.Stmt.Statement | None:
    for stmt in reversed(block.statements):
        if not isinstance(stmt, ailment.Stmt.Label):
            return stmt
    return None


def last_node(node: BaseNode) -> BaseNode | ailment.Block | None:
    """
    Get the last node in a sequence or code node.
    """
    if isinstance(node, CodeNode):
        return last_node(node.node)
    if isinstance(node, SequenceNode):
        if not node.nodes:
            return None
        return last_node(node.nodes[-1])
    return node


def first_nonlabel_node(seq: SequenceNode) -> BaseNode | ailment.Block | None:
    for node in seq.nodes:
        inner_node = node.node if isinstance(node, CodeNode) else node
        if isinstance(inner_node, ailment.Block) and not has_nonlabel_statements(inner_node):
            continue
        return node
    return None


def first_nonlabel_nonphi_node(seq: SequenceNode) -> BaseNode | ailment.Block | None:
    for node in seq.nodes:
        inner_node = node.node if isinstance(node, CodeNode) else node
        if isinstance(inner_node, ailment.Block) and not has_nonlabel_nonphi_statements(inner_node):
            continue
        return node
    return None


def remove_labels(graph: networkx.DiGraph):
    new_graph = networkx.DiGraph()
    nodes_map = {}
    for node in graph:
        node_copy = node.copy()
        node_copy.statements = [stmt for stmt in node_copy.statements if not isinstance(stmt, ailment.Stmt.Label)]
        nodes_map[node] = node_copy

    for old_node in graph.nodes:
        new_graph.add_node(nodes_map[old_node])

    for src, dst, data in graph.edges(data=True):
        new_graph.add_edge(nodes_map[src], nodes_map[dst], **data)

    return new_graph


def add_labels(graph: networkx.DiGraph):
    new_graph = networkx.DiGraph()
    nodes_map = {}
    for node in graph:
        lbl = ailment.Stmt.Label(None, f"LABEL_{node.addr:x}", ins_addr=node.addr, block_idx=node.idx)
        node_copy = node.copy()
        node_copy.statements = [lbl, *node_copy.statements]
        nodes_map[node] = node_copy

    for old_node in graph.nodes:
        new_graph.add_node(nodes_map[old_node])

    for src, dst in graph.edges:
        new_graph.add_edge(nodes_map[src], nodes_map[dst])

    return new_graph


def update_labels(graph: networkx.DiGraph):
    """
    A utility function to recreate the labels for every node in an AIL graph. This useful when you are working with
    a graph where only _some_ of the nodes have labels.
    """
    return add_labels(remove_labels(graph))


def _flatten_structured_node(packed_node: SequenceNode | MultiNode) -> list[ailment.Block]:
    if not packed_node or not packed_node.nodes:
        return []

    blocks = []
    if packed_node.nodes is not None:
        for _node in packed_node.nodes:
            if isinstance(_node, (SequenceNode, MultiNode)):
                blocks += _flatten_structured_node(_node)
            else:
                blocks.append(_node)

    return blocks


def _find_node_in_graph(node: ailment.Block, graph: networkx.DiGraph) -> ailment.Block | None:
    for bb in graph:
        if isinstance(bb, ailment.Block) and bb.addr == node.addr and bb.idx == node.idx:
            return bb
    return None


def structured_node_has_multi_predecessors(
    node: SequenceNode | MultiNode | ailment.Block, graph: networkx.DiGraph
) -> bool:
    if graph is None:
        return False

    first_block = None
    if isinstance(node, (SequenceNode, MultiNode)) and node.nodes:
        flat_blocks = _flatten_structured_node(node)
        node = flat_blocks[0]

    if isinstance(node, ailment.Block):
        first_block = node

    if first_block is not None:
        graph_node = _find_node_in_graph(first_block, graph)
        if graph_node is not None:
            return len(list(graph.predecessors(graph_node))) > 1

    return False


def structured_node_is_simple_return(
    node: SequenceNode | MultiNode, graph: networkx.DiGraph, use_packed_successors=False
) -> bool:
    """
    Will check if a "simple return" is contained within the node a simple returns looks like this:
    if (cond) {
      // simple return
      ...
      return 0;
    }
    ...

    Returns true on any block ending in linear statements and a return.
    """
    # sanity check: we need a graph to understand returning blocks
    if graph is None:
        return False

    last_block = None
    if isinstance(node, (SequenceNode, MultiNode)) and node.nodes:
        flat_blocks = _flatten_structured_node(node)
        if all(isinstance(block, ailment.Block) for block in flat_blocks):
            last_block = flat_blocks[-1]
    elif isinstance(node, ailment.Block):
        last_block = node

    valid_last_stmt = last_block is not None
    if last_block is not None and last_block.statements:
        valid_last_stmt = not isinstance(last_block.statements[-1], (ailment.Stmt.ConditionalJump, ailment.Stmt.Jump))

    if use_packed_successors:
        last_block = node

    if valid_last_stmt:
        # note that the block may not be the same block in the AIL graph post dephication. we must find the block again
        # in the graph.
        assert isinstance(last_block, ailment.Block)
        last_graph_block = _find_node_in_graph(last_block, graph)
        if last_graph_block is not None:
            succs = list(graph.successors(last_graph_block))
            return not succs or succs == [last_graph_block]
    return False


def structured_node_is_simple_return_strict(node: BaseNode | SequenceNode | MultiNode | ailment.Block) -> bool:
    """
    Returns True iff the node exclusively contains a return statement.
    """
    if isinstance(node, (SequenceNode, MultiNode)) and node.nodes:
        flat_blocks = _flatten_structured_node(node)
        if len(flat_blocks) != 1:
            return False
        node = flat_blocks[-1]

    return (
        isinstance(node, ailment.Block)
        and len(node.statements) == 1
        and isinstance(node.statements[0], ailment.Stmt.Return)
    )


def is_statement_terminating(stmt: ailment.statement.Statement, functions) -> bool:
    if isinstance(stmt, ailment.Stmt.Return):
        return True
    if isinstance(stmt, ailment.Stmt.Call) and isinstance(stmt.target, ailment.Expr.Const):
        # is it calling a non-returning function?
        target_func_addr = stmt.target.value
        try:
            func = functions.get_by_addr(target_func_addr)
            return func.returning is False
        except KeyError:
            pass
    return False


class _PeepholeExprsWalker(ailment.AILBlockRewriter):
    """
    Walker to apply peephole optimizers (1)
    """

    def __init__(self, *args, expr_opts: list[PeepholeOptimizationExprBase], **kwargs):
        self.expr_opts = expr_opts
        self.any_update = False

        super().__init__(*args, **kwargs)

    def _handle_expr(
        self, expr_idx: int, expr: ailment.Expr.Expression, stmt_idx: int, stmt: ailment.Stmt.Statement | None, block
    ) -> ailment.Expression:
        # process the expr
        expr = super()._handle_expr(expr_idx, expr, stmt_idx, stmt, block)
        old_expr = expr

        redo = True
        while redo:
            redo = False
            for expr_opt in self.expr_opts:
                if isinstance(expr, expr_opt.expr_classes):
                    r = expr_opt.optimize(expr, stmt_idx=stmt_idx, block=block)
                    if r is not None and r is not expr:
                        expr = r
                        redo = True
                        break

        if expr is not old_expr:
            self.any_update = True

        return expr


def peephole_optimize_exprs(block, expr_opts):
    # run expression optimizers
    walker = _PeepholeExprsWalker(expr_opts=expr_opts)
    walker.walk(block)
    return walker.any_update


class _PeepholeExprWalker(ailment.AILBlockRewriter):
    """
    Walker to apply peephole optimizers (2)
    """

    def __init__(self, *args, expr_opts: list[PeepholeOptimizationExprBase], **kwargs):
        self.expr_opts = expr_opts

        super().__init__(*args, **kwargs)

    def _handle_expr(
        self, expr_idx: int, expr: ailment.Expr.Expression, stmt_idx: int, stmt: ailment.Stmt.Statement | None, block
    ) -> ailment.Expression:
