Use toposort instead of topological_op_nodes for DAG reconstruction

crates/circuit/src/dag_circuit.rs

@@ -8245,6 +8245,30 @@ impl DAGCircuit {
         }
         Ok(())
     }
+
+    /// Build a new dag copy by iterating over this dag and building up nodes in the new dag from
+    /// them.
+    ///
+    /// This function will build a new output dag builder with the same capacity and iterate over
+    /// the nodes in this dag in a topological order. The given callback is called at each operation
+    /// node and a mutable reference to the [`DAGCircuitBuilder`] which is where the new dag is built
+    /// and a [`PackedInstruction`] of the current op node in the dag. The caller is responsible for
+    /// adding any nodes to the dag as part of the callback.
+    pub fn rebuild_dag_with<F, E>(&self, mut callback: F) -> Result<Self, E>
+    where
+        F: FnMut(&mut DAGCircuitBuilder, &PackedInstruction) -> Result<(), E>,
+    {
+        let new_dag = self.copy_empty_like_with_same_capacity(VarsMode::Alike, BlocksMode::Keep);
+        let mut builder = new_dag.into_builder();
+        for node in
+            petgraph::algo::toposort(&self.dag, None).expect("DAGCircuit can't have a cycle")
+        {
+            if let NodeType::Operation(ref inst) = self.dag[node] {
+                callback(&mut builder, inst)?;
+            }
+        }
+        Ok(builder.build())
+    }
 }
 
 struct NodesOnWireIter<'a> {

crates/transpiler/src/passes/basis_translator/mod.rs

@@ -20,21 +20,24 @@ use errors::BasisTranslatorError;
 use hashbrown::{HashMap, HashSet};
 use pyo3::prelude::*;
 use qiskit_util::{IndexMap, IndexSet};
+use rustworkx_core::petgraph::algo::toposort;
 
 mod basis_search;
 mod compose_transforms;
 mod errors;
 
-use qiskit_circuit::dag_circuit::{DAGCircuit, DAGCircuitBuilder};
 use qiskit_circuit::instruction::Parameters;
-use qiskit_circuit::operations::{Operation, OperationRef, Param, PauliBased, PythonOperation};
 use qiskit_circuit::packed_instruction::{PackedInstruction, PackedOperation};
 use qiskit_circuit::parameter::parameter_expression::ParameterError;
 use qiskit_circuit::parameter::parameter_expression::ParameterExpression;
 use qiskit_circuit::parameter::symbol_expr::Symbol;
 use qiskit_circuit::parameter::symbol_expr::SymbolExpr;
 use qiskit_circuit::parameter::symbol_expr::Value;
-use qiskit_circuit::{BlocksMode, Clbit, PhysicalQubit, Qubit, VarsMode};
+use qiskit_circuit::{Clbit, PhysicalQubit, Qubit};
+use qiskit_circuit::{
+    dag_circuit::{DAGCircuit, DAGCircuitBuilder, NodeType},
+    operations::{Operation, OperationRef, Param, PauliBased, PythonOperation},
+};
 use smallvec::SmallVec;
 
 use crate::equivalence::EquivalenceLibrary;
@@ -331,10 +334,8 @@ fn apply_translation(
     qargs_with_non_global_operation: &AhashIndexMap<Qargs, AhashIndexSet<&str>>,
     qarg_mapping: Option<&HashMap<Qubit, Qubit>>,
 ) -> Result<DAGCircuit, BasisTranslatorError> {
-    let out_dag = dag.copy_empty_like(VarsMode::Alike, BlocksMode::Keep);
-    let mut out_dag_builder = out_dag.into_builder();
-    for node in dag.topological_op_nodes(false) {
-        let node_obj = dag[node].unwrap_operation();
+    let rebuilder_callback = |out_dag_builder: &mut DAGCircuitBuilder,
+                              node_obj: &PackedInstruction| {
         let node_qarg = dag.get_qargs(node_obj.qubits);
         let node_carg = dag.get_cargs(node_obj.clbits);
         let qubit_set: AhashIndexSet<Qubit> = AhashIndexSet::from_iter(node_qarg.iter().copied());
@@ -401,7 +402,7 @@ fn apply_translation(
                         )
                     })?;
             }
-            continue;
+            return Ok(());
         }
         // Map to the absolute indices when provided to avoid mistakenly tracking
         // the operation as global.
@@ -431,7 +432,7 @@ fn apply_translation(
                         "Error applying operation to DAGCircuit".to_string(),
                     )
                 })?;
