Add Expr support to DAGCircuit.replace_block_with_node

qiskit/dagcircuit/dagcircuit.py

@@ -1012,9 +1012,12 @@ def replace_block_with_op(self, node_block, op, wire_pos_map, cycle_check=True):
                 node block to be replaced
             op (qiskit.circuit.Operation): The operation to replace the
                 block with
-            wire_pos_map (Dict[Qubit, int]): The dictionary mapping the qarg to
-                the position. This is necessary to reconstruct the qarg order
-                over multiple gates in the combined single op node.
+            wire_pos_map (Dict[Bit, int]): The dictionary mapping the bits to their positions in the
+                output ``qargs`` or ``cargs``. This is necessary to reconstruct the arg order over
+                multiple gates in the combined single op node.  If a :class:`.Bit` is not in the
+                dictionary, it will not be added to the args; this can be useful when dealing with
+                control-flow operations that have inherent bits in their ``condition`` or ``target``
+                fields.
             cycle_check (bool): When set to True this method will check that
                 replacing the provided ``node_block`` with a single node
                 would introduce a cycle (which would invalidate the
@@ -1044,16 +1047,21 @@ 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)
             block_cargs |= set(nd.cargs)
-            cond = getattr(nd.op, "condition", None)
-            if cond is not None:
-                block_cargs.update(condition_resources(cond).clbits)
-
-        # Create replacement node
-        new_node = DAGOpNode(
-            op,
-            sorted(block_qargs, key=lambda x: wire_pos_map[x]),
-            sorted(block_cargs, key=lambda x: wire_pos_map[x]),
-        )
+            if (condition := getattr(nd.op, "condition", None)) is not None:
+                block_cargs.update(condition_resources(condition).clbits)
+            elif isinstance(nd.op, SwitchCaseOp):
+                if isinstance(nd.op.target, Clbit):
+                    block_cargs.add(nd.op.target)
+                elif isinstance(nd.op.target, ClassicalRegister):
+                    block_cargs.update(nd.op.target)
+                else:
+                    block_cargs.update(node_resources(nd.op.target).clbits)
+
+        block_qargs = [bit for bit in block_qargs if bit in wire_pos_map]
+        block_qargs.sort(key=wire_pos_map.get)
+        block_cargs = [bit for bit in block_cargs if bit in wire_pos_map]
+        block_cargs.sort(key=wire_pos_map.get)
+        new_node = DAGOpNode(op, block_qargs, block_cargs)
 
         try:
             new_node._node_id = self._multi_graph.contract_nodes(

test/python/dagcircuit/test_dagcircuit.py

@@ -2197,6 +2197,43 @@ def test_single_node_block(self):
         self.assertEqual(expected_dag.count_ops(), dag.count_ops())
         self.assertIsInstance(new_node.op, XGate)
 
+    def test_replace_control_flow_block(self):
+        """Test that we can replace a block of control-flow nodes with a single one."""
+        body = QuantumCircuit(1)
+        body.x(0)
+
+        qr = QuantumRegister(1)
+        cr1 = ClassicalRegister(3)
+        cr2 = ClassicalRegister(3)
+        dag = DAGCircuit()
+        dag.add_qreg(qr)
+        dag.add_creg(cr1)
+        dag.add_creg(cr2)
+        nodes = [
+            dag.apply_operation_back(IfElseOp(expr.logic_not(cr1[0]), body.copy(), None), qr, []),
+            dag.apply_operation_back(
+                SwitchCaseOp(expr.lift(cr2), [((0, 1, 2, 3), body.copy())]), qr, []
+            ),
+        ]
+
+        dag.replace_block_with_op(
+            nodes,
+            IfElseOp(expr.logic_or(expr.logic_not(cr1[0]), expr.less(cr2, 4)), body.copy(), None),
+            {qr[0]: 0},
+        )
+
+        expected = DAGCircuit()
+        expected.add_qreg(qr)
+        expected.add_creg(cr1)
+        expected.add_creg(cr2)
+        expected.apply_operation_back(
+            IfElseOp(expr.logic_or(expr.logic_not(cr1[0]), expr.less(cr2, 4)), body.copy(), None),
+            qr,
+            [],
+        )
+
+        self.assertEqual(dag, expected)
+
 
 class TestDagProperties(QiskitTestCase):
     """Test the DAG properties."""