Fix logic for handling repetition_ids in CircuitOperation.replace (#6984)

* Fix logic for handling `repetition_ids` in `CircuitOperation.replace`

Previously, passing `repetition_ids` to `replace` would enable
`use_repetition_ids`, even if the value passed was `None`. This fixes
things so we only enable `use_repetition_ids` if a non-None value is
passed, which is the same as what is done in the `CircuitOperation`
constructor.

* fmt

Author: maffoo

cirq-core/cirq/circuits/circuit_operation.py

@@ -266,7 +266,9 @@ def replace(self, **changes) -> 'cirq.CircuitOperation':
             'repetition_ids': self.repetition_ids,
             'parent_path': self.parent_path,
             'extern_keys': self._extern_keys,
-            'use_repetition_ids': True if 'repetition_ids' in changes else self.use_repetition_ids,
+            'use_repetition_ids': (
+                True if changes.get('repetition_ids') is not None else self.use_repetition_ids
+            ),
             'repeat_until': self.repeat_until,
             **changes,
         }

cirq-core/cirq/circuits/circuit_operation_test.py

@@ -332,6 +332,28 @@ def test_repeat(add_measurements: bool, use_default_ids_for_initial_rep: bool) -
     assert op_base.repeat(2.99999999999).repetitions == 3
 
 
+def test_replace_repetition_ids() -> None:
+    a, b = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(cirq.H(a), cirq.CX(a, b), cirq.M(b, key='mb'), cirq.M(a, key='ma'))
+    op = cirq.CircuitOperation(circuit.freeze())
+    assert op.repetitions == 1
+    assert not op.use_repetition_ids
+
+    op2 = op.replace(repetitions=2)
+    assert op2.repetitions == 2
+    assert not op2.use_repetition_ids
+
+    op3 = op.replace(repetitions=3, repetition_ids=None)
+    assert op3.repetitions == 3
+    assert not op3.use_repetition_ids
+
+    # Passing `repetition_ids` will also enable `use_repetition_ids`
+    op4 = op.replace(repetitions=4, repetition_ids=['a', 'b', 'c', 'd'])
+    assert op4.repetitions == 4
+    assert op4.use_repetition_ids
+    assert op4.repetition_ids == ['a', 'b', 'c', 'd']
+
+
 @pytest.mark.parametrize('add_measurements', [True, False])
 @pytest.mark.parametrize('use_repetition_ids', [True, False])
 @pytest.mark.parametrize('initial_reps', [0, 1, 2, 3])