Fixing BlockCollapser with Clbits (#9823)

* Bug fix: collapsing blocks with clbits

Co-authored-by: chriseclectic <cjwood@us.ibm.com>

* release notes

* similar fix to DAGDependency

* additional fixes to handle classical bits in conditions

* fixing lint issue

* Fixing handling of conditions over full registers, following review comments

---------

Co-authored-by: chriseclectic <cjwood@us.ibm.com>

qiskit/dagcircuit/collect_blocks.py

@@ -14,7 +14,8 @@
 """Various ways to divide a DAG into blocks of nodes, to split blocks of nodes
 into smaller sub-blocks, and to consolidate blocks."""
 
-from qiskit.circuit import QuantumCircuit, CircuitInstruction
+from qiskit.circuit import QuantumCircuit, CircuitInstruction, ClassicalRegister
+from qiskit.circuit.controlflow.condition import condition_bits
 from . import DAGOpNode, DAGCircuit, DAGDependency
 from .exceptions import DAGCircuitError
 
@@ -254,22 +255,46 @@ def collapse_to_operation(self, blocks, collapse_fn):
         then uses collapse_fn to collapse this circuit into a single operation.
         """
         global_index_map = {wire: idx for idx, wire in enumerate(self.dag.qubits)}
+        global_index_map.update({wire: idx for idx, wire in enumerate(self.dag.clbits)})
+
         for block in blocks:
-            # Find the set of qubits used in this block (which might be much smaller than
-            # the set of all qubits).
+            # Find the sets of qubits/clbits used in this block (which might be much smaller
+            # than the set of all qubits/clbits).
             cur_qubits = set()
+            cur_clbits = set()
+
+            # Additionally, find the set of classical registers used in conditions over full registers
+            # (in such a case, we need to add that register to the block circuit, not just its clbits).
+            cur_clregs = []
+
             for node in block:
                 cur_qubits.update(node.qargs)
-
-            # For reproducibility, order these qubits compatibly with the global order.
+                cur_clbits.update(node.cargs)
+                cond = getattr(node.op, "condition", None)
+                if cond is not None:
+                    cur_clbits.update(condition_bits(cond))
+                    if isinstance(cond[0], ClassicalRegister):
+                        cur_clregs.append(cond[0])
+
+            # For reproducibility, order these qubits/clbits compatibly with the global order.
             sorted_qubits = sorted(cur_qubits, key=lambda x: global_index_map[x])
+            sorted_clbits = sorted(cur_clbits, key=lambda x: global_index_map[x])
+
+            qc = QuantumCircuit(sorted_qubits, sorted_clbits)
+
+            # Add classical registers used in conditions over registers
+            for reg in cur_clregs:
+                qc.add_register(reg)
 
             # Construct a quantum circuit from the nodes in the block, remapping the qubits.
             wire_pos_map = {qb: ix for ix, qb in enumerate(sorted_qubits)}
-            qc = QuantumCircuit(len(cur_qubits))
+            wire_pos_map.update({qb: ix for ix, qb in enumerate(sorted_clbits)})
+
             for node in block:
-                remapped_qubits = [wire_pos_map[qarg] for qarg in node.qargs]
-                qc.append(CircuitInstruction(node.op, remapped_qubits, node.cargs))
+                instructions = qc.append(CircuitInstruction(node.op, node.qargs, node.cargs))
+                cond = getattr(node.op, "condition", None)
+                if cond is not None:
+                    instructions.c_if(*cond)
 
             # Collapse this quantum circuit into an operation.
             op = collapse_fn(qc)

qiskit/dagcircuit/dagcircuit.py

@@ -30,6 +30,7 @@
 import rustworkx as rx
 
 from qiskit.circuit import ControlFlowOp, ForLoopOp, IfElseOp, WhileLoopOp, SwitchCaseOp
+from qiskit.circuit.controlflow.condition import condition_bits
 from qiskit.circuit.exceptions import CircuitError
 from qiskit.circuit.quantumregister import QuantumRegister, Qubit
 from qiskit.circuit.classicalregister import ClassicalRegister, Clbit
@@ -1129,8 +1130,10 @@ def replace_block_with_op(self, node_block, op, wire_pos_map, cycle_check=True):
 
         for nd in node_block:
             block_qargs |= set(nd.qargs)
-            if isinstance(nd, DAGOpNode) and getattr(nd.op, "condition", None):
-                block_cargs |= set(nd.cargs)
+            block_cargs |= set(nd.cargs)
+            cond = getattr(nd.op, "condition", None)
+            if cond is not None:
+                block_cargs.update(condition_bits(cond))
 
         # Create replacement node
         new_node = DAGOpNode(

qiskit/dagcircuit/dagdependency.py

@@ -19,6 +19,7 @@
 
 import rustworkx as rx
 
+from qiskit.circuit.controlflow.condition import condition_bits
 from qiskit.circuit.quantumregister import QuantumRegister, Qubit
 from qiskit.circuit.classicalregister import ClassicalRegister, Clbit
 from qiskit.dagcircuit.exceptions import DAGDependencyError
@@ -394,11 +395,8 @@ def _create_op_node(self, operation, qargs, cargs):
                 #   (1) cindices_list are specific to template optimization and should not be computed
                 #       in this place.
                 #   (2) Template optimization pass needs currently does not handle general conditions.
-                if isinstance(operation.condition[0], Clbit):
-                    condition_bits = [operation.condition[0]]
-                else:
-                    condition_bits = operation.condition[0]
-                cindices_list = [self.clbits.index(clbit) for clbit in condition_bits]
+                cond_bits = condition_bits(operation.condition)
+                cindices_list = [self.clbits.index(clbit) for clbit in cond_bits]
             else:
                 cindices_list = []
         else:
@@ -591,8 +589,10 @@ def replace_block_with_op(self, node_block, op, wire_pos_map, cycle_check=True):
 
         for nd in node_block:
             block_qargs |= set(nd.qargs)
-            if nd.op.condition:
-                block_cargs |= set(nd.cargs)
+            block_cargs |= set(nd.cargs)
+            cond = getattr(nd.op, "condition", None)
+            if cond is not None:
+                block_cargs.update(condition_bits(cond))
 
         # Create replacement node
         new_node = self._create_op_node(

releasenotes/notes/fix-collapse-with-clbits-e14766353303d442.yaml

@@ -0,0 +1,9 @@
+---
+fixes:
+  - |
+    Fixed a bug in :class:`~BlockCollapser` where classical bits were ignored when collapsing
+    a block of nodes.
+  - |
+    Fixed a bug in :meth:`~qiskit.dagcircuit.DAGCircuit.replace_block_with_op` and
+    :meth:`~qiskit.dagcircuit.DAGDependency.replace_block_with_op`
+    that led to ignoring classical bits.