BugFix in cirq.map_operations (#6155)

* BugFix in cirq.map_operations * Remove stale comment
quantumlib · Jun 22, 2023 · 3b699a4 · 3b699a4
1 parent b337659
commit 3b699a4
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/cirq-core/cirq/transformers/transformer_primitives.py b/cirq-core/cirq/transformers/transformer_primitives.py
@@ -158,7 +158,8 @@ def apply_map(op: ops.Operation, idx: int) -> ops.OP_TREE:
 
     return map_moments(
         circuit,
-        lambda m, i: [circuits.Moment(apply_map(op, i) for op in m.operations)],
+        lambda m, i: circuits.Circuit(apply_map(op, i) for op in m.operations).moments
+        or [circuits.Moment()],
         deep=deep,
         tags_to_ignore=tags_to_ignore,
     )

diff --git a/cirq-core/cirq/transformers/transformer_primitives_test.py b/cirq-core/cirq/transformers/transformer_primitives_test.py
@@ -414,6 +414,17 @@ def test_map_operations_can_add_qubits_if_flag_false():
     cirq.testing.assert_same_circuits(c_mapped, cirq.Circuit(cirq.CNOT(q[0], q[1])))
 
 
+def test_map_operations_maps_different_ops_from_same_moment_to_shared_qubits():
+    q = cirq.LineQubit.range(3)
+    c = cirq.Circuit(cirq.H.on_each(q[:2]))
+    c_mapped = cirq.map_operations(
+        c, lambda op, _: op.controlled_by(q[2]), raise_if_add_qubits=False
+    )
+    cirq.testing.assert_same_circuits(
+        c_mapped, cirq.Circuit(cirq.H(q[0]).controlled_by(q[2]), cirq.H(q[1]).controlled_by(q[2]))
+    )
+
+
 def test_map_operations_can_drop_operations():
     q = cirq.LineQubit.range(2)
     c = cirq.Circuit(cirq.X(q[0]), cirq.Y(q[1]), cirq.X(q[1]), cirq.Y(q[0]))