Fix exponent location on classically controlled diagrams (#6897)

* Fix exponent location on classically controlled diagrams

The original implementation seems to have assumed the classical conditions were prepended to the wire symbols, not appended.

* Differentiate None and 0

cirq-core/cirq/ops/classically_controlled_operation.py

@@ -169,15 +169,12 @@ def _circuit_diagram_info_(
                 + ', '.join(str(c) for c in self._conditions)
                 + '])',
             ) + 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_label_count
-        elif sub_info.exponent is not None:
-            exponent_qubit_index = control_label_count
+        exp_index = sub_info.exponent_qubit_index
+        if exp_index is None:
+            # None means at bottom, which means the last of the original wire symbols
+            exp_index = len(sub_info.wire_symbols) - 1
         return protocols.CircuitDiagramInfo(
-            wire_symbols=wire_symbols,
-            exponent=sub_info.exponent,
-            exponent_qubit_index=exponent_qubit_index,
+            wire_symbols=wire_symbols, exponent=sub_info.exponent, exponent_qubit_index=exp_index
         )
 
     def _json_dict_(self) -> Dict[str, Any]:

cirq-core/cirq/ops/classically_controlled_operation_test.py

@@ -1033,3 +1033,64 @@ def test_moment_diagram():
   │
     """.strip()
     )
+
+
+def test_diagram_exponents():
+    q0, q1 = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='m'), (cirq.X(q1) ** 0.5).with_classical_controls('m')
+    )
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+0: ───M───────────
+      ║
+1: ───╫───X^0.5───
+      ║   ║
+m: ═══@═══^═══════
+""",
+    )
+
+
+def test_diagram_exponents_cx():
+    q0, q1, q2 = cirq.LineQubit.range(3)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='m'), (cirq.CX(q2, q1) ** 0.5).with_classical_controls('m')
+    )
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+0: ───M───────────
+      ║
+1: ───╫───X^0.5───
+      ║   ║
+2: ───╫───@───────
+      ║   ║
+m: ═══@═══^═══════
+""",
+    )
+
+
+def test_diagram_exponents_multiple_keys():
+    q0, q1, q2 = cirq.LineQubit.range(3)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='m0'),
+        cirq.measure(q1, key='m1'),
+        (cirq.X(q2) ** 0.5).with_classical_controls('m0', 'm1'),
+    )
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+       ┌──┐
+0: ─────M─────────────
+        ║
+1: ─────╫M────────────
+        ║║
+2: ─────╫╫────X^0.5───
+        ║║    ║
+m0: ════@╬════^═══════
+         ║    ║
+m1: ═════@════^═══════
+       └──┘
+""",
+    )