Fix: Correct the behavior of compose_transforms.

-  Use `PackedOperation.control_flow` instead of `CONTROL_FLOW_OP_NAMES` to check if an instructuion is a control flow operation.
- Remove panic statement when checking for example gates.
- Use qreg when adding an instruction to the mock_dag in `compose_transforms`.
- Add missing comparison check in for loop to compare the mapped instructions.
- Use borrowed `DAGCircuit` instances in the recursion of `get_example_gates`, do not use extract.
- Fix behavior of `get_example_gates` to return `PackedInstruction` instances instead of `NodeIndex`.

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

@@ -11,31 +11,23 @@
 // that they have been altered from the originals.
 
 use hashbrown::{HashMap, HashSet};
-use once_cell::sync::Lazy;
-use pyo3::types::PyTuple;
 use pyo3::{exceptions::PyTypeError, prelude::*};
 use qiskit_circuit::imports::{
     CIRCUIT_TO_DAG, GATE, PARAMETER_VECTOR, QUANTUM_CIRCUIT, QUANTUM_REGISTER,
 };
 use qiskit_circuit::operations::OperationRef;
+use qiskit_circuit::packed_instruction::PackedInstruction;
 use qiskit_circuit::parameter_table::ParameterUuid;
 use qiskit_circuit::{
     circuit_data::CircuitData,
     dag_circuit::{DAGCircuit, NodeType},
-    dag_node::DAGOpNode,
     operations::{Operation, Param},
 };
-use rustworkx_core::petgraph::graph::NodeIndex;
 use smallvec::SmallVec;
 
 // Custom types
 // TODO: Remove these and use the version from `basis_search`
 type BasisTransforms = Vec<(String, u32, SmallVec<[Param; 3]>, CircuitRep)>;
-static CONTROL_FLOW_OP_NAMES: Lazy<HashSet<&'static str>> = Lazy::new(|| {
-    ["for_loop", "while_loop", "if_else", "switch_case"]
-        .into_iter()
-        .collect()
-});
 // TODO: Remove these and use the version from `EquivalenceLibrary`
 
 /// Representation of QuantumCircuit which the original circuit object + an
@@ -101,15 +93,7 @@ pub(super) fn compose_transforms<'a>(
     for (gate_name, gate_num_qubits) in source_basis.iter().cloned() {
         // Need to grab a gate instance to find num_qubits and num_params.
         // Can be removed following https://github.com/Qiskit/qiskit-terra/pull/3947 .
-        let Some(NodeType::Operation(example_gate)) = source_dag
-            .dag
-            .node_weight(example_gates[&(gate_name.clone(), gate_num_qubits)])
-        else {
-            panic!(
-                "Nodeindex {:?} should be in the source_dag",
-                example_gates[&(gate_name, gate_num_qubits)]
-            )
-        };
+        let example_gate = &example_gates[&(gate_name.clone(), gate_num_qubits)];
         let num_params = example_gate
             .params
             .as_ref()
@@ -135,31 +119,31 @@ pub(super) fn compose_transforms<'a>(
                 .collect::<SmallVec<[Param; 3]>>(),
         ))?;
 
-        dag.py_apply_operation_back(
-            py,
-            gate,
-            Some(PyTuple::new_bound(py, 0..gate_num_qubits).extract()?),
-            None,
-            true,
-        )?;
+        dag.py_apply_operation_back(py, gate, Some(qubits.extract()?), None, false)?;
         mapped_instructions.insert(
             (gate_name.clone(), gate_num_qubits),
             (placeholder_params, dag),
         );
 
-        for (_gate_name, _gate_num_qubitss, equiv_params, equiv) in basis_transforms {
+        for (gate_name, gate_num_qubits, equiv_params, equiv) in basis_transforms {
             for ((_mapped_instr_name, _), (_dag_params, dag)) in &mut mapped_instructions {
                 let doomed_nodes = dag
                     .op_nodes(true)
                     .filter_map(|node| {
                         if let Some(NodeType::Operation(op)) = dag.dag.node_weight(node) {
-                            Some((
-                                node,
-                                op.params_view()
-                                    .iter()
-                                    .map(|x| x.clone_ref(py))
-                                    .collect::<SmallVec<[Param; 3]>>(),
-                            ))
+                            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
                         }
@@ -178,11 +162,11 @@ pub(super) fn compose_transforms<'a>(
                     replacement
                         .0
                         .assign_parameters_from_mapping(py, param_mapping)?;
-                    let replace_dag: DAGCircuit = CIRCUIT_TO_DAG
+                    let replace_dag: PyRef<DAGCircuit> = CIRCUIT_TO_DAG
                         .get_bound(py)
                         .call1((replacement,))?
                         .downcast_into::<DAGCircuit>()?
-                        .extract()?;
+                        .borrow();
                     let op_node = dag.get_node(py, node)?;
                     dag.substitute_node_with_dag(py, op_node.bind(py), &replace_dag, None, true)?;
                 }
@@ -195,34 +179,21 @@ pub(super) fn compose_transforms<'a>(
 fn get_example_gates(
     py: Python,
     dag: &DAGCircuit,
-    example_gates: Option<Box<HashMap<(String, u32), NodeIndex>>>,
-) -> PyResult<Box<HashMap<(String, u32), NodeIndex>>> {
-    /*
-    def _get_example_gates(source_dag):
-        def recurse(dag, example_gates=None):
-            example_gates = example_gates or {}
-            for node in dag.op_nodes():
-                example_gates[(node.name, node.num_qubits)] = node
-                if node.name in CONTROL_FLOW_OP_NAMES:
-                    for block in node.op.blocks:
-                        example_gates = recurse(circuit_to_dag(block), example_gates)
-            return example_gates
-
-        return recurse(source_dag)
-    */
+    example_gates: Option<Box<HashMap<(String, u32), PackedInstruction>>>,
+) -> PyResult<Box<HashMap<(String, u32), PackedInstruction>>> {
     let mut example_gates = example_gates.unwrap_or_default();
     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()), node);
-            if CONTROL_FLOW_OP_NAMES.contains(op.op.name()) {
+            example_gates.insert((op.op.name().to_string(), op.op.num_qubits()), op.clone());
+            if op.op.control_flow() {
                 let OperationRef::Instruction(inst) = op.op.view() else {
                     panic!("Control flow op can only be an instruction")
                 };
                 let inst_bound = inst.instruction.bind(py);
                 let blocks = inst_bound.getattr("blocks")?;
                 for block in blocks.iter()? {
                     let block_as_dag = CIRCUIT_TO_DAG.get_bound(py).call1((block?,))?;
-                    let block_as_dag = block_as_dag.downcast_into::<DAGCircuit>()?.extract()?;
+                    let block_as_dag = block_as_dag.downcast_into::<DAGCircuit>()?.borrow();
                     example_gates = get_example_gates(py, &block_as_dag, Some(example_gates))?;
                 }
             }