diff --git a/pylingual/control_flow_reconstruction/templates/Conditional.py b/pylingual/control_flow_reconstruction/templates/Conditional.py index 1e00c31..a17cb1a 100644 --- a/pylingual/control_flow_reconstruction/templates/Conditional.py +++ b/pylingual/control_flow_reconstruction/templates/Conditional.py @@ -1,5 +1,5 @@ from ..cft import ControlFlowTemplate, EdgeKind, register_template -from ..utils import T, N, defer_source_to, run_is, has_no_lines, with_instructions, has_instval, starting_instructions, to_indented_source, make_try_match, without_top_level_instructions +from ..utils import T, N, defer_source_to, run_is, has_no_lines, with_instructions, exact_instructions, without_instructions, has_instval, starting_instructions, to_indented_source, make_try_match, without_top_level_instructions from .Loop import BreakTemplate, ContinueTemplate @@ -51,7 +51,7 @@ class IfElseLoop(ControlFlowTemplate): class IfThen(ControlFlowTemplate): template = T( if_header=~N("if_body", "tail").with_cond(without_top_level_instructions("WITH_EXCEPT_START", "CHECK_EXC_MATCH", "FOR_ITER", "JUMP_IF_NOT_EXC_MATCH")), - if_body=~N(None).with_in_deg(1).of_type(BreakTemplate, ContinueTemplate) | ~N("tail").with_in_deg(1) | ~N("tail.").with_in_deg(1).with_cond(run_is(2)), + if_body=~N(None).with_in_deg(1).of_type(BreakTemplate, ContinueTemplate) | ~N("tail.").with_in_deg(1).with_cond(run_is(2)) | ~N.tail().with_in_deg(1).with_cond(exact_instructions("LOAD_CONST","RETURN_VALUE")), tail=N.tail(), ) diff --git a/pylingual/control_flow_reconstruction/templates/Loop.py b/pylingual/control_flow_reconstruction/templates/Loop.py index 0188718..b01595e 100644 --- a/pylingual/control_flow_reconstruction/templates/Loop.py +++ b/pylingual/control_flow_reconstruction/templates/Loop.py @@ -11,6 +11,9 @@ from ..utils import ( no_back_edges, versions_below, versions_from, + with_instructions, + exact_instructions, + ending_instructions, has_no_lines, condense_mapping, defer_source_to, @@ -42,6 +45,22 @@ class ForLoop(ControlFlowTemplate): {for_body} """ +@register_template(0, 2) +class LoopedReturn(ControlFlowTemplate): + template = T( + for_iter=~N("for_body", "tail").with_cond(ending_instructions("FOR_ITER")), + for_body=~N.tail().with_in_deg(1).with_cond(ending_instructions("RETURN_CONST"),ending_instructions("RETURN_VALUE")), + tail=N.tail(), + ) + + try_match = make_try_match({EdgeKind.Fall: "tail"}, "for_iter", "for_body") + + @to_indented_source + def to_indented_source(): + """ + {for_iter} + {for_body} + """ @register_template(0, 2, *versions_below(3, 10)) class SelfLoop3_6(ControlFlowTemplate):