angr 9.2.117__py3-none-manylinux2014_x86_64.whl → 9.2.119__py3-none-manylinux2014_x86_64.whl

angr/analyses/decompiler/optimization_passes/lowered_switch_simplifier.py

@@ -1,4 +1,5 @@
-from typing import DefaultDict, TYPE_CHECKING
+from __future__ import annotations
+from typing import TYPE_CHECKING
 from collections import defaultdict, OrderedDict
 import logging
 
@@ -9,7 +10,7 @@ from ailment.statement import ConditionalJump, Label, Assignment, Jump
 from ailment.expression import Expression, BinaryOp, Const, Load
 
 from angr.utils.graph import GraphUtils
-from ..utils import first_nonlabel_statement, remove_last_statement
+from ..utils import first_nonlabel_nonphi_statement, remove_last_statement
 from ..structuring.structurer_nodes import IncompleteSwitchCaseHeadStatement, SequenceNode, MultiNode
 from .optimization_pass import MultipleBlocksException, StructuringOptimizationPass
 from ..region_simplifiers.switch_cluster_simplifier import SwitchClusterFinder
@@ -60,7 +61,7 @@ class Case:
         if self.value == "default":
             return f"Case default@{self.target:#x}{'' if self.target_idx is None else '.' + str(self.target_idx)}"
         return (
-            f"Case {repr(self.original_node)}@{self.target:#x}"
+            f"Case {self.original_node!r}@{self.target:#x}"
             f"{'' if self.target_idx is None else '.' + str(self.target_idx)}: {self.expr} == {self.value}"
         )
 
@@ -119,14 +120,14 @@ class StableVarExprHasher(AILBlockWalkerBase):
         self._hash_lst.append(expr.op)
         super()._handle_BinaryOp(expr_idx, expr, stmt_idx, stmt, block)
 
-    def _handle_UnaryOp(self, expr_idx: int, expr: "UnaryOp", stmt_idx: int, stmt, block: Block | None):
+    def _handle_UnaryOp(self, expr_idx: int, expr: UnaryOp, stmt_idx: int, stmt, block: Block | None):
         self._hash_lst.append(expr.op)
         super()._handle_UnaryOp(expr_idx, expr, stmt_idx, stmt, block)
 
     def _handle_Const(self, expr_idx: int, expr: Const, stmt_idx: int, stmt, block: Block | None):
         self._hash_lst.append((expr.value, expr.bits))
 
-    def _handle_Convert(self, expr_idx: int, expr: "Convert", stmt_idx: int, stmt, block: Block | None):
+    def _handle_Convert(self, expr_idx: int, expr: Convert, stmt_idx: int, stmt, block: Block | None):
         self._hash_lst.append(expr.to_bits)
         super()._handle_Convert(expr_idx, expr, stmt_idx, stmt, block)
 
@@ -193,8 +194,7 @@ class LoweredSwitchSimplifier(StructuringOptimizationPass):
                 current_len = 1
 
         # Final check to include the last sequence
-        max_len = max(max_len, current_len)
-        return max_len
+        return max(max_len, current_len)
 
     @staticmethod
     def _count_distinct_cases(cases: list[Case]) -> int:
@@ -369,18 +369,18 @@ class LoweredSwitchSimplifier(StructuringOptimizationPass):
         for node in sorted_nodes:
             r = self._find_switch_variable_comparison_type_a(node)
             if r is not None:
-                variable_comparisons[node] = ("a",) + r
+                variable_comparisons[node] = ("a", *r)
                 continue
             r = self._find_switch_variable_comparison_type_b(node)
             if r is not None:
-                variable_comparisons[node] = ("b",) + r
+                variable_comparisons[node] = ("b", *r)
                 continue
             r = self._find_switch_variable_comparison_type_c(node)
             if r is not None:
-                variable_comparisons[node] = ("c",) + r
+                variable_comparisons[node] = ("c", *r)
                 continue
 
-        varhash_to_caselists: DefaultDict[int, list[tuple[list[Case], list]]] = defaultdict(list)
+        varhash_to_caselists: defaultdict[int, list[tuple[list[Case], list]]] = defaultdict(list)
         used_nodes = set()
 
         for head in variable_comparisons:
