# pylint:disable=unused-argument,arguments-differ
from collections import defaultdict
from typing import Dict, List
import ailment
from ..sequence_walker import SequenceWalker
from ..structuring.structurer_nodes import (
SequenceNode,
CodeNode,
MultiNode,
LoopNode,
ConditionNode,
ContinueNode,
CascadingConditionNode,
)
[docs]class LoopSimplifier(SequenceWalker):
"""
Simplifies loops.
"""
[docs] 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.continue_preludes: Dict[LoopNode, List[ailment.Block]] = defaultdict(list)
self.walk(node)
@staticmethod
def _control_transferring_statement(stmt: ailment.Stmt.Statement) -> bool:
return isinstance(
stmt, (ailment.Stmt.Call, ailment.Stmt.Return, ailment.Stmt.Jump, ailment.Stmt.ConditionalJump)
)
def _handle_sequencenode(self, node, predecessor=None, successor=None, loop=None, loop_successor=None, **kwargs):
for n0, n1, n2 in zip(node.nodes, node.nodes[1:] + [successor], [predecessor] + node.nodes[:-1]):
self._handle(n0, predecessor=n2, successor=n1, loop=loop, loop_successor=loop_successor)
def _handle_codenode(self, node, predecessor=None, successor=None, loop=None, loop_successor=None, **kwargs):
self._handle(node.node, predecessor=predecessor, successor=successor, loop=loop, loop_successor=loop_successor)
def _handle_conditionnode(self, node, predecessor=None, successor=None, loop=None, loop_successor=None, **kwargs):
if node.true_node is not None:
self._handle(
node.true_node, predecessor=predecessor, successor=successor, loop=loop, loop_successor=loop_successor
)
if node.false_node is not None:
self._handle(
node.false_node, predecessor=predecessor, successor=successor, loop=loop, loop_successor=loop_successor
)
def _handle_cascadingconditionnode(
self, node: CascadingConditionNode, predecessor=None, successor=None, loop=None, loop_successor=None, **kwargs
):
for _, child_node in node.condition_and_nodes:
self._handle(
child_node, predecessor=predecessor, successor=successor, loop=loop, loop_successor=loop_successor
)
if node.else_node is not None:
self._handle(
node.else_node, predecessor=predecessor, successor=successor, loop=loop, loop_successor=loop_successor
)
def _handle_loopnode(
self, node: LoopNode, predecessor=None, successor=None, loop=None, loop_successor=None, **kwargs
):
self._handle(
node.sequence_node, predecessor=predecessor, successor=successor, loop=node, loop_successor=successor
)
# find for-loop iterators
if (
node.sort == "while"
and self.continue_preludes[node]
and (
(node.condition is not None and not isinstance(node.condition, ailment.Expr.Const))
or len(self.continue_preludes[node]) > 1
)
):
if (
all(block.statements for block in self.continue_preludes[node])
and all(
not self._control_transferring_statement(block.statements[-1])
for block in self.continue_preludes[node]
)
and all(
block.statements[-1] == self.continue_preludes[node][0].statements[-1]
for block in self.continue_preludes[node]
)
):
node.sort = "for"
node.iterator = self.continue_preludes[node][0].statements[-1]
for block in self.continue_preludes[node]:
block.statements = block.statements[:-1]
# find for-loop initializers
if isinstance(predecessor, MultiNode):
predecessor = predecessor.nodes[-1]
if (
node.sort == "for"
and isinstance(predecessor, ailment.Block)
and predecessor.statements
and isinstance(predecessor.statements[-1], (ailment.Stmt.Assignment, ailment.Stmt.Store))
):
node.initializer = predecessor.statements[-1]
predecessor.statements = predecessor.statements[:-1]
def _handle_multinode(self, node, predecessor=None, successor=None, loop=None, loop_successor=None, **kwargs):
for n0, n1, n2 in zip(node.nodes, node.nodes[1:] + [successor], [predecessor] + node.nodes[:-1]):
self._handle(n0, predecessor=n2, successor=n1, loop=loop, loop_successor=loop_successor)
def _handle_block(
self, block, predecessor=None, successor=None, loop=None, loop_successor=None, **kwargs
): # pylint:disable=no-self-use
if isinstance(successor, ContinueNode) or successor is loop_successor:
self.continue_preludes[loop].append(block)