# pylint:disable=unused-argument,arguments-differ
import logging
import ailment
from ..sequence_walker import SequenceWalker
from ..structuring.structurer_nodes import (
SequenceNode,
CodeNode,
MultiNode,
LoopNode,
ConditionNode,
CascadingConditionNode,
)
l = logging.getLogger(name=__name__)
[docs]class IfSimplifier(SequenceWalker):
"""
Remove unnecessary jump or conditional jump statements if they jump to the successor right afterwards.
"""
def __init__(self, node):
handlers = {
SequenceNode: self._handle_sequencenode,
CodeNode: self._handle_codenode,
MultiNode: self._handle_multinode,
LoopNode: self._handle_loopnode,
ConditionNode: self._handle_conditionnode,
CascadingConditionNode: self._handle_cascadingconditionnode,
ailment.Block: self._handle_block,
}
super().__init__(handlers)
self.walk(node)
def _handle_sequencenode(self, node, successor=None, **kwargs):
"""
:param SequenceNode node:
:return:
"""
for n0, n1 in zip(node.nodes, node.nodes[1:] + [successor]):
self._handle(n0, successor=n1)
def _handle_codenode(self, node, successor=None, **kwargs):
"""
:param CodeNode node:
:return:
"""
self._handle(node.node, successor=successor)
def _handle_conditionnode(self, node, successor=None, **kwargs):
"""
:param ConditionNode node:
:param successor:
:return:
"""
if node.true_node is not None:
self._handle(node.true_node, successor=successor)
if node.false_node is not None:
self._handle(node.false_node, successor=successor)
def _handle_cascadingconditionnode(self, node: CascadingConditionNode, successor=None, **kwargs):
for _, child_node in node.condition_and_nodes:
self._handle(child_node, successor=successor)
if node.else_node is not None:
self._handle(node.else_node, successor=successor)
def _handle_loopnode(self, node, successor=None, **kwargs):
"""
:param LoopNode node:
:param successor:
:return:
"""
self._handle(node.sequence_node, successor=successor)
def _handle_multinode(self, node, successor=None, **kwargs):
"""
:param MultiNode node:
:return:
"""
for n0, n1 in zip(node.nodes, node.nodes[1:] + [successor]):
self._handle(n0, successor=n1)
def _handle_block(self, block, successor=None, **kwargs): # pylint:disable=no-self-use
"""
Remove unnecessary jump or conditional jump statements if they jump to the successor right afterwards.
:param ailment.Block block:
:return:
"""
if block.statements and isinstance(block.statements[-1], ailment.Stmt.ConditionalJump):
cond_stmt = block.statements[-1] # ailment.Stmt.ConditionalJump
if isinstance(successor, ConditionNode):
true_cond = False
if cond_stmt.true_target is not None and successor.true_node is not None:
# True branch exists. Test if the true target is the address
if (
isinstance(cond_stmt.true_target, ailment.Expr.Const)
and cond_stmt.true_target.value == successor.true_node.addr
):
true_cond = True
if cond_stmt.true_target is not None and successor.false_node is not None:
# True branch exists. Test if the true target is the address
if (
isinstance(cond_stmt.true_target, ailment.Expr.Const)
and cond_stmt.true_target.value == successor.false_node.addr
):
true_cond = True
false_cond = False
if cond_stmt.false_target is not None and successor.false_node is not None:
# False branch exists. Test if the false target is the address
if (
isinstance(cond_stmt.true_target, ailment.Expr.Const)
and cond_stmt.false_target.value == successor.false_node.addr
):
false_cond = True
if cond_stmt.false_target is not None and successor.true_node is not None:
# True branch exists. Test if the true target is the address
if (
isinstance(cond_stmt.true_target, ailment.Expr.Const)
and cond_stmt.false_target.value == successor.true_node.addr
):
false_cond = True
if true_cond or false_cond:
# We can safely remove this statement
block.statements = block.statements[:-1]