@@ -405,10 +405,7 @@ class LoweredSwitchSimplifier(StructuringOptimizationPass):
                     next_addr,
                     next_addr_idx,
                 ) = variable_comparisons[comp]
-                if cases:
-                    last_varhash = cases[-1].variable_hash
-                else:
-                    last_varhash = None
+                last_varhash = cases[-1].variable_hash if cases else None
 
                 if op == "eq":
                     # eq always indicates a new case
@@ -423,20 +420,18 @@ class LoweredSwitchSimplifier(StructuringOptimizationPass):
                         if value in {0xFFFF_FFFF, 0xFFFF_FFFF_FFFF_FFFF}:
                             break
 
-                        if comp is not head:
-                            # non-head node has at most one predecessor
-                            if self._graph.in_degree[comp] > 1:
-                                break
+                        # non-head node has at most one predecessor
+                        if comp is not head and self._graph.in_degree[comp] > 1:
+                            break
 
                         cases.append(Case(comp, comp_type, variable_hash, expr, value, target, target_idx, next_addr))
                         used_nodes.add(comp)
                     else:
                         # new variable!
-                        if last_comp is not None:
-                            if comp.addr not in default_case_candidates:
-                                default_case_candidates[comp.addr] = Case(
-                                    last_comp, None, last_varhash, None, "default", comp.addr, comp.idx, None
-                                )
+                        if last_comp is not None and comp.addr not in default_case_candidates:
+                            default_case_candidates[comp.addr] = Case(
+                                last_comp, None, last_varhash, None, "default", comp.addr, comp.idx, None
+                            )
                         break
 
                     successors = [succ for succ in self._graph.successors(comp) if succ is not comp]
@@ -522,11 +517,10 @@ class LoweredSwitchSimplifier(StructuringOptimizationPass):
                                     default_case_candidates[le_addr] = Case(
                                         comp, None, variable_hash, expr, "default", le_addr, le_idx, None
                                     )
-                            elif not gt_added:
-                                if gt_addr not in default_case_candidates:
-                                    default_case_candidates[gt_addr] = Case(
-                                        comp, None, variable_hash, expr, "default", gt_addr, gt_idx, None
-                                    )
+                            elif not gt_added and gt_addr not in default_case_candidates:
+                                default_case_candidates[gt_addr] = Case(
+                                    comp, None, variable_hash, expr, "default", gt_addr, gt_idx, None
+                                )
                             extra_cmp_nodes.append(comp)
                             used_nodes.add(comp)
                         else:
@@ -608,39 +602,41 @@ class LoweredSwitchSimplifier(StructuringOptimizationPass):
 
         if isinstance(node, Block) and node.statements:
             stmt = node.statements[-1]
-            if stmt is not None and stmt is not first_nonlabel_statement(node):
-                if (
+            if (
+                stmt is not None
+                and stmt is not first_nonlabel_nonphi_statement(node)
+                and (
                     isinstance(stmt, ConditionalJump)
                     and isinstance(stmt.true_target, Const)
                     and isinstance(stmt.false_target, Const)
-                ):
-                    cond = stmt.condition
-                    if isinstance(cond, BinaryOp):
-                        if isinstance(cond.operands[1], Const):
-                            variable_hash = StableVarExprHasher(cond.operands[0]).hash
-                            value = cond.operands[1].value
-                            if cond.op == "CmpEQ":
-                                target = stmt.true_target.value
-                                target_idx = stmt.true_target_idx
-                                next_node_addr = stmt.false_target.value
-                                next_node_idx = stmt.false_target_idx
-                            elif cond.op == "CmpNE":
-                                target = stmt.false_target.value
-                                target_idx = stmt.false_target_idx
-                                next_node_addr = stmt.true_target.value
-                                next_node_idx = stmt.true_target_idx
-                            else:
-                                return None
-                            return (
-                                variable_hash,
-                                "eq",
-                                cond.operands[0],
-                                value,
-                                target,
-                                target_idx,
-                                next_node_addr,
-                                next_node_idx,
-                            )
+                )
+            ):
+                cond = stmt.condition
+                if isinstance(cond, BinaryOp) and isinstance(cond.operands[1], Const):
+                    variable_hash = StableVarExprHasher(cond.operands[0]).hash
+                    value = cond.operands[1].value
+                    if cond.op == "CmpEQ":
+                        target = stmt.true_target.value
+                        target_idx = stmt.true_target_idx
+                        next_node_addr = stmt.false_target.value
+                        next_node_idx = stmt.false_target_idx
+                    elif cond.op == "CmpNE":
+                        target = stmt.false_target.value
+                        target_idx = stmt.false_target_idx
+                        next_node_addr = stmt.true_target.value
+                        next_node_idx = stmt.true_target_idx
+                    else:
+                        return None
+                    return (
+                        variable_hash,
+                        "eq",
+                        cond.operands[0],
+                        value,
+                        target,
+                        target_idx,
+                        next_node_addr,
+                        next_node_idx,
+                    )
 
         return None
 
