diff --git a/pylingual/control_flow_reconstruction/templates/Block.py b/pylingual/control_flow_reconstruction/templates/Block.py index 7f1ae49..371b2cb 100644 --- a/pylingual/control_flow_reconstruction/templates/Block.py +++ b/pylingual/control_flow_reconstruction/templates/Block.py @@ -130,6 +130,7 @@ class BlockTemplate(ControlFlowTemplate): @override @classmethod def try_match(cls, cfg, node) -> ControlFlowTemplate | None: + from .Loop import LoopElse members: list[ControlFlowTemplate] = [] out = out_edge_dict(cfg, node) exc = out[EdgeCategory.Exception] @@ -141,6 +142,8 @@ class BlockTemplate(ControlFlowTemplate): break if current in members: break + if isinstance(current, LoopElse): + break members.append(current) next = out[EdgeCategory.Natural] if next is None: diff --git a/pylingual/control_flow_reconstruction/templates/Loop.py b/pylingual/control_flow_reconstruction/templates/Loop.py index 21735a3..4310cbd 100644 --- a/pylingual/control_flow_reconstruction/templates/Loop.py +++ b/pylingual/control_flow_reconstruction/templates/Loop.py @@ -8,11 +8,10 @@ from ..cft import ControlFlowTemplate, EdgeKind, register_template from ..utils import ( T, N, + is_not_type, no_back_edges, versions_below, versions_from, - with_instructions, - exact_instructions, ending_instructions, has_no_lines, condense_mapping, @@ -28,12 +27,24 @@ if TYPE_CHECKING: from pylingual.control_flow_reconstruction.cfg import CFG +class LoopElse(ControlFlowTemplate): + @classmethod + def try_match(cls, cfg, node): + if has_no_lines(cfg, node): + return None + else: + return condense_mapping(cls, cfg, {"child": node}, "child") + + def to_indented_source(self, source): + return self.child.to_indented_source(source) + + @register_template(0, 1) class ForLoop(ControlFlowTemplate): template = T( for_iter=~N("for_body", "tail"), for_body=~N("for_iter").with_in_deg(1), - tail=N.tail(), + tail=N.tail().with_cond(is_not_type(LoopElse)), ) try_match = make_try_match({EdgeKind.Fall: "tail"}, "for_iter", "for_body") @@ -45,6 +56,28 @@ class ForLoop(ControlFlowTemplate): {for_body} """ + +@register_template(0, 1) +class ForElseLoop(ControlFlowTemplate): + template = T( + for_iter=~N("for_body", "else_body"), + for_body=~N("for_iter").with_in_deg(1), + else_body=~N("tail.").of_type(LoopElse), + tail=N.tail(), + ) + + try_match = make_try_match({EdgeKind.Fall: "tail"}, "for_iter", "for_body", "else_body") + + @to_indented_source + def to_indented_source(): + """ + {for_iter} + {for_body} + else: + {else_body} + """ + + @register_template(0, 2) class LoopedReturn(ControlFlowTemplate): template = T( @@ -259,10 +292,12 @@ class FixLoop(ControlFlowTemplate): edges_to_remove = [] # Find the candidate end that break connects to + false_edge = None candidate_end = None for succ in cfg.successors(node): if cfg.get_edge_data(node, succ).get("kind") == EdgeKind.FalseJump and not any(n == node for n in cfg.successors(succ)): candidate_end = succ + false_edge = succ # Candidate end is a buffer node if cfg.in_degree(candidate_end) == 1: @@ -307,6 +342,8 @@ class FixLoop(ControlFlowTemplate): break_node = BreakTemplate.try_match(cfg, pred) if break_node is not None: cfg.remove_edge(break_node, succ) + if succ != false_edge: + LoopElse.try_match(cfg, false_edge) cfg.iterate() return diff --git a/pylingual/control_flow_reconstruction/utils.py b/pylingual/control_flow_reconstruction/utils.py index dc02c5c..86ee251 100644 --- a/pylingual/control_flow_reconstruction/utils.py +++ b/pylingual/control_flow_reconstruction/utils.py @@ -116,6 +116,13 @@ def has_type(*template_type: type[ControlFlowTemplate]): return check_type +def is_not_type(*template_type: type[ControlFlowTemplate]): + def check_type(cfg: CFG, node: ControlFlowTemplate | None) -> bool: + return not isinstance(node, template_type) + + return check_type + + def no_back_edges(cfg: CFG, node: ControlFlowTemplate | None) -> bool: return node is None or not any(cfg.dominates(succ, node) for succ in cfg.successors(node))