angr 9.2.192__cp311-cp311-macosx_10_12_x86_64.whl

angr/analyses/decompiler/optimization_passes/code_motion.py

@@ -0,0 +1,362 @@
+from __future__ import annotations
+import itertools
+import logging
+
+from angr.ailment import Block
+from angr.ailment.statement import Jump, ConditionalJump, Statement, DirtyStatement
+import networkx as nx
+
+from angr.analyses.decompiler.optimization_passes.optimization_pass import OptimizationPass, OptimizationPassStage
+from angr.analyses.decompiler.block_similarity import is_similar, index_of_similar_stmts
+from angr.analyses.decompiler.block_io_finder import BlockIOFinder
+from angr.analyses.decompiler.utils import to_ail_supergraph, remove_labels, add_labels
+
+_l = logging.getLogger(name=__name__)
+
+
+class CodeMotionOptimization(OptimizationPass):
+    """
+    Moves common statements out of blocks that share the same predecessors or the same
+    successors. This is done to reduce the number of statements in a block and to make the
+    blocks more similar to each other.
+
+    As an example:
+    if (x) {
+        b = 2;
+        a = 1;
+        c = 3;
+    } else {
+        b = 2;
+        c = 3;
+    }
+
+    Will be turned into:
+    if (x) {
+        a = 1;
+    }
+    b = 2;
+    c = 3;
+
+    Current limitations (for very conservative operations):
+    - moving statements above conditional jumps is not supported
+    - only immediate children and parents are considered for moving statements
+    - when moving statements down, a block is only considered if already has a matching statement at the end
+    """
+
+    ARCHES = None
+    PLATFORMS = None
+    NAME = "Merge common statements in sub-scopes"
+    STAGE = OptimizationPassStage.AFTER_GLOBAL_SIMPLIFICATION
+    DESCRIPTION = __doc__
+
+    def __init__(self, func, *args, max_iters=10, node_idx_start: int = 0, **kwargs):
+        super().__init__(func, *args, **kwargs)
+        self._node_idx_start = node_idx_start
+        self._max_optimization_runs = max_iters
+        self.analyze()
+
+    def _check(self):
+        return True, None
+
+    def _analyze(self, cache=None):
+        _l.warning("CodeMotionOptimization is likely broken right now, use with caution!")
+        optimization_runs = 0
+        graph_copy = remove_labels(nx.DiGraph(self._graph))
+        updates = True
+        graph_changed = False
+        while optimization_runs < self._max_optimization_runs and updates:
+            optimization_runs += 1
+            super_graph = to_ail_supergraph(graph_copy)
+            updates, updated_blocks = self._move_common_code(super_graph)
+            if updates:
+                critical_fail = self.update_graph_with_super_edits(graph_copy, super_graph, updated_blocks)
+                if critical_fail:
+                    _l.error("Critical failure in updating graph with super edits, aborting")
+                    break
+                graph_changed = True
+
+        if graph_changed:
+            self.out_graph = add_labels(graph_copy)
+
+    @staticmethod
+    def update_graph_with_super_edits(
+        original_graph: nx.DiGraph, super_graph: nx.DiGraph, updated_blocks: dict[Block, Block]
+    ) -> bool:
+        """
+        This function updates an graph when doing block edits on a supergraph version of that same graph.
+        The updated blocks must be provided as a dictionary where the keys are original block in the supergraph and
+        the values are the new blocks that should replace them.
+
+        The supergraph MUST be generated using the to_ail_supergraph function, since it stores the original nodes
+        each super node represents. This is necessary to update the original graph with the new super nodes.
+        """
+        og_to_super = {}
+        for old_super, new_super in updated_blocks.items():
+            original_blocks = super_graph.nodes[old_super]["original_nodes"]
+            for original_block in original_blocks:
+                og_to_super[original_block] = new_super
+
+        for old_super, new_super in updated_blocks.items():
+            original_blocks = super_graph.nodes[old_super]["original_nodes"]
+            first_node_preds = []
+            last_node_preds = []
+            for original_block in original_blocks:
+                if original_block not in original_graph.nodes:
+                    return True
+
+                external_preds = [
+                    pred for pred in original_graph.predecessors(original_block) if pred not in original_blocks
+                ]
+                external_succs = [
+                    succ for succ in original_graph.successors(original_block) if succ not in original_blocks
+                ]
+                if external_preds:
+                    first_node_preds = external_preds
+                if external_succs:
+                    last_node_preds = external_succs
+
+            original_graph.remove_nodes_from(original_blocks)
+            original_graph.add_node(new_super)
+            for pred in first_node_preds:
+                original_graph.add_edge(og_to_super.get(pred, pred), new_super)
+            for succ in last_node_preds:
+                original_graph.add_edge(new_super, og_to_super.get(succ, succ))
+
+        return False
+
+    def _move_common_code(self, graph) -> tuple[bool, dict[Block, Block] | None]:
+        """
+        Does two things at a high level:
+        1. rearrange code in blocks to maximize the number of similar statements at the end of the block
+        2. move common code out of blocks
+
+        To understand the limitations of this approach, see the TODOs.
+        """
+        # TODO: how can you handle an odd-numbered switch case? or many blocks with the same child?
+        for b0, b1 in itertools.combinations(graph.nodes, 2):
+            if (
+                b0 is b1
+                or not b0.statements
+                or not b1.statements
+                or any(isinstance(stmt, DirtyStatement) for stmt in b0.statements + b1.statements)
+                or is_similar(b0, b1)
+            ):
+                continue
+
+            # TODO: add support for moving code to a shared parent block, which requires that we figure out how to
+            #   move code above conditional jumps. Hard since you need to know if the condition executes code.
+            # TODO: also, how do you deal with short-circuiting, which is a region parent, not just a block?
+
+            # target any blocks that have a shared child and move common code to the child
+            b0_succs = list(graph.successors(b0))
+            b1_succs = list(graph.successors(b1))
+            if (len(b0_succs) == len(b1_succs) == 1) and b0_succs[0] == b1_succs[0]:
+                common_succ = b0_succs[0]
+                common_succ_preds = list(graph.predecessors(common_succ))
+                # you can only safely move code to a child if all the common_succ's preds are the ones
+                # we are moving code from (2 nodes).
+                if all(csp in (b0, b1) for csp in common_succ_preds):
+                    success, updated_blocks = self._move_common_code_to_child(b0, b1, common_succ)
+                    if success:
+                        return True, updated_blocks
+
+        return False, None
+
+    def _move_common_code_to_parent(self, b0: Block, b1: Block, parent: Block):
+        # TODO: this function does not work yet because you need to figure out if you can move a stmt above
+        #   a conditional jump, which requires cross-block analysis
+        changed, new_b0, new_b1 = self._make_stmts_end_similar(b0, b1, up=True)
+        if not changed:
+            return False, None
+
+        # move the longest common suffix to the parent
+        new_b0_stmts = new_b0.statements
+        new_b1_stmts = new_b1.statements
+        common_len = 0
+        for idx, new_b0_stmt in enumerate(new_b0_stmts):
+            if not new_b0_stmt.likes(new_b1_stmts[idx]):
+                break
+            common_len += 1
+
+        if not common_len:
+            raise ValueError("No common statements found, this is unexpected")
+
+        common_stmts = [new_b0_stmts.pop(0) for _ in range(common_len)]
+        for _ in range(common_len):
+            new_b1_stmts.pop(0)
+
+        parent_stmts = parent.statements.copy() or []
+        if isinstance(parent_stmts[-1], (ConditionalJump, Jump)):
+            parent_stmts = parent_stmts[:-1] + common_stmts + [parent_stmts[-1]]
+        new_parent = parent.copy(statements=parent_stmts)
+
+        return True, {b0: new_b0, b1: new_b1, parent: new_parent}
+
+    def _move_common_code_to_child(self, b0: Block, b1: Block, child: Block):
+        changed, new_b0, new_b1 = self._make_stmts_end_similar(b0, b1, down=True)
+        if not changed:
+            return False, None
+
+        # move the longest common suffix to the parent
+        new_b0_stmts = new_b0.statements
+        new_b1_stmts = new_b1.statements
+        common_len = 0
+        # start from the -1 index and go backwards
+        for idx in range(-1, -(min(len(new_b0_stmts), len(new_b1_stmts))) - 1, -1):
+            if not new_b0_stmts[idx].likes(new_b1_stmts[idx]):
+                break
+            common_len += 1
+
+        if not common_len:
+            raise ValueError("No common statements found, this is unexpected")
+
+        common_stmts = [new_b0_stmts.pop() for _ in range(common_len)]
+        for _ in range(common_len):
+            new_b1_stmts.pop()
+
+        child_stmts = child.statements.copy() or []
+        new_child = child.copy(statements=common_stmts[::-1] + child_stmts)
+
+        return True, {b0: new_b0, b1: new_b1, child: new_child}
+
+    def _make_stmts_end_similar(
+        self, b0: Block, b1: Block, up=False, down=False
+    ) -> tuple[bool, Block | None, Block | None]:
+        """
+        This algorithm attempts to rearrange two blocks to have the longest common sequence of statements
+        at either ends of the blocks. It is flawed in that it currently only attempts to do this rearrangement
+        if the blocks have at least one matching statement at the end.
+
+        This algorithm iteratively removes statements from the ends of the blocks and then attempts to match
+        the ends of the blocks. It will only do this if one of the two ends has a matching statement in the other.
+        """
+        self._assert_up_or_down(up, down)
+        # copy the statements while filtering out statements that are not needed in the specific
+        # movement case (up or down)
+        curr_stmts = {}
+        for blk in (b0, b1):
+            new_stmts = blk.statements.copy()
+            if down:
+                last_stmt = new_stmts[-1]
+                if isinstance(last_stmt, Jump):
+                    new_stmts.pop()
+                elif isinstance(last_stmt, ConditionalJump):
+                    _l.warning("ConditionalJump at the end of block %s, this should never happen!", blk)
+                    return False, None, None
+
+            curr_stmts[blk] = new_stmts
+        if not curr_stmts[b0] or not curr_stmts[b1]:
+            return False, None, None
+
+        # attempt to do a swapping algorithm to maximize the number of similar statements at the end
+        changed = True
+        stmts_updated = False
+        matched_stmts = {b0: [], b1: []}
+        max_iters = len(curr_stmts[b0]) * len(curr_stmts[b1])
+        curr_iters = 0
+        while changed and curr_iters < max_iters:
+            changed = False
+            try_next_swap = False
+            for tgt0, tgt1 in ((b0, b1), (b1, b0)):
+                t0_stmts = curr_stmts[tgt0]
+                t1_stmts = curr_stmts[tgt1]
+                if not t0_stmts or not t1_stmts:
+                    break
+
+                if up:
+                    # maximize up
+                    while t0_stmts and t1_stmts:
+                        if t0_stmts[0].likes(t1_stmts[0]):
+                            matched_stmts[b0].append((t0_stmts.pop(0), 0))
+                            matched_stmts[b1].append((t1_stmts.pop(0), 0))
+                            try_next_swap = True
+                        else:
+                            break
+                    if not t0_stmts or not t1_stmts:
+                        break
+                elif down:
+                    # maximize down
+                    while t0_stmts and t1_stmts:
+                        if t0_stmts[-1].likes(t1_stmts[-1]):
+                            matched_stmts[b0].append((t0_stmts.pop(), -1))
+                            matched_stmts[b1].append((t1_stmts.pop(), -1))
+                            try_next_swap = True
+                        else:
+                            break
+                    if not t0_stmts or not t1_stmts:
+                        break
+
+                if not try_next_swap:
+                    continue
+
+                stmts_updated = True
+                swap_occurred, new_stmts = self._maximize_ends(t0_stmts, t1_stmts, up=up, down=down)
+                if swap_occurred:
+                    changed = True
+                    curr_stmts[b0], curr_stmts[b1] = new_stmts
+                    break
+
+                try_next_swap = True
+
+            curr_iters += 1
+            if curr_iters > max_iters:
+                raise ValueError("Exceeded max iterations, likely stuck in infinite loop")
+
+        # did any changes occur?
+        if not stmts_updated:
+            return False, None, None
+
+        # reconstruct the blocks and return them
+        new_blks = {}
+        for blk in (b0, b1):
+            new_stmts = curr_stmts[blk]
+            for stmt, idx in matched_stmts[blk][::-1]:
+                if idx == -1:
+                    new_stmts.append(stmt)
+                else:
+                    new_stmts.insert(idx, stmt)
+
+            new_blks[blk] = blk.copy(statements=new_stmts)
+
+        return True, new_blks[b0], new_blks[b1]
+
+    def _maximize_ends(
+        self, b0_stmts, b1_stmts, up=False, down=False
+    ) -> tuple[bool, tuple[list[Statement], list[Statement]]]:
+        self._assert_up_or_down(up, down)
+
+        similar_stmt = b0_stmts[0] if up else b0_stmts[-1]
+        idx_similar = index_of_similar_stmts([similar_stmt], b1_stmts)
+        if idx_similar is None or len(b1_stmts) == 1:
+            return False, (b0_stmts, b1_stmts)
+
+        target_stmt = b1_stmts[idx_similar]
+        success, new_b1_stmts = self._move_to_end(target_stmt, b1_stmts, up=up, down=down)
+        return (success and (b1_stmts != new_b1_stmts)), (b0_stmts, new_b1_stmts)
+
+    def _move_to_end(self, stmt, stmts, up=False, down=False) -> tuple[bool, list[Statement]]:
+        """
+        Attempts to move a stmt to either the top or the bottom of stmts.
+        It does this by attempting to swap, 1 by 1, in either direction it is targeting.
+        """
+        new_stmts = stmts.copy()
+        stmt_idx = new_stmts.index(stmt)
+        swap_offset = -1 if up else 1
+        swap_order = range(stmt_idx + 1, len(new_stmts)) if down else range(stmt_idx - 1, -1, -1)
+        io_finder = BlockIOFinder(new_stmts, self.project)
+        for swap_pos in swap_order:
+            src_stmt = new_stmts[stmt_idx]
+            if io_finder.can_swap(src_stmt, new_stmts, 1 if down else -1):
+                new_stmts[stmt_idx], new_stmts[swap_pos] = new_stmts[swap_pos], new_stmts[stmt_idx]
+                stmt_idx += swap_offset
+            else:
+                return False, stmts
+
+        return True, new_stmts
+
+    @staticmethod
+    def _assert_up_or_down(up, down):
+        if up and down:
+            raise ValueError("Cannot maximize both up and down")
+        if not up and not down:
+            raise ValueError("Must maximize either up or down")

