Support OpenQASM end statements

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,41 @@ 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)
+        if isinstance(final_statement, qasm3_ast.BranchingStatement):
+            return (
+                bool(final_statement.else_block)
+                and cls._ends_with_end_statement(final_statement.if_block)
+                and cls._ends_with_end_statement(final_statement.else_block)
+            )
+        return False
+
     def visit_statement(
         self, statement: qasm3_ast.Statement | qasm3_ast.Pragma
     ) -> list[qasm3_ast.Statement]:
@@ -3554,7 +3604,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]:

tests/qasm3/test_end.py

@@ -0,0 +1,211 @@
+# 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)
+
+
+def test_end_in_both_runtime_branches_terminates_program():
+    """Later statements are unreachable when every runtime branch terminates."""
+    module = loads("""
+        OPENQASM 3.0;
+        include "stdgates.inc";
+        bit[1] c;
+        qubit q;
+        if (c[0]) {
+            h q;
+            end;
+        } else {
+            z q;
+            end;
+        }
+        x q;
+        qubit unreachable;
+        """)
+
+    module.unroll()
+    output = dumps(module)
+
+    assert "h q[0];" in output
+    assert "z q[0];" in output
+    assert output.count("end;") == 2
+    assert "x q[0];" not in output
+    assert "unreachable" not in output
+    assert module.num_qubits == 1