Fix classically controlled op Moment diagram (#5777)

Fixes #5776 by accounting for the fact that `args.label_map` is null when coming from a Moment renderer, and handling appropriately.

Author: daxfohl

cirq-core/cirq/ops/classically_controlled_operation.py

@@ -149,9 +149,13 @@ def _circuit_diagram_info_(
         sub_info = protocols.circuit_diagram_info(self._sub_operation, sub_args, None)
         if sub_info is None:
             return NotImplemented  # coverage: ignore
-        control_count = len({k for c in self._conditions for k in c.keys})
-        wire_symbols = sub_info.wire_symbols + ('^',) * control_count
-        if any(not isinstance(c, value.KeyCondition) for c in self._conditions):
+        control_label_count = 0
+        if args.label_map is not None:
+            control_label_count = len({k for c in self._conditions for k in c.keys})
+        wire_symbols = sub_info.wire_symbols + ('^',) * control_label_count
+        if control_label_count == 0 or any(
+            not isinstance(c, value.KeyCondition) for c in self._conditions
+        ):
             wire_symbols = (
                 wire_symbols[0]
                 + '(conditions=['
@@ -160,9 +164,9 @@ def _circuit_diagram_info_(
             ) + wire_symbols[1:]
         exponent_qubit_index = None
         if sub_info.exponent_qubit_index is not None:
-            exponent_qubit_index = sub_info.exponent_qubit_index + control_count
+            exponent_qubit_index = sub_info.exponent_qubit_index + control_label_count
         elif sub_info.exponent is not None:
-            exponent_qubit_index = control_count
+            exponent_qubit_index = control_label_count
         return protocols.CircuitDiagramInfo(
             wire_symbols=wire_symbols,
             exponent=sub_info.exponent,

cirq-core/cirq/ops/classically_controlled_operation_test.py

@@ -1017,3 +1017,19 @@ def test_commutes():
     assert not cirq.commutes(
         cirq.X(q0).with_classical_controls('a'), cirq.H(q0).with_classical_controls('a')
     )
+
+
+def test_moment_diagram():
+    a, _, c, d = cirq.GridQubit.rect(2, 2)
+    m = cirq.Moment(cirq.CZ(a, d), cirq.X(c).with_classical_controls('m'))
+    assert (
+        str(m).strip()
+        == """
+  ╷ 0                 1
+╶─┼─────────────────────
+0 │ @─────────────────┐
+  │                   │
+1 │ X(conditions=[m]) @
+  │
+    """.strip()
+    )