angr/analyses/decompiler/optimization_passes/condition_constprop.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from collections import defaultdict
+
+import networkx
+
+from angr.ailment import AILBlockRewriter, Block, Expression
+from angr.ailment.statement import ConditionalJump, Statement, Assignment
+from angr.ailment.expression import Const, BinaryOp, VirtualVariable
+
+from angr.analyses.decompiler.utils import first_nonlabel_nonphi_statement
+from angr.utils.graph import dominates
+from angr.utils.timing import timethis
+from .optimization_pass import OptimizationPass, OptimizationPassStage
+
+if TYPE_CHECKING:
+    from angr.analyses.s_reaching_definitions import SRDAModel
+
+
+class ConstantCondition:
+    """
+    Describes an opportunity for replacing a vvar with a constant value.
+    """
+
+    def __init__(self, vvar_id: int, value: Const, block_addr: int, block_idx: int | None):
+        self.vvar_id = vvar_id
+        self.value = value
+        self.block_addr = block_addr
+        self.block_idx = block_idx
+
+    def __repr__(self):
+        return f"<ConstCond vvar_{self.vvar_id} == {self.value} since {self.block_addr:#x}-{self.block_idx}>"
+
+
+class CCondPropBlockWalker(AILBlockRewriter):
+    """
+    Block walker for ConditionConstantPropagation to replace vvars with constant values.
+    """
+
+    def __init__(self, vvar_id: int, const_value: Const):
+        super().__init__()
+        self._new_block: Block | None = None  # output
+        self.vvar_id = vvar_id
+        self.const_value = const_value
+        self.abort = False
+
+    def _handle_stmt(self, stmt_idx: int, stmt: Statement, block: Block):  # type: ignore
+        if self.abort:
+            return stmt
+
+        return super()._handle_stmt(stmt_idx, stmt, block)
+
+    def _handle_Assignment(self, stmt_idx: int, stmt: Assignment, block: Block | None) -> Statement:
+        if isinstance(stmt.dst, VirtualVariable) and stmt.dst.varid == self.vvar_id:
+            # we see the assignment of this virtual variable; this is the original block that creates this variable
+            # and checks if this variable is equal to a constant value. as such, we stop processing this block.
+            # an example appears in binary 1de5cda760f9ed80bb6f4a35edcebc86ccec14c49cf4775ddf2ffc3e05ff35f4, function
+            # 0x4657C0, blocks 0x465bd6 and 0x465a5c
+            self.abort = True
+            return stmt
+        return super()._handle_Assignment(stmt_idx, stmt, block)
+
+    def _handle_VirtualVariable(  # type: ignore
+        self, expr_idx: int, expr: VirtualVariable, stmt_idx: int, stmt: Statement, block: Block | None
+    ) -> Expression:
+        if expr.varid == self.vvar_id and not (
+            isinstance(stmt, Assignment) and isinstance(stmt.dst, VirtualVariable) and stmt.dst.varid == self.vvar_id
+        ):
+            return Const(expr.idx, None, self.const_value.value, self.const_value.bits, **expr.tags)
+        return expr
+
+
+class ConditionConstantPropagation(OptimizationPass):
+    """
+    Reason about constant propagation opportunities from conditionals and propagate constants in the graph accordingly.
+    """
+
+    ARCHES = None
+    PLATFORMS = None
+    STAGE = OptimizationPassStage.AFTER_SINGLE_BLOCK_SIMPLIFICATION
+    NAME = "Propagate constants using information deduced from conditionals."
+    DESCRIPTION = __doc__.strip()  # type: ignore
+
+    def __init__(self, func, **kwargs):
+        super().__init__(func, **kwargs)
+        self.analyze()
+
+    def _check(self):
+        cconds = self._find_const_conditions()
+
+        if not cconds:
+            return False, None
+
+        # group cconds according to their sources
+        cconds_by_src: dict[tuple[int, int | None], list[ConstantCondition]] = {}
+        for ccond in cconds:
+            src = ccond.block_addr, ccond.block_idx
+            if src not in cconds_by_src:
+                cconds_by_src[src] = []
+            cconds_by_src[src].append(ccond)
+
+        # eliminate sources with more than one in-edges; this is because the condition may not hold on all in-edges!
+        for src in list(cconds_by_src):
+            block = self._get_block(src[0], idx=src[1])
+            if block is not None and block in self._graph and self._graph.in_degree[block] > 1:
+                del cconds_by_src[src]
+
+        # eliminate conflicting conditions
+        for src in list(cconds_by_src):
+            cconds = cconds_by_src[src]
+            vvar_id_to_values = defaultdict(set)
+            ccond_dict = {}  # keyed by vvar_id; used for deduplication
+            for ccond in cconds:
+                vvar_id_to_values[ccond.vvar_id].add(ccond.value)
+                ccond_dict[ccond.vvar_id] = ccond
+            new_cconds = []
+            for vid, vvalues in vvar_id_to_values.items():
+                if len(vvalues) == 1:
+                    new_cconds.append(ccond_dict[vid])
+            if new_cconds:
+                cconds_by_src[src] = new_cconds
+            else:
+                del cconds_by_src[src]
+
+        if not cconds_by_src:
+            return False, None
+        return True, {"cconds_by_src": cconds_by_src}
+
+    @timethis
+    def _analyze(self, cache=None):
+        if not cache or cache.get("cconds_by_src", None) is None:
+            return
+        cconds_by_src = cache["cconds_by_src"]
+
+        if not cconds_by_src:
+            return
+
+        # calculate a dominance frontier for each block
+        entry_node_addr, entry_node_idx = self.entry_node_addr
+        entry_node = self._get_block(entry_node_addr, idx=entry_node_idx)
+        idoms = networkx.algorithms.immediate_dominators(self._graph, entry_node)
+        rda: SRDAModel = self.project.analyses.SReachingDefinitions(self._func, func_graph=self._graph).model
+
+        for src, cconds in cconds_by_src.items():
+            head_block = self._get_block(src[0], idx=src[1])
+            if head_block is None:
+                continue
+
+            for ccond in cconds:
+                for _, loc in rda.all_vvar_uses[ccond.vvar_id]:
+                    loc_block = self._get_block(loc.block_addr, idx=loc.block_idx)
+                    if loc_block is None:
+                        continue
+                    if dominates(idoms, head_block, loc_block):
+                        # the constant condition dominates the use site
+                        walker = CCondPropBlockWalker(ccond.vvar_id, ccond.value)
+                        new_block = walker.walk(loc_block)
+                        if new_block is not None:
+                            self._update_block(loc_block, new_block)
+
+    @timethis
+    def _find_const_conditions(self) -> list[ConstantCondition]:
+        cconds = []
+
+        for block in self._graph:
+            if block.statements:
+                last_stmt = block.statements[-1]
+                if (
+                    isinstance(last_stmt, ConditionalJump)
+                    and isinstance(last_stmt.true_target, Const)
+                    and isinstance(last_stmt.false_target, Const)
+                ):
+                    self._extract_const_condition_from_stmt(last_stmt, cconds)
+                else:
+                    # also check the first non-phi statement; rep stos may generate blocks whose conditional checks
+                    # are at the beginning of the block
+
+                    # we could have used is_head_controlled_loop_block, but at this point the block is simplified enough
+                    # that the first non-label, non-phi statement must be a ConditionalJump that controls the execution
+                    # of the loop body, so the following logic should work fine.
+
+                    first_stmt = first_nonlabel_nonphi_statement(block)
+                    if (
+                        first_stmt is not last_stmt
+                        and isinstance(first_stmt, ConditionalJump)
+                        and isinstance(first_stmt.true_target, Const)
+                        and isinstance(first_stmt.false_target, Const)
+                    ):
+                        self._extract_const_condition_from_stmt(first_stmt, cconds)
+
+        return cconds
+
+    @staticmethod
+    def _extract_const_condition_from_stmt(stmt: ConditionalJump, cconds: list[ConstantCondition]) -> None:
+        if isinstance(stmt.condition, BinaryOp):
+            cond = stmt.condition
+            op = cond.op
+            op0, op1 = cond.operands
+            if isinstance(op0, Const):
+                op0, op1 = op1, op0
+            if isinstance(op0, VirtualVariable) and isinstance(op1, Const) and op1.is_int:
+                if op == "CmpEQ":
+                    ccond = ConstantCondition(
+                        op0.varid, op1, stmt.true_target.value, stmt.true_target_idx  # type: ignore
+                    )
+                    cconds.append(ccond)
+                elif op == "CmpNE":
+                    ccond = ConstantCondition(
+                        op0.varid, op1, stmt.false_target.value, stmt.false_target_idx  # type: ignore
+                    )
+                    cconds.append(ccond)