Include &Packedinstruction in DAGCircuit::op_nodes iteration (Qiskit#13681)

Uses of the `DAGCircuit::op_nodes` iterator are almost invariably
followed by indexing into the DAG to retrieve the node weight, which is
dynamically asserted to be `NodeType::Operation`.  Since the iterator
already verifies this, we can just include the unwrapped
`&PackedInstruction` in the iterator return to simply calling code and
reduce the number of `unreachable!()`, `panic!()`, etcs.

A new `op_node_indices` is provided for ease for cases where the
iterator needs to be consumed to avoid being lifetime-tied to the DAG,
such as when the DAG is going to be mutated based on the nodes.

`topological_op_nodes` is left for a follow-up; it might be slightly
more challenging because of the interaction with rustworkx, but it can
also be done separately anyway.

Author: jakelishman

crates/accelerate/src/barrier_before_final_measurement.rs

@@ -30,24 +30,21 @@ pub fn barrier_before_final_measurements(
     dag: &mut DAGCircuit,
     label: Option<String>,
 ) -> PyResult<()> {
+    let is_exactly_final = |inst: &PackedInstruction| FINAL_OP_NAMES.contains(&inst.op.name());
     let final_ops: HashSet<NodeIndex> = dag
         .op_nodes(true)
-        .filter(|node| {
-            let NodeType::Operation(ref inst) = dag.dag()[*node] else {
-                unreachable!();
-            };
-            if !FINAL_OP_NAMES.contains(&inst.op.name()) {
-                return false;
+        .filter_map(|(node, inst)| {
+            if !is_exactly_final(inst) {
+                return None;
             }
-            let is_final_op = dag.bfs_successors(*node).all(|(_, child_successors)| {
-                !child_successors.iter().any(|suc| match dag.dag()[*suc] {
-                    NodeType::Operation(ref suc_inst) => {
-                        !FINAL_OP_NAMES.contains(&suc_inst.op.name())
-                    }
-                    _ => false,
+            dag.bfs_successors(node)
+                .all(|(_, child_successors)| {
+                    child_successors.iter().all(|suc| match dag.dag()[*suc] {
+                        NodeType::Operation(ref suc_inst) => is_exactly_final(suc_inst),
+                        _ => true,
+                    })
                 })
-            });
-            is_final_op
+                .then_some(node)
         })
         .collect();
     if final_ops.is_empty() {

crates/accelerate/src/basis/basis_translator/compose_transforms.rs

@@ -18,7 +18,7 @@ use qiskit_circuit::parameter_table::ParameterUuid;
 use qiskit_circuit::Qubit;
 use qiskit_circuit::{
     circuit_data::CircuitData,
-    dag_circuit::{DAGCircuit, NodeType},
+    dag_circuit::DAGCircuit,
     operations::{Operation, Param},
 };
 use smallvec::SmallVec;
@@ -83,24 +83,18 @@ pub(super) fn compose_transforms<'a>(
             for (_, dag) in &mut mapped_instructions.values_mut() {
                 let nodes_to_replace = dag
                     .op_nodes(true)
-                    .filter_map(|node| {
-                        if let Some(NodeType::Operation(op)) = dag.dag().node_weight(node) {
-                            if (gate_name.as_str(), *gate_num_qubits)
-                                == (op.op.name(), op.op.num_qubits())
-                            {
-                                Some((
-                                    node,
-                                    op.params_view()
-                                        .iter()
-                                        .map(|x| x.clone_ref(py))
-                                        .collect::<SmallVec<[Param; 3]>>(),
-                                ))
-                            } else {
-                                None
-                            }
-                        } else {
-                            None
-                        }
+                    .filter(|(_, op)| {
+                        (op.op.num_qubits() == *gate_num_qubits)
+                            && (op.op.name() == gate_name.as_str())
+                    })
+                    .map(|(node, op)| {
+                        (
+                            node,
+                            op.params_view()
+                                .iter()
+                                .map(|x| x.clone_ref(py))
+                                .collect::<SmallVec<[Param; 3]>>(),
+                        )
                     })
                     .collect::<Vec<_>>();
                 for (node, params) in nodes_to_replace {
@@ -141,17 +135,15 @@ fn get_gates_num_params(
     dag: &DAGCircuit,
     example_gates: &mut HashMap<GateIdentifier, usize>,
 ) -> PyResult<()> {
-    for node in dag.op_nodes(true) {
-        if let Some(NodeType::Operation(op)) = dag.dag().node_weight(node) {
-            example_gates.insert(
-                (op.op.name().to_string(), op.op.num_qubits()),
-                op.params_view().len(),
-            );
-            if op.op.control_flow() {
-                let blocks = op.op.blocks();
-                for block in blocks {
-                    get_gates_num_params_circuit(&block, example_gates)?;
-                }
+    for (_, inst) in dag.op_nodes(true) {
+        example_gates.insert(
+            (inst.op.name().to_string(), inst.op.num_qubits()),
+            inst.params_view().len(),
+        );
+        if inst.op.control_flow() {
+            let blocks = inst.op.blocks();
+            for block in blocks {
+                get_gates_num_params_circuit(&block, example_gates)?;
             }
         }
     }

crates/accelerate/src/basis/basis_translator/mod.rs

@@ -212,8 +212,7 @@ fn extract_basis(
         basis: &mut HashSet<GateIdentifier>,
         min_qubits: usize,
     ) -> PyResult<()> {
-        for node in circuit.op_nodes(true) {
-            let operation: &PackedInstruction = circuit.dag()[node].unwrap_operation();
+        for (node, operation) in circuit.op_nodes(true) {
             if !circuit.has_calibration_for_index(py, node)?
                 && circuit.get_qargs(operation.qubits).len() >= min_qubits
             {
@@ -279,8 +278,7 @@ fn extract_basis_target(
     min_qubits: usize,
     qargs_with_non_global_operation: &HashMap<Option<Qargs>, HashSet<String>>,
 ) -> PyResult<()> {
-    for node in dag.op_nodes(true) {
-        let node_obj: &PackedInstruction = dag.dag()[node].unwrap_operation();
+    for (node, node_obj) in dag.op_nodes(true) {
         let qargs: &[Qubit] = dag.get_qargs(node_obj.qubits);
         if dag.has_calibration_for_index(py, node)? || qargs.len() < min_qubits {
             continue;

crates/accelerate/src/check_map.rs

@@ -16,7 +16,7 @@ use pyo3::prelude::*;
 use pyo3::wrap_pyfunction;
 
 use qiskit_circuit::circuit_data::CircuitData;
-use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType};
+use qiskit_circuit::dag_circuit::DAGCircuit;
 use qiskit_circuit::imports::CIRCUIT_TO_DAG;
 use qiskit_circuit::operations::{Operation, OperationRef};
 use qiskit_circuit::Qubit;
@@ -36,45 +36,43 @@ fn recurse<'py>(
             None => edge_set.contains(&[qubits[0].into(), qubits[1].into()]),
         }
     };
-    for node in dag.op_nodes(false) {
-        if let NodeType::Operation(inst) = &dag.dag()[node] {
-            let qubits = dag.get_qargs(inst.qubits);
-            if inst.op.control_flow() {
-                if let OperationRef::Instruction(py_inst) = inst.op.view() {
-                    let raw_blocks = py_inst.instruction.getattr(py, "blocks")?;
-                    let circuit_to_dag = CIRCUIT_TO_DAG.get_bound(py);
-                    for raw_block in raw_blocks.bind(py).iter().unwrap() {
-                        let block_obj = raw_block?;
-                        let block = block_obj
-                            .getattr(intern!(py, "_data"))?
-                            .downcast::<CircuitData>()?
-                            .borrow();
-                        let new_dag: DAGCircuit =
-                            circuit_to_dag.call1((block_obj.clone(),))?.extract()?;
-                        let wire_map = (0..block.num_qubits())
-                            .map(|inner| {
-                                let outer = qubits[inner];
-                                match wire_map {
-                                    Some(wire_map) => wire_map[outer.index()],
-                                    None => outer,
-                                }
-                            })
-                            .collect::<Vec<_>>();
-                        let res = recurse(py, &new_dag, edge_set, Some(&wire_map))?;
-                        if res.is_some() {
-                            return Ok(res);
-                        }
+    for (node, inst) in dag.op_nodes(false) {
+        let qubits = dag.get_qargs(inst.qubits);
+        if inst.op.control_flow() {
+            if let OperationRef::Instruction(py_inst) = inst.op.view() {
+                let raw_blocks = py_inst.instruction.getattr(py, "blocks")?;
+                let circuit_to_dag = CIRCUIT_TO_DAG.get_bound(py);
+                for raw_block in raw_blocks.bind(py).iter().unwrap() {
+                    let block_obj = raw_block?;
+                    let block = block_obj
+                        .getattr(intern!(py, "_data"))?
+                        .downcast::<CircuitData>()?
+                        .borrow();
+                    let new_dag: DAGCircuit =
+                        circuit_to_dag.call1((block_obj.clone(),))?.extract()?;
+                    let wire_map = (0..block.num_qubits())
+                        .map(|inner| {
+                            let outer = qubits[inner];
+                            match wire_map {
+                                Some(wire_map) => wire_map[outer.index()],
+                                None => outer,
+                            }
+                        })
+                        .collect::<Vec<_>>();
+                    let res = recurse(py, &new_dag, edge_set, Some(&wire_map))?;
+                    if res.is_some() {
+                        return Ok(res);
                     }
                 }
-            } else if qubits.len() == 2
-                && (dag.calibrations_empty() || !dag.has_calibration_for_index(py, node)?)
-                && !check_qubits(qubits)
-            {
-                return Ok(Some((
-                    inst.op.name().to_string(),
-                    [qubits[0].0, qubits[1].0],
-                )));
             }
+        } else if qubits.len() == 2
+            && (dag.calibrations_empty() || !dag.has_calibration_for_index(py, node)?)
+            && !check_qubits(qubits)
+        {
+            return Ok(Some((
+                inst.op.name().to_string(),
+                [qubits[0].0, qubits[1].0],
+            )));
         }
     }
     Ok(None)

crates/accelerate/src/filter_op_nodes.rs

@@ -29,7 +29,7 @@ pub fn py_filter_op_nodes(
         predicate.call1((dag_op_node,))?.extract()
     };
     let mut remove_nodes: Vec<NodeIndex> = Vec::new();
-    for node in dag.op_nodes(true) {
+    for node in dag.op_node_indices(true) {
         if !callable(node)? {
             remove_nodes.push(node);
         }

crates/accelerate/src/gate_direction.rs

@@ -23,7 +23,7 @@ use qiskit_circuit::{
     circuit_instruction::CircuitInstruction,
     circuit_instruction::ExtraInstructionAttributes,
     converters::{circuit_to_dag, QuantumCircuitData},
-    dag_circuit::{DAGCircuit, NodeType},
+    dag_circuit::DAGCircuit,
     dag_node::{DAGNode, DAGOpNode},
     imports,
     imports::get_std_gate_class,
@@ -105,11 +105,7 @@ fn check_gate_direction<T>(
 where
     T: Fn(&PackedInstruction, &[Qubit]) -> bool,
 {
-    for node in dag.op_nodes(false) {
-        let NodeType::Operation(packed_inst) = &dag.dag()[node] else {
-            panic!("PackedInstruction is expected");
-        };
-
+    for (_, packed_inst) in dag.op_nodes(false) {
         let inst_qargs = dag.get_qargs(packed_inst.qubits);
 
         if let OperationRef::Instruction(py_inst) = packed_inst.op.view() {
@@ -254,9 +250,7 @@ where
     let mut nodes_to_replace: Vec<(NodeIndex, DAGCircuit)> = Vec::new();
     let mut ops_to_replace: Vec<(NodeIndex, Vec<Bound<PyAny>>)> = Vec::new();
 
-    for node in dag.op_nodes(false) {
-        let packed_inst = dag.dag()[node].unwrap_operation();
-
+    for (node, packed_inst) in dag.op_nodes(false) {
         let op_args = dag.get_qargs(packed_inst.qubits);
 
         if let OperationRef::Instruction(py_inst) = packed_inst.op.view() {

crates/accelerate/src/gates_in_basis.rs

@@ -79,8 +79,7 @@ fn any_gate_missing_from_target(dag: &DAGCircuit, target: &Target) -> PyResult<b
     );
 
     // Process the DAG.
-    for gate in dag.op_nodes(true) {
-        let gate = dag.dag()[gate].unwrap_operation();
+    for (_, gate) in dag.op_nodes(true) {
         if is_universal(gate) {
             continue;
         }

crates/accelerate/src/remove_diagonal_gates_before_measure.rs

@@ -47,12 +47,7 @@ fn run_remove_diagonal_before_measure(dag: &mut DAGCircuit) -> PyResult<()> {
     static DIAGONAL_3Q_GATES: [StandardGate; 1] = [StandardGate::CCZGate];
 
     let mut nodes_to_remove = Vec::new();
-    for index in dag.op_nodes(true) {
-        let node = &dag.dag()[index];
-        let NodeType::Operation(inst) = node else {
-            panic!()
-        };
-
+    for (index, inst) in dag.op_nodes(true) {
         if inst.op.name() == "measure" {
             let predecessor = (dag.quantum_predecessors(index))
                 .next()

crates/accelerate/src/remove_identity_equiv.rs

@@ -74,8 +74,7 @@ fn remove_identity_equiv(
         }
     };
 
-    for op_node in dag.op_nodes(false) {
-        let inst = dag.dag()[op_node].unwrap_operation();
+    for (op_node, inst) in dag.op_nodes(false) {
         match inst.op.view() {
             OperationRef::Standard(gate) => {
                 let (dim, trace) = match gate {

crates/accelerate/src/split_2q_unitaries.rs

@@ -31,7 +31,7 @@ pub fn split_2q_unitaries(
     if !dag.get_op_counts().contains_key("unitary") {
         return Ok(());
     }
-    let nodes: Vec<NodeIndex> = dag.op_nodes(false).collect();
+    let nodes: Vec<NodeIndex> = dag.op_node_indices(false).collect();
 
     for node in nodes {
         if let NodeType::Operation(inst) = &dag.dag()[node] {