Add double-line diagrams for control keys (#4627)

Adds double-line wires to diagrams between measurement gates and their associated control gates.

* Cbits are represented similarly as qubits with a double-line axis.
* The measurement key name is the label of the cbit (and the key label from `M` is removed).
* `@` is the symbol for writing to the cbit.
  * Currently we only allow write-once, but in the future we may want to allow overwrite, xor, bit-and, sum, etc.
* `^` is the symbol for reading from / executing conditionally upon a cbit.
* `X` below is hard-coded in the test case, but will be the conditional operation applied.
* Measurement keys that are not used in control operations will not have a line associated, to reduce clutter.

For example a circuit
```
cirq.Circuit(
    cirq.measure(q0, key='a')
    ControlOp(qubits=[q1], keys=['a'])
)
```
would have the following representation.

```
0: ───M───────
      ║
1: ───╫───X───
      ║   ║
a: ═══@═══^═══
```

See `circuit_test` for more examples

Parts 6 of https://tinyurl.com/cirq-feedforward.

cirq-core/cirq/circuits/_box_drawing_character_data.py

@@ -282,6 +282,62 @@ def parts_with(val: int) -> List[str]:
 )
 
 
+NORMAL_THEN_DOUBLED_MIXED_BOX_CHARS = MixedBoxDrawCharacterSet(
+    first_char_set=NORMAL_BOX_CHARS,
+    second_char_set=DOUBLED_BOX_CHARS,
+    top_then_bottom=' ',
+    top_then_left='╛',
+    top_then_right='╘',
+    top_then_bottom_left=' ',
+    top_then_bottom_right=' ',
+    top_then_left_right='╧',
+    top_then_bottom_left_right=' ',
+    bottom_then_top=' ',
+    bottom_then_left='╕',
+    bottom_then_right='╒',
+    bottom_then_top_left=' ',
+    bottom_then_top_right=' ',
+    bottom_then_left_right='╤',
+    bottom_then_top_left_right=' ',
+    left_then_top='╜',
+    left_then_bottom='╖',
+    left_then_right=' ',
+    left_then_top_bottom='╢',
+    left_then_bottom_right=' ',
+    left_then_top_right=' ',
+    left_then_top_bottom_right=' ',
+    right_then_top='╙',
+    right_then_bottom='╓',
+    right_then_left=' ',
+    right_then_top_bottom='╟',
+    right_then_top_left=' ',
+    right_then_bottom_left=' ',
+    right_then_top_bottom_left=' ',
+    top_bottom_then_left='╡',
+    top_bottom_then_right='╞',
+    top_bottom_then_left_right='╪',
+    top_left_then_bottom=' ',
+    top_left_then_right=' ',
+    top_left_then_bottom_right=' ',
+    top_right_then_bottom=' ',
+    top_right_then_left=' ',
+    top_right_then_bottom_left=' ',
+    bottom_left_then_top=' ',
+    bottom_left_then_right=' ',
+    bottom_left_then_top_right=' ',
+    bottom_right_then_top=' ',
+    bottom_right_then_left=' ',
+    bottom_right_then_top_left=' ',
+    left_right_then_top='╨',
+    left_right_then_bottom='╥',
+    left_right_then_top_bottom='╫',
+    top_bottom_left_then_right=' ',
+    top_bottom_right_then_left=' ',
+    top_left_right_then_bottom=' ',
+    bottom_left_right_then_top=' ',
+)
+
+
 def box_draw_character(
     first: Optional[BoxDrawCharacterSet],
     second: BoxDrawCharacterSet,
@@ -327,6 +383,11 @@ def box_draw_character(
     if first is BOLD_BOX_CHARS and second is NORMAL_BOX_CHARS:
         combo = NORMAL_THEN_BOLD_MIXED_BOX_CHARS
         sign = -1
+    if first is NORMAL_BOX_CHARS and second is DOUBLED_BOX_CHARS:
+        combo = NORMAL_THEN_DOUBLED_MIXED_BOX_CHARS
+    if first is DOUBLED_BOX_CHARS and second is NORMAL_BOX_CHARS:
+        combo = NORMAL_THEN_DOUBLED_MIXED_BOX_CHARS
+        sign = -1
 
     if combo is None:
         choice = second if +1 in [top, bottom, left, right] else first

cirq-core/cirq/circuits/_box_drawing_character_data_test.py

@@ -17,6 +17,7 @@
     NORMAL_BOX_CHARS,
     NORMAL_THEN_BOLD_MIXED_BOX_CHARS,
     BOLD_BOX_CHARS,
+    DOUBLED_BOX_CHARS,
 )
 
 
@@ -31,3 +32,5 @@ def test_chars():
     assert box_draw_character(None, NORMAL_BOX_CHARS) is None
     assert box_draw_character(NORMAL_BOX_CHARS, BOLD_BOX_CHARS, top=-1, bottom=+1) == '╽'
     assert box_draw_character(BOLD_BOX_CHARS, NORMAL_BOX_CHARS, top=-1, bottom=+1) == '╿'
+    assert box_draw_character(DOUBLED_BOX_CHARS, NORMAL_BOX_CHARS, left=-1, bottom=+1) == '╕'
+    assert box_draw_character(NORMAL_BOX_CHARS, DOUBLED_BOX_CHARS, left=-1, bottom=+1) == '╖'

cirq-core/cirq/circuits/circuit.py

@@ -1150,7 +1150,15 @@ def to_text_diagram_drawer(
         Returns:
             The TextDiagramDrawer instance.
         """
-        qubits = ops.QubitOrder.as_qubit_order(qubit_order).order_for(self.all_qubits())
+        qubits: Tuple[Any, ...] = ops.QubitOrder.as_qubit_order(qubit_order).order_for(
+            self.all_qubits()
+        )
+        cbits = tuple(
+            sorted(
+                (key for op in self.all_operations() for key in protocols.control_keys(op)), key=str
+            )
+        )
+        qubits = qubits + cbits
         qubit_map = {qubits[i]: i for i in range(len(qubits))}
 
         if qubit_namer is None:
@@ -1181,7 +1189,7 @@ def to_text_diagram_drawer(
 
         w = diagram.width()
         for i in qubit_map.values():
-            diagram.horizontal_line(i, 0, w)
+            diagram.horizontal_line(i, 0, w, doubled=not isinstance(qubits[i], ops.Qid))
 
         if moment_groups and draw_moment_groups:
             _draw_moment_groups_in_diagram(moment_groups, use_unicode_characters, diagram)
@@ -2389,7 +2397,12 @@ def _draw_moment_in_diagram(
 
     max_x = x0
     for op in non_global_ops:
-        indices = [qubit_map[q] for q in op.qubits]
+        qubits: Tuple[Any, ...] = tuple(op.qubits)
+        cbits = tuple(
+            (protocols.measurement_key_objs(op) | protocols.control_keys(op)) & qubit_map.keys()
+        )
+        qubits += cbits
+        indices = [qubit_map[q] for q in qubits]
         y1 = min(indices)
         y2 = max(indices)
 
@@ -2411,16 +2424,14 @@ def _draw_moment_in_diagram(
 
         # Draw vertical line linking the gate's qubits.
         if y2 > y1 and info.connected:
-            out_diagram.vertical_line(x, y1, y2)
+            out_diagram.vertical_line(x, y1, y2, doubled=len(cbits) != 0)
 
         # Print gate qubit labels.
         symbols = info._wire_symbols_including_formatted_exponent(
             args,
-            preferred_exponent_index=max(
-                range(len(op.qubits)), key=lambda i: qubit_map[op.qubits[i]]
-            ),
+            preferred_exponent_index=max(range(len(qubits)), key=lambda i: qubit_map[qubits[i]]),
         )
-        for s, q in zip(symbols, op.qubits):
+        for s, q in zip(symbols, qubits):
             out_diagram.write(x, qubit_map[q], s)
 
         if x > max_x:

cirq-core/cirq/circuits/circuit_test.py

@@ -84,19 +84,26 @@ def validate_moment(self, moment):
 
 
 class ControlOp(cirq.Operation):
-    def __init__(self, keys):
-        self._keys = keys
+    def __init__(self, keys, qubits=None):
+        self._keys = [cirq.MeasurementKey(k) if isinstance(k, str) else k for k in keys]
+        self._qubits = qubits or []
 
     def with_qubits(self, *new_qids):
         pass  # coverage: ignore
 
     @property
     def qubits(self):
-        return []  # coverage: ignore
+        return self._qubits
 
     def _control_keys_(self):
         return self._keys
 
+    def _circuit_diagram_info_(
+        self, args: 'cirq.CircuitDiagramInfoArgs'
+    ) -> 'cirq.CircuitDiagramInfo':
+        symbols = ['X'] * len(self._qubits) + ['^'] * len(self._keys)
+        return cirq.CircuitDiagramInfo(symbols)
+
 
 def test_alignment():
     assert repr(cirq.Alignment.LEFT) == 'cirq.Alignment.LEFT'
@@ -241,19 +248,191 @@ def test_append_single():
 
 def test_append_control_key():
     q = cirq.LineQubit(0)
-
     c = cirq.Circuit()
     c.append(cirq.measure(q, key='a'))
-    c.append(ControlOp([cirq.MeasurementKey('a')]))
+    c.append(ControlOp(['a']))
     assert len(c) == 2
 
     c = cirq.Circuit()
     c.append(cirq.measure(q, key='a'))
-    c.append(ControlOp([cirq.MeasurementKey('b')]))
-    c.append(ControlOp([cirq.MeasurementKey('b')]))
+    c.append(ControlOp(['b']))
+    c.append(ControlOp(['b']))
     assert len(c) == 1
 
 
+def test_control_key_diagram():
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(cirq.measure(q0, key='a'), ControlOp(qubits=[q1], keys=['a']))
+
+    cirq.testing.assert_has_diagram(
+        c,
+        """
+0: ───M───────
+      ║
+1: ───╫───X───
+      ║   ║
+a: ═══@═══^═══
+""",
+        use_unicode_characters=True,
+    )
+
+
+def test_control_key_diagram_pauli():
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(
+        cirq.measure_single_paulistring(cirq.X(q0), key='a'), ControlOp(qubits=[q1], keys=['a'])
+    )
+
+    cirq.testing.assert_has_diagram(
+        c,
+        """
+0: ───M(X)───────
+      ║
+1: ───╫──────X───
+      ║      ║
+a: ═══@══════^═══
+""",
+        use_unicode_characters=True,
+    )
+
+
+def test_control_key_diagram_extra_measurements():
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(
+        cirq.measure(q0, key='a'), cirq.measure(q0, key='b'), ControlOp(qubits=[q1], keys=['a'])
+    )
+
+    cirq.testing.assert_has_diagram(
+        c,
+        """
+0: ───M───M('b')───
+      ║
+1: ───╫───X────────
+      ║   ║
+a: ═══@═══^════════
+""",
+        use_unicode_characters=True,
+    )
+
+
+def test_control_key_diagram_extra_controlled_bits():
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(cirq.measure(q0, key='a'), ControlOp(qubits=[q0, q1], keys=['a']))
+
+    cirq.testing.assert_has_diagram(
+        c,
+        """
+0: ───M───X───
+      ║   ║
+1: ───╫───X───
+      ║   ║
+a: ═══@═══^═══
+""",
+        use_unicode_characters=True,
+    )
+
+
+def test_control_key_diagram_extra_control_bits():
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.measure(q0, key='b'),
+        ControlOp(qubits=[q1], keys=['a', 'b']),
+    )
+
+    cirq.testing.assert_has_diagram(
+        c,
+        """
+0: ───M───M───────
+      ║   ║
+1: ───╫───╫───X───
+      ║   ║   ║
+a: ═══@═══╬═══^═══
+          ║   ║
+b: ═══════@═══^═══
+""",
+        use_unicode_characters=True,
+    )
+
+
+def test_control_key_diagram_multiple_ops_single_moment():
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.measure(q1, key='b'),
+        ControlOp(qubits=[q0], keys=['a']),
+        ControlOp(qubits=[q1], keys=['b']),
+    )
+
+    cirq.testing.assert_has_diagram(
+        c,
+        """
+      ┌──┐   ┌──┐
+0: ────M──────X─────
+       ║      ║
+1: ────╫M─────╫X────
+       ║║     ║║
+a: ════@╬═════^╬════
+        ║      ║
+b: ═════@══════^════
+      └──┘   └──┘
+""",
+        use_unicode_characters=True,
+    )
+
+
+def test_control_key_diagram_subcircuit():
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(
+        cirq.CircuitOperation(
+            cirq.FrozenCircuit(cirq.measure(q0, key='a'), ControlOp(qubits=[q1], keys=['a']))
+        )
+    )
+
+    cirq.testing.assert_has_diagram(
+        c,
+        """
+      Circuit_0xfba37d11898c0e81:
+      [ 0: ───M───────          ]
+0: ───[       ║                 ]───
+      [ 1: ───╫───X───          ]
+      [       ║   ║             ]
+      [ a: ═══@═══^═══          ]
+      │
+1: ───#2────────────────────────────
+""",
+        use_unicode_characters=True,
+    )
+
+
+def test_control_key_diagram_subcircuit_layered():
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.CircuitOperation(
+            cirq.FrozenCircuit(cirq.measure(q0, key='a'), ControlOp(qubits=[q1], keys=['a'])),
+        ),
+        ControlOp(qubits=[q1], keys=['a']),
+    )
+
+    cirq.testing.assert_has_diagram(
+        c,
+        """
+          Circuit_0xa3bc42bd21c25cca:
+          [ 0: ───M───────          ]
+0: ───M───[       ║                 ]───────
+      ║   [ 1: ───╫───X───          ]
+      ║   [       ║   ║             ]
+      ║   [ a: ═══@═══^═══          ]
+      ║   ║
+1: ───╫───#2────────────────────────────X───
+      ║   ║                             ║
+a: ═══@═══╩═════════════════════════════^═══
+""",
+        use_unicode_characters=True,
+    )
+
+
 def test_append_multiple():
     a = cirq.NamedQubit('a')
     b = cirq.NamedQubit('b')