fix: walk nested box and if bodies in module transformation passes

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))

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

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

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()

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;

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;