Fix measurement and control key commutativity bug in insertion_sort_transformer. (#7822)

Check measurement and control key commutativity before swapping two
operations.

Fixes #7472

---------

Co-authored-by: Pavol Juhas <juhas@google.com>

Author: codrut3

cirq-core/cirq/transformers/insertion_sort.py

@@ -42,22 +42,41 @@ def insertion_sort_transformer(
         q: idx for idx, q in enumerate(sorted(circuit.all_qubits()))
     }
     cached_qubit_indices: dict[int, list[int]] = {}
+    cached_measurement_keys: dict[int, frozenset[cirq.MeasurementKey]] = {}
+    cached_control_keys: dict[int, frozenset[cirq.MeasurementKey]] = {}
     for pos, op in enumerate(circuit.all_operations()):
         # here `pos` is at the append position of final_operations
-        if (op_qubit_indices := cached_qubit_indices.get(id(op))) is None:
-            op_qubit_indices = cached_qubit_indices[id(op)] = sorted(
+        op_id = id(op)
+        if (op_qubit_indices := cached_qubit_indices.get(op_id)) is None:
+            op_qubit_indices = cached_qubit_indices[op_id] = sorted(
                 qubit_index[q] for q in op.qubits
             )
+        if (op_measurement_keys := cached_measurement_keys.get(op_id)) is None:
+            op_measurement_keys = cached_measurement_keys[op_id] = protocols.measurement_key_objs(
+                op
+            )
+        if (op_control_keys := cached_control_keys.get(op_id)) is None:
+            op_control_keys = cached_control_keys[op_id] = protocols.control_keys(op)
+
         for tail_op in reversed(final_operations):
-            tail_qubit_indices = cached_qubit_indices[id(tail_op)]
-            if op_qubit_indices < tail_qubit_indices and (
-                # special case for zero-qubit gates
-                not op_qubit_indices
-                # check if two sorted sequences are disjoint
-                or op_qubit_indices[-1] < tail_qubit_indices[0]
-                or set(op_qubit_indices).isdisjoint(tail_qubit_indices)
-                # fallback to more expensive commutation check
-                or protocols.commutes(op, tail_op, default=False)
+            tail_id = id(tail_op)
+            tail_qubit_indices = cached_qubit_indices[tail_id]
+            tail_measurement_keys = cached_measurement_keys[tail_id]
+            tail_control_keys = cached_control_keys[tail_id]
+            if (
+                op_qubit_indices < tail_qubit_indices
+                and op_measurement_keys.isdisjoint(tail_measurement_keys)
+                and op_control_keys.isdisjoint(tail_measurement_keys)
+                and tail_control_keys.isdisjoint(op_measurement_keys)
+                and (
+                    # special case for zero-qubit gates
+                    not op_qubit_indices
+                    # check if two sorted sequences are disjoint
+                    or op_qubit_indices[-1] < tail_qubit_indices[0]
+                    or set(op_qubit_indices).isdisjoint(tail_qubit_indices)
+                    # fallback to more expensive commutation check
+                    or protocols.commutes(op, tail_op, default=False)
+                )
             ):
                 pos -= 1
                 continue

cirq-core/cirq/transformers/insertion_sort_test.py

@@ -27,10 +27,55 @@ def test_insertion_sort() -> None:
         cirq.GlobalPhaseGate(1j).on(),
     )
     sorted_circuit = cirq.transformers.insertion_sort_transformer(c)
-    assert sorted_circuit == cirq.Circuit(
-        cirq.GlobalPhaseGate(1j).on(),
-        cirq.CZ(cirq.q(0), cirq.q(1)),
-        cirq.CZ(cirq.q(2), cirq.q(1)),
-        cirq.CZ(cirq.q(2), cirq.q(1)),
-        cirq.CZ(cirq.q(2), cirq.q(4)),
+    cirq.testing.assert_same_circuits(
+        sorted_circuit,
+        cirq.Circuit(
+            cirq.GlobalPhaseGate(1j).on(),
+            cirq.CZ(cirq.q(0), cirq.q(1)),
+            cirq.CZ(cirq.q(2), cirq.q(1)),
+            cirq.CZ(cirq.q(2), cirq.q(1)),
+            cirq.CZ(cirq.q(2), cirq.q(4)),
+        ),
+    )
+
+
+def test_insertion_sort_same_measurement_key() -> None:
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(cirq.measure(q1, key='k'), cirq.measure(q0, key='k'))
+    cirq.testing.assert_same_circuits(cirq.transformers.insertion_sort_transformer(c), c)
+
+
+def test_insertion_sort_measurement_and_control_key_conflict() -> None:
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(cirq.measure(q1, key='k'), cirq.X(q0).with_classical_controls('k'))
+    # Second operation depends on the first so they don't commute.
+    cirq.testing.assert_same_circuits(cirq.transformers.insertion_sort_transformer(c), c)
+
+
+def test_insertion_sort_measurement_and_control_key_conflict_other_way_around() -> None:
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(
+        cirq.measure(q0, key='k'),
+        cirq.X(q1).with_classical_controls('k'),
+        cirq.measure(q0, key='k'),
+    )
+    cirq.testing.assert_same_circuits(cirq.transformers.insertion_sort_transformer(c), c)
+
+
+def test_insertion_sort_distinct_measurement_keys() -> None:
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(cirq.measure(q1, key='k1'), cirq.measure(q0, key='k0'))
+    # Measurement keys are distinct, so the measurements commute.
+    expected = cirq.Circuit(cirq.measure(q0, key='k0'), cirq.measure(q1, key='k1'))
+    assert cirq.transformers.insertion_sort_transformer(c)[0].operations == expected[0].operations
+
+
+def test_insertion_sort_shared_control_key() -> None:
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(
+        cirq.X(q1).with_classical_controls('k'), cirq.X(q0).with_classical_controls('k')
+    )
+    expected = cirq.Circuit(
+        cirq.X(q0).with_classical_controls('k'), cirq.X(q1).with_classical_controls('k')
     )
+    assert cirq.transformers.insertion_sort_transformer(c)[0].operations == expected[0].operations