@@ -652,40 +648,42 @@ class LoweredSwitchSimplifier(StructuringOptimizationPass):
         # there is only one non-label statement
 
         if isinstance(node, Block):
-            stmt = first_nonlabel_statement(node)
-            if stmt is not None and stmt is node.statements[-1]:
-                if (
+            stmt = first_nonlabel_nonphi_statement(node)
+            if (
+                stmt is not None
+                and stmt is node.statements[-1]
+                and (
                     isinstance(stmt, ConditionalJump)
                     and isinstance(stmt.true_target, Const)
                     and isinstance(stmt.false_target, Const)
-                ):
-                    cond = stmt.condition
-                    if isinstance(cond, BinaryOp):
-                        if isinstance(cond.operands[1], Const):
-                            variable_hash = StableVarExprHasher(cond.operands[0]).hash
-                            value = cond.operands[1].value
-                            if cond.op == "CmpEQ":
-                                target = stmt.true_target.value
-                                target_idx = stmt.true_target_idx
-                                next_node_addr = stmt.false_target.value
-                                next_node_idx = stmt.false_target_idx
-                            elif cond.op == "CmpNE":
-                                target = stmt.false_target.value
-                                target_idx = stmt.false_target_idx
-                                next_node_addr = stmt.true_target.value
-                                next_node_idx = stmt.true_target_idx
-                            else:
-                                return None
-                            return (
-                                variable_hash,
-                                "eq",
-                                cond.operands[0],
-                                value,
-                                target,
-                                target_idx,
-                                next_node_addr,
-                                next_node_idx,
-                            )
+                )
+            ):
+                cond = stmt.condition
+                if isinstance(cond, BinaryOp) and isinstance(cond.operands[1], Const):
+                    variable_hash = StableVarExprHasher(cond.operands[0]).hash
+                    value = cond.operands[1].value
+                    if cond.op == "CmpEQ":
+                        target = stmt.true_target.value
+                        target_idx = stmt.true_target_idx
+                        next_node_addr = stmt.false_target.value
+                        next_node_idx = stmt.false_target_idx
+                    elif cond.op == "CmpNE":
+                        target = stmt.false_target.value
+                        target_idx = stmt.false_target_idx
+                        next_node_addr = stmt.true_target.value
+                        next_node_idx = stmt.true_target_idx
+                    else:
+                        return None
+                    return (
+                        variable_hash,
+                        "eq",
+                        cond.operands[0],
+                        value,
+                        target,
+                        target_idx,
+                        next_node_addr,
+                        next_node_idx,
+                    )
 
         return None
 
@@ -697,59 +695,61 @@ class LoweredSwitchSimplifier(StructuringOptimizationPass):
         # there is only one non-label statement
 
         if isinstance(node, Block):
