fix: walk loop and switch bodies in module statement passes

[ryanhill1]

Fixes #354

has_measurements() / remove_measurements() / has_barriers() / remove_barriers() fall back to self._statements when the module has not been unrolled, and the statement walker did not descend into for / while / switch bodies — so an occurrence inside one was invisible and removal was a no-op.

iter_quantum_statements and drop_statements (the walker #345 introduced for box/if) now descend into loop blocks and switch cases/default.

Two consequences, both raised in review:

• The pre-unroll answer is an over-approximation. A measurement or barrier inside a zero-iteration loop, an unmatched switch case or an if (false) block is now reported even though the unrolled program never runs it. That is the safe direction, and it is documented on iter_quantum_statements, has_measurements and has_barriers.
• drop_statements no longer mutates its input. It previously filtered nested bodies in place, so an in-place removal reached through the shared AST — _statements is _original_program.statements — and stripped nested statements out of the module's own original_program while leaving the top level intact. Containers whose bodies are filtered are now shallow-copied, which also closes the same problem for box and if from fix: walk nested box and if bodies in module transformation passes #345.

Note that loops and switches can reach _unrolled_ast: a pre-unroll remove_barriers() assigns the filtered raw statements straight to it. That is not exploitable today, because every operand-rewriting caller runs unroll() first, but the invariant is now recorded on the walker rather than left implicit.

[argus-eye]

Argus review

Auto-review is off for this repo. Tick the box below to run a review on this PR.

• Trigger Argus review

Estimated cost

• Files changed: 4
• Diff lines (±): 93
• Historical avg: ~318.9k tokens · ~$1.35 · across last 6 review(s)

Tip: you can also comment @argus-eye review at any time.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: ce5a0528-3e2d-4719-b3c3-0ba955c91f06

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codecov-commenter]

Codecov Report

❌ Patch coverage is 90.90909% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/pyqasm/modules/base.py	90.90%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[ryanhill1]

@Argus-Eye review

[argus-eye]

🔎 Argus · 10/10 — Correctly extends statement passes into loop and switch bodies

🔍 PR intent vs diff (LLM analysis)

Argus read the diff against the stated intent. This is not an execution log — reviewer still needs to test behavior.

Goal: Make module statement passes detect and remove measurements and barriers inside for, while, and switch bodies before unrolling.
Not in scope:

• The unrolled path is unchanged; loops and switches are fully expanded by unroll() and do not reach _unrolled_ast with their bodies intact.
  Stated acceptance criteria (from PR/issue — not independently verified):
• iter_quantum_statements descends into for and while loop blocks and switch cases/default.
• drop_statements descends into for and while loop blocks and switch cases/default.
• has_measurements(), remove_measurements(), has_barriers(), and remove_barriers() handle occurrences inside non-unrolled loop and switch bodies.

✅ Intent delivered

Verdict: This PR adds the intended traversal and removal support for measurements and barriers in non-unrolled loop and switch bodies. It is ready to merge.

💡 2 P2 · 4 files reviewed

Architecture: The recursive traversal keeps module-level statement passes consistent with nested AST structure; broader nested-control-flow coverage would further protect this behavior.

2 findings · 2 inline · 0 folded

_{🔢 118.9k tokens · $0.2737 total}

Stage	Tokens	Cost
Intent	2.6k	$0.0000
Triage	2.8k	$0.0000
Lead agent	1.5k	$0.0000
Review · bug_hunter	20.9k	$0.0674
Review · security	21.0k	$0.0665
Review · architecture	20.8k	$0.0707
Review · regression	21.1k	$0.0692
Review	24.9k	$0.0000
Acceptance	1.2k	$0.0000
Scoring	1.0k	$0.0000
Synthesis	932	$0.0000

---

<sub>Contract: production/full · checked: bug_hunter, security, architecture, regression · review took 1m32s

Dashboard → · React 👎 to dismiss · Reply to any inline comment or use @argus-eye help to chat</sub>

[TheGupta2012]

Verdict

Approve with comments. The fix is correct, minimal, and lands exactly where issue #354 said it hurt. Every behaviour the PR claims was reproduced on this branch and confirmed absent on origin/main. Nothing here blocks merge. Four follow-up items are raised inline; three are Low and one is a Medium worth a short docstring or a small change before this pattern spreads further.

