diff --git a/src/pyqasm/visitor.py b/src/pyqasm/visitor.py index bd45db75..dec9e3e4 100644 --- a/src/pyqasm/visitor.py +++ b/src/pyqasm/visitor.py @@ -200,6 +200,7 @@ def _construct_visit_map(self): qasm3_ast.IODeclaration: lambda x: [], qasm3_ast.BreakStatement: self._visit_break, qasm3_ast.ContinueStatement: self._visit_continue, + qasm3_ast.EndStatement: self._visit_end_statement, qasm3_ast.DelayInstruction: self._visit_delay_statement, qasm3_ast.Box: self._visit_box_statement, qasm3_ast.Pragma: self._visit_pragma, @@ -2545,9 +2546,10 @@ def _visit_forin_loop(self, statement: qasm3_ast.ForInLoop) -> list[qasm3_ast.St if statement_block != statement.block: statement_block = copy.deepcopy(statement.block) - result.extend(self.visit_basic_block(statement_block)) + iteration_statements = self.visit_basic_block(statement_block) else: - result.extend(self.visit_basic_block(statement.block)) + iteration_statements = self.visit_basic_block(statement.block) + result.extend(iteration_statements) # scope not persistent between loop iterations self._scope_manager.pop_scope() @@ -2557,6 +2559,8 @@ def _visit_forin_loop(self, statement: qasm3_ast.ForInLoop) -> list[qasm3_ast.St # not runtime errors, we can break here if self._check_only: return [] + if self._ends_with_end_statement(iteration_statements): + break return result def _visit_subroutine_definition( @@ -2717,9 +2721,12 @@ def _visit_function_call( return_statement = copy.copy(function_op) break try: - result.extend(self.visit_statement(copy.copy(function_op))) + function_statements = self.visit_statement(copy.copy(function_op)) except (TypeError, copy.Error): - result.extend(self.visit_statement(copy.deepcopy(function_op))) + function_statements = self.visit_statement(copy.deepcopy(function_op)) + result.extend(function_statements) + if self._ends_with_end_statement(function_statements): + break if return_statement: return_value, stmts = Qasm3ExprEvaluator.evaluate_expression( @@ -2779,8 +2786,10 @@ def _visit_while_loop(self, statement: qasm3_ast.WhileLoop) -> list[qasm3_ast.St self._scope_manager.push_context(Context.BLOCK) self._scope_manager.push_scope({}) + loop_statements = [] try: - result.extend(self.visit_basic_block(statement.block)) + loop_statements = self.visit_basic_block(statement.block) + result.extend(loop_statements) except LoopControlSignal as lcs: self._scope_manager.pop_scope() self._scope_manager.restore_context() @@ -2792,6 +2801,9 @@ def _visit_while_loop(self, statement: qasm3_ast.WhileLoop) -> list[qasm3_ast.St self._scope_manager.pop_scope() self._scope_manager.restore_context() + if self._ends_with_end_statement(loop_statements): + break + loop_counter += 1 if loop_counter >= max_iterations: raise_qasm3_error( @@ -2968,7 +2980,10 @@ def _evaluate_case(statements): result = [] for stmt in statements: Qasm3Validator.validate_statement_type(SWITCH_BLACKLIST_STMTS, stmt, "switch") - result.extend(self.visit_statement(stmt)) + case_statements = self.visit_statement(stmt) + result.extend(case_statements) + if self._ends_with_end_statement(case_statements): + break self._scope_manager.pop_scope() self._scope_manager.restore_context() @@ -3505,6 +3520,35 @@ def _visit_include(self, include: qasm3_ast.Include) -> list[qasm3_ast.Statement return [include] + def _visit_end_statement( + self, statement: qasm3_ast.EndStatement + ) -> list[qasm3_ast.EndStatement]: + """Visit a statement that terminates the program. + + Args: + statement (EndStatement): The terminating statement to visit. + + Returns: + list[EndStatement]: The statement in a list, or an empty list if + self._check_only is True. + """ + if self._check_only: + return [] + return [statement] + + @classmethod + def _ends_with_end_statement(cls, statements: Sequence[qasm3_ast.Statement]) -> bool: + """Return whether a statement sequence ends with unconditional termination.""" + if not statements: + return False + + final_statement = statements[-1] + if isinstance(final_statement, qasm3_ast.EndStatement): + return True + if isinstance(final_statement, qasm3_ast.Box): + return cls._ends_with_end_statement(final_statement.body) + return False + def visit_statement( self, statement: qasm3_ast.Statement | qasm3_ast.Pragma ) -> list[qasm3_ast.Statement]: @@ -3554,7 +3598,10 @@ def visit_basic_block( """ result = [] for stmt in stmt_list: - result.extend(self.visit_statement(stmt)) + statements = self.visit_statement(stmt) + result.extend(statements) + if self._ends_with_end_statement(statements): + break return result def finalize(self, unrolled_stmts: list[qasm3_ast.Statement]) -> list[qasm3_ast.Statement]: diff --git a/tests/qasm3/test_end.py b/tests/qasm3/test_end.py new file mode 100644 index 00000000..5494e575 --- /dev/null +++ b/tests/qasm3/test_end.py @@ -0,0 +1,182 @@ +# Copyright 2026 qBraid +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for terminating an OpenQASM 3 program with ``end``.""" + +import pytest + +from pyqasm.entrypoint import dumps, loads +from tests.utils import check_unrolled_qasm + + +def test_end_stops_global_unrolling_and_bookkeeping(): + """Statements after a global ``end`` are unreachable.""" + module = loads(""" + OPENQASM 3.0; + include "stdgates.inc"; + qubit q; + h q; + end; + x q; + qubit[2] unreachable; + """) + + module.validate() + module.unroll() + + expected = """ + OPENQASM 3.0; + include "stdgates.inc"; + qubit[1] q; + h q[0]; + end; + """ + check_unrolled_qasm(dumps(module), expected) + assert module.num_qubits == 1 + assert module.depth() == 1 + + round_tripped = loads(dumps(module)) + round_tripped.unroll() + check_unrolled_qasm(dumps(round_tripped), expected) + + +@pytest.mark.parametrize( + "control_flow", + [ + "if (true) { h q; end; x q; }", + "if (false) { x q; } else { h q; end; x q; }", + "for int i in [0:2] { h q; end; x q; }", + "int i = 1; switch (i) { case 1 { h q; end; x q; } default { x q; } }", + ], +) +def test_end_propagates_from_static_control_flow(control_flow): + """A reachable ``end`` in static control flow terminates the program.""" + module = loads(f""" + OPENQASM 3.0; + include "stdgates.inc"; + qubit q; + {control_flow} + x q; + """) + + module.unroll() + output = dumps(module) + + assert output.count("h q[0];") == 1 + assert output.count("end;") == 1 + assert "x q[0];" not in output + assert output.rstrip().endswith("end;") + + +def test_end_stops_while_loop_and_program(): + """A terminating while-loop body is not expanded again.""" + module = loads(""" + OPENQASM 3.0; + include "stdgates.inc"; + qubit q; + int i = 0; + while (i < 2) { + h q; + end; + i += 1; + } + x q; + """) + + module.unroll(max_loop_iters=2) + output = dumps(module) + + assert output.count("h q[0];") == 1 + assert output.count("end;") == 1 + assert "x q[0];" not in output + + +def test_end_propagates_from_inlined_subroutine(): + """An ``end`` reached in an inlined subroutine terminates its caller.""" + module = loads(""" + OPENQASM 3.0; + include "stdgates.inc"; + def stop(qubit q) { + h q; + end; + x q; + } + qubit q; + stop(q); + x q; + """) + + module.unroll() + output = dumps(module) + + assert output.count("h q[0];") == 1 + assert output.count("end;") == 1 + assert "x q[0];" not in output + + +def test_end_propagates_from_box(): + """A box preserves its ``end`` and terminates the surrounding block.""" + module = loads(""" + OPENQASM 3.0; + include "stdgates.inc"; + qubit q; + box { + h q; + end; + x q; + } + x q; + """) + + module.unroll() + + expected = """ + OPENQASM 3.0; + include "stdgates.inc"; + qubit[1] q; + box { + h q[0]; + end; + } + """ + check_unrolled_qasm(dumps(module), expected) + + +def test_runtime_conditional_end_remains_conditional(): + """A runtime-dependent ``end`` does not truncate the surrounding block.""" + module = loads(""" + OPENQASM 3.0; + include "stdgates.inc"; + qubit q; + bit[1] c; + if (c[0]) { + end; + x q; + } + h q; + """) + + module.unroll() + + expected = """ + OPENQASM 3.0; + include "stdgates.inc"; + qubit[1] q; + bit[1] c; + if (c[0] == true) { + end; + } + h q[0]; + """ + check_unrolled_qasm(dumps(module), expected)