-            stmt = first_nonlabel_statement(node)
-            if stmt is not None and stmt is node.statements[-1]:
-                if (
+            stmt = first_nonlabel_nonphi_statement(node)
+            if (
+                stmt is not None
+                and stmt is node.statements[-1]
+                and (
                     isinstance(stmt, ConditionalJump)
                     and isinstance(stmt.true_target, Const)
                     and isinstance(stmt.false_target, Const)
-                ):
-                    cond = stmt.condition
-                    if isinstance(cond, BinaryOp):
-                        if isinstance(cond.operands[1], Const):
-                            variable_hash = StableVarExprHasher(cond.operands[0]).hash
-                            value = cond.operands[1].value
-                            op = cond.op
-                            if stmt.true_target.value == stmt.false_target.value:
-                                return None
-                            if op == "CmpGT":
-                                op_str = "gt"
-                                gt_node_addr = stmt.true_target.value
-                                gt_node_idx = stmt.true_target_idx
-                                le_node_addr = stmt.false_target.value
-                                le_node_idx = stmt.false_target_idx
-                            elif op == "CmpGE":
-                                op_str = "gt"
-                                value += 1
-                                gt_node_addr = stmt.true_target.value
-                                gt_node_idx = stmt.true_target_idx
-                                le_node_addr = stmt.false_target.value
-                                le_node_idx = stmt.false_target_idx
-                            elif op == "CmpLT":
-                                op_str = "gt"
-                                value -= 1
-                                gt_node_addr = stmt.false_target.value
-                                gt_node_idx = stmt.false_target_idx
-                                le_node_addr = stmt.true_target.value
-                                le_node_idx = stmt.true_target_idx
-                            elif op == "CmpLE":
-                                op_str = "gt"
-                                gt_node_addr = stmt.false_target.value
-                                gt_node_idx = stmt.false_target_idx
-                                le_node_addr = stmt.true_target.value
-                                le_node_idx = stmt.true_target_idx
-                            else:
-                                return None
-                            return (
-                                variable_hash,
-                                op_str,
-                                cond.operands[0],
-                                value,
-                                gt_node_addr,
-                                gt_node_idx,
-                                le_node_addr,
-                                le_node_idx,
-                            )
+                )
+            ):
+                cond = stmt.condition
+                if isinstance(cond, BinaryOp) and isinstance(cond.operands[1], Const):
+                    variable_hash = StableVarExprHasher(cond.operands[0]).hash
+                    value = cond.operands[1].value
+                    op = cond.op
+                    if stmt.true_target.value == stmt.false_target.value:
+                        return None
+                    if op == "CmpGT":
+                        op_str = "gt"
+                        gt_node_addr = stmt.true_target.value
+                        gt_node_idx = stmt.true_target_idx
+                        le_node_addr = stmt.false_target.value
+                        le_node_idx = stmt.false_target_idx
+                    elif op == "CmpGE":
+                        op_str = "gt"
+                        value += 1
+                        gt_node_addr = stmt.true_target.value
+                        gt_node_idx = stmt.true_target_idx
+                        le_node_addr = stmt.false_target.value
+                        le_node_idx = stmt.false_target_idx
+                    elif op == "CmpLT":
+                        op_str = "gt"
+                        value -= 1
+                        gt_node_addr = stmt.false_target.value
+                        gt_node_idx = stmt.false_target_idx
+                        le_node_addr = stmt.true_target.value
+                        le_node_idx = stmt.true_target_idx
+                    elif op == "CmpLE":
+                        op_str = "gt"
+                        gt_node_addr = stmt.false_target.value
+                        gt_node_idx = stmt.false_target_idx
+                        le_node_addr = stmt.true_target.value
+                        le_node_idx = stmt.true_target_idx
+                    else:
+                        return None
+                    return (
+                        variable_hash,
+                        op_str,
+                        cond.operands[0],
+                        value,
+                        gt_node_addr,
+                        gt_node_idx,
+                        le_node_addr,
+                        le_node_idx,
+                    )
 
         return None
 
@@ -784,7 +784,7 @@ class LoweredSwitchSimplifier(StructuringOptimizationPass):
             onode, value, target, target_idx, next_node_addr = ca_others[next_node_addr]
             onode: Block
 
-            if first_nonlabel_statement(onode) is not onode.statements[-1]:
+            if first_nonlabel_nonphi_statement(onode) is not onode.statements[-1]:
                 onode = onode.copy(statements=[onode.statements[-1]])
 
             graph.add_edge(last_node, onode)
