Fix CCO related nits in cirq.Operation and cirq.TaggedOperation (quantumlib#5390)

* Improve CCO support in cirq.Operation and cirq.TaggedOperation

* Fix failing tests

* Update docstrings

* Clarify without_classical_controls docstring

* Reword docstrings

Author: tanujkhattar

cirq/ops/classically_controlled_operation.py

@@ -134,7 +134,7 @@ def _resolve_parameters_(
         self, resolver: 'cirq.ParamResolver', recursive: bool
     ) -> 'ClassicallyControlledOperation':
         new_sub_op = protocols.resolve_parameters(self._sub_operation, resolver, recursive)
-        return new_sub_op.with_classical_controls(*self._conditions)
+        return ClassicallyControlledOperation(new_sub_op, self._conditions)
 
     def _circuit_diagram_info_(
         self, args: 'cirq.CircuitDiagramInfoArgs'

cirq/ops/classically_controlled_operation_test.py

@@ -385,7 +385,7 @@ def test_condition_removal():
     op = op.without_classical_controls()
     assert not cirq.control_keys(op)
     assert not op.classical_controls
-    assert set(map(str, op.tags)) == {'t1'}
+    assert not op.tags
 
 
 def test_qubit_mapping():

cirq/ops/raw_types.py

@@ -622,27 +622,33 @@ def classical_controls(self) -> FrozenSet['cirq.Condition']:
 
     def with_classical_controls(
         self, *conditions: Union[str, 'cirq.MeasurementKey', 'cirq.Condition', sympy.Expr]
-    ) -> 'cirq.ClassicallyControlledOperation':
+    ) -> 'cirq.Operation':
         """Returns a classically controlled version of this operation.
 
         An operation that is classically controlled is executed iff all
         conditions evaluate to True. Currently the only condition type is a
         measurement key. A measurement key evaluates to True iff any qubit in
         the corresponding measurement operation evaluated to a non-zero value.
 
-        The classical control will hide any tags on the existing operation,
-        since tags are considered a local attribute.
+        If no conditions are specified, returns self.
+
+        The classical control will remove any tags on the existing operation,
+        since tags are fragile, and we always opt to get rid of the tags when
+        the underlying operation is changed.
 
         Args:
             *conditions: A list of measurement keys, strings that can be parsed
                 into measurement keys, or sympy expressions where the free
                 symbols are measurement key strings.
 
         Returns:
-            A `ClassicallyControlledOperation` wrapping the operation.
+            A `ClassicallyControlledOperation` wrapping the operation. If no conditions
+           are specified, returns self.
         """
         from cirq.ops.classically_controlled_operation import ClassicallyControlledOperation
 
+        if not conditions:
+            return self
         return ClassicallyControlledOperation(self, conditions)
 
     def without_classical_controls(self) -> 'cirq.Operation':
@@ -655,10 +661,10 @@ def without_classical_controls(self) -> 'cirq.Operation':
         If there are no classical controls on the operation, it will return
         `self`.
 
-        Since tags are considered local, this will also remove any tags from
-        the operation (unless there are no classical controls on it). If a
-        `TaggedOperation` is under all the classical control layers, that
-        `TaggedOperation` will be returned from this function.
+        Since tags are fragile, this will also remove any tags from the operation,
+        when called on `TaggedOperation` (unless there are no classical controls on it).
+        If a `TaggedOperation` is under all the classical control layers,
+        that `TaggedOperation` will be returned from this function.
 
         Returns:
             The operation with all classical controls removed.
@@ -712,6 +718,8 @@ def controlled_by(
         *control_qubits: 'cirq.Qid',
         control_values: Optional[Sequence[Union[int, Collection[int]]]] = None,
     ) -> 'cirq.Operation':
+        if len(control_qubits) == 0:
+            return self
         return self.sub_operation.controlled_by(*control_qubits, control_values=control_values)
 
     @property
@@ -864,6 +872,13 @@ def without_classical_controls(self) -> 'cirq.Operation':
         new_sub_operation = self.sub_operation.without_classical_controls()
         return self if new_sub_operation is self.sub_operation else new_sub_operation
 
+    def with_classical_controls(
+        self, *conditions: Union[str, 'cirq.MeasurementKey', 'cirq.Condition', sympy.Expr]
+    ) -> 'cirq.Operation':
+        if not conditions:
+            return self
+        return self.sub_operation.with_classical_controls(*conditions)
+
     def _control_keys_(self) -> AbstractSet['cirq.MeasurementKey']:
         return protocols.control_keys(self.sub_operation)
 

cirq/ops/raw_types_test.py

@@ -627,6 +627,10 @@ def test_tagged_operation_forwards_protocols():
     assert controlled_y.qubits == (q2, q1)
     assert isinstance(controlled_y, cirq.Operation)
     assert not isinstance(controlled_y, cirq.TaggedOperation)
+    classically_controlled_y = tagged_y.with_classical_controls("a")
+    assert classically_controlled_y == y.with_classical_controls("a")
+    assert isinstance(classically_controlled_y, cirq.Operation)
+    assert not isinstance(classically_controlled_y, cirq.TaggedOperation)
 
     clifford_x = cirq.SingleQubitCliffordGate.X(q1)
     tagged_x = cirq.SingleQubitCliffordGate.X(q1).with_tags(tag)
@@ -932,3 +936,12 @@ def __iter__(self):
             raise NotImplementedError()
 
     assert cirq.H.on_each(QidIter())[0] == cirq.H.on(QidIter())
+
+
+@pytest.mark.parametrize(
+    'op', [cirq.X(cirq.NamedQubit("q")), cirq.X(cirq.NamedQubit("q")).with_tags("tagged_op")]
+)
+def test_with_methods_return_self_on_empty_conditions(op):
+    assert op is op.with_tags(*[])
+    assert op is op.with_classical_controls(*[])
+    assert op is op.controlled_by(*[])