-            continue;
+            return Ok(());
         }
 
         // Map the unique qargs with the absolute indices as well
@@ -445,21 +446,22 @@ fn apply_translation(
         };
         if extra_inst_map.contains_key(&unique_qargs) {
             replace_node(
-                &mut out_dag_builder,
+                out_dag_builder,
                 node_obj.clone(),
                 &extra_inst_map[&unique_qargs],
             )?;
         } else if instr_map
             .contains_key(&(node_obj.op.name().to_string(), node_obj.op.num_qubits()))
         {
-            replace_node(&mut out_dag_builder, node_obj.clone(), instr_map)?;
+            replace_node(out_dag_builder, node_obj.clone(), instr_map)?;
         } else {
             return Err(BasisTranslatorError::ApplyTranslationMappingError(
                 node_obj.op.name().to_string(),
             ));
         }
-    }
-    Ok(out_dag_builder.build())
+        Ok(())
+    };
+    dag.rebuild_dag_with(rebuilder_callback)
 }
 
 fn replace_node(
@@ -480,8 +482,11 @@ fn replace_node(
         });
     }
     if params_view.is_empty() {
-        for inner_index in target_dag.topological_op_nodes(false) {
-            let inner_node = &target_dag[inner_index].unwrap_operation();
+        for inner_index in toposort(target_dag.dag(), None).expect("DAGCircuit can't have a cycle")
+        {
+            let NodeType::Operation(ref inner_node) = target_dag[inner_index] else {
+                continue;
+            };
             let old_qargs = dag.qargs_interner().get(node.qubits);
             let old_cargs = dag.cargs_interner().get(node.clbits);
             let new_qubits: Vec<Qubit> = target_dag
@@ -544,8 +549,11 @@ fn replace_node(
                     _ => None,
                 }),
         );
-        for inner_index in target_dag.topological_op_nodes(false) {
-            let inner_node = &target_dag[inner_index].unwrap_operation();
+        for inner_index in toposort(target_dag.dag(), None).expect("DAGCircuit can't have a cycle")
+        {
+            let NodeType::Operation(ref inner_node) = target_dag[inner_index] else {
+                continue;
+            };
             let old_qargs = dag.qargs_interner().get(node.qubits);
             let old_cargs = dag.cargs_interner().get(node.clbits);
             let new_qubits: Vec<Qubit> = target_dag

crates/transpiler/src/passes/disjoint_layout.rs

@@ -20,14 +20,15 @@ use pyo3::prelude::*;
 use pyo3::types::PyList;
 use rustworkx_core::connectivity::connected_components;
 use rustworkx_core::petgraph::EdgeType;
+use rustworkx_core::petgraph::algo::toposort;
 use rustworkx_core::petgraph::prelude::*;
 use rustworkx_core::petgraph::visit::{IntoEdgeReferences, IntoNodeReferences, NodeFiltered};
 use uuid::Uuid;
 
 use crate::TranspilerError;
 use crate::target::{Qargs, Target};
 use qiskit_circuit::bit::ShareableQubit;
-use qiskit_circuit::dag_circuit::DAGCircuit;
+use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType};
 use qiskit_circuit::operations::{Operation, OperationRef, StandardInstruction};
 use qiskit_circuit::packed_instruction::PackedOperation;
 use qiskit_circuit::{
@@ -474,8 +475,10 @@ fn separate_dag(dag: &mut DAGCircuit) -> PyResult<Vec<DAGCircuit>> {
             new_dag.set_global_phase_f64(0.);
             let old_qubits = dag.qubits();
             let mut block_map = BlockMapper::new();
-            for index in dag.topological_op_nodes(false) {
-                let node = dag[index].unwrap_operation();
+            for index in toposort(dag.dag(), None).expect("DAGCircuit can't have a cycle") {
+                let NodeType::Operation(ref node) = dag.dag()[index] else {
+                    continue;
+                };
                 let qargs: HashSet<Qubit> = dag.get_qargs(node.qubits).iter().copied().collect();
                 if dag_qubits.is_superset(&qargs) {
                     let qargs = dag.get_qargs(node.qubits);

crates/transpiler/src/passes/split_2q_unitaries.rs

@@ -18,10 +18,10 @@ use pyo3::prelude::*;
 use rustworkx_core::petgraph::stable_graph::NodeIndex;
 use smallvec::{SmallVec, smallvec};
 
-use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType, Wire};
+use qiskit_circuit::Qubit;
+use qiskit_circuit::dag_circuit::{DAGCircuit, DAGCircuitBuilder, NodeType, Wire};
 use qiskit_circuit::operations::{ArrayType, Operation, OperationRef, Param, UnitaryGate};
-use qiskit_circuit::packed_instruction::PackedOperation;
-use qiskit_circuit::{BlocksMode, Qubit, VarsMode};
+use qiskit_circuit::packed_instruction::{PackedInstruction, PackedOperation};
 
 use qiskit_synthesis::two_qubit_decompose::{Specialization, TwoQubitWeylDecomposition};
 
@@ -99,12 +99,7 @@ pub fn run_split_2q_unitaries(
     // We have swap-like unitaries, so we create a new DAG in a manner similar to
     // The Elide Permutations pass, while also splitting the unitaries to 1-qubit gates
     let mut mapping: Vec<usize> = (0..dag.num_qubits()).collect();
-    let new_dag = dag.copy_empty_like(VarsMode::Alike, BlocksMode::Keep);
-    let mut new_dag = new_dag.into_builder();
-    for node in dag.topological_op_nodes(false) {
-        let NodeType::Operation(inst) = &dag.dag()[node] else {
-            unreachable!("Op nodes contain a non-operation");
-        };
+    let rebuilder_callback = |new_dag: &mut DAGCircuitBuilder, inst: &PackedInstruction| {
         if let OperationRef::Unitary(unitary_gate) = inst.op.view() {
             if unitary_gate.num_qubits() == 2 {
                 let decomp = TwoQubitWeylDecomposition::new_inner(
@@ -156,7 +151,7 @@ pub fn run_split_2q_unitaries(
                         None,
                     )?;
                     new_dag.add_global_phase(&Param::Float(decomp.global_phase + PI4))?;
-                    continue; // skip the general instruction handling code
+                    return Ok(());
                 }
             }
         }
@@ -177,8 +172,10 @@ pub fn run_split_2q_unitaries(
             #[cfg(feature = "cache_pygates")]
             inst.py_op.get().cloned(),
         )?;
-    }
-    Ok(Some((new_dag.build(), mapping)))
+        Ok(())
+    };
+    dag.rebuild_dag_with(rebuilder_callback)
+        .map(|x| Some((x, mapping)))
 }
 
 pub fn split_2q_unitaries_mod(m: &Bound<PyModule>) -> PyResult<()> {

crates/transpiler/src/passes/unitary_synthesis/mod.rs

@@ -33,7 +33,7 @@ use qiskit_circuit::dag_circuit::{DAGCircuit, DAGCircuitBuilder};
 use qiskit_circuit::instruction::Parameters;
 use qiskit_circuit::operations::{Operation, OperationRef, Param, PythonOperation, StandardGate};
 use qiskit_circuit::packed_instruction::{PackedInstruction, PackedOperation};
-use qiskit_circuit::{BlocksMode, PhysicalQubit, Qubit, VarsMode};
+use qiskit_circuit::{PhysicalQubit, Qubit};
 use qiskit_synthesis::euler_one_qubit_decomposer::unitary_to_gate_sequence_inner;
 use qiskit_synthesis::qsd::quantum_shannon_decomposition;
 use qiskit_synthesis::two_qubit_decompose::TwoQubitGateSequence;
@@ -342,18 +342,14 @@ pub fn run_unitary_synthesis(
         return Ok(None);
     }
 
-    let mut out = dag
-        .copy_empty_like(VarsMode::Alike, BlocksMode::Drop)
-        .into_builder();
-    for node in dag.topological_op_nodes(false) {
-        let inst = dag[node].unwrap_operation();
+    let rebuilder_callback = |out: &mut DAGCircuitBuilder, inst: &PackedInstruction| {
         let Some(cf) = dag.try_view_control_flow(inst) else {
             // Handle regular instructions - this path is where we end up most of the time.
-            if !synthesize_onto(&mut out, state, inst)? {
+            if !synthesize_onto(out, state, inst)? {
                 // No synthesis was necessary, so reinstate the operation.
                 out.push_back(inst.clone())?;
             }
-            continue;
+            return Ok(());
         };
         // If we make it here, we've got control flow and have to set ourselves up to recurse.
         let blocks = cf
@@ -389,8 +385,9 @@ pub fn run_unitary_synthesis(
             inst.clbits,
             inst.label.as_deref().cloned(),
         ))?;
-    }
-    Ok(Some(out.build()))
+        Ok(())
+    };
+    Some(dag.rebuild_dag_with(rebuilder_callback)).transpose()
 }
 
 /// Synthesise a matrix onto the DAG.