fix: remove_idle_qubits KeyError on shared operand nodes

CHANGELOG.md

@@ -25,6 +25,7 @@ Types of changes:
 ### Removed
 
 ### Fixed
+- 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))
 - Fixed `reset` on a physical qubit rewriting the operand to the internal pulse register, e.g. `reset $2;` unrolled to `reset __PYQASM_QUBITS__[2];`. That names a register the program never declares, so the unrolled output did not round-trip through `dumps()`/`loads()`, and the qubit was never registered (a program whose only operation was `reset $3;` reported `num_qubits == 0`). Physical qubits are now kept as-is in plain QASM programs, matching how gate and measurement operands already treat them; the rename is still applied for OpenPulse programs, where the pulse visitor expects it. ([#325](https://github.com/qBraid/pyqasm/pull/325))
 - Fixed `measure` and `reset` on a user register whose name merely starts with the reserved internal register name (e.g. `__PYQASM_QUBITS__foo`) being mistaken for the internal register itself. Such statements were short-circuited out of unrolling and emitted verbatim, so `c = measure __PYQASM_QUBITS__foo;` was never expanded per qubit and `reset __PYQASM_QUBITS__foo;` silently reset nothing. The register is now matched on its exact name, or on the name followed by an index or slice. ([#325](https://github.com/qBraid/pyqasm/pull/325))

src/pyqasm/modules/base.py

@@ -377,14 +377,20 @@ def _remap_qubits(self, reg_name: str, size: int, idle_indices: list[int]):
             del self._qubit_depths[(reg_name, idx)]
 
         # update the operations that use the qubits
+        # 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:
-                        old_idx = bit.indices[0][0].value  # type: ignore[union-attr,index]
-                        bit.indices[0][0].value = idx_map[old_idx]  # type: ignore[union-attr,index]
+                        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

tests/qasm3/test_transformations.py

@@ -45,6 +45,25 @@ def test_remove_idle_qubits_qasm3_small():
     check_unrolled_qasm(dumps(module), expected_qasm3_str)
 
 
+def test_remove_idle_qubits_qasm3_shared_operand_nodes():
+    """Test remove_idle_qubits when a gate decomposition reuses operand nodes"""
+    qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[3] q;
+    crz(0.5) q[1], q[2];
+    """
+    module = loads(qasm3_str)
+    module.unroll()
+    assert module.num_qubits == 3
+    module.remove_idle_qubits()
+    assert module.num_qubits == 2
+
+    unrolled_qasm = dumps(module)
+    assert "q[2]" not in unrolled_qasm
+    assert "cx q[0], q[1];" in unrolled_qasm
+
+
 def test_remove_idle_qubits_qasm3():
     """Test conversion of qasm3 to compressed contiguous qasm3"""
     qasm3_str = """