@@ -837,7 +837,4 @@ class LoweredSwitchSimplifier(StructuringOptimizationPass):
 
         if len(cases_0) > len(cases_1):
             return False
-        for case in cases_0:
-            if case not in cases_1:
-                return False
-        return True
+        return all(case in cases_1 for case in cases_0)

angr/analyses/decompiler/optimization_passes/mod_simplifier.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import logging
 
 from ailment import Expr
@@ -16,29 +17,32 @@ class ModSimplifierAILEngine(SimplifierAILEngine):
         operand_1 = self._expr(expr.operands[1])
 
         x_0, c_0, x_1, c_1 = None, None, None, None
-        if isinstance(operand_1, Expr.BinaryOp) and isinstance(operand_1.operands[1], Expr.Const):
-            if operand_1.op == "Mul":
-                if (
-                    isinstance(operand_1.operands[0], Expr.BinaryOp)
-                    and isinstance(operand_1.operands[0].operands[1], Expr.Const)
-                    and operand_1.operands[0].op in ["Div", "Mod"]
-                ):
-                    x_0 = operand_1.operands[0].operands[0]
-                    x_1 = operand_0
-                    c_0 = operand_1.operands[1]
-                    c_1 = operand_1.operands[0].operands[1]
-                elif (
-                    isinstance(operand_1.operands[0], Expr.Convert)
-                    and isinstance(operand_1.operands[0].operand, Expr.BinaryOp)
-                    and operand_1.operands[0].operand.op in ["Div", "Mod"]
-                ):
-                    x_0 = operand_1.operands[0].operand.operands[0]
-                    x_1 = operand_0
-                    c_0 = operand_1.operands[1]
-                    c_1 = operand_1.operands[0].operand.operands[1]
-
-                if x_0 is not None and x_1 is not None and x_0.likes(x_1) and c_0.value == c_1.value:
-                    return Expr.BinaryOp(expr.idx, "Mod", [x_0, c_0], expr.signed, **expr.tags)
+        if (
+            isinstance(operand_1, Expr.BinaryOp)
+            and isinstance(operand_1.operands[1], Expr.Const)
+            and operand_1.op == "Mul"
+        ):
+            if (
+                isinstance(operand_1.operands[0], Expr.BinaryOp)
+                and isinstance(operand_1.operands[0].operands[1], Expr.Const)
+                and operand_1.operands[0].op in ["Div", "Mod"]
+            ):
+                x_0 = operand_1.operands[0].operands[0]
+                x_1 = operand_0
+                c_0 = operand_1.operands[1]
+                c_1 = operand_1.operands[0].operands[1]
+            elif (
+                isinstance(operand_1.operands[0], Expr.Convert)
+                and isinstance(operand_1.operands[0].operand, Expr.BinaryOp)
+                and operand_1.operands[0].operand.op in ["Div", "Mod"]
+            ):
+                x_0 = operand_1.operands[0].operand.operands[0]
+                x_1 = operand_0
+                c_0 = operand_1.operands[1]
+                c_1 = operand_1.operands[0].operand.operands[1]
+
+            if x_0 is not None and x_1 is not None and x_0.likes(x_1) and c_0.value == c_1.value:
+                return Expr.BinaryOp(expr.idx, "Mod", [x_0, c_0], expr.signed, **expr.tags)
 
         if (operand_0, operand_1) != (expr.operands[0], expr.operands[1]):
             return Expr.BinaryOp(expr.idx, "Sub", [operand_0, operand_1], expr.signed, **expr.tags)

angr/analyses/decompiler/optimization_passes/multi_simplifier.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import logging
 
 from ailment import Expr
@@ -20,11 +21,13 @@ class MultiSimplifierAILEngine(SimplifierAILEngine):
         operand_1 = self._expr(expr.operands[1])
 
         # x + x = 2*x
