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

decompiler/pipeline/controlflowanalysis/readability_based_refinement.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Dict, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from decompiler.pipeline.stage import PipelineStage
 from decompiler.structures.ast.ast_nodes import (
@@ -19,7 +19,7 @@
 )
 from decompiler.structures.ast.syntaxtree import AbstractSyntaxTree
 from decompiler.structures.logic.logic_condition import LogicCondition
-from decompiler.structures.pseudo import Assignment, Condition, Expression, Operation, Variable
+from decompiler.structures.pseudo import Assignment, BinaryOperation, Condition, Constant, Expression, Operation, OperationType, Variable
 from decompiler.structures.visitors.assignment_visitor import AssignmentVisitor
 from decompiler.task import DecompilerTask
 from decompiler.util.options import Options
@@ -277,6 +277,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      
@@ -287,9 +288,6 @@ def run(self):
             or self._invalid_simple_for_loop_condition_type(loop_node.condition):
                 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
 
@@ -300,6 +298,10 @@ def run(self):
                     continue
                 if not self._force_for_loops and continuation.instruction.complexity > self._modification_max_complexity:
                     continue
+                if (equalizable_continue_nodes := self._get_continue_nodes_with_equalizable_definition(loop_node, continuation, variable_init)) is None:
+                    break
+                for node in equalizable_continue_nodes:
+                    self._remove_continuation_from_last_definition(node, continuation, variable_init)
                 self._replace_with_for_loop(loop_node, continuation, variable_init)
                 break
 
@@ -315,6 +317,63 @@ def run(self):
                     )
                 )
 
+    def _get_continue_nodes_with_equalizable_definition(self, loop_node: WhileLoopNode, continuation: AstInstruction, variable_init: AstInstruction) -> List[CodeNode]:
+        """
+        Finds code nodes of a while loop containing continue statements and a definition of the continuation instruction, which can be easily equalized.
+
+        :param loop_node: While-loop to search in
+        :param continuation: Instruction defining the for-loops modification
+        :param variable_init: Instruction defining the for-loops declaration
+        :return: List of continue code nodes that has equalizable definitions, None if at least one continue node does not match the requirements
+        """
+        equalizable_nodes = []
+        for code_node in loop_node.body.get_descendant_code_nodes_interrupting_ancestor_loop():
+            if not code_node.does_end_with_continue:
+                continue
+            if not self._is_expression_simple_binary_operation(continuation.instruction.value):
+                return None
+            if (last_definition_index := _get_last_definition_index_of(code_node, variable_init.instruction.destination)) == -1:
+                return None
+            if not (self._is_expression_simple_binary_operation(code_node.instructions[last_definition_index].value) \
+                or isinstance(code_node.instructions[last_definition_index].value, Constant)):
+                return None
+            equalizable_nodes.append(code_node)
+        return equalizable_nodes
+
+    def _is_expression_simple_binary_operation(self, expression: Expression) -> bool:
+        """Checks if an expression is a binary operation with plus or minus and a constant."""
+        if not isinstance(expression, BinaryOperation):
+            return False
+        if not expression.operation in [OperationType.plus, OperationType.minus]:
+            return False
+        if not any(isinstance(operand, Constant) for operand in expression.operands):
+            return False
+        return True
+
+    def _remove_continuation_from_last_definition(self, code_node: CodeNode, continuation: AstInstruction, variable_init: AstInstruction):
+        """
+        Substracts the value of the continuation instruction from the last definition, which must be a simple binary operation or a constant.
+
+        :param code_node: Code node whose last instruction is to be changed
+        :param continuation: Instruction defining the for-loops modification
+        :param variable_init: Instruction defining the for-loops declaration
+        """
+        last_definition = code_node.instructions[_get_last_definition_index_of(code_node, variable_init.instruction.destination)]
+        continuation_operand = continuation.instruction.value.right if isinstance(continuation.instruction.value.right, Constant) else continuation.instruction.value.left
+
+        if isinstance(last_definition.value, BinaryOperation):
+            last_definition_operand = last_definition.value.right if isinstance(last_definition.value.right, Constant) else last_definition.value.left
+
+            if continuation.instruction.value.operation == last_definition.value.operation:
+                last_definition_operand.value -= continuation_operand.value
+            else:
+                last_definition_operand.value += continuation_operand.value
+        else:
+            if continuation.instruction.value.operation == OperationType.plus:
+                last_definition.value.value -= continuation_operand.value
+            else:
+                last_definition.value.value += continuation_operand.value
+
     def _replace_with_for_loop(self, loop_node: WhileLoopNode, continuation: AstInstruction, init: AstInstruction):
         """
         Replaces a given WhileLoopNode with a ForLoopNode.
@@ -347,7 +406,7 @@ def _replace_with_for_loop(self, loop_node: WhileLoopNode, continuation: AstInst
         )
         continuation.node.instructions.remove(continuation.instruction)
         self._ast.clean_up()
-       
+
     def _invalid_simple_for_loop_condition_type(self, logic_condition) -> bool:
         """ Checks if a logic condition is only a symbol, if true checks condition type of symbol for forbidden ones"""
         if not logic_condition.is_symbol or not self._forbidden_condition_types:

tests/pipeline/controlflowanalysis/test_readability_based_refinement.py

@@ -7,6 +7,7 @@
     WhileLoopReplacer,
     WhileLoopVariableRenamer,
     _find_continuation_instruction,
+    _get_last_definition_index_of,
     _has_deep_requirement,
     _initialization_reaches_loop_node,
 )
@@ -2080,8 +2081,10 @@ def test_declaration_listop(self, ast_call_for_loop):
             if isinstance(node, ForLoopNode):
                 assert node.declaration.destination.operands[0].name == "i"
 
-    def test_skip_for_loop_recovery_if_continue_in_while(self):
+    def test_for_loop_recovery_if_continue_in_while_1(self):
         """
