Fix break/continue mishandling in for/while loops (#386)

CHANGELOG.md

@@ -26,6 +26,7 @@ Types of changes:
 - Fixed `pyqasm validate` wrapping its diagnostics at the console width, which split a file path longer than the width across lines mid-token and left it neither copyable nor clickable. The error console now uses `soft_wrap`, keeping one diagnostic per line.
 - Fixed an indirect cycle between gate definitions exhausting the Python stack: `gate a q { b q; }` with `gate b q { a q; }` raised a bare `RecursionError` naming nothing, while the direct case was already reported cleanly. The guard compared the body's gate name against one name, so it saw only a cycle of length one. It now tests membership of the whole expansion chain, and names the path: `Recursive definitions not allowed for gate 'a' (a -> b -> a)`. A gate reached twice down separate paths is a diamond, not a cycle, and still expands. ([#369](https://github.com/qBraid/pyqasm/issues/369))
 - Fixed a nested external custom gate counting the depth of the decomposition it skipped, the shape the [#352](https://github.com/qBraid/pyqasm/issues/352) fix did not reach: `unroll(external_gates=["outer"])` on a gate whose body calls another custom gate emitted one statement but reported `depth() == 13`. The suppression flag was assigned and cleared without save-restore, so the inner gate clobbered the outer gate's state in both directions. It is now saved and restored, and the depth is recorded once, from the outermost external gate. ([#367](https://github.com/qBraid/pyqasm/issues/367))
+- Fixed `break` and `continue` in loops. A `for` loop had no handler, so the internal `BreakSignal` / `ContinueSignal` escaped `validate()` and `unroll()` to the caller. A `while` loop caught the signal but discarded the statements the interrupted iteration had already emitted, so `while (i < 3) { h q[0]; i += 1; break; }` unrolled to nothing instead of one `h q[0]`. Both followed from `visit_basic_block` losing its accumulated result when a nested statement raised; the signal now carries those statements back to the enclosing loop. `break` / `continue` outside any loop now raise a `ValidationError`. Two `while` iteration-limit bugs are fixed alongside: an iteration cut short by `continue` was not counted, so `while (cond) { continue; }` never tripped the guard and unrolled forever; and the guard fired one iteration early, allowing only `max_loop_iters - 1` iterations to complete where a `for` loop allows the full count. ([#386](https://github.com/qBraid/pyqasm/issues/386))
 
 ### Dependencies
 

src/pyqasm/exceptions.py

@@ -58,29 +58,50 @@ def __init__(self, message: str = "Loop limit exceeded."):
 
 class LoopControlSignal(Exception):
     """Base class for loop control signals like break and continue.
-    This class is used to signal control flow changes within loops during AST traversal."""
 
-    def __init__(self, signal_type: str):
+    Signals a control flow change within a loop during AST traversal.
+    ``partial_result`` carries the statements the interrupted iteration emitted
+    before the signal, so the enclosing loop handler can still emit them.
+    """
+
+    signal_type: str
+    partial_result: list
+
+    def __init__(self, signal_type: str, msg: Optional[str] = None) -> None:
+        """Initialize the signal.
+
+        Args:
+            signal_type: Either ``"break"`` or ``"continue"``.
+            msg: Message for the base ``Exception``. Defaults to ``signal_type``.
+        """
         assert signal_type in ("break", "continue")
         self.signal_type = signal_type
+        self.partial_result = []
+        super().__init__(msg if msg is not None else signal_type)
 
 
 class BreakSignal(LoopControlSignal):
     """Signal to break out of a loop during AST traversal."""
 
-    def __init__(self, msg: Optional[str] = None):
-        if msg is None:
-            msg = "break"
-        super().__init__(msg)
+    def __init__(self, msg: Optional[str] = None) -> None:
+        """Initialize the signal.
+
+        Args:
+            msg: Message for the base ``Exception``. Defaults to ``"break"``.
+        """
+        super().__init__("break", msg)
 
 
 class ContinueSignal(LoopControlSignal):
     """Signal to continue to the next iteration of a loop during AST traversal."""
 
-    def __init__(self, msg: Optional[str] = None):
-        if msg is None:
-            msg = "continue"
-        super().__init__("continue")
+    def __init__(self, msg: Optional[str] = None) -> None:
+        """Initialize the signal.
+
+        Args:
+            msg: Message for the base ``Exception``. Defaults to ``"continue"``.
+        """
+        super().__init__("continue", msg)
 
 
 def raise_qasm3_error(

src/pyqasm/visitor.py

@@ -135,6 +135,7 @@ def __init__(  # pylint: disable=too-many-arguments
         self._unroll_barriers: bool = unroll_barriers
         self._recording_ext_gate_depth = False
         self._in_branching_statement: int = 0
+        self._loop_depth: int = 0
         self._is_branch_qubits: set[tuple[str, int]] = set()
         self._is_branch_clbits: set[tuple[str, int]] = set()
         self._measurement_set: set[str] = set()
@@ -1276,12 +1277,42 @@ def _visit_basic_gate_operation(
         return result
 
     def _visit_break(self, statement: qasm3_ast.BreakStatement) -> None:
+        """Visit a break statement by raising the signal its enclosing loop catches.
+
+        Args:
+            statement (BreakStatement): The break statement to visit.
+
+        Raises:
+            BreakSignal: Always, when inside a loop.
+            ValidationError: If the statement appears outside any loop.
+        """
+        if self._loop_depth <= 0:
+            raise_qasm3_error(
+                "'break' statement outside of loop",
+                error_node=statement,
+                span=statement.span,
+            )
         raise_qasm3_error(
             err_type=BreakSignal,
             error_node=statement,
         )
 
     def _visit_continue(self, statement: qasm3_ast.ContinueStatement) -> None:
+        """Visit a continue statement by raising the signal its enclosing loop catches.
+
+        Args:
+            statement (ContinueStatement): The continue statement to visit.
+
+        Raises:
+            ContinueSignal: Always, when inside a loop.
+            ValidationError: If the statement appears outside any loop.
+        """
+        if self._loop_depth <= 0:
+            raise_qasm3_error(
+                "'continue' statement outside of loop",
+                error_node=statement,
+                span=statement.span,
+            )
         raise_qasm3_error(
             err_type=ContinueSignal,
             error_node=statement,
@@ -2333,6 +2364,27 @@ def _visit_branching_statement(
         self._scope_manager.increment_scope_level()
         self._in_branching_statement += 1
 
+        try:
+            return self._visit_branching_statement_body(statement)
+        except LoopControlSignal:
+            # Undo what this frame pushed, then re-raise for the enclosing loop.
+            # visit_basic_block has already attached the branch body's statements.
+            self._scope_manager.decrement_scope_level()
+            self._scope_manager.pop_scope()
+            self._scope_manager.restore_context()
+            self._in_branching_statement -= 1
+            if not self._in_branching_statement:
+                self._update_branching_gate_depths()
+            raise
+
+    def _visit_branching_statement_body(
+        self, statement: qasm3_ast.BranchingStatement
+    ) -> list[qasm3_ast.Statement]:
+        """Body of :meth:`_visit_branching_statement`.
+
+        Split out so the caller can wrap it in a ``try/except LoopControlSignal``
+        that still pops the scope it pushed when a ``break``/``continue`` raises.
+        """
         result = []
         condition = statement.condition
 
@@ -2499,39 +2551,56 @@ def _visit_forin_loop(self, statement: qasm3_ast.ForInLoop) -> list[qasm3_ast.St
 
         result = []
         statement_block = None
-        for ival in irange:
-            self._scope_manager.push_context(Context.BLOCK)
-            self._scope_manager.push_scope({})
-
-            # Initialize loop variable in loop scope
-            # need to re-declare as we discard the block scope in subsequent
-            # iterations of the loop
-            result.extend(
-                self._visit_classical_declaration(
-                    qasm3_ast.ClassicalDeclaration(statement.type, statement.identifier, init_exp)
+        self._loop_depth += 1
+        try:
+            for ival in irange:
+                self._scope_manager.push_context(Context.BLOCK)
+                self._scope_manager.push_scope({})
+
+                # Initialize loop variable in loop scope
+                # need to re-declare as we discard the block scope in subsequent
+                # iterations of the loop
+                result.extend(
+                    self._visit_classical_declaration(
+                        qasm3_ast.ClassicalDeclaration(
+                            statement.type, statement.identifier, init_exp
+                        )
+                    )
                 )
-            )
-            i = self._scope_manager.get_from_visible_scope(statement.identifier.name)
+                i = self._scope_manager.get_from_visible_scope(statement.identifier.name)
 
-            # Update scope with current value of loop Variable
-            if i is not None:
-                i.value = ival
-                self._scope_manager.update_var_in_scope(i)
+                # Update scope with current value of loop Variable
+                if i is not None:
+                    i.value = ival
+                    self._scope_manager.update_var_in_scope(i)
 
-            if statement_block != statement.block:
-                statement_block = copy.deepcopy(statement.block)
-                result.extend(self.visit_basic_block(statement_block))
-            else:
-                result.extend(self.visit_basic_block(statement.block))
+                try:
+                    if statement_block != statement.block:
+                        statement_block = copy.deepcopy(statement.block)
+                        result.extend(self.visit_basic_block(statement_block))
+                    else:
+                        result.extend(self.visit_basic_block(statement.block))
+                except LoopControlSignal as lcs:
+                    # Keep what this iteration emitted before the break/continue.
+                    result.extend(lcs.partial_result)
+                    self._scope_manager.pop_scope()
+                    self._scope_manager.restore_context()
+                    if self._check_only:
+                        return []
+                    if lcs.signal_type == "break":
+                        break
+                    continue
 
-            # scope not persistent between loop iterations
-            self._scope_manager.pop_scope()
-            self._scope_manager.restore_context()
+                # scope not persistent between loop iterations
+                self._scope_manager.pop_scope()
+                self._scope_manager.restore_context()
 
-            # as we are only checking compile time errors
-            # not runtime errors, we can break here
-            if self._check_only:
-                return []
+                # as we are only checking compile time errors
+                # not runtime errors, we can break here
+                if self._check_only:
+                    return []
+        finally:
+            self._loop_depth -= 1
         return result
 
     def _visit_subroutine_definition(
@@ -2746,35 +2815,46 @@ def _visit_while_loop(self, statement: qasm3_ast.WhileLoop) -> list[qasm3_ast.St
                 span=statement.span,
             )
 
-        while True:
-            cond_value = Qasm3ExprEvaluator.evaluate_expression(statement.while_condition)[0]
-            if not cond_value:
-                break
+        self._loop_depth += 1
+        try:
+            while True:
+                cond_value = Qasm3ExprEvaluator.evaluate_expression(statement.while_condition)[0]
+                if not cond_value:
+                    break
 
-            self._scope_manager.push_context(Context.BLOCK)
-            self._scope_manager.push_scope({})
+                # Checked once the condition is known true, so exactly
+                # `max_iterations` iterations may complete - matching the bound
+                # `_visit_forin_loop` applies to its range length.
+                if loop_counter >= max_iterations:
+                    raise_qasm3_error(
+                        "Loop exceeded max allowed iterations",
+                        err_type=LoopLimitExceededError,
+                        error_node=statement,
+                        span=statement.span,
+                    )
+
+                self._scope_manager.push_context(Context.BLOCK)
+                self._scope_manager.push_scope({})
+
+                signal_type: Optional[str] = None
+                try:
+                    result.extend(self.visit_basic_block(statement.block))
+                except LoopControlSignal as lcs:
+                    # Keep what this iteration emitted before the break/continue.
+                    result.extend(lcs.partial_result)
+                    signal_type = lcs.signal_type
 
-            try:
-                result.extend(self.visit_basic_block(statement.block))
-            except LoopControlSignal as lcs:
                 self._scope_manager.pop_scope()
                 self._scope_manager.restore_context()
-                if lcs.signal_type == "break":
-                    break
-                if lcs.signal_type == "continue":
-                    continue
 
-            self._scope_manager.pop_scope()
-            self._scope_manager.restore_context()
+                if signal_type == "break":
+                    break
 
-            loop_counter += 1
-            if loop_counter >= max_iterations:
-                raise_qasm3_error(
-                    "Loop exceeded max allowed iterations",
-                    err_type=LoopLimitExceededError,
-                    error_node=statement,
-                    span=statement.span,
-                )
+                # Counted on the `continue` path too, or `while (cond) { continue; }`
+                # never reaches the limit.
+                loop_counter += 1
+        finally:
+            self._loop_depth -= 1
 
         return result
 
@@ -2935,15 +3015,36 @@ def _visit_switch_statement(  # type: ignore[return]
         #    each element in the list of the values
         #    should be of const int type and no duplicates should be present
 
-        def _evaluate_case(statements):
+        def _evaluate_case(
+            statements: list[qasm3_ast.Statement],
+        ) -> list[qasm3_ast.Statement]:
+            """Visit one case body in its own scope.
+
+            Args:
+                statements: The statements making up the case body.
+
+            Returns:
+                list[Statement]: The unrolled statements of the case body.
+
+            Raises:
+                LoopControlSignal: Re-raised on a ``break``/``continue`` bound for an
+                    enclosing loop, carrying the statements emitted so far.
+            """
             # can not put 'context' outside
             # BECAUSE the case expression CAN CONTAIN VARS from global scope
             self._scope_manager.push_context(Context.BLOCK)
             self._scope_manager.push_scope({})
             result = []
-            for stmt in statements:
-                Qasm3Validator.validate_statement_type(SWITCH_BLACKLIST_STMTS, stmt, "switch")
-                result.extend(self.visit_statement(stmt))
+            try:
+                for stmt in statements:
+                    Qasm3Validator.validate_statement_type(SWITCH_BLACKLIST_STMTS, stmt, "switch")
+                    result.extend(self.visit_statement(stmt))
+            except LoopControlSignal as lcs:
+                # Carry this case body's statements up to the enclosing loop.
+                lcs.partial_result = result + list(lcs.partial_result)
+                self._scope_manager.pop_scope()
+                self._scope_manager.restore_context()
+                raise
 
             self._scope_manager.pop_scope()
             self._scope_manager.restore_context()
@@ -3526,10 +3627,19 @@ def visit_basic_block(
 
         Returns:
             list[Statement]: The list of unrolled statements.
+
+        Raises:
+            LoopControlSignal: Re-raised on a nested ``break``/``continue``, with
+                this block's statements attached to ``partial_result``.
         """
-        result = []
+        result: list[qasm3_ast.Statement] = []
         for stmt in stmt_list:
-            result.extend(self.visit_statement(stmt))
+            try:
+                result.extend(self.visit_statement(stmt))
+            except LoopControlSignal as lcs:
+                # This block's statements precede those from deeper blocks.
+                lcs.partial_result = result + list(lcs.partial_result)
+                raise
         return result
 
     def finalize(self, unrolled_stmts: list[qasm3_ast.Statement]) -> list[qasm3_ast.Statement]: