diff --git a/pylingual/control_flow_reconstruction/templates/Conditional.py b/pylingual/control_flow_reconstruction/templates/Conditional.py index 05f2239..a920de0 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, with_instructions, starting_instructions, to_indented_source, make_try_match, without_top_level_instructions +from ..utils import T, N, defer_source_to, run_is, with_instructions, has_instval, starting_instructions, to_indented_source, make_try_match, without_top_level_instructions @register_template(1, 40) @@ -46,7 +46,7 @@ class IfThen(ControlFlowTemplate): class Assertion(ControlFlowTemplate): template = T( assertion=~N("fail", "tail"), - fail=+N().with_cond(starting_instructions("LOAD_ASSERTION_ERROR"), with_instructions("LOAD_GLOBAL", "RAISE_VARARGS")), + fail=+N().with_cond(starting_instructions("LOAD_ASSERTION_ERROR"), has_instval("LOAD_GLOBAL", argval = "AssertionError")), tail=N.tail(), ) diff --git a/pylingual/control_flow_reconstruction/utils.py b/pylingual/control_flow_reconstruction/utils.py index 1eab7ab..8724bfd 100644 --- a/pylingual/control_flow_reconstruction/utils.py +++ b/pylingual/control_flow_reconstruction/utils.py @@ -118,6 +118,18 @@ def has_incoming_edge_of_categories(*categories: str): return False return check + +def has_instval(opname: str, argval : Any): + def check_instructions(cfg: CFG, node: ControlFlowTemplate | None) -> bool: + ops = {x.opname : x.argval for x in node.get_instructions()} + for op in ops: + if ops[op] == argval: + return True + return False + + return check_instructions + + def run_is(n: int): def check_run(cfg: CFG, node: ControlFlowTemplate | None) -> bool: return cfg.run == n