One correction to the PR description, and it is the most useful thing in this review: the description states "the unrolled path is unchanged: loops and switches are fully expanded by unroll() and never reach _unrolled_ast." That is not quite true. remove_measurements() / remove_barriers() end by assigning curr_module._unrolled_ast.statements = <filtered raw statements>, so calling either on a non-unrolled module writes loop nodes straight into _unrolled_ast. Measured directly after remove_barriers() on a module holding a for loop:

_unrolled_ast statement types: ['Include', 'QubitDeclaration', 'ForInLoop']

That matters because iter_quantum_statements has seven call sites, not four. Besides the four has_* / remove_* methods, it is also used by the operand-rewriting passes _remap_qubits (base.py:446) and reverse_qubit_order (base.py:602 and base.py:619), both of which read _unrolled_ast.statements. Those passes now descend into loop and switch bodies too. This was traced end-to-end and found safe today — both passes call unroll() first, which rebuilds _unrolled_ast from _statements with loops expanded. reverse_qubit_order(), remove_idle_qubits(), populate_idle_qubits() and depth() were each exercised after a pre-unroll remove_barriers() and none regressed. The residual hazard is noted inline.

Findings checklist

#	Type	Severity	Item
1	Maintenance	Medium	In-place removal leaves original_program half-stripped
2	Implementation	Low	Emptied loop/switch bodies are kept, while an emptied box is dropped
3	Implementation	Low	Broadened walker now reaches loop bodies from the operand-rewriting passes
4	Maintenance	Low	Docstring should record the new pre-unroll over-approximation

Security: no findings. The change adds pure AST traversal, introduces no I/O, no external input handling, no new dependency, and no eval/exec surface. Recursion depth follows program nesting depth, which the parser already bounds.

Minor polish, no action strictly required: the new switch test in tests/qasm3/test_measurement.py places the measurement in case 1 only, so the stmt.default is not None branch is walked but never actually removes anything. A second assertion with the measurement in default would close that gap; the equivalent was verified manually and passes.

What worked — verified, not merely read

Every row below was run on the PR worktree at b58c87c9 and again on a read-only origin/main worktree, with openqasm3 1.0.1.

AST shape confirmed against openqasm3 1.0.1 (a wrong attribute name would only raise on the untested path):

• ForInLoop.block, WhileLoop.block — present.
• SwitchStatement.cases — list[tuple[list[Expression], CompoundStatement]], so the for _, case_block in stmt.cases / case_block.statements unpacking is right.
• SwitchStatement.default — CompoundStatement | None, so the is not None guard is right.
• QUANTUM_STATEMENTS is (QuantumGate, QuantumBarrier, QuantumReset, QuantumMeasurementStatement). None of those is also a container type, so the new elif branches cannot shadow the final elif isinstance(stmt, QUANTUM_STATEMENTS). Ordering is safe.

The #354 repro, all four methods, base vs PR. Target placed inside a for, a while, a switch case, and a switch default:

Placement	has_* on main	has_* on PR	remove_* on main	remove_* on PR
for body	False	True	no-op	removed
while body	False	True	no-op	removed
switch case	False	True	no-op	removed
switch default	False	True	no-op	removed

The old behaviour was worse than a plain no-op: remove_measurements() set _has_measurements = False unconditionally, so has_measurements() reported False while dumps() still emitted measure. That silent lie is gone.

Nesting combinations — all reached, box handling from #345 still intact: for inside if; if inside a switch case; while inside a box; for inside for. In each case detection flipped False -> True and removal became effective. (The while-inside-box sample fails unroll() with Global variable 'i' must be a constant to use it in a local scope — reproduced identically on origin/main, unrelated to this PR.)

Round-trip after removal. Every filtered program was re-serialized with dumps() and re-parsed with loads(); all succeeded, including the degenerate cases where a body was emptied. See finding 2 for the cosmetic residue.

Pre-unroll / post-unroll consistency. For all eight sample programs, has_measurements() before unroll() and after unroll() now agree, where they disagreed on origin/main. The one remaining class of disagreement is a safe over-approximation and is covered by finding 4.

Mutation semantics. remove_barriers(in_place=False) was confirmed to leave a caller-held Program completely untouched — the copy()-first ordering introduced by #345 holds for the new loop and switch branches too. remove_barriers() followed by unroll() produces the correct expanded program and does not double-apply.

How it was tested

