From ca97bf1ac059411abda3a0c942fa9695457af2aa Mon Sep 17 00:00:00 2001
From: Ryan Hill <ryanjh88@gmail.com>
Date: Fri, 7 Aug 2026 08:14:07 -0500
Subject: [PATCH 1/2] fix: walk loop and switch bodies in module statement
 passes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

has_measurements(), remove_measurements(), has_barriers() and
remove_barriers() fall back to _statements when the module has not been
unrolled, and that list can still hold for/while/switch bodies the
walker did not descend into — so an occurrence inside a loop or switch
was invisible and removal was a no-op. iter_quantum_statements and
drop_statements now descend into loop blocks and switch cases.

Fixes #354
---
 CHANGELOG.md                    |  1 +
 src/pyqasm/modules/base.py      | 18 ++++++++++-
 tests/qasm3/test_barrier.py     | 17 ++++++++++
 tests/qasm3/test_measurement.py | 57 +++++++++++++++++++++++++++++++++
 4 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bcc32f0..abb00ab 100644
--- a/CHANGELOG.md
+++ b/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))
diff --git a/src/pyqasm/modules/base.py b/src/pyqasm/modules/base.py
index e87e242..499b5aa 100644
--- a/src/pyqasm/modules/base.py
+++ b/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)
diff --git a/tests/qasm3/test_barrier.py b/tests/qasm3/test_barrier.py
index 293e76f..3f51db6 100644
--- a/tests/qasm3/test_barrier.py
+++ b/tests/qasm3/test_barrier.py
@@ -143,6 +143,23 @@ 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()
+    assert "barrier" not in dumps(module)
+
+
 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;
diff --git a/tests/qasm3/test_measurement.py b/tests/qasm3/test_measurement.py
index 6a4f514..5ab445a 100644
--- a/tests/qasm3/test_measurement.py
+++ b/tests/qasm3/test_measurement.py
@@ -167,6 +167,63 @@ 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
+
+
+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 = """

From b58c87c910a03d7b7537e23e740ae5de6601756d Mon Sep 17 00:00:00 2001
From: Ryan Hill <ryanjh88@gmail.com>
Date: Fri, 7 Aug 2026 11:23:06 -0500
Subject: [PATCH 2/2] test: strengthen loop-body removal tests (Argus P2s)

- while-loop measurement removal now asserted, not just detection
- barrier removal test asserts the loop's gates survive the pass
---
 tests/qasm3/test_barrier.py     | 5 ++++-
 tests/qasm3/test_measurement.py | 6 ++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/tests/qasm3/test_barrier.py b/tests/qasm3/test_barrier.py
index 3f51db6..9f33eef 100644
--- a/tests/qasm3/test_barrier.py
+++ b/tests/qasm3/test_barrier.py
@@ -157,7 +157,10 @@ def test_has_and_remove_barriers_inside_loop_before_unroll():
     module.remove_barriers()
     assert module.has_barriers() is False
     module.unroll()
-    assert "barrier" not in dumps(module)
+    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():
diff --git a/tests/qasm3/test_measurement.py b/tests/qasm3/test_measurement.py
index 5ab445a..50829fc 100644
--- a/tests/qasm3/test_measurement.py
+++ b/tests/qasm3/test_measurement.py
@@ -196,6 +196,12 @@ def test_has_and_remove_measurements_inside_loop_before_unroll():
     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)."""