From 60647d0b86aa9783257038f0a18dac4aab3b3b95 Mon Sep 17 00:00:00 2001
From: Ryan Hill <ryanjh88@gmail.com>
Date: Tue, 4 Aug 2026 12:30:14 -0500
Subject: [PATCH 1/3] fix: walk nested box and if bodies in module
 transformation passes

remove_idle_qubits and reverse_qubit_order only looked at top-level
statements, so operands inside box and if blocks kept stale indices --
silently addressing the wrong qubit, or naming a qubit outside the
shrunken register, which no longer parses.

Both passes now walk nested bodies, as do the measurement and barrier
removal passes. A qubit touched only inside a branch is recorded as used
so it is not mistaken for idle, and operands that are physical qubits are
skipped by the remap instead of tripping an assert.
---
 CHANGELOG.md                        |   1 +
 src/pyqasm/elements.py              |  11 ++-
 src/pyqasm/modules/base.py          | 136 +++++++++++++++++++---------
 src/pyqasm/visitor.py               |  31 +++++--
 tests/qasm3/test_barrier.py         |  34 +++++++
 tests/qasm3/test_measurement.py     |  34 +++++++
 tests/qasm3/test_transformations.py | 118 ++++++++++++++++++++++++
 7 files changed, 314 insertions(+), 51 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fc340e8..6b00ac1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,7 @@ Types of changes:
 ### Removed
 
 ### Fixed