• Detached worktree at PR head b58c87c9, read-only baseline worktree at origin/main (7c31308), same interpreter and same openqasm3 1.0.1 for both.
• Purpose-built probe scripts, run against both worktrees, covering: the four-method x four-placement matrix, eight nesting combinations, empty-body round-tripping, caller-AST mutation for both in_place values, _unrolled_ast contents after a pre-unroll removal, and the downstream passes reverse_qubit_order(), remove_idle_qubits(), populate_idle_qubits(), depth(), num_qubits.
• Full suite on the PR worktree: 717 passed, 4 skipped, 2 failed. Both failures are tests/cli/test_cli_commands.py::test_validate_qasm_with_invalid_file and ::test_validate_command_with_invalid_file; both fail identically on the origin/main baseline worktree, so they are environmental and not attributable to this PR.
• black --check src tests -> clean, 93 files unchanged.
• gh pr checks 360 -> Changelog pass, CodeRabbit skipped.

Could not be verified

• pylint and isort are not installed in the available environment, so the repo's full lint gate was not reproduced locally. black was.
• No test workflow appears in gh pr checks 360 — only the changelog check ran on CI, so CI has not independently exercised the new tests.
• Qiskit interop was not exercised (qiskit not installed); no qiskit-facing surface is touched by this diff.

Next steps

1. Rebase on main. The PR is currently CONFLICTING, but only in CHANGELOG.md. There is no code conflict — src/pyqasm/modules/base.py merges cleanly. A rebase and a re-stated changelog entry clears it.
2. Consider finding 1 before merge, or file it as a follow-up covering box and if as well, since the root cause predates this PR.
3. Findings 2, 3 and 4 are safe to defer.

Nice, tightly scoped fix — the walker was already the right abstraction, and extending it beat adding a parallel traversal.

[TheGupta2012]

Type: Maintenance
Severity: Low

Rationale: Walking loop and switch bodies makes the pre-unroll answer conservative rather than exact, and the docstrings elsewhere still claim the unrolled AST is simply "a better indicator". The pre-unroll check now reports a measurement that the unrolled program provably never executes. Measured on this branch:

Program	pre-unroll has_measurements()	post-unroll
measure in a non-taken switch case (const int i = 5; case 1 { measure })	True	False
measure in a zero-iteration loop (for int i in [0:-1])	True	False
measure in if (false) { ... }	True	False

