From 2c0aad55e7b51ae5676915a6a9b9159e1a74ca9e Mon Sep 17 00:00:00 2001
From: Matthew Treinish <mtreinish@kortar.org>
Date: Fri, 23 Aug 2024 16:54:50 -0400
Subject: [PATCH] Tweak control_flow_op_nodes() method to avoid dag traversal
 when not necessary

---
 crates/circuit/src/dag_circuit.rs             | 38 ++++++++++++-------
 .../transpiler/passes/utils/control_flow.py   | 10 +++--
 2 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs
index 59eb01728dca..5013877a4047 100644
--- a/crates/circuit/src/dag_circuit.rs
+++ b/crates/circuit/src/dag_circuit.rs
@@ -3923,21 +3923,31 @@ def _format(operand):
     /// Get a list of "op" nodes in the dag that contain control flow instructions.
     ///
     /// Returns:
-    ///     list[DAGOpNode]: The list of dag nodes containing control flow ops.
-    fn control_flow_op_nodes(&self, py: Python) -> PyResult<Vec<Py<PyAny>>> {
-        self.dag
-            .node_references()
-            .filter_map(|(node_index, node_type)| match node_type {
-                NodeType::Operation(ref node) => {
-                    if node.op.control_flow() {
-                        Some(self.unpack_into(py, node_index, node_type))
-                    } else {
-                        None
+    ///     list[DAGOpNode] | None: The list of dag nodes containing control flow ops. If there
+    ///         are no control flow nodes None is returned
+    fn control_flow_op_nodes(&self, py: Python) -> PyResult<Option<Vec<Py<PyAny>>>> {
+        if CONTROL_FLOW_OP_NAMES
+            .iter()
+            .any(|name| self.op_names.contains_key(*name))
+        {
+            let result: PyResult<Vec<Py<PyAny>>> = self
+                .dag
+                .node_references()
+                .filter_map(|(node_index, node_type)| match node_type {
+                    NodeType::Operation(ref node) => {
+                        if node.op.control_flow() {
+                            Some(self.unpack_into(py, node_index, node_type))
+                        } else {
+                            None
+                        }
                     }
-                }
-                _ => None,
-            })
-            .collect()
+                    _ => None,
+                })
+                .collect();
+            Ok(Some(result?))
+        } else {
+            Ok(None)
+        }
     }
 
     /// Get the list of gate nodes in the dag.
diff --git a/qiskit/transpiler/passes/utils/control_flow.py b/qiskit/transpiler/passes/utils/control_flow.py
index 9b471aadf928..4fc95587e782 100644
--- a/qiskit/transpiler/passes/utils/control_flow.py
+++ b/qiskit/transpiler/passes/utils/control_flow.py
@@ -54,10 +54,12 @@ def out(self, dag):
         def bound_wrapped_method(dag):
             return out(self, dag)
 
-        for node in dag.control_flow_op_nodes():
-            dag.substitute_node(
-                node, map_blocks(bound_wrapped_method, node.op), propagate_condition=False
-            )
+        control_flow_nodes = dag.control_flow_op_nodes()
+        if control_flow_nodes is not None:
+            for node in control_flow_nodes:
+                dag.substitute_node(
+                    node, map_blocks(bound_wrapped_method, node.op), propagate_condition=False
+                )
         return method(self, dag)
 
     return out