-        if type(operand_0) in [Expr.Convert, Expr.Register]:
-            if isinstance(operand_1, (Expr.Convert, Expr.Register)):
-                if operand_0 == operand_1:
-                    count = Expr.Const(expr.idx, None, 2, operand_1.bits)
-                    return Expr.BinaryOp(expr.idx, "Mul", [operand_1, count], expr.signed, **expr.tags)
+        if (
+            type(operand_0) in [Expr.Convert, Expr.VirtualVariable]
+            and isinstance(operand_1, (Expr.Convert, Expr.VirtualVariable))
+            and operand_0 == operand_1
+        ):
+            count = Expr.Const(expr.idx, None, 2, operand_1.bits)
+            return Expr.BinaryOp(expr.idx, "Mul", [operand_1, count], expr.signed, **expr.tags)
         # 2*x + x = 3*x
         if Expr.BinaryOp in [type(operand_0), type(operand_1)]:
             if (
@@ -41,8 +44,7 @@ class MultiSimplifierAILEngine(SimplifierAILEngine):
                 const_x1 = operand_1.operands[1 - x1_index]
                 if x0 == x1:
                     new_const = Expr.Const(const_x1.idx, None, const_x1.value + 1, const_x1.bits)
-                    new_expr = Expr.BinaryOp(expr.idx, "Mul", [x0, new_const], expr.signed, **expr.tags)
-                    return new_expr
+                    return Expr.BinaryOp(expr.idx, "Mul", [x0, new_const], expr.signed, **expr.tags)
             elif (
                 isinstance(operand_0, Expr.BinaryOp)
                 and operand_0.op == "Mul"
@@ -57,8 +59,7 @@ class MultiSimplifierAILEngine(SimplifierAILEngine):
                 const_x0 = operand_0.operands[1 - x0_index]
                 if x0 == x1:
                     new_const = Expr.Const(const_x0.idx, None, const_x0.value + 1, const_x0.bits)
-                    new_expr = Expr.BinaryOp(expr.idx, "Mul", [x1, new_const], expr.signed, **expr.tags)
-                    return new_expr
+                    return Expr.BinaryOp(expr.idx, "Mul", [x1, new_const], expr.signed, **expr.tags)
             # 2*x + 3*x = 5*x
             elif (
                 isinstance(operand_0, Expr.BinaryOp)
@@ -79,8 +80,7 @@ class MultiSimplifierAILEngine(SimplifierAILEngine):
                     const_x1 = operand_1.operands[1 - x1_index]
                     if x0 == x1:
                         new_const = Expr.Const(const_x1.idx, None, const_x1.value + const_x0.value, const_x1.bits)
-                        new_expr = Expr.BinaryOp(expr.idx, "Mul", [x0, new_const], expr.signed, **expr.tags)
-                        return new_expr
+                        return Expr.BinaryOp(expr.idx, "Mul", [x0, new_const], expr.signed, **expr.tags)
 
         if (operand_0, operand_1) != (expr.operands[0], expr.operands[1]):
             return Expr.BinaryOp(expr.idx, "Add", [operand_0, operand_1], expr.signed, **expr.tags)
@@ -91,12 +91,13 @@ class MultiSimplifierAILEngine(SimplifierAILEngine):
         operand_1 = self._expr(expr.operands[1])
 
         # x + x = 2*x
-        if type(operand_0) in [Expr.Convert, Expr.Register]:
-            if isinstance(operand_1, (Expr.Convert, Expr.Register)):
-                if operand_0 == operand_1:
-                    count = Expr.Const(expr.idx, None, 0, 8)
-                    new_expr = Expr.BinaryOp(expr.idx, "Mul", [operand_1, count], expr.signed, **expr.tags)
-                    return new_expr
+        if (
+            type(operand_0) in [Expr.Convert, Expr.VirtualVariable]
+            and isinstance(operand_1, (Expr.Convert, Expr.VirtualVariable))
+            and operand_0 == operand_1
+        ):
+            count = Expr.Const(expr.idx, None, 0, 8)
+            return Expr.BinaryOp(expr.idx, "Mul", [operand_1, count], expr.signed, **expr.tags)
 
         # 2*x - x = x
         if Expr.BinaryOp in [type(operand_0), type(operand_1)]:
@@ -114,8 +115,7 @@ class MultiSimplifierAILEngine(SimplifierAILEngine):
                 const_x1 = operand_1.operands[1 - x1_index]
                 if x0 == x1:
                     new_const = Expr.Const(const_x1.idx, None, const_x1.value - 1, const_x1.bits)
-                    new_expr = Expr.BinaryOp(expr.idx, "Mul", [x0, new_const], expr.signed, **expr.tags)
-                    return new_expr
+                    return Expr.BinaryOp(expr.idx, "Mul", [x0, new_const], expr.signed, **expr.tags)
             elif (
                 isinstance(operand_0, Expr.BinaryOp)
                 and operand_0.op == "Mul"
@@ -130,8 +130,7 @@ class MultiSimplifierAILEngine(SimplifierAILEngine):
                 const_x0 = operand_0.operands[1 - x0_index]
                 if x0 == x1:
                     new_const = Expr.Const(const_x0.idx, None, const_x0.value - 1, const_x0.bits)
-                    new_expr = Expr.BinaryOp(expr.idx, "Mul", [x1, new_const], expr.signed, **expr.tags)
-                    return new_expr
+                    return Expr.BinaryOp(expr.idx, "Mul", [x1, new_const], expr.signed, **expr.tags)
             # 3*x - 2*x = x
             elif (
                 isinstance(operand_0, Expr.BinaryOp)
@@ -152,8 +151,7 @@ class MultiSimplifierAILEngine(SimplifierAILEngine):
                     const_x1 = operand_1.operands[1 - x1_index]
                     if x0 == x1:
                         new_const = Expr.Const(const_x1.idx, None, const_x0.value - const_x1.value, const_x1.bits)
-                        new_expr = Expr.BinaryOp(expr.idx, "Mul", [x0, new_const], expr.signed, **expr.tags)
-                        return new_expr
+                        return Expr.BinaryOp(expr.idx, "Mul", [x0, new_const], expr.signed, **expr.tags)
 
         if (operand_0, operand_1) != (expr.operands[0], expr.operands[1]):
             return Expr.BinaryOp(expr.idx, "Sub", [operand_0, operand_1], expr.signed, **expr.tags)
@@ -175,17 +173,15 @@ class MultiSimplifierAILEngine(SimplifierAILEngine):
         operand_0 = self._expr(expr.operands[0])
         operand_1 = self._expr(expr.operands[1])
 
-        if Expr.Const in [type(operand_0), type(operand_1)]:
-            if Expr.BinaryOp in [type(operand_0), type(operand_1)]:
-                const_, x0 = (operand_0, operand_1) if isinstance(operand_0, Expr.Const) else (operand_1, operand_0)
-                if x0.op == "Mul" and Expr.Const in [type(x0.operands[0]), type(x0.operands[1])]:
-                    if isinstance(x0.operands[0], Expr.Const):
-                        const_x0, x = x0.operands[0], x0.operands[1]
-                    else:
-                        const_x0, x = x0.operands[1], x0.operands[0]
-                    new_const = Expr.Const(const_.idx, None, const_.value * const_x0.value, const_.bits)
-                    new_expr = Expr.BinaryOp(expr.idx, "Mul", [x, new_const], expr.signed, **expr.tags)
-                    return new_expr
+        if Expr.Const in [type(operand_0), type(operand_1)] and Expr.BinaryOp in [type(operand_0), type(operand_1)]:
+            const_, x0 = (operand_0, operand_1) if isinstance(operand_0, Expr.Const) else (operand_1, operand_0)
+            if x0.op == "Mul" and Expr.Const in [type(x0.operands[0]), type(x0.operands[1])]:
+                if isinstance(x0.operands[0], Expr.Const):
+                    const_x0, x = x0.operands[0], x0.operands[1]
+                else:
+                    const_x0, x = x0.operands[1], x0.operands[0]
+                new_const = Expr.Const(const_.idx, None, const_.value * const_x0.value, const_.bits)
+                return Expr.BinaryOp(expr.idx, "Mul", [x, new_const], expr.signed, **expr.tags)
 
         if (operand_0, operand_1) != (expr.operands[0], expr.operands[1]):
             return Expr.BinaryOp(expr.idx, "Mul", [operand_0, operand_1], expr.signed, **expr.tags)