The if (false) row predates this PR (#345); the loop and switch rows are new. The direction is the safe one — a false positive is far better than the false negative that #354 reported — but a caller that branches on has_measurements() to decide whether to append its own measurements will now take the wrong branch on these programs, and nothing in the code says so.

Change Requested: Extend this docstring, and ideally the has_measurements / has_barriers docstrings, to state that the pre-unroll answer is an over-approximation: it reports occurrences that are syntactically present even when the branch or iteration never executes. One sentence is enough; no behaviour change is wanted here.

[ryanhill1]

Reproduced your table exactly — pre=True, post=False for the zero-iteration loop, the unmatched switch case and if (false). Documented in 6e6cf28 on iter_quantum_statements and on both has_measurements and has_barriers, with the note that a caller needing an exact answer should unroll first. No behaviour change, as you asked.

[TheGupta2012]

Type: Implementation
Severity: Low

Rationale: This helper has seven call sites, not four. Beyond the has_* / remove_* methods it also backs _remap_qubits (base.py:446) and reverse_qubit_order (base.py:602, base.py:619), which rewrite qubit operands in _unrolled_ast.statements. Those passes now descend into loop and switch bodies as well.

The PR description says loops "never reach _unrolled_ast". They can. remove_barriers() / remove_measurements() end with curr_module._unrolled_ast.statements = <filtered raw statements>, so on a non-unrolled module the raw loop nodes land in _unrolled_ast. Verified directly after remove_barriers() on a module holding a for loop:

_unrolled_ast statement types: ['Include', 'QubitDeclaration', 'ForInLoop']

This was traced end-to-end and is not exploitable today: _remap_qubits and reverse_qubit_order both call unroll() first, which rebuilds _unrolled_ast from _statements with loops expanded. reverse_qubit_order(), remove_idle_qubits(), populate_idle_qubits() and depth() were each run after a pre-unroll remove_barriers() on a loop-bearing program and all behaved correctly.

The hazard is latent rather than live. If any future operand pass reads _unrolled_ast without unrolling first, it will now reach a q[i] inside a for body, where the index node is an Identifier and bit.indices[0][0].value raises AttributeError. Before this PR the loop was invisible, so that pass silently skipped it.

Change Requested: No code change is required. Please correct the "never reach _unrolled_ast" claim in the PR description, and consider a short comment here noting that the operand-rewriting callers rely on unroll() having expanded loops first — that invariant is currently implicit and easy to break.

[ryanhill1]

You are right on both counts and the PR description was wrong — corrected now. Confirmed ['Include', 'QubitDeclaration', 'ForInLoop'] in _unrolled_ast after a pre-unroll remove_barriers().

Added the invariant as a comment on iter_quantum_statements in 6e6cf28: the operand-rewriting callers depend on unroll() having expanded loops away, and reading _unrolled_ast without that can reach a q[i] whose index is an Identifier. Agreed it is latent rather than live, but it was implicit and easy to break.

[TheGupta2012]

Type: Implementation
Severity: Low

Rationale: The Box branch a few lines above drops a box whose body was emptied (if not stmt.body: continue), but the loop and switch branches keep an emptied container. Removing the only statement from a loop or a case leaves dead scaffolding in the output:

// after remove_barriers() on: for int i in [0:1] { barrier q; }
for int i in [0:1] {
}

// after remove_barriers() on: switch (i) { case 1 { barrier q; } default { x q; } }
switch (i) {
  case 1 {
  }
  default {
    x q;
  }
}

This was checked carefully and is cosmetic, not a correctness bug: both forms re-serialize with dumps(), re-parse with loads(), and unroll() cleanly to nothing. Severity stays Low for that reason.

Change Requested: Consider mirroring the Box treatment so an emptied ForInLoop / WhileLoop / case body is dropped, which would also let the docstring above stop singling out box. Deliberately no suggestion block: dropping an emptied while changes runtime behaviour (an emptied non-terminating while becomes a no-op rather than a hang), so the call belongs to the author. Leaving this as-is and noting it in the docstring is an equally acceptable resolution.

[ryanhill1]

Taking the second option — leaving the behaviour and documenting it. Your while point is the deciding one for me: the box drop exists because pyqasm rejects an empty box, which is a forcing reason, and there is no equivalent for loops or cases. Dropping an emptied while would quietly convert a non-terminating loop into a no-op, and I would rather not make that call as a side effect of removing barriers. An emptied for could be dropped safely, but splitting the two would be a rule nobody would remember.

Documented on drop_statements in 6e6cf28 so the asymmetry with box is explained rather than looking like an oversight. Happy to revisit if the dead scaffolding turns out to bother anyone downstream.

[TheGupta2012]

Type: Maintenance
Severity: Medium

Rationale: These branches filter nested bodies in place, so an in-place removal now mutates the module's own _original_program, which stays reachable through the public original_program property. The top level is rebuilt into a fresh list, but the nested bodies are edited on the shared nodes, so the exposed "original" ends up being neither the original nor the filtered result. Reproduced on this branch:

m = loads('OPENQASM 3.0; include "stdgates.inc"; qubit[2] q; barrier q; for int i in [0:1] { h q[i]; barrier q; }')
m.remove_barriers()

# module.original_program  -> keeps `barrier q;` at top level, but the barrier
#                             INSIDE the for body is gone
# dumps(module)            -> both barriers gone (correct)

The design comes from #345 and already applies to box and if; this PR widens its reach to loops and switches, which is where it becomes easy to hit. To be precise about the blast radius: __str__ is not affected after a removal, because _unrolled_ast.statements is non-empty by then and __str__ takes the unrolled branch. The exposure is the original_program property, plus any caller that constructed the module from a Program it still holds.

The in_place=False path is sound — copy() runs before filtering, and a caller-held Program was confirmed byte-identical afterwards. Only in_place=True (the default) is affected.

Change Requested: Either deep-copy the nested bodies before filtering them so _original_program stays pristine, or document on remove_measurements / remove_barriers that in_place=True also rewrites nested bodies of the original program and that original_program is not a faithful pre-removal snapshot afterwards. Since the root cause predates this PR, splitting this into a follow-up that covers box and if at the same time is a perfectly reasonable answer.

[ryanhill1]

Confirmed, and fixed properly rather than documented — 6e6cf28.

Your repro is exact: _statements is _original_program.statements (same list object, base.py:147), so the top level survives because drop_statements returns a fresh list, while the nested bodies were edited on shared nodes.

drop_statements now shallow-copies any container whose body it filters, so it no longer mutates its input at all. Kept statements are still shared, only the containers are copied. Verified original_program comes back byte-identical for for, box, if and switch while the removal still lands in the output, and pinned that as a parametrized test.

Since it works at the walker rather than the call site, it closes the box/if case from #345 at the same time — so no follow-up needed. Only two callers outside the recursion (remove_measurements, remove_barriers), both removal paths, so the blast radius was small enough to do here.