Fix issues with ControlFlowOp displays (#10842)

* Fix override bug and testing

* Preliminary changes

* First attempts

* Work on x_index

* Early testing

* Cleanup getattr

* Complete adding exprs

* First fix for cregbundle control flow

* Finish basic cregbundle mods

* Finish all cregbundle except wire_map bug

* Move carg check to circuit_vis and change text tests to circuit_drawer

* Move mpl tests to use circuit_drawer and update images for cregbundle changes

* Change mpl tests to circuit_drawer

* Fix flow_wire_map and tranpile bug

* Final cleanup

* Missing outputs on mpl tests

* Fix output

* Restore tests and move carg check to text and mpl and make circuit

* Prepare final cregbundle fix

* Update ref images

* Fixing refs

* Update refs again

* Final mods

* Incorporate mtreinish changes

* Add reno

* Fix if_else_op test

* Lint

* Fix target == expr bug in text

* Review fixes and outer_circuit, wire_map fixes

* Remove png

* Minor changes

* Fix for_loop measure placement bug

---------

Co-authored-by: Jake Lishman <jake.lishman@ibm.com>
(cherry picked from commit 404df43)

# Conflicts:
#	qiskit/visualization/circuit/text.py
#	test/python/visualization/test_circuit_text_drawer.py
#	test/visual/mpl/circuit/test_circuit_matplotlib_drawer.py

qiskit/visualization/circuit/text.py

@@ -17,7 +17,6 @@
 from warnings import warn
 from shutil import get_terminal_size
 import collections
-import itertools
 import sys
 
 from qiskit.circuit import Qubit, Clbit, ClassicalRegister
@@ -630,19 +629,27 @@ def __init__(
         self.vertical_compression = vertical_compression
         self._wire_map = {}
 
-        for node in itertools.chain.from_iterable(self.nodes):
-            if node.cargs and node.op.name != "measure":
-                if cregbundle:
-                    warn(
-                        "Cregbundle set to False since an instruction needs to refer"
-                        " to individual classical wire",
-                        RuntimeWarning,
-                        2,
-                    )
-                self.cregbundle = False
-                break
-        else:
-            self.cregbundle = True if cregbundle is None else cregbundle
+        def check_clbit_in_inst(circuit, cregbundle):
+            if cregbundle is False:
+                return False
+            for inst in circuit.data:
+                if isinstance(inst.operation, ControlFlowOp):
+                    for block in inst.operation.blocks:
+                        if check_clbit_in_inst(block, cregbundle) is False:
+                            return False
+                elif inst.clbits and not isinstance(inst.operation, Measure):
+                    if cregbundle is not False:
+                        warn(
+                            "Cregbundle set to False since an instruction needs to refer"
+                            " to individual classical wire",
+                            RuntimeWarning,
+                            3,
+                        )
+                    return False
+
+            return True
+
+        self.cregbundle = check_clbit_in_inst(circuit, cregbundle)
 
         if encoding:
             self.encoding = encoding
@@ -1138,8 +1145,18 @@ def build_layers(self):
         layers = [InputWire.fillup_layer(wire_names)]
 
         for node_layer in self.nodes:
+<<<<<<< HEAD
             layer = Layer(self.qubits, self.clbits, self.cregbundle, self._circuit)
 
+=======
+            layer = Layer(
+                self.qubits,
+                self.clbits,
+                self.cregbundle,
+                self._circuit,
+                self._wire_map,
+            )
+>>>>>>> 404df4327 (Fix issues with ControlFlowOp displays (#10842))
             for node in node_layer:
                 layer, current_cons, current_cons_cond, connection_label = self._node_to_gate(
                     node, layer
@@ -1152,15 +1169,175 @@ def build_layers(self):
 
         return layers
 
+<<<<<<< HEAD
+=======
+    def add_control_flow(self, node, layers, wire_map):
+        """Add control flow ops to the circuit drawing."""
+
+        # # Draw a left box such as If, While, For, and Switch
+        flow_layer = self.draw_flow_box(node, wire_map, CF_LEFT)
+        layers.append(flow_layer.full_layer)
+
+        # Get the list of circuits in the ControlFlowOp from the node blocks
+        circuit_list = list(node.op.blocks)
+
+        if isinstance(node.op, SwitchCaseOp):
+            # Create an empty circuit at the head of the circuit_list if a Switch box
+            circuit_list.insert(0, list(node.op.cases_specifier())[0][1].copy_empty_like())
+
+        for circ_num, circuit in enumerate(circuit_list):
+            # Update the wire_map with the qubits and clbits from the inner circuit
+            flow_wire_map = wire_map.copy()
+            flow_wire_map.update(
+                {inner: wire_map[outer] for outer, inner in zip(node.qargs, circuit.qubits)}
+            )
+            for outer, inner in zip(node.cargs, circuit.clbits):
+                if self.cregbundle and (
+                    (in_reg := get_bit_register(self._circuit, inner)) is not None
+                ):
+                    out_reg = get_bit_register(self._circuit, outer)
+                    flow_wire_map.update({in_reg: wire_map[out_reg]})
+                else:
+                    flow_wire_map.update({inner: wire_map[outer]})
+
+            if circ_num > 0:
+                # Draw a middle box such as Else and Case
+                flow_layer = self.draw_flow_box(node, flow_wire_map, CF_MID, circ_num - 1)
+                layers.append(flow_layer.full_layer)
+
+            _, _, nodes = _get_layered_instructions(circuit, wire_map=flow_wire_map)
+            for layer_nodes in nodes:
+                # Limit qubits sent to only ones from main circuit, so qubit_layer is correct length
+                flow_layer2 = Layer(
+                    self.qubits, self.clbits, self.cregbundle, self._circuit, flow_wire_map
+                )
+                for layer_node in layer_nodes:
+                    if isinstance(layer_node.op, ControlFlowOp):
+                        # Recurse on this function if nested ControlFlowOps
+                        self._nest_depth += 1
+                        self.add_control_flow(layer_node, layers, flow_wire_map)
+                        self._nest_depth -= 1
+                    else:
+                        (
+                            flow_layer2,
+                            current_cons,
+                            current_cons_cond,
+                            connection_label,
+                        ) = self._node_to_gate(layer_node, flow_layer2, flow_wire_map)
+                        flow_layer2.connections.append((connection_label, current_cons))
+                        flow_layer2.connections.append((None, current_cons_cond))
+
+                flow_layer2.connect_with("│")
+                layers.append(flow_layer2.full_layer)
+
+        # Draw the right box for End
+        flow_layer = self.draw_flow_box(node, flow_wire_map, CF_RIGHT)
+        layers.append(flow_layer.full_layer)
+
+    def draw_flow_box(self, node, flow_wire_map, section, circ_num=0):
+        """Draw the left, middle, or right of a control flow box"""
+
+        conditional = section == CF_LEFT and not isinstance(node.op, ForLoopOp)
+        depth = str(self._nest_depth)
+        if section == CF_LEFT:
+            if isinstance(node.op, IfElseOp):
+                label = "If-" + depth
+            elif isinstance(node.op, WhileLoopOp):
+                label = "While-" + depth
+            elif isinstance(node.op, ForLoopOp):
+                indexset = node.op.params[0]
+                # If tuple of values instead of range, cut it off at 4 items
+                if "range" not in str(indexset) and len(indexset) > 4:
+                    index_str = str(indexset[:4])
+                    index_str = index_str[:-1] + ", ...)"
+                else:
+                    index_str = str(indexset)
+                label = "For-" + depth + " " + index_str
+            else:
+                label = "Switch-" + depth
+        elif section == CF_MID:
+            if isinstance(node.op, IfElseOp):
+                label = "Else-" + depth
+            else:
+                jump_list = []
+                for jump_values, _ in list(node.op.cases_specifier()):
+                    jump_list.append(jump_values)
+
+                if "default" in str(jump_list[circ_num][0]):
+                    jump_str = "default"
+                else:
+                    jump_str = str(jump_list[circ_num]).replace(",)", ")")
+                label = "Case-" + depth + " " + jump_str
+
+        else:
+            label = "End-" + depth
+
+        flow_layer = Layer(
+            self.qubits,
+            self.clbits,
+            self.cregbundle,
+            self._circuit,
+            flow_wire_map,
+        )
+        # If only 1 qubit, draw basic 1 qubit box
+        if len(node.qargs) == 1:
+            flow_layer.set_qubit(
+                self.qubits[flow_wire_map[node.qargs[0]]],
+                FlowOnQuWire(section, label=label, conditional=conditional),
+            )
+        else:
+            # If multiple qubits, must use wire_map to handle wire_order changes.
+            idx_list = [flow_wire_map[qarg] for qarg in node.qargs]
+            min_idx = min(idx_list)
+            max_idx = max(idx_list)
+            box_height = max_idx - min_idx + 1
+
+            flow_layer.set_qubit(
+                self.qubits[min_idx], FlowOnQuWireTop(section, label=label, wire_label="")
+            )
+            for order, i in enumerate(range(min_idx + 1, max_idx)):
+                flow_layer.set_qubit(
+                    self.qubits[i],
+                    FlowOnQuWireMid(
+                        section,
+                        label=label,
+                        input_length=box_height,
+                        order=order,
+                        wire_label="",
+                    ),
+                )
+            flow_layer.set_qubit(
+                self.qubits[max_idx],
+                FlowOnQuWireBot(
+                    section,
+                    label=label,
+                    input_length=box_height,
+                    conditional=conditional,
+                    wire_label="",
+                ),
+            )
+        if conditional:
+            if isinstance(node.op, SwitchCaseOp):
+                if isinstance(node.op.target, expr.Expr):
+                    condition = node.op.target
+                elif isinstance(node.op.target, Clbit):
+                    condition = (node.op.target, 1)
+                else:
+                    condition = (node.op.target, 2 ** (node.op.target.size) - 1)
+            else:
+                condition = node.op.condition
+            _ = flow_layer.set_cl_multibox(condition, flow_wire_map, top_connect="╨")
+
+        return flow_layer
+
+>>>>>>> 404df4327 (Fix issues with ControlFlowOp displays (#10842))
 
 class Layer:
     """A layer is the "column" of the circuit."""
 
     def __init__(self, qubits, clbits, cregbundle, circuit=None):
         self.qubits = qubits
-        self.clbits_raw = clbits  # list of clbits ignoring cregbundle change below
         self._circuit = circuit
-
         if cregbundle:
             self.clbits = []
             previous_creg = None
@@ -1350,14 +1527,16 @@ def set_cl_multibox(self, condition, top_connect="┴"):
         if isinstance(condition, expr.Expr):
             # If fixing this, please update the docstrings of `QuantumCircuit.draw` and
             # `visualization.circuit_drawer` to remove warnings.
-            label = "<expression>"
+            label = "[expr]"
             out = []
             condition_bits = node_resources(condition).clbits
             registers = collections.defaultdict(list)
             for bit in condition_bits:
                 registers[get_bit_register(self._circuit, bit)].append(bit)
             if registerless := registers.pop(None, ()):
-                out.extend(self.set_cond_bullets(label, ["1"] * len(registerless), registerless))
+                out.extend(
+                    self.set_cond_bullets(label, ["1"] * len(registerless), registerless, wire_map)
+                )
             if self.cregbundle:
                 # It's hard to do something properly sensible here without more major rewrites, so
                 # as a minimum to *not crash* we'll just treat a condition that touches part of a
@@ -1366,7 +1545,7 @@ def set_cl_multibox(self, condition, top_connect="┴"):
                     self.set_clbit(register[0], BoxOnClWire(label=label, top_connect=top_connect))
             else:
                 for register, bits in registers.items():
-                    out.extend(self.set_cond_bullets(label, ["1"] * len(bits), bits))
+                    out.extend(self.set_cond_bullets(label, ["1"] * len(bits), bits, wire_map))
             return out
         label, val_bits = get_condition_label_val(condition, self._circuit, self.cregbundle)
         if isinstance(condition[0], ClassicalRegister):

releasenotes/notes/fix-flow-cregbundle-transpiled-482979d57382b1dc.yaml

@@ -0,0 +1,16 @@
+---
+fixes:
+  - |
+    Fixed an issue with the matplotlib based visualization in the
+    :meth:`.QuantumCircuit.draw` method and the :func:`~.circuit_drawer`
+    function when visualizing circuits that had control flow instructions.
+    Previously in some situations, especially with a layout set, the output
+    visualization could have put gates inside a control flow block on the
+    wrong wires in the visualization.
+    Fixed `#10601 <https://github.com/Qiskit/qiskit-terra/issues/10601>`__
+  - |
+    Fixed an issue with the matplotlib based visualization in the
+    :meth:`.QuantumCircuit.draw` method and the :func:`~.circuit_drawer`
+    function when visualizing circuits that had control flow instructions.
+    Previously when the `cregbundle` option was set to None or True, the
+    drawer did not properly display the `cregs` as bundled.