[For-Loop Reconstruction] Follow up: Not handling "continue" instructions correctly

decompiler/pipeline/controlflowanalysis/loop_utility_methods.py

@@ -1,12 +1,30 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Dict, Optional
-
-from decompiler.structures.ast.ast_nodes import AbstractSyntaxTreeNode, CaseNode, CodeNode, ConditionNode, LoopNode, SeqNode, SwitchNode
+from typing import Dict, List, Optional
+
+from decompiler.structures.ast.ast_nodes import (
+    AbstractSyntaxTreeNode,
+    CaseNode,
+    CodeNode,
+    ConditionNode,
+    LoopNode,
+    SeqNode,
+    SwitchNode,
+    WhileLoopNode,
+)
 from decompiler.structures.ast.syntaxtree import AbstractSyntaxTree
 from decompiler.structures.logic.logic_condition import LogicCondition
-from decompiler.structures.pseudo import Assignment, Condition, Variable
+from decompiler.structures.pseudo import (
+    Assignment,
+    BinaryOperation,
+    Condition,
+    Constant,
+    Expression,
+    OperationType,
+    UnaryOperation,
+    Variable,
+)
 from decompiler.structures.visitors.assignment_visitor import AssignmentVisitor
 
 
@@ -210,3 +228,85 @@ def _requirement_without_reinitialization(ast: AbstractSyntaxTree, node: Abstrac
                 return True
             elif variable in assignment.requirements:
                 return True
+
+
+def _get_equalizable_last_definitions(loop_node: WhileLoopNode, continuation: AstInstruction) -> List[CodeNode]:
+    """
+    Finds equalizable last definitions of the continuation instruction in the code nodes of a while loop containing continue statements.
+
+    :param loop_node: While-loop to search in
+    :param continuation: Instruction defining the for-loops modification
+    :return: List of equalizable last definitions, None if at least one continue node does not match the requirements
+    """
+    equalizable_nodes = []
+    for code_node in (
+        node for node in loop_node.body.get_descendant_code_nodes_interrupting_ancestor_loop() if node.does_end_with_continue
+    ):
+        if not _is_expression_simple_binary_operation(continuation.instruction.value):
+            return None
+        if (last_definition_index := _get_last_definition_index_of(code_node, continuation.instruction.destination)) == -1:
+            return None
+        else:
+            last_definition = code_node.instructions[last_definition_index]
+        if not (
+            isinstance(last_definition.value, Constant)
+            or _is_expression_simple_binary_operation(last_definition.value)
+            and continuation.instruction.destination
+            == _get_variable_in_binary_operation(continuation.instruction.value)
+            == _get_variable_in_binary_operation(last_definition.value)
+        ):
+            return None
+
+        _unify_binary_operation_in_assignment(continuation.instruction)
+        equalizable_nodes.append(last_definition)
+    return equalizable_nodes
+
+
+def _is_expression_simple_binary_operation(expression: Expression) -> bool:
+    """Checks if an expression is a simple binary operation. Meaning it includes a variable and a constant and uses plus or minus as operation type."""
+    return (
+        isinstance(expression, BinaryOperation)
+        and expression.operation in {OperationType.plus, OperationType.minus}
+        and any(isinstance(operand, Constant) or _is_negated_constant_variable(operand, Constant) for operand in expression.operands)
+        and any(isinstance(operand, Variable) or _is_negated_constant_variable(operand, Variable) for operand in expression.operands)
+    )
+
+
+def _is_negated_constant_variable(operand: Expression, expression: Constant | Variable) -> bool:
+    """Checks if an operand (constant or variable) is negated."""
+    return isinstance(operand, UnaryOperation) and operand.operation == OperationType.negate and isinstance(operand.operand, expression)
+
+
+def _get_variable_in_binary_operation(binaryoperation: BinaryOperation) -> Variable:
+    """Returns the used variable of a binary operation if available."""
+    for operand in binaryoperation.operands:
+        if isinstance(operand, Variable):
+            return operand
+        if _is_negated_constant_variable(operand, Variable):
+            return operand.operand
+    return None
+
+
+def _unify_binary_operation_in_assignment(assignment: Assignment):
+    """Brings a simple binary operation of an assignment into a unified representation like 'var = -var + const' instead of 'var = const - var'."""
+    if not assignment.value.operation == OperationType.plus:
+        assignment.substitute(assignment.value, BinaryOperation(OperationType.plus, [assignment.value.left, assignment.value.right]))
+        assignment.value.substitute(assignment.value.right, UnaryOperation(OperationType.negate, [assignment.value.right]))
+
+    if any(isinstance(operand, Constant) for operand in assignment.value.left.subexpressions()):
+        assignment.substitute(assignment.value, BinaryOperation(OperationType.plus, [assignment.value.right, assignment.value.left]))
+
+
+def _substract_continuation_from_last_definition(last_definition: Assignment, continuation: AstInstruction):
+    """
+    Substracts the value of the continuation instruction from the last definition, which must be a simple binary operation or a constant,
+    defining the same value as the continuation instruction in the given code node.
+
+    :param last_definition: Last definition that is to be changed
+    :param continuation: Instruction defining the for-loops modification
+    """
+    last_definition.substitute(
+        last_definition.value, BinaryOperation(OperationType.minus, [last_definition.value, continuation.instruction.value.right])
+    )
+    if _is_negated_constant_variable(continuation.instruction.value.left, Variable):
+        last_definition.substitute(last_definition.value, UnaryOperation(OperationType.negate, [last_definition.value]))

decompiler/pipeline/controlflowanalysis/readability_based_refinement.py

@@ -6,10 +6,12 @@
 from decompiler.pipeline.controlflowanalysis.loop_utility_methods import (
     AstInstruction,
     _find_continuation_instruction,
+    _get_equalizable_last_definitions,
     _get_variable_initialisation,
     _initialization_reaches_loop_node,
     _is_single_instruction_loop_node,
     _single_defininition_reaches_node,
+    _substract_continuation_from_last_definition,
 )
 from decompiler.pipeline.stage import PipelineStage
 from decompiler.structures.ast.ast_nodes import ConditionNode, DoWhileLoopNode, ForLoopNode, WhileLoopNode
@@ -76,6 +78,7 @@ def run(self):
             -> loop condition complexity < condition complexity
             -> possible modification complexity < modification complexity
             -> if condition is only a symbol: check condition type for allowed one
+            -> has a continue statement which must and can be equalized
 
         If 'force_for_loops' is enabled, the complexity options are ignored and every while loop after the
         initial transformation will be forced into a for loop with an empty declaration/modification
@@ -90,9 +93,6 @@ def run(self):
             ):
                 continue
 
-            if any(node.does_end_with_continue for node in loop_node.body.get_descendant_code_nodes_interrupting_ancestor_loop()):
-                continue
-
             if not self._force_for_loops and loop_node.condition.get_complexity(self._ast.condition_map) > self._condition_max_complexity:
                 continue
 
@@ -103,6 +103,10 @@ def run(self):
                     continue
                 if not self._force_for_loops and continuation.instruction.complexity > self._modification_max_complexity:
                     continue
+                if (equalizable_last_definitions := _get_equalizable_last_definitions(loop_node, continuation)) is None:
+                    continue
+                for last_definition in equalizable_last_definitions:
+                    _substract_continuation_from_last_definition(last_definition, continuation)
                 self._replace_with_for_loop(loop_node, continuation, variable_init)
                 break