+        Test for loop recovery if a continue occurs in a while loop and the last definition is a simple binary operation
+
         a = 0
         while(a < 10) {
             if(a == 2) {
@@ -2127,9 +2130,68 @@ def test_skip_for_loop_recovery_if_continue_in_while(self):
         )
 
         WhileLoopReplacer(ast, _generate_options()).run()
-        assert not any(isinstance(loop_node, ForLoopNode) for loop_node in list(ast.get_loop_nodes_post_order()))
+        assert all(isinstance(loop_node, ForLoopNode) for loop_node in list(ast.get_loop_nodes_post_order()))
+
+        condition_nodes = list(ast.get_condition_nodes_post_order())
+        last_definition = condition_nodes[0].true_branch_child.instructions[_get_last_definition_index_of(condition_nodes[0].true_branch_child, Variable("a"))]
+        assert last_definition.value.right.value == 1
+
+    def test_for_loop_recovery_if_continue_in_while_2(self):
+        """
+        Test for loop recovery if a continue occurs in a while loop and the last definition is a contant assignment
+
+        a = 0
+        while(a < 10) {
+            if(a == 2) {
+                a = 4
+                continue
+            }
+            a = a + 1
+        }
+        """
+        true_value = LogicCondition.initialize_true(context := LogicCondition.generate_new_context())
+        ast = AbstractSyntaxTree(
+            root := SeqNode(true_value),
+            condition_map={
+                logic_cond("x1", context): Condition(OperationType.less, [Variable("a"), Constant(10)]),
+                logic_cond("x2", context): Condition(OperationType.equal, [Variable("a"), Constant(2)])
+            }
+        )
+
+        true_branch = ast._add_code_node(
+            [
+                Assignment(Variable("a"), Constant(4)),
+                Continue()
+            ]
+        )
+        if_condition = ast._add_condition_node_with(logic_cond("x2", context), true_branch)
+
+        init_code_node = ast._add_code_node([Assignment(Variable("a"), Constant(0))])
+
+        while_loop = ast.factory.create_while_loop_node(logic_cond("x1", context))
+        while_loop_body = ast.factory.create_seq_node()
+        while_loop_iteration = ast._add_code_node([Assignment(Variable("a"), BinaryOperation(OperationType.plus, [Variable("a"), Constant(1)]))])
+        ast._add_node(while_loop)
+        ast._add_node(while_loop_body)
+
+        ast._add_edges_from(
+            [
+                (root, init_code_node),
+                (root, while_loop),
+                (while_loop, while_loop_body),
+                (while_loop_body, if_condition),
+                (while_loop_body, while_loop_iteration)
+            ]
+        )
 
-    def test_skip_for_loop_recovery_if_continue_in_nested_while(self):
+        WhileLoopReplacer(ast, _generate_options()).run()
+        assert all(isinstance(loop_node, ForLoopNode) for loop_node in list(ast.get_loop_nodes_post_order()))
+
+        condition_nodes = list(ast.get_condition_nodes_post_order())
+        last_definition = condition_nodes[0].true_branch_child.instructions[_get_last_definition_index_of(condition_nodes[0].true_branch_child, Variable("a"))]
+        assert last_definition.value.value == 3
+
+    def test_for_loop_recovery_if_continue_in_nested_while(self):
         """
         while(a < 5) {
             a = a + b
@@ -2188,5 +2250,156 @@ def test_skip_for_loop_recovery_if_continue_in_nested_while(self):
         )
 
         WhileLoopReplacer(ast, _generate_options()).run()
-        loop_nodes = list(ast.get_loop_nodes_post_order())
-        assert not isinstance(loop_nodes[0], ForLoopNode) and isinstance(loop_nodes[1], ForLoopNode)
+        assert all(isinstance(loop_node, ForLoopNode) for loop_node in list(ast.get_loop_nodes_post_order()))
+
+        condition_nodes = list(ast.get_condition_nodes_post_order())
+        last_definition = condition_nodes[0].true_branch_child.instructions[_get_last_definition_index_of(condition_nodes[0].true_branch_child, Variable("b"))]
+        assert last_definition.value.right.value == 1
+
+    def test_skip_for_loop_recovery_if_continue_in_while_1(self):
+        """
+        Test skip of for loop recovery if a continue occurs in a while loop, because the continuation instruction is no simple binary operation
+
+        a = 0
+        while(a < 10) {
+            if(a == 2) {
+                a = a + 2
+                continue
+            }
+            a = a * 2
+        }
+        """
+        true_value = LogicCondition.initialize_true(context := LogicCondition.generate_new_context())
+        ast = AbstractSyntaxTree(
+            root := SeqNode(true_value),
+            condition_map={
+                logic_cond("x1", context): Condition(OperationType.less, [Variable("a"), Constant(10)]),
+                logic_cond("x2", context): Condition(OperationType.equal, [Variable("a"), Constant(2)])
+            }
+        )
+
+        true_branch = ast._add_code_node(
+            [
+                Assignment(Variable("a"), BinaryOperation(OperationType.plus, [Variable("a"), Constant(2)])),
+                Continue()
+            ]
+        )
+        if_condition = ast._add_condition_node_with(logic_cond("x2", context), true_branch)
+
+        init_code_node = ast._add_code_node([Assignment(Variable("a"), Constant(0))])
+
+        while_loop = ast.factory.create_while_loop_node(logic_cond("x1", context))
+        while_loop_body = ast.factory.create_seq_node()
+        while_loop_iteration = ast._add_code_node([Assignment(Variable("a"), BinaryOperation(OperationType.multiply, [Variable("a"), Constant(2)]))])
+        ast._add_node(while_loop)
+        ast._add_node(while_loop_body)
+
+        ast._add_edges_from(
+            [
+                (root, init_code_node),
+                (root, while_loop),
+                (while_loop, while_loop_body),
+                (while_loop_body, if_condition),
+                (while_loop_body, while_loop_iteration)
+            ]
+        )
+
+        WhileLoopReplacer(ast, _generate_options()).run()
+        assert not any(isinstance(loop_node, ForLoopNode) for loop_node in list(ast.get_loop_nodes_post_order()))
+
+
+    def test_skip_for_loop_recovery_if_continue_in_while_2(self):
+        """
+        Test skip of for loop recovery if a continue occurs in a while loop, because the last definition is no simple binary operation
+
+        a = 0
+        while(a < 10) {
+            if(a == 2) {
+                a = a * 2
+                continue
+            }
+            a = a + 1
+        }
+        """
+        true_value = LogicCondition.initialize_true(context := LogicCondition.generate_new_context())
+        ast = AbstractSyntaxTree(
+            root := SeqNode(true_value),
+            condition_map={
+                logic_cond("x1", context): Condition(OperationType.less, [Variable("a"), Constant(10)]),
+                logic_cond("x2", context): Condition(OperationType.equal, [Variable("a"), Constant(2)])
+            }
+        )
+
+        true_branch = ast._add_code_node(
+            [
+                Assignment(Variable("a"), BinaryOperation(OperationType.multiply, [Variable("a"), Constant(2)])),
+                Continue()
+            ]
+        )
+        if_condition = ast._add_condition_node_with(logic_cond("x2", context), true_branch)
+
+        init_code_node = ast._add_code_node([Assignment(Variable("a"), Constant(0))])
+
+        while_loop = ast.factory.create_while_loop_node(logic_cond("x1", context))
+        while_loop_body = ast.factory.create_seq_node()
+        while_loop_iteration = ast._add_code_node([Assignment(Variable("a"), BinaryOperation(OperationType.plus, [Variable("a"), Constant(1)]))])
+        ast._add_node(while_loop)
+        ast._add_node(while_loop_body)
+
+        ast._add_edges_from(
+            [
+                (root, init_code_node),
+                (root, while_loop),
+                (while_loop, while_loop_body),
+                (while_loop_body, if_condition),
+                (while_loop_body, while_loop_iteration)
+            ]
+        )
+
+        WhileLoopReplacer(ast, _generate_options()).run()
+        assert not any(isinstance(loop_node, ForLoopNode) for loop_node in list(ast.get_loop_nodes_post_order()))
+
+    def test_skip_for_loop_recovery_if_continue_in_while_3(self):
+        """
+        Test skip of for loop recovery if a continue occurs in a while loop, because no last definition exists
+
+        a = 0
+        while(a < 10) {
+            if(a == 2) {
+                continue
+            }
+            a = a + 1
+        }
+        """
+        true_value = LogicCondition.initialize_true(context := LogicCondition.generate_new_context())
+        ast = AbstractSyntaxTree(
+            root := SeqNode(true_value),
+            condition_map={
+                logic_cond("x1", context): Condition(OperationType.less, [Variable("a"), Constant(10)]),
+                logic_cond("x2", context): Condition(OperationType.equal, [Variable("a"), Constant(2)])
+            }
+        )
+
+        true_branch = ast._add_code_node([Continue()])
+        if_condition = ast._add_condition_node_with(logic_cond("x2", context), true_branch)
+
+        init_code_node = ast._add_code_node([Assignment(Variable("a"), Constant(0))])
+
+        while_loop = ast.factory.create_while_loop_node(logic_cond("x1", context))
+        while_loop_body = ast.factory.create_seq_node()
+        while_loop_iteration = ast._add_code_node([Assignment(Variable("a"), BinaryOperation(OperationType.plus, [Variable("a"), Constant(1)]))])
+        ast._add_node(while_loop)
+        ast._add_node(while_loop_body)
+
+        ast._add_edges_from(
+            [
+                (root, init_code_node),
+                (root, while_loop),
+                (while_loop, while_loop_body),
+                (while_loop_body, if_condition),
+                (while_loop_body, while_loop_iteration)
+            ]
+        )
+
+        WhileLoopReplacer(ast, _generate_options()).run()
+        assert not any(isinstance(loop_node, ForLoopNode) for loop_node in list(ast.get_loop_nodes_post_order()))