fix: walk loop and switch bodies in module statement passes

CHANGELOG.md

@@ -26,6 +26,7 @@ Types of changes:
 ### Removed
 
 ### Fixed
+- Fixed `has_measurements()` / `remove_measurements()` and `has_barriers()` / `remove_barriers()` missing occurrences inside `for` / `while` / `switch` bodies on a module that has not been unrolled — e.g. a measurement inside a `for` loop was invisible and `remove_measurements()` was a no-op. The statement walker now descends into loop and switch bodies. ([#354](https://github.com/qBraid/pyqasm/issues/354))
 - Fixed `remove_idle_qubits()` and `reverse_qubit_order()` ignoring statements nested inside `box` and `if` blocks. Top-level operands were rewritten while nested ones kept their old indices, so the result silently addressed the wrong qubits — and when a nested index fell outside the shrunken register, the output was not a loadable program at all. Both passes now walk nested bodies, as do `has_measurements()` / `remove_measurements()` and `has_barriers()` / `remove_barriers()`; a box left empty by a removal is dropped, since pyqasm rejects a box with no statements. Two consequences of the same blind spot are fixed alongside: a qubit operated on only inside an `if` block no longer counts as idle, and `remove_idle_qubits()` no longer raises `AssertionError` on a program that mixes physical qubits with declared registers. ([#345](https://github.com/qBraid/pyqasm/pull/345))
 - Fixed `unroll(consolidate_qubits=True)` raising `AttributeError: 'str' object has no attribute 'name'` for any gate applied to a physical qubit, e.g. `h $1;`. Consolidation assumed every gate operand was an `IndexedIdentifier`, but a physical qubit survives unrolling as `Identifier("$1")`. Physical qubits are absolute hardware indices belonging to no declared register, so they are now left as written — matching how `measure`, `reset` and `barrier` already treat them. ([#344](https://github.com/qBraid/pyqasm/pull/344))
 - Fixed `unroll()` and `rebase()` emitting statements that share operand AST nodes: gate decompositions passed the same `IndexedIdentifier` objects into every statement they emitted, so transformations that rewrite qubit indices in place mutated a shared node once per referencing statement. This crashed `reverse_qubit_order()` (`KeyError: -1`) and `remove_idle_qubits()` (`KeyError`, [#331](https://github.com/qBraid/pyqasm/issues/331)) on any decomposed gate (e.g. `crz`) whenever the remap was not the identity. Statement constructors in `maps/gates.py` and `Decomposer` now copy their qubit operands so every emitted statement owns its nodes. ([#333](https://github.com/qBraid/pyqasm/issues/333))

src/pyqasm/modules/base.py

@@ -44,7 +44,9 @@ def iter_quantum_statements(
     """Yield the quantum statements in a statement list, nested ones included.
 
     ``box`` and ``if`` blocks survive unrolling with their bodies intact, so a pass that
-    rewrites qubit operands has to reach the statements inside them too.
+    rewrites qubit operands has to reach the statements inside them too. Loop and switch
+    bodies only exist before unrolling, but walking them lets the same pass work on a
+    module that has not been unrolled (issue #354).
 
     Args:
         statements (Sequence[qasm3_ast.QASMNode]): The statements to walk.
@@ -58,6 +60,13 @@ def iter_quantum_statements(
         elif isinstance(stmt, BranchingStatement):
             yield from iter_quantum_statements(stmt.if_block)
             yield from iter_quantum_statements(stmt.else_block)
+        elif isinstance(stmt, (qasm3_ast.ForInLoop, qasm3_ast.WhileLoop)):
+            yield from iter_quantum_statements(stmt.block)
+        elif isinstance(stmt, qasm3_ast.SwitchStatement):
+            for _, case_block in stmt.cases:
+                yield from iter_quantum_statements(case_block.statements)
+            if stmt.default is not None:
+                yield from iter_quantum_statements(stmt.default.statements)
         elif isinstance(stmt, QUANTUM_STATEMENTS):
             yield stmt
 
@@ -84,6 +93,13 @@ def drop_statements(statements: list[StatementT], unwanted: type) -> list[Statem
         elif isinstance(stmt, BranchingStatement):
             stmt.if_block = drop_statements(stmt.if_block, unwanted)
             stmt.else_block = drop_statements(stmt.else_block, unwanted)
+        elif isinstance(stmt, (qasm3_ast.ForInLoop, qasm3_ast.WhileLoop)):
+            stmt.block = drop_statements(stmt.block, unwanted)
+        elif isinstance(stmt, qasm3_ast.SwitchStatement):
+            for _, case_block in stmt.cases:
+                case_block.statements = drop_statements(case_block.statements, unwanted)
+            if stmt.default is not None:
+                stmt.default.statements = drop_statements(stmt.default.statements, unwanted)
         elif isinstance(stmt, unwanted):
             continue
         kept.append(stmt)

tests/qasm3/test_barrier.py

@@ -143,6 +143,26 @@ def test_remove_barriers_inside_box_and_branch():
     check_unrolled_qasm(dumps(module), expected_qasm)
 
 
+def test_has_and_remove_barriers_inside_loop_before_unroll():
+    """Barriers inside a loop or switch must be visible before unroll() (issue #354)."""
+    qasm_str = """OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    for int i in [0:1] { h q[i]; barrier q; }
+    """
+    module = loads(qasm_str)
+    assert module.has_barriers() is True
+
+    module = loads(qasm_str)
+    module.remove_barriers()
+    assert module.has_barriers() is False
+    module.unroll()
+    unrolled_qasm = dumps(module)
+    assert "barrier" not in unrolled_qasm
+    # the loop's gates must survive the removal pass
+    assert unrolled_qasm.count("h q[") == 2
+
+
 def test_remove_barriers_not_in_place_leaves_the_original_alone():
     """Filtering rewrites nested bodies in place, so it must run on the returned copy."""
     qasm_str = """OPENQASM 3.0;

tests/qasm3/test_measurement.py

@@ -167,6 +167,69 @@ def test_remove_measurement_inside_box_and_branch():
     check_unrolled_qasm(dumps(module), expected_qasm)
 
 
+def test_has_and_remove_measurements_inside_loop_before_unroll():
+    """Measurements inside a for/while loop must be visible before unroll() (issue #354)."""
+    qasm3_string = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    bit[2] c;
+    for int i in [0:1] { c[i] = measure q[i]; }
+    """
+    module = loads(qasm3_string)
+    assert module.has_measurements() is True
+
+    module = loads(qasm3_string)
+    module.remove_measurements()
+    assert module.has_measurements() is False
+    module.unroll()
+    assert "measure" not in dumps(module)
+
+    while_string = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[1] q;
+    bit[1] c;
+    int i = 0;
+    while (i < 1) { c[0] = measure q[0]; i += 1; }
+    """
+    module = loads(while_string)
+    assert module.has_measurements() is True
+
+    module = loads(while_string)
+    module.remove_measurements()
+    assert module.has_measurements() is False
+    module.unroll()
+    assert "measure" not in dumps(module)
+
+
+def test_has_and_remove_measurements_inside_switch_before_unroll():
+    """Measurements inside a switch case must be visible before unroll() (issue #354)."""
+    qasm3_string = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    const int i = 1;
+    qubit[1] q;
+    bit[1] c;
+    switch(i) {
+    case 1 {
+        c[0] = measure q[0];
+    }
+    default {
+        x q;
+    }
+    }
+    """
+    module = loads(qasm3_string)
+    assert module.has_measurements() is True
+
+    module = loads(qasm3_string)
+    module.remove_measurements()
+    assert module.has_measurements() is False
+    module.unroll()
+    assert "measure" not in dumps(module)
+
+
 def test_remove_measurement_not_in_place_leaves_the_original_alone():
     """Filtering rewrites nested bodies in place, so it must run on the returned copy."""
     qasm3_string = """