+- 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 `remove_idle_qubits(in_place=False)` updating the qubit count on the wrong module: the original module's `num_qubits` was decremented while the returned copy kept the stale pre-removal count. The copy's AST was already correct; only the counters were swapped. ([#336](https://github.com/qBraid/pyqasm/pull/336))
 - Fixed `remove_idle_qubits()` raising `KeyError` when the unrolled AST contains operand nodes shared across multiple statements (e.g. the `crz` decomposition) and an idle lower-indexed qubit shifts the register indices. `_remap_qubits` now remaps each operand node exactly once instead of once per statement that references it. ([#332](https://github.com/qBraid/pyqasm/pull/332))
 - Fixed `box` duration validation summing `delay` durations across all qubits instead of tracking each qubit's timeline. Delays on disjoint qubits run in parallel, so `box[300ns] { delay[200ns] q[0]; delay[200ns] q[1]; }` was rejected while the identical schedule written as a broadcast delay (`delay[200ns] q;`) was accepted. Delays are now accumulated per qubit and the box is validated against the busiest single timeline; the error message names the offending qubit. Nested boxes now also contribute their declared duration to the enclosing box's timelines (previously the accumulator was reset when an inner box closed, dropping all inner delay accounting). ([#330](https://github.com/qBraid/pyqasm/pull/330))
diff --git a/src/pyqasm/elements.py b/src/pyqasm/elements.py
index b75c740..e0445d0 100644
--- a/src/pyqasm/elements.py
+++ b/src/pyqasm/elements.py
@@ -86,9 +86,18 @@ class QubitDepthNode(DepthNode):
     num_measurements: int = 0
     num_gates: int = 0
     num_barriers: int = 0
+    # Operations applied inside an if/else block. Counted apart from the others because
+    # their depth is only settled once the branch closes, but they still make the qubit used.
+    num_branch_ops: int = 0
 
     def _total_ops(self) -> int:
-        return self.num_resets + self.num_measurements + self.num_gates + self.num_barriers
+        return (
+            self.num_resets
+            + self.num_measurements
+            + self.num_gates
+            + self.num_barriers
+            + self.num_branch_ops
+        )
 
     def is_idle(self) -> bool:
         return self._total_ops() == 0
diff --git a/src/pyqasm/modules/base.py b/src/pyqasm/modules/base.py
index ea2ac84..70a1d14 100644
--- a/src/pyqasm/modules/base.py
+++ b/src/pyqasm/modules/base.py
@@ -22,7 +22,7 @@
 from abc import ABC, abstractmethod
 from collections import Counter
 from copy import deepcopy
-from typing import Optional
+from typing import Iterator, Optional, Sequence, TypeVar
 
 import openqasm3.ast as qasm3_ast
 from openqasm3.ast import BranchingStatement, Program, QuantumGate
@@ -35,6 +35,60 @@
 from pyqasm.maps.decomposition_rules import DECOMPOSITION_RULES
 from pyqasm.visitor import QasmVisitor, ScopeManager
 
+StatementT = TypeVar("StatementT", bound=qasm3_ast.QASMNode)
+
+
+def iter_quantum_statements(
+    statements: Sequence[qasm3_ast.QASMNode],
+) -> Iterator[qasm3_ast.Statement]:
+    """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.
+
+    Args:
+        statements (Sequence[qasm3_ast.QASMNode]): The statements to walk.
+
+    Yields:
+        qasm3_ast.Statement: Each quantum statement, in program order.
+    """
+    for stmt in statements:
+        if isinstance(stmt, qasm3_ast.Box):
+            yield from iter_quantum_statements(stmt.body)
+        elif isinstance(stmt, BranchingStatement):
+            yield from iter_quantum_statements(stmt.if_block)
+            yield from iter_quantum_statements(stmt.else_block)
+        elif isinstance(stmt, QUANTUM_STATEMENTS):
+            yield stmt
+
+
+def drop_statements(statements: list[StatementT], unwanted: type) -> list[StatementT]:
+    """Return the statements without those of the unwanted type, nested ones included.
+
+    Bodies of ``box`` and ``if`` blocks are filtered in place. A box left with an empty
+    body is dropped, since pyqasm rejects a box that holds no statement.
+
+    Args:
+        statements (list[StatementT]): The statements to filter.
+        unwanted (type): The statement type to remove.
+
+    Returns:
+        list[StatementT]: The remaining statements.
+    """
+    kept: list[StatementT] = []
+    for stmt in statements:
+        if isinstance(stmt, qasm3_ast.Box):
+            stmt.body = drop_statements(stmt.body, unwanted)
+            if not stmt.body:
+                continue
+        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, unwanted):
+            continue
+        kept.append(stmt)
+    return kept
+
 
 def track_user_operation(func):
     """Decorator to track user operations on a QasmModule."""
@@ -183,7 +237,7 @@ def has_measurements(self) -> bool:
                 if len(self._unrolled_ast.statements) > 0
                 else self._statements
             )
-            for stmt in stmts_to_check:
+            for stmt in iter_quantum_statements(stmts_to_check):
                 if isinstance(stmt, qasm3_ast.QuantumMeasurementStatement):
                     self._has_measurements = True
                     break
@@ -204,11 +258,7 @@ def remove_measurements(self, in_place: bool = True) -> Optional["QasmModule"]:
             if len(self._unrolled_ast.statements) == 0
             else self._unrolled_ast.statements
         )
-        stmts_without_meas = [
-            stmt
-            for stmt in stmt_list
-            if not isinstance(stmt, qasm3_ast.QuantumMeasurementStatement)
-        ]
+        stmts_without_meas = drop_statements(stmt_list, qasm3_ast.QuantumMeasurementStatement)
         curr_module = self
 
         if not in_place:
@@ -243,7 +293,7 @@ def has_barriers(self) -> bool:
                 if len(self._unrolled_ast.statements) > 0
                 else self._statements
             )
-            for stmt in stmts_to_check:
+            for stmt in iter_quantum_statements(stmts_to_check):
                 if isinstance(stmt, qasm3_ast.QuantumBarrier):
                     self._has_barriers = True
                     break
@@ -264,9 +314,7 @@ def remove_barriers(self, in_place: bool = True) -> Optional["QasmModule"]:
             if len(self._unrolled_ast.statements) == 0
             else self._unrolled_ast.statements
         )
-        stmts_without_barriers = [
-            stmt for stmt in stmt_list if not isinstance(stmt, qasm3_ast.QuantumBarrier)
-        ]
+        stmts_without_barriers = drop_statements(stmt_list, qasm3_ast.QuantumBarrier)
         curr_module = self
         if not in_place:
             curr_module = self.copy()
@@ -380,17 +428,17 @@ def _remap_qubits(self, reg_name: str, size: int, idle_indices: list[int]):
         # gate decompositions can reuse the same operand node across multiple statements,
         # so track visited nodes to avoid remapping a shared node more than once
         visited_node_ids = set()
-        for operation in self._unrolled_ast.statements:
-            if isinstance(operation, QUANTUM_STATEMENTS):
-                bit_list = Qasm3Analyzer.get_op_bit_list(operation)
-                for bit in bit_list:
-                    assert isinstance(bit, qasm3_ast.IndexedIdentifier)
-                    if bit.name.name == reg_name:
-                        index_node = bit.indices[0][0]  # type: ignore[index]
-                        if id(index_node) in visited_node_ids:
-                            continue
-                        visited_node_ids.add(id(index_node))
-                        index_node.value = idx_map[index_node.value]  # type: ignore[union-attr]
+        for operation in iter_quantum_statements(self._unrolled_ast.statements):
+            bit_list = Qasm3Analyzer.get_op_bit_list(operation)
+            for bit in bit_list:
+                if not isinstance(bit, qasm3_ast.IndexedIdentifier):
+                    continue  # physical qubit ("$n"): not part of any register
+                if bit.name.name == reg_name:
+                    index_node = bit.indices[0][0]  # type: ignore[index]
+                    if id(index_node) in visited_node_ids:
+                        continue
+                    visited_node_ids.add(id(index_node))
+                    index_node.value = idx_map[index_node.value]  # type: ignore[union-attr]
 
     def _get_idle_qubit_indices(self) -> dict[str, list[int]]:
         """Get the indices of the idle qubits in the module
@@ -536,29 +584,31 @@ def reverse_qubit_order(self, in_place=True):
         #    the depth maps here
 
         # 2. replace each qubit index in the Quantum Operations with the new index
-        for operation in qasm_module._unrolled_ast.statements:
-            if isinstance(operation, QUANTUM_STATEMENTS):
-                bit_list = Qasm3Analyzer.get_op_bit_list(operation)
-                for bit in bit_list:
-                    curr_reg_name = bit.name.name
-                    curr_reg_idx = bit.indices[0][0].value
-                    new_reg_idx = new_qubit_mappings[curr_reg_name][curr_reg_idx]
-
-                    # make the idx -ve so that this is not touched
-                    # while updating the same index later
-
-                    # idx -> -1 * idx - 1 as we also have to look at index 0
-                    # which will remain 0 if we just multiply by -1
-                    bit.indices[0][0].value = -1 * new_reg_idx - 1
+        for operation in iter_quantum_statements(qasm_module._unrolled_ast.statements):
+            bit_list = Qasm3Analyzer.get_op_bit_list(operation)
+            for bit in bit_list:
+                if not isinstance(bit, qasm3_ast.IndexedIdentifier):
+                    continue  # physical qubit ("$n"): not part of any register
+                curr_reg_name = bit.name.name
+                curr_reg_idx = bit.indices[0][0].value
+                new_reg_idx = new_qubit_mappings[curr_reg_name][curr_reg_idx]
+
+                # make the idx -ve so that this is not touched
+                # while updating the same index later
+
+                # idx -> -1 * idx - 1 as we also have to look at index 0
+                # which will remain 0 if we just multiply by -1
+                bit.indices[0][0].value = -1 * new_reg_idx - 1
 
         # remove the -ve marker
-        for operation in qasm_module._unrolled_ast.statements:
-            if isinstance(operation, QUANTUM_STATEMENTS):
-                bit_list = Qasm3Analyzer.get_op_bit_list(operation)
-                for bit in bit_list:
-                    if bit.indices[0][0].value < 0:
-                        bit.indices[0][0].value += 1
-                        bit.indices[0][0].value *= -1
+        for operation in iter_quantum_statements(qasm_module._unrolled_ast.statements):
+            bit_list = Qasm3Analyzer.get_op_bit_list(operation)
+            for bit in bit_list:
+                if not isinstance(bit, qasm3_ast.IndexedIdentifier):
+                    continue
+                if bit.indices[0][0].value < 0:
+                    bit.indices[0][0].value += 1
+                    bit.indices[0][0].value *= -1
 
         # 3. update the original AST with the unrolled AST
         qasm_module._statements = qasm_module._unrolled_ast.statements
diff --git a/src/pyqasm/visitor.py b/src/pyqasm/visitor.py
index fd4d0dd..c4d3d69 100644
--- a/src/pyqasm/visitor.py
+++ b/src/pyqasm/visitor.py
@@ -628,11 +628,13 @@ def _visit_measurement(  # pylint: disable=too-many-locals,too-many-branches,too
                     )
                 )
                 # if measurement gate is not in branching statement
+                src_name, src_idx = QasmVisitor._get_qubit_name_and_id(src_id)
                 if not self._in_branching_statement:
-                    src_name, src_id = src_id.name.name, src_id.indices[0][0].value  # type: ignore
-                    qubit_node = self._module._qubit_depths[(src_name, src_id)]
+                    qubit_node = self._module._qubit_depths[(src_name, src_idx)]
                     qubit_node.depth += 1
                     qubit_node.num_measurements += 1
+                else:
+                    self._mark_branch_qubit(src_name, src_idx)
         else:
             target_name: str = (
                 target.name if isinstance(target, qasm3_ast.Identifier) else target.name.name
@@ -661,12 +663,12 @@ def _visit_measurement(  # pylint: disable=too-many-locals,too-many-branches,too
                     target=tgt_id if target else None,
                 )
                 # if measurement gate is not in branching statement
+                src_name, src_idx = QasmVisitor._get_qubit_name_and_id(src_id)
                 if not self._in_branching_statement:
-                    src_name, src_id = src_id.name.name, src_id.indices[0][0].value  # type: ignore
                     tgt_name, tgt_id = tgt_id.name.name, tgt_id.indices[0][0].value  # type: ignore
 
                     qubit_node, clbit_node = (
-                        self._module._qubit_depths[(src_name, src_id)],
+                        self._module._qubit_depths[(src_name, src_idx)],
                         self._module._clbit_depths[(tgt_name, tgt_id)],
                     )
                     qubit_node.depth += 1
@@ -682,6 +684,8 @@ def _visit_measurement(  # pylint: disable=too-many-locals,too-many-branches,too
                         self._measurement_set.add(target.name)
                     elif isinstance(target, qasm3_ast.IndexedIdentifier):
                         self._measurement_set.add(target.name.name)
+                else:
+                    self._mark_branch_qubit(src_name, src_idx)
 
                 unrolled_measurements.append(unrolled_measure)
 
@@ -769,7 +773,7 @@ def _visit_reset(self, statement: qasm3_ast.QuantumReset) -> list[qasm3_ast.Quan
                 qubit_node.depth += 1
                 qubit_node.num_resets += 1
             else:
-                self._is_branch_qubits.add((qubit_name, qubit_id))
+                self._mark_branch_qubit(qubit_name, qubit_id)
 
             unrolled_resets.append(unrolled_reset)
 
@@ -1042,6 +1046,19 @@ def _register_physical_qubit(self, name: str) -> int:
             self._module.num_qubits = max(self._module.num_qubits, phys_idx + 1)
         return phys_idx
 
+    def _mark_branch_qubit(self, qubit_name: str, qubit_idx: int) -> None:
+        """Record a qubit operated on inside an if/else block.
+
+        The depth of such a qubit is settled once the branch closes, but it counts as
+        used right away so that it is not mistaken for an idle qubit.
+
+        Args:
+            qubit_name: The register name (or physical qubit identifier).
+            qubit_idx: The index of the qubit within the register.
+        """
+        self._is_branch_qubits.add((qubit_name, qubit_idx))
+        self._module._qubit_depths[(qubit_name, qubit_idx)].num_branch_ops += 1
+
     @staticmethod
     def _get_qubit_name_and_id(
         qubit: qasm3_ast.IndexedIdentifier | qasm3_ast.Identifier,
@@ -1182,7 +1199,7 @@ def _visit_basic_gate_operation(
                     for qubit_subset in unrolled_targets + [ctrls]:
                         for qubit in qubit_subset:
                             qubit_name, qubit_idx = QasmVisitor._get_qubit_name_and_id(qubit)
-                            self._is_branch_qubits.add((qubit_name, qubit_idx))
+                            self._mark_branch_qubit(qubit_name, qubit_idx)
 
         # check for duplicate bits
         for final_gate in result:
@@ -1299,7 +1316,7 @@ def _visit_custom_gate_operation(
                 for qubit_subset in [op_qubits] + [ctrls]:
                     for qubit in qubit_subset:
                         qubit_name, qubit_idx = QasmVisitor._get_qubit_name_and_id(qubit)
-                        self._is_branch_qubits.add((qubit_name, qubit_idx))
+                        self._mark_branch_qubit(qubit_name, qubit_idx)
 
         self._scope_manager.pop_scope()
         self._scope_manager.restore_context()
diff --git a/tests/qasm3/test_barrier.py b/tests/qasm3/test_barrier.py
index 6798777..b405254 100644
--- a/tests/qasm3/test_barrier.py
+++ b/tests/qasm3/test_barrier.py
@@ -109,6 +109,40 @@ def test_remove_barriers():
     check_unrolled_qasm(dumps(module), expected_qasm)
 
 
+def test_remove_barriers_inside_box_and_branch():
+    """Barriers nested in a box or an if block must be found and removed too (see #342)."""
+    qasm_str = """OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    bit c;
+    h q[0];
+    box {
+        barrier q;
+        x q[1];
+    }
+    if (c == 1) {
+        barrier q;
+    }
+    """
+    expected_qasm = """OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    bit[1] c;
+    h q[0];
+    box {
+        x q[1];
+    }
+    if (c[0] == true) {
+    }
+    """
+    module = loads(qasm_str)
+    module.unroll()
+    assert module.has_barriers() is True
+    module.remove_barriers()
+    assert module.has_barriers() is False
+    check_unrolled_qasm(dumps(module), expected_qasm)
+
+
 def test_unroll_barrier():
     qasm_str = """
     OPENQASM 3.0;
diff --git a/tests/qasm3/test_measurement.py b/tests/qasm3/test_measurement.py
index aba814a..cd91f41 100644
--- a/tests/qasm3/test_measurement.py
+++ b/tests/qasm3/test_measurement.py
@@ -133,6 +133,40 @@ def test_remove_measurement():
     check_unrolled_qasm(dumps(module), expected_qasm)
 
 
+def test_remove_measurement_inside_box_and_branch():
+    """Measurements nested in a box or an if block must be found and removed too (see #342)."""
+    qasm3_string = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    bit[2] c;
+    h q[0];
+    box {
+        c[0] = measure q[0];
+    }
+    if (c[0] == 1) {
+        x q[1];
+        c[1] = measure q[1];
+    }
+    """
+    expected_qasm = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    bit[2] c;
+    h q[0];
+    if (c[0] == true) {
+        x q[1];
+    }
+    """
+    module = loads(qasm3_string)
+    module.unroll()
+    assert module.has_measurements() is True
+    module.remove_measurements()
+    # the box held nothing but the measurement, and an empty box is not a valid program
+    check_unrolled_qasm(dumps(module), expected_qasm)
+
+
 def test_init_measure():
     qasm3_string = """
     OPENQASM 3.0;
diff --git a/tests/qasm3/test_transformations.py b/tests/qasm3/test_transformations.py
index 97e2d76..2250d21 100644
--- a/tests/qasm3/test_transformations.py
+++ b/tests/qasm3/test_transformations.py
@@ -239,3 +239,121 @@ def test_populate_idle_qubits_increases_depth_by_one():
     original_depth = module.depth()
     module.populate_idle_qubits()
     assert module.depth() == original_depth + 1
+
+
+def test_remove_idle_qubits_inside_box():
+    """Operands nested in a box must be remapped along with the top-level ones (see #342)."""
+    qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[3] q;
+    bit c;
+    h q[2];
+    box {
+        cx q[2], q[1];
+    }
+    c = measure q[1];
+    """
+    expected_qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    bit[1] c;
+    h q[1];
+    box {
+        cx q[1], q[0];
+    }
+    c[0] = measure q[0];
+    """
+    module = loads(qasm3_str)
+    module.remove_idle_qubits()
+    assert module.num_qubits == 2
+    check_unrolled_qasm(dumps(module), expected_qasm3_str)
+
+
+def test_remove_idle_qubits_keeps_qubits_used_in_a_branch():
+    """A qubit touched only inside an if block is in use, so it must not be removed."""
+    qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[3] q;
+    bit c;
+    h q[2];
+    c = measure q[2];
+    if (c == 1) {
+        x q[1];
+    }
+    """
+    expected_qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    bit[1] c;
+    h q[1];
+    c[0] = measure q[1];
+    if (c[0] == true) {
+        x q[0];
+    }
+    """
+    module = loads(qasm3_str)
+    module.remove_idle_qubits()
+    assert module.num_qubits == 2
+    check_unrolled_qasm(dumps(module), expected_qasm3_str)
+    # the result must still be a valid program
+    loads(dumps(module)).validate()
+
+
+def test_remove_idle_qubits_with_physical_qubits():
+    """Physical qubits belong to no register, so the remap must skip over them."""
+    qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    h $1;
+    h q[0];
+    """
+    expected_qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[1] q;
+    h $1;
+    h q[0];
+    """
+    module = loads(qasm3_str)
+    module.unroll()
+    module.remove_idle_qubits()
+    check_unrolled_qasm(dumps(module), expected_qasm3_str)
+
+
+def test_reverse_qubit_order_inside_box_and_branch():
+    qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[3] q;
+    bit c;
+    h q[2];
+    box {
+        cx q[2], q[1];
+    }
+    c = measure q[0];
+    if (c == 1) {
+        x q[1];
+    }
+    """
+    expected_qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[3] q;
+    bit[1] c;
+    h q[0];
+    box {
+        cx q[0], q[1];
+    }
+    c[0] = measure q[2];
+    if (c[0] == true) {
+        x q[1];
+    }
+    """
+    module = loads(qasm3_str)
+    module.reverse_qubit_order()
+    check_unrolled_qasm(dumps(module), expected_qasm3_str)

From ea621d633d96b9cacc11540cc23233184baf5583 Mon Sep 17 00:00:00 2001
From: Ryan Hill <ryanjh88@gmail.com>
Date: Tue, 4 Aug 2026 16:26:31 -0500
Subject: [PATCH 2/3] review: copy before filtering, flag branch use instead of
 counting it

- remove_measurements/remove_barriers(in_place=False) took the copy after
  filtering, and the nested filter rewrites box and if bodies in place,
  so the original lost its nested measurements and barriers too
- branch operations are now flagged rather than counted: a branch body
  can be visited more than once and a custom gate marks both its own
  operands and those of its expansion, so the count was never meaningful
- reverse-order test used a branch operand that maps to itself, which
  could not tell a walked branch from a skipped one
---
 src/pyqasm/elements.py              | 10 ++++++----
 src/pyqasm/modules/base.py          | 24 +++++++++++-------------
 src/pyqasm/visitor.py               |  2 +-
 tests/qasm3/test_barrier.py         | 19 +++++++++++++++++++
 tests/qasm3/test_measurement.py     | 21 +++++++++++++++++++++
 tests/qasm3/test_transformations.py | 23 +++++++++++++++++++++--
 6 files changed, 79 insertions(+), 20 deletions(-)

diff --git a/src/pyqasm/elements.py b/src/pyqasm/elements.py
index e0445d0..fda5319 100644
--- a/src/pyqasm/elements.py
+++ b/src/pyqasm/elements.py
@@ -86,9 +86,11 @@ class QubitDepthNode(DepthNode):
     num_measurements: int = 0
     num_gates: int = 0
     num_barriers: int = 0
-    # Operations applied inside an if/else block. Counted apart from the others because
-    # their depth is only settled once the branch closes, but they still make the qubit used.
-    num_branch_ops: int = 0
+    # Set when the qubit is operated on inside an if/else block. Those operations are
+    # flagged rather than counted: their depth is only settled once the branch closes,
+    # and a branch body may be visited more than once, so only the fact that the qubit
+    # is used can be relied on.
+    used_in_branch: bool = False
 
     def _total_ops(self) -> int:
         return (
@@ -96,7 +98,7 @@ def _total_ops(self) -> int:
             + self.num_measurements
             + self.num_gates
             + self.num_barriers
-            + self.num_branch_ops
+            + int(self.used_in_branch)
         )
 
     def is_idle(self) -> bool:
diff --git a/src/pyqasm/modules/base.py b/src/pyqasm/modules/base.py
index 70a1d14..7f0b5b6 100644
--- a/src/pyqasm/modules/base.py
+++ b/src/pyqasm/modules/base.py
@@ -253,16 +253,15 @@ def remove_measurements(self, in_place: bool = True) -> Optional["QasmModule"]:
         Returns:
             QasmModule: The module with the measurements removed if in_place is False
         """
+        # copy first: the filtering rewrites nested box and if bodies in place, so it
+        # has to run on the module that is being returned
+        curr_module = self if in_place else self.copy()
         stmt_list = (
-            self._statements
-            if len(self._unrolled_ast.statements) == 0
-            else self._unrolled_ast.statements
+            curr_module._statements
+            if len(curr_module._unrolled_ast.statements) == 0
+            else curr_module._unrolled_ast.statements
         )
         stmts_without_meas = drop_statements(stmt_list, qasm3_ast.QuantumMeasurementStatement)
-        curr_module = self
-
-        if not in_place:
-            curr_module = self.copy()
 
         for qubit in curr_module._qubit_depths.values():
             qubit.num_measurements = 0
@@ -309,15 +308,14 @@ def remove_barriers(self, in_place: bool = True) -> Optional["QasmModule"]:
         Returns:
             QasmModule: The module with the barriers removed if in_place is False
         """
+        # copy first, as in remove_measurements: nested bodies are filtered in place
+        curr_module = self if in_place else self.copy()
         stmt_list = (
-            self._statements
-            if len(self._unrolled_ast.statements) == 0
-            else self._unrolled_ast.statements
+            curr_module._statements
+            if len(curr_module._unrolled_ast.statements) == 0
+            else curr_module._unrolled_ast.statements
         )
         stmts_without_barriers = drop_statements(stmt_list, qasm3_ast.QuantumBarrier)
-        curr_module = self
-        if not in_place:
-            curr_module = self.copy()
 
         for qubit in curr_module._qubit_depths.values():
             qubit.num_barriers = 0
diff --git a/src/pyqasm/visitor.py b/src/pyqasm/visitor.py
index c4d3d69..9edbaa5 100644
--- a/src/pyqasm/visitor.py
+++ b/src/pyqasm/visitor.py
@@ -1057,7 +1057,7 @@ def _mark_branch_qubit(self, qubit_name: str, qubit_idx: int) -> None:
             qubit_idx: The index of the qubit within the register.
         """
         self._is_branch_qubits.add((qubit_name, qubit_idx))
-        self._module._qubit_depths[(qubit_name, qubit_idx)].num_branch_ops += 1
+        self._module._qubit_depths[(qubit_name, qubit_idx)].used_in_branch = True
 
     @staticmethod
     def _get_qubit_name_and_id(
diff --git a/tests/qasm3/test_barrier.py b/tests/qasm3/test_barrier.py
index b405254..293e76f 100644
--- a/tests/qasm3/test_barrier.py
+++ b/tests/qasm3/test_barrier.py
@@ -143,6 +143,25 @@ def test_remove_barriers_inside_box_and_branch():
     check_unrolled_qasm(dumps(module), expected_qasm)
 
 
+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;
+    include "stdgates.inc";
+    qubit[2] q;
+    h q[0];
+    box {
+        barrier q;
+        x q[1];
+    }
+    """
+    module = loads(qasm_str)
+    module.unroll()
+    new_module = module.remove_barriers(in_place=False)
+
+    assert "barrier" not in dumps(new_module)
+    assert "barrier" in dumps(module)
+
+
 def test_unroll_barrier():
     qasm_str = """
     OPENQASM 3.0;
diff --git a/tests/qasm3/test_measurement.py b/tests/qasm3/test_measurement.py
index cd91f41..6a4f514 100644
--- a/tests/qasm3/test_measurement.py
+++ b/tests/qasm3/test_measurement.py
@@ -167,6 +167,27 @@ def test_remove_measurement_inside_box_and_branch():
     check_unrolled_qasm(dumps(module), expected_qasm)
 
 
+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 = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    bit[2] c;
+    h q[0];
+    box {
+        c[0] = measure q[0];
+        x q[1];
+    }
+    """
+    module = loads(qasm3_string)
+    module.unroll()
+    new_module = module.remove_measurements(in_place=False)
+
+    assert "measure" not in dumps(new_module)
+    assert "measure" in dumps(module)
+
+
 def test_init_measure():
     qasm3_string = """
     OPENQASM 3.0;
diff --git a/tests/qasm3/test_transformations.py b/tests/qasm3/test_transformations.py
index 2250d21..cb633ca 100644
--- a/tests/qasm3/test_transformations.py
+++ b/tests/qasm3/test_transformations.py
@@ -337,7 +337,7 @@ def test_reverse_qubit_order_inside_box_and_branch():
     }
     c = measure q[0];
     if (c == 1) {
-        x q[1];
+        x q[0];
     }
     """
     expected_qasm3_str = """
@@ -351,9 +351,28 @@ def test_reverse_qubit_order_inside_box_and_branch():
     }
     c[0] = measure q[2];
     if (c[0] == true) {
-        x q[1];
+        x q[2];
     }
     """
     module = loads(qasm3_str)
     module.reverse_qubit_order()
     check_unrolled_qasm(dumps(module), expected_qasm3_str)
+
+
+def test_remove_idle_qubits_keeps_qubits_used_by_a_custom_gate_in_a_branch():
+    """The expansion of a custom gate in a branch marks its qubits as used."""
+    qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    bit c;
+    gate my_gate p, r { x p; cx p, r; }
+    h q[0];
+    c = measure q[0];
+    if (c == 1) {
+        my_gate q[0], q[1];
+    }
+    """
+    module = loads(qasm3_str)
+    module.remove_idle_qubits()
+    assert module.num_qubits == 2

From 8b1f2ac9dc1d006a8fff0d52e3175db3a528d037 Mon Sep 17 00:00:00 2001
From: Ryan Hill <ryanjh88@gmail.com>
Date: Wed, 5 Aug 2026 08:59:40 -0500
Subject: [PATCH 3/3] review: keep the branch flag out of the operation count

_total_ops() means 'operations on this qubit' again -- byte-identical to main --
and the branch flag is expressed in is_idle(), the only thing it was introduced
to affect. printer.py reads _total_ops() for max_depth, so the term did not
belong in the arithmetic (L1).
---
 src/pyqasm/elements.py | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/src/pyqasm/elements.py b/src/pyqasm/elements.py
index fda5319..b7dd303 100644
--- a/src/pyqasm/elements.py
+++ b/src/pyqasm/elements.py
@@ -93,16 +93,10 @@ class QubitDepthNode(DepthNode):
     used_in_branch: bool = False
 
     def _total_ops(self) -> int:
-        return (
-            self.num_resets
-            + self.num_measurements
-            + self.num_gates
-            + self.num_barriers
-            + int(self.used_in_branch)
-        )
+        return self.num_resets + self.num_measurements + self.num_gates + self.num_barriers
 
     def is_idle(self) -> bool:
-        return self._total_ops() == 0
+        return self._total_ops() == 0 and not self.used_in_branch
 
 
 @dataclass(slots=True)