angr 9.2.116__py3-none-manylinux2014_x86_64.whl → 9.2.118__py3-none-manylinux2014_x86_64.whl

angr/analyses/decompiler/condition_processor.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from collections import defaultdict, OrderedDict
 from typing import Any
 from collections.abc import Generator
@@ -32,7 +33,7 @@ from .structuring.structurer_nodes import (
     IncompleteSwitchCaseNode,
 )
 from .graph_region import GraphRegion
-from .utils import first_nonlabel_statement, peephole_optimize_expr
+from .utils import first_nonlabel_nonphi_statement, peephole_optimize_expr
 
 if is_pyinstaller():
     # PyInstaller is not happy with lazy import
@@ -63,7 +64,7 @@ _UNIFIABLE_COMPARISONS = {
 def _op_with_unified_size(op, conv, operand0, operand1):
     # ensure operand1 is of the same size as operand0
     if isinstance(operand1, ailment.Expr.Const):
-        # amazing - we do the eazy thing here
+        # amazing - we do the easy thing here
         return op(conv(operand0, nobool=True), operand1.value)
     if operand1.bits == operand0.bits:
         return op(conv(operand0, nobool=True), conv(operand1))
@@ -74,13 +75,13 @@ def _op_with_unified_size(op, conv, operand0, operand1):
 
 
 def _dummy_bvs(condition, condition_mapping, name_suffix=""):
-    var = claripy.BVS(f"ailexpr_{repr(condition)}{name_suffix}", condition.bits, explicit_name=True)
+    var = claripy.BVS(f"ailexpr_{condition!r}{name_suffix}", condition.bits, explicit_name=True)
     condition_mapping[var.args[0]] = condition
     return var
 
 
 def _dummy_bools(condition, condition_mapping, name_suffix=""):
-    var = claripy.BoolS(f"ailexpr_{repr(condition)}{name_suffix}", explicit_name=True)
+    var = claripy.BoolS(f"ailexpr_{condition!r}{name_suffix}", explicit_name=True)
     condition_mapping[var.args[0]] = condition
     return var
 
@@ -98,6 +99,16 @@ _ail2claripy_op_mapping = {
     "CmpGEs": lambda expr, conv, _: claripy.SGE(conv(expr.operands[0]), conv(expr.operands[1])),
     "CmpGT": lambda expr, conv, _: conv(expr.operands[0]) > conv(expr.operands[1]),
     "CmpGTs": lambda expr, conv, _: claripy.SGT(conv(expr.operands[0]), conv(expr.operands[1])),
+    "CasCmpEQ": lambda expr, conv, _: conv(expr.operands[0]) == conv(expr.operands[1]),
+    "CasCmpNE": lambda expr, conv, _: conv(expr.operands[0]) != conv(expr.operands[1]),
+    "CasCmpLE": lambda expr, conv, _: conv(expr.operands[0]) <= conv(expr.operands[1]),
+    "CasCmpLEs": lambda expr, conv, _: claripy.SLE(conv(expr.operands[0]), conv(expr.operands[1])),
+    "CasCmpLT": lambda expr, conv, _: conv(expr.operands[0]) < conv(expr.operands[1]),
+    "CasCmpLTs": lambda expr, conv, _: claripy.SLT(conv(expr.operands[0]), conv(expr.operands[1])),
+    "CasCmpGE": lambda expr, conv, _: conv(expr.operands[0]) >= conv(expr.operands[1]),
+    "CasCmpGEs": lambda expr, conv, _: claripy.SGE(conv(expr.operands[0]), conv(expr.operands[1])),
+    "CasCmpGT": lambda expr, conv, _: conv(expr.operands[0]) > conv(expr.operands[1]),
+    "CasCmpGTs": lambda expr, conv, _: claripy.SGT(conv(expr.operands[0]), conv(expr.operands[1])),
     "Add": lambda expr, conv, _: conv(expr.operands[0], nobool=True) + conv(expr.operands[1], nobool=True),
     "Sub": lambda expr, conv, _: conv(expr.operands[0], nobool=True) - conv(expr.operands[1], nobool=True),
     "Mul": lambda expr, conv, _: conv(expr.operands[0], nobool=True) * conv(expr.operands[1], nobool=True),
@@ -207,7 +218,7 @@ class ConditionProcessor:
 
         if graph:
             _g = graph
-            head = [node for node in graph.nodes if graph.in_degree(node) == 0][0]
+            head = next(node for node in graph.nodes if graph.in_degree(node) == 0)
         else:
             if with_successors and region.graph_with_successors is not None:
                 _g = region.graph_with_successors
@@ -227,13 +238,12 @@ class ConditionProcessor:
         terminating_nodes = []
         for node in sorted_nodes:
             # create special conditions for all nodes that are jump table entries
-            if case_entry_to_switch_head:
-                if node.addr in case_entry_to_switch_head:
-                    jump_target_var = self.create_jump_target_var(case_entry_to_switch_head[node.addr])
-                    cond = jump_target_var == claripy.BVV(node.addr, self.arch.bits)
-                    reaching_conditions[node] = cond
-                    self.jump_table_conds[case_entry_to_switch_head[node.addr]].add(cond)
-                    continue
+            if case_entry_to_switch_head and node.addr in case_entry_to_switch_head:
+                jump_target_var = self.create_jump_target_var(case_entry_to_switch_head[node.addr])
+                cond = jump_target_var == claripy.BVV(node.addr, self.arch.bits)
+                reaching_conditions[node] = cond
+                self.jump_table_conds[case_entry_to_switch_head[node.addr]].add(cond)
+                continue
 
             preds = _g.predecessors(node)
             reaching_condition = None
@@ -317,19 +327,17 @@ class ConditionProcessor:
             for n in node.nodes:
                 new_node = self.remove_claripy_bool_asts(n, memo=memo)
                 new_nodes.append(new_node)
-            new_seq_node = SequenceNode(node.addr, new_nodes)
-            return new_seq_node
+            return SequenceNode(node.addr, new_nodes)
 
-        elif isinstance(node, MultiNode):
+        if isinstance(node, MultiNode):
             new_nodes = []
             for n in node.nodes:
                 new_node = self.remove_claripy_bool_asts(n, memo=memo)
                 new_nodes.append(new_node)
-            new_multinode = MultiNode(nodes=new_nodes)
-            return new_multinode
+            return MultiNode(nodes=new_nodes)
 
-        elif isinstance(node, CodeNode):
-            node = CodeNode(
+        if isinstance(node, CodeNode):
+            return CodeNode(
                 self.remove_claripy_bool_asts(node.node, memo=memo),
                 (
                     None
@@ -337,16 +345,15 @@ class ConditionProcessor:
                     else self.convert_claripy_bool_ast(node.reaching_condition, memo=memo)
                 ),
             )
-            return node
 
-        elif isinstance(node, ConditionalBreakNode):
+        if isinstance(node, ConditionalBreakNode):
             return ConditionalBreakNode(
                 node.addr,
                 self.convert_claripy_bool_ast(node.condition, memo=memo),
                 node.target,
             )
 
-        elif isinstance(node, ConditionNode):
+        if isinstance(node, ConditionNode):
             return ConditionNode(
                 node.addr,
                 (
@@ -359,7 +366,7 @@ class ConditionProcessor:
                 self.remove_claripy_bool_asts(node.false_node, memo=memo),
             )
 
-        elif isinstance(node, CascadingConditionNode):
+        if isinstance(node, CascadingConditionNode):
             cond_and_nodes = []
             for cond, child_node in node.condition_and_nodes:
                 cond_and_nodes.append(
@@ -375,7 +382,7 @@ class ConditionProcessor:
                 else_node=else_node,
             )
 
-        elif isinstance(node, LoopNode):
+        if isinstance(node, LoopNode):
             result = node.copy()
             result.condition = (
                 self.convert_claripy_bool_ast(node.condition, memo=memo) if node.condition is not None else None
@@ -383,7 +390,7 @@ class ConditionProcessor:
             result.sequence_node = self.remove_claripy_bool_asts(node.sequence_node, memo=memo)
             return result
 
-        elif isinstance(node, SwitchCaseNode):
+        if isinstance(node, SwitchCaseNode):
             return SwitchCaseNode(
                 self.convert_claripy_bool_ast(node.switch_expr, memo=memo),
                 OrderedDict(
@@ -393,15 +400,14 @@ class ConditionProcessor:
                 addr=node.addr,
             )
 
-        elif isinstance(node, IncompleteSwitchCaseNode):
+        if isinstance(node, IncompleteSwitchCaseNode):
             return IncompleteSwitchCaseNode(
                 node.addr,
                 self.remove_claripy_bool_asts(node.head, memo=memo),
                 [self.remove_claripy_bool_asts(case, memo=memo) for case in node.cases],
             )
 
-        else:
-            return node
+        return node
 
     @classmethod
     def get_last_statement(cls, block):
@@ -413,26 +419,25 @@ class ConditionProcessor:
         if type(block) is SequenceNode:
             if block.nodes:
                 return cls.get_last_statement(block.nodes[-1])
-            raise EmptyBlockNotice()
+            raise EmptyBlockNotice
         if type(block) is CodeNode:
             return cls.get_last_statement(block.node)
         if type(block) is ailment.Block:
             if not block.statements:
-                raise EmptyBlockNotice()
+                raise EmptyBlockNotice
             return block.statements[-1]
         if type(block) is Block:
-            raise NotImplementedError()
+            raise NotImplementedError
         if type(block) is BlockNode:
-            raise NotImplementedError()
+            raise NotImplementedError
         if type(block) is MultiNode:
             # get the last node
             for the_block in reversed(block.nodes):
                 try:
-                    last_stmt = cls.get_last_statement(the_block)
-                    return last_stmt
+                    return cls.get_last_statement(the_block)
                 except EmptyBlockNotice:
                     continue
-            raise EmptyBlockNotice()
+            raise EmptyBlockNotice
         if type(block) is LoopNode:
             return cls.get_last_statement(block.sequence_node)
         if type(block) is ConditionalBreakNode:
@@ -469,43 +474,41 @@ class ConditionProcessor:
             # normally this should not happen. however, we have test cases that trigger this case.
             return None
 
-        raise NotImplementedError()
+        raise NotImplementedError
 
     @classmethod
     def get_last_statements(cls, block) -> list[ailment.Stmt.Statement | None]:
         if type(block) is SequenceNode:
             for last_node in reversed(block.nodes):
                 try:
-                    last_stmts = cls.get_last_statements(last_node)
-                    return last_stmts
+                    return cls.get_last_statements(last_node)
                 except EmptyBlockNotice:
                     # the node is empty. try the next one
                     continue
 
-            raise EmptyBlockNotice()
+            raise EmptyBlockNotice
 
         if type(block) is CodeNode:
             return cls.get_last_statements(block.node)
         if type(block) is ailment.Block:
             if not block.statements:
-                raise EmptyBlockNotice()
+                raise EmptyBlockNotice
             return [block.statements[-1]]
         if type(block) is Block:
-            raise NotImplementedError()
+            raise NotImplementedError
         if type(block) is BlockNode:
-            raise NotImplementedError()
+            raise NotImplementedError
         if type(block) is MultiNode:
             # get the last node
             for the_block in reversed(block.nodes):
                 try:
-                    last_stmts = cls.get_last_statements(the_block)
-                    return last_stmts
+                    return cls.get_last_statements(the_block)
                 except EmptyBlockNotice:
                     continue
-            raise EmptyBlockNotice()
+            raise EmptyBlockNotice
         if type(block) is LoopNode:
             if block.sequence_node is None:
-                raise EmptyBlockNotice()
+                raise EmptyBlockNotice
             return cls.get_last_statements(block.sequence_node)
         if type(block) is ConditionalBreakNode:
             return [block]
@@ -561,7 +564,7 @@ class ConditionProcessor:
             # normally this should not happen. however, we have test cases that trigger this case.
             return []
 
-        raise NotImplementedError()
+        raise NotImplementedError
 
     #
     # Path predicate
@@ -590,8 +593,7 @@ class ConditionProcessor:
             bool_var = src_block.condition
             if src_block.target == dst_block.addr:
                 return bool_var
-            else:
-                return claripy.Not(bool_var)
+            return claripy.Not(bool_var)
 
         if type(src_block) is GraphRegion:
             return claripy.true
@@ -600,9 +602,9 @@ class ConditionProcessor:
         if (
             isinstance(src_block, ailment.Block)
             and src_block.statements
-            and isinstance(first_nonlabel_statement(src_block), ailment.Stmt.ConditionalJump)
+            and isinstance(first_nonlabel_nonphi_statement(src_block), ailment.Stmt.ConditionalJump)
         ):
-            last_stmt = first_nonlabel_statement(src_block)
+            last_stmt = first_nonlabel_nonphi_statement(src_block)
         else:
             last_stmt = self.get_last_statement(src_block)
 
@@ -618,8 +620,7 @@ class ConditionProcessor:
             bool_var = self.claripy_ast_from_ail_condition(last_stmt.condition)
             if isinstance(last_stmt.true_target, ailment.Expr.Const) and last_stmt.true_target.value == dst_block.addr:
                 return bool_var
-            else:
-                return claripy.Not(bool_var)
+            return claripy.Not(bool_var)
 
         return claripy.true
 
@@ -737,7 +738,7 @@ class ConditionProcessor:
                 cond_tags = {}
             return _mapping[cond.op](cond, cond_tags)
         raise NotImplementedError(
-            ("Condition variable %s has an unsupported operator %s. Consider implementing.") % (cond, cond.op)
+            f"Condition variable {cond} has an unsupported operator {cond.op}. Consider implementing."
         )
 
     def claripy_ast_from_ail_condition(self, condition, nobool: bool = False) -> claripy.ast.Bool:
@@ -750,19 +751,19 @@ class ConditionProcessor:
             (ailment.Expr.DirtyExpression, ailment.Expr.BasePointerOffset, ailment.Expr.ITE),
         ):
             return _dummy_bvs(condition, self._condition_mapping)
-        elif isinstance(condition, ailment.Stmt.Call):
+        if isinstance(condition, ailment.Stmt.Call):
             return _dummy_bvs(condition, self._condition_mapping, name_suffix=hex(condition.tags.get("ins_addr", 0)))
-        elif isinstance(condition, (ailment.Expr.Load, ailment.Expr.Register)):
+        if isinstance(condition, (ailment.Expr.Load, ailment.Expr.Register, ailment.Expr.VirtualVariable)):
             # does it have a variable associated?
             if condition.variable is not None:
                 var = claripy.BVS(
-                    f"ailexpr_{repr(condition)}-{condition.variable.ident}", condition.bits, explicit_name=True
+                    f"ailexpr_{condition!r}-{condition.variable.ident}", condition.bits, explicit_name=True
                 )
             else:
-                var = claripy.BVS(f"ailexpr_{repr(condition)}-{condition.idx}", condition.bits, explicit_name=True)
+                var = claripy.BVS(f"ailexpr_{condition!r}-{condition.idx}", condition.bits, explicit_name=True)
             self._condition_mapping[var.args[0]] = condition
             return var
-        elif isinstance(condition, ailment.Expr.Convert):
+        if isinstance(condition, ailment.Expr.Convert):
             # convert is special. if it generates a 1-bit variable, it should be treated as a BoolS
             if condition.to_bits == 1:
                 var_ = self.claripy_ast_from_ail_condition(condition.operands[0])
@@ -774,7 +775,7 @@ class ConditionProcessor:
                 var = claripy.BVS(name, condition.to_bits, explicit_name=True)
             self._condition_mapping[var.args[0]] = condition
             return var
-        elif isinstance(condition, ailment.Expr.Const):
+        if isinstance(condition, ailment.Expr.Const):
             if condition.value is True or condition.value is False:
                 var = claripy.BoolV(condition.value)
             else:
@@ -782,7 +783,7 @@ class ConditionProcessor:
             if isinstance(var, claripy.Bits) and var.size() == 1:
                 var = claripy.true if var.concrete_value == 1 else claripy.false
             return var
-        elif isinstance(condition, ailment.Expr.Tmp):
+        if isinstance(condition, ailment.Expr.Tmp):
             l.warning("Left-over ailment.Tmp variable %s.", condition)
             if condition.bits == 1:
                 var = claripy.BoolS("ailtmp_%d" % condition.tmp_idx, explicit_name=True)
@@ -790,7 +791,7 @@ class ConditionProcessor:
                 var = claripy.BVS("ailtmp_%d" % condition.tmp_idx, condition.bits, explicit_name=True)
             self._condition_mapping[var.args[0]] = condition
             return var
-        elif isinstance(condition, ailment.Expr.MultiStatementExpression):
+        if isinstance(condition, ailment.Expr.MultiStatementExpression):
             # just cache it
             if condition.bits == 1:
                 var = claripy.BoolS("mstmtexpr_%d" % hash(condition), explicit_name=True)
@@ -805,20 +806,19 @@ class ConditionProcessor:
             lambda_expr = _ail2claripy_op_mapping.get(condition.op, None)
         if lambda_expr is None:
             raise NotImplementedError(
-                "Unsupported AIL expression operation %s or %s. Consider implementing."
-                % (condition.op, condition.verbose_op)
+                f"Unsupported AIL expression operation {condition.op} or {condition.verbose_op}. Consider implementing."
             )
         r = lambda_expr(condition, self.claripy_ast_from_ail_condition, self._condition_mapping)
 
         if isinstance(r, claripy.ast.Bool) and nobool:
-            r = claripy.BVS("ailexpr_from_bool_%r" % r, 1, explicit_name=True)
+            r = claripy.BVS(f"ailexpr_from_bool_{r!r}", 1, explicit_name=True)
             self._condition_mapping[r.args[0]] = condition
 
         if r is NotImplemented:
             if condition.bits == 1:
-                r = claripy.BoolS("ailexpr_%r" % condition, explicit_name=True)
+                r = claripy.BoolS(f"ailexpr_{condition!r}", explicit_name=True)
             else:
-                r = claripy.BVS("ailexpr_%r" % condition, condition.bits, explicit_name=True)
+                r = claripy.BVS(f"ailexpr_{condition!r}", condition.bits, explicit_name=True)
             self._condition_mapping[r.args[0]] = condition
         # don't lose tags
         self._ast2annotations[r] = condition.tags
@@ -871,8 +871,7 @@ class ConditionProcessor:
         if cond.depth > depth_limit or len(cond.variables) > variables_limit:
             return cond
         sympy_expr = ConditionProcessor.claripy_ast_to_sympy_expr(cond, memo=memo)
-        r = ConditionProcessor.sympy_expr_to_claripy_ast(sympy.simplify_logic(sympy_expr, deep=False), memo)
-        return r
+        return ConditionProcessor.sympy_expr_to_claripy_ast(sympy.simplify_logic(sympy_expr, deep=False), memo)
 
     @staticmethod
     def simplify_condition_deprecated(cond):
@@ -893,8 +892,7 @@ class ConditionProcessor:
         # cond = simplified if simplified is not None else cond
         # in the end, use claripy's simplification to handle really easy cases again
         simplified = ConditionProcessor._simplify_trivial_cases(cond)
-        cond = simplified if simplified is not None else cond
-        return cond
+        return simplified if simplified is not None else cond
 
     @staticmethod
     def _simplify_trivial_cases(cond):
@@ -947,10 +945,8 @@ class ConditionProcessor:
             a = claripy.Not(not_a)
             if len(cond.args) <= 2:
                 return claripy.Not(claripy.And(a, b))
-            else:
-                return claripy.Or(claripy.Not(claripy.And(a, b)), *cond.args[2:])
-        else:
-            return cond
+            return claripy.Or(claripy.Not(claripy.And(a, b)), *cond.args[2:])
+        return cond
 
     @staticmethod
     def _fold_double_negations(cond):
@@ -967,30 +963,24 @@ class ConditionProcessor:
         if cond.args[0].op == "And" and len(cond.args[0].args) == 2:
             and_0, and_1 = cond.args[0].args
             if and_0.op == "Not" and and_1.op == "Not":
-                expr = claripy.Or(and_0.args[0], and_1.args[0])
-                return expr
+                return claripy.Or(and_0.args[0], and_1.args[0])
 
             if and_0.op == "Not":  # and_1.op != "Not"
-                expr = claripy.Or(and_0.args[0], ConditionProcessor.simplify_condition(claripy.Not(and_1)))
-                return expr
+                return claripy.Or(and_0.args[0], ConditionProcessor.simplify_condition(claripy.Not(and_1)))
 
         if cond.args[0].op == "Or" and len(cond.args[0].args) == 2:
             or_0, or_1 = cond.args[0].args
-            expr = claripy.And(
+            return claripy.And(
                 ConditionProcessor.simplify_condition(claripy.Not(or_0)),
                 ConditionProcessor.simplify_condition(claripy.Not(or_1)),
             )
-            return expr
 
         return None
 
     @staticmethod
     def _extract_common_subexpressions(cond):
         def _expr_inside_collection(expr_, coll_) -> bool:
-            for ex_ in coll_:
-                if expr_ is ex_:
-                    return True
-            return False
+            return any(expr_ is ex_ for ex_ in coll_)
 
         # (A && B) || (A && C) => A && (B || C)
         if cond.op == "And":
@@ -1037,11 +1027,8 @@ class ConditionProcessor:
         return None
 
     @staticmethod
-    def _extract_terms(ast: claripy.ast.Bool) -> Generator[claripy.ast.Bool, None, None]:
-        if ast.op == "And":
-            for arg in ast.args:
-                yield from ConditionProcessor._extract_terms(arg)
-        elif ast.op == "Or":
+    def _extract_terms(ast: claripy.ast.Bool) -> Generator[claripy.ast.Bool]:
+        if ast.op == "And" or ast.op == "Or":
             for arg in ast.args:
                 yield from ConditionProcessor._extract_terms(arg)
         elif ast.op == "Not":
@@ -1061,20 +1048,19 @@ class ConditionProcessor:
             return ast.make_like(
                 "And", (ConditionProcessor._replace_term_in_ast(arg, r0, r0_with, r1, r1_with) for arg in ast.args)
             )
-        elif ast.op == "Or":
+        if ast.op == "Or":
             return ast.make_like(
                 "Or", (ConditionProcessor._replace_term_in_ast(arg, r0, r0_with, r1, r1_with) for arg in ast.args)
             )
-        elif ast.op == "Not":
+        if ast.op == "Not":
             return ast.make_like(
                 "Not", (ConditionProcessor._replace_term_in_ast(ast.args[0], r0, r0_with, r1, r1_with),)
             )
-        else:
-            if ast is r0:
-                return r0_with
-            if ast is r1:
-                return r1_with
-            return ast
+        if ast is r0:
+            return r0_with
+        if ast is r1:
+            return r1_with
+        return ast
 
     @staticmethod
     def _remove_redundant_terms(cond):
@@ -1104,7 +1090,7 @@ class ConditionProcessor:
 
         solver = claripy.SolverCacheless()
         for term in all_terms_without_negs:
-            neg = negations.get(term, None)
+            neg = negations.get(term)
 
             replaced_with_true = ConditionProcessor._replace_term_in_ast(cond, term, claripy.true, neg, claripy.false)
             sat0 = solver.satisfiable(
@@ -1187,4 +1173,4 @@ class ConditionProcessor:
     #
 
     def create_jump_target_var(self, jumptable_head_addr: int):
-        return claripy.BVS("jump_table_%x" % jumptable_head_addr, self.arch.bits, explicit_name=True)
+        return claripy.BVS(f"jump_table_{jumptable_head_addr:x}", self.arch.bits, explicit_name=True)

angr/analyses/decompiler/counters/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+from .boolean_counter import BooleanCounter
+from .call_counter import AILBlockCallCounter
+from .seq_cf_structure_counter import ControlFlowStructureCounter
+from .expression_counters import SingleExpressionCounter, RegisterExpressionCounter, OperatorCounter

angr/analyses/decompiler/counters/boolean_counter.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+import typing
+
+from ailment import AILBlockWalkerBase
+
+if typing.TYPE_CHECKING:
+    from ailment.expression import BinaryOp
+    from ailment.statement import Statement
+    from ailment.block import Block
+
+
+class BooleanCounter(AILBlockWalkerBase):
+    """
+    This class counts the number of Boolean operators an expression has.
+    In the case of: `if (a || (b && c))`, it will count 2 Boolean operators.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.boolean_cnt = 0
+
+    def _handle_BinaryOp(self, expr_idx: int, expr: BinaryOp, stmt_idx: int, stmt: Statement, block: Block | None):
+        if expr.op in {"LogicalAnd", "LogicalOr"}:
+            self.boolean_cnt += 1
+
+        self._handle_expr(0, expr.operands[0], stmt_idx, stmt, block)
+        self._handle_expr(1, expr.operands[1], stmt_idx, stmt, block)

angr/analyses/decompiler/{call_counter.py → counters/call_counter.py}

@@ -1,9 +1,10 @@
-from typing import Optional, TYPE_CHECKING
+from __future__ import annotations
+from typing import TYPE_CHECKING
 
 from ailment import Block
 from ailment.block_walker import AILBlockWalkerBase
 
-from .sequence_walker import SequenceWalker
+from angr.analyses.decompiler.sequence_walker import SequenceWalker
 
 if TYPE_CHECKING:
     from ailment.statement import Call
@@ -16,11 +17,11 @@ class AILBlockCallCounter(AILBlockWalkerBase):
 
     calls = 0
 
-    def _handle_CallExpr(self, expr_idx: int, expr: "Call", stmt_idx: int, stmt, block: Optional["Block"]):
+    def _handle_CallExpr(self, expr_idx: int, expr: Call, stmt_idx: int, stmt, block: Block | None):
         self.calls += 1
         super()._handle_CallExpr(expr_idx, expr, stmt_idx, stmt, block)
 
-    def _handle_Call(self, stmt_idx: int, stmt: "Call", block: Optional["Block"]):
+    def _handle_Call(self, stmt_idx: int, stmt: Call, block: Block | None):
         self.calls += 1
         super()._handle_Call(stmt_idx, stmt, block)
 

angr/analyses/decompiler/{expression_counters.py → counters/expression_counters.py}

@@ -1,4 +1,5 @@
-from typing import Any, DefaultDict, TYPE_CHECKING
+from __future__ import annotations
+from typing import Any, TYPE_CHECKING
 from collections.abc import Iterable
 from collections import defaultdict
 
@@ -37,7 +38,7 @@ class RegisterExpressionCounter(AILBlockWalkerBase):
 
     def __init__(self, expr_or_stmt: Expression | Statement):
         super().__init__()
-        self.counts: DefaultDict[tuple[int, int], int] = defaultdict(int)
+        self.counts: defaultdict[tuple[int, int], int] = defaultdict(int)
         if isinstance(expr_or_stmt, Expression):
             self.walk_expression(expr_or_stmt)
         elif isinstance(expr_or_stmt, Statement):
@@ -65,12 +66,12 @@ class OperatorCounter(AILBlockWalkerBase):
         else:
             raise TypeError(f"Unsupported argument type {type(expr_or_stmt)}")
 
-    def _handle_BinaryOp(self, expr_idx: int, expr: "BinaryOp", stmt_idx: int, stmt: Statement, block: Block | None):
+    def _handle_BinaryOp(self, expr_idx: int, expr: BinaryOp, stmt_idx: int, stmt: Statement, block: Block | None):
         if expr.op in self.operators:
             self.count += 1
         return super()._handle_BinaryOp(expr_idx, expr, stmt_idx, stmt, block)
 
-    def _handle_UnaryOp(self, expr_idx: int, expr: "UnaryOp", stmt_idx: int, stmt: Statement, block: Block | None):
+    def _handle_UnaryOp(self, expr_idx: int, expr: UnaryOp, stmt_idx: int, stmt: Statement, block: Block | None):
         if expr.op in self.operators:
             self.count += 1
         return super()._handle_UnaryOp(expr_idx, expr, stmt_idx, stmt, block)

angr/analyses/decompiler/counters/seq_cf_structure_counter.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+from collections import defaultdict
+
+import ailment
+
+from angr.analyses.decompiler.sequence_walker import SequenceWalker
+from angr.analyses.decompiler.structuring.structurer_nodes import LoopNode
+
+
+class ControlFlowStructureCounter(SequenceWalker):
+    """
+    Counts the number of different types of control flow structures found in a sequence of nodes.
+    This should be used after the sequence has been simplified.
+    """
+
+    def __init__(self, node):
+        handlers = {
+            LoopNode: self._handle_Loop,
+            ailment.Block: self._handle_Block,
+        }
+        super().__init__(handlers)
+
+        self.while_loops = 0
+        self.do_while_loops = 0
+        self.for_loops = 0
+        self.ordered_labels = []
+        self.goto_targets = defaultdict(int)
+
+        self.walk(node)
+
+        # eliminate gotos without labels
+        self.goto_targets = {k: v for k, v in self.goto_targets.items() if k in self.ordered_labels}
+        # correct labels that are not used
+        self.ordered_labels = [lbl for lbl in self.ordered_labels if lbl in self.goto_targets]
+
+    # pylint: disable=unused-argument
+    def _handle_Block(self, node: ailment.Block, **kwargs):
+        if not node.statements:
+            return
+
+        for stmt in node.statements:
+            # labels found in the block at this point will be labels in the output
+            if isinstance(stmt, ailment.statement.Label):
+                label_addr = stmt.ins_addr
+                if label_addr is not None:
+                    self.ordered_labels.append(label_addr)
+
+            # goto targets found in the block at this point will be goto targets in the output
+            if isinstance(stmt, ailment.statement.Jump):
+                target = stmt.target
+                target_value = target.value if target is not None and isinstance(target, ailment.Expr.Const) else None
+                if target_value is not None:
+                    self.goto_targets[target_value] += 1
+
+    def _handle_Loop(self, node: LoopNode, **kwargs):
+        if node.sort == "while":
+            self.while_loops += 1
+        elif node.sort == "do-while":
+            self.do_while_loops += 1
+        elif node.sort == "for":
+            self.for_loops += 1
+
+        return super()._handle_Loop(node, **kwargs)