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

Author: ryanhill1

src/pyqasm/elements.py

@@ -86,17 +86,19 @@ 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 (
             self.num_resets
             + self.num_measurements
             + self.num_gates
             + self.num_barriers
-            + self.num_branch_ops
+            + int(self.used_in_branch)
         )
 
     def is_idle(self) -> bool:

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

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(

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;

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;

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