From 473ca0396f6a2b37c869f09dc77d5b6020d02abf Mon Sep 17 00:00:00 2001
From: Raynel Sanchez <raynelfss@hotmail.com>
Date: Tue, 10 Dec 2024 12:45:53 -0500
Subject: [PATCH] Fix: Make `py_op` private. - Add views for `py_op`. - Remove
 stale saching being done in `apply_operation_*` in `DAGCircuit`. - Adapt the
 code to new changes. - Add missing docstring

---
 .../src/barrier_before_final_measurement.rs   | 31 +++--------
 .../basis_translator/compose_transforms.rs    |  2 -
 .../src/basis/basis_translator/mod.rs         | 12 -----
 crates/accelerate/src/elide_permutations.rs   |  2 -
 crates/accelerate/src/gate_direction.rs       |  4 +-
 crates/accelerate/src/twirling.rs             |  2 -
 crates/accelerate/src/unitary_synthesis.rs    |  4 --
 crates/circuit/src/circuit_data.rs            | 13 ++---
 crates/circuit/src/converters.rs              |  3 --
 crates/circuit/src/dag_circuit.rs             | 51 ++-----------------
 crates/circuit/src/packed_instruction.rs      | 16 +++++-
 11 files changed, 32 insertions(+), 108 deletions(-)

diff --git a/crates/accelerate/src/barrier_before_final_measurement.rs b/crates/accelerate/src/barrier_before_final_measurement.rs
index c3c9ee7809b3..36e3fd59b691 100644
--- a/crates/accelerate/src/barrier_before_final_measurement.rs
+++ b/crates/accelerate/src/barrier_before_final_measurement.rs
@@ -84,29 +84,14 @@ pub fn barrier_before_final_measurements(
         instruction: new_barrier.unbind(),
     };
     let qargs: Vec<Qubit> = (0..dag.num_qubits() as u32).map(Qubit).collect();
-    #[cfg(feature = "cache_pygates")]
-    {
-        dag.apply_operation_back(
-            py,
-            PackedOperation::from_instruction(Box::new(new_barrier_py_inst)),
-            qargs.as_slice(),
-            &[],
-            None,
-            ExtraInstructionAttributes::new(label, None, None, None),
-            Some(new_barrier.unbind()),
-        )?;
-    }
-    #[cfg(not(feature = "cache_pygates"))]
-    {
-        dag.apply_operation_back(
-            py,
-            PackedOperation::from_instruction(Box::new(new_barrier_py_inst)),
-            qargs.as_slice(),
-            &[],
-            None,
-            ExtraInstructionAttributes::new(label, None, None, None),
-        )?;
-    }
+    dag.apply_operation_back(
+        py,
+        PackedOperation::from_instruction(Box::new(new_barrier_py_inst)),
+        qargs.as_slice(),
+        &[],
+        None,
+        ExtraInstructionAttributes::new(label, None, None, None),
+    )?;
     for inst in final_packed_ops {
         dag.push_back(py, inst)?;
     }
diff --git a/crates/accelerate/src/basis/basis_translator/compose_transforms.rs b/crates/accelerate/src/basis/basis_translator/compose_transforms.rs
index fa07b224de08..79efdac90f6c 100644
--- a/crates/accelerate/src/basis/basis_translator/compose_transforms.rs
+++ b/crates/accelerate/src/basis/basis_translator/compose_transforms.rs
@@ -74,8 +74,6 @@ pub(super) fn compose_transforms<'a>(
                 Some(gate_obj.params)
             },
             gate_obj.extra_attrs,
-            #[cfg(feature = "cache_pygates")]
-            Some(gate.into()),
         )?;
         mapped_instructions.insert((gate_name, gate_num_qubits), (placeholder_params, dag));
 
diff --git a/crates/accelerate/src/basis/basis_translator/mod.rs b/crates/accelerate/src/basis/basis_translator/mod.rs
index f10b0adbdeb4..c36aa73a7486 100644
--- a/crates/accelerate/src/basis/basis_translator/mod.rs
+++ b/crates/accelerate/src/basis/basis_translator/mod.rs
@@ -481,8 +481,6 @@ fn apply_translation(
                     node_carg,
                     (!new_op.params.is_empty()).then_some(new_op.params),
                     new_op.extra_attrs,
-                    #[cfg(feature = "cache_pygates")]
-                    None,
                 )?;
             } else {
                 out_dag.apply_operation_back(
@@ -498,8 +496,6 @@ fn apply_translation(
                             .collect(),
                     ),
                     node_obj.extra_attrs().clone(),
-                    #[cfg(feature = "cache_pygates")]
-                    None,
                 )?;
             }
             continue;
@@ -523,8 +519,6 @@ fn apply_translation(
                         .collect(),
                 ),
                 node_obj.extra_attrs().clone(),
-                #[cfg(feature = "cache_pygates")]
-                None,
             )?;
             continue;
         }
@@ -543,8 +537,6 @@ fn apply_translation(
                         .collect(),
                 ),
                 node_obj.extra_attrs().clone(),
-                #[cfg(feature = "cache_pygates")]
-                None,
             )?;
             continue;
         }
@@ -642,8 +634,6 @@ fn replace_node(
                 &new_clbits,
                 (!new_params.is_empty()).then_some(new_params),
                 new_extra_props,
-                #[cfg(feature = "cache_pygates")]
-                None,
             )?;
         }
         dag.add_global_phase(py, target_dag.global_phase())?;
@@ -742,8 +732,6 @@ fn replace_node(
                 &new_clbits,
                 (!new_params.is_empty()).then_some(new_params),
                 inner_node.extra_attrs().clone(),
-                #[cfg(feature = "cache_pygates")]
-                None,
             )?;
         }
 
diff --git a/crates/accelerate/src/elide_permutations.rs b/crates/accelerate/src/elide_permutations.rs
index 511fb54761b1..f74857d02db5 100644
--- a/crates/accelerate/src/elide_permutations.rs
+++ b/crates/accelerate/src/elide_permutations.rs
@@ -89,8 +89,6 @@ fn run(py: Python, dag: &mut DAGCircuit) -> PyResult<Option<(DAGCircuit, Vec<usi
                         cargs,
                         (!inst.params_view().is_empty()).then_some(inst.params_view().into()),
                         inst.extra_attrs().clone(),
-                        #[cfg(feature = "cache_pygates")]
-                        None,
                     )?;
                 }
             }
diff --git a/crates/accelerate/src/gate_direction.rs b/crates/accelerate/src/gate_direction.rs
index 6d4128ac3044..4dcb5241a088 100755
--- a/crates/accelerate/src/gate_direction.rs
+++ b/crates/accelerate/src/gate_direction.rs
@@ -406,7 +406,7 @@ fn has_calibration_for_op_node(
                         .collect(),
                     extra_attrs: packed_inst.extra_attrs().clone(),
                     #[cfg(feature = "cache_pygates")]
-                    py_op: packed_inst.py_op.clone(),
+                    py_op: packed_inst.py_op().clone(),
                 },
                 sort_key: "".into_py(py),
             },
@@ -479,8 +479,6 @@ fn apply_operation_back(
         &[],
         param,
         ExtraInstructionAttributes::default(),
-        #[cfg(feature = "cache_pygates")]
-        None,
     )?;
 
     Ok(())
diff --git a/crates/accelerate/src/twirling.rs b/crates/accelerate/src/twirling.rs
index 6870d9e662d6..c38b604b703a 100644
--- a/crates/accelerate/src/twirling.rs
+++ b/crates/accelerate/src/twirling.rs
@@ -353,8 +353,6 @@ fn generate_twirled_circuit(
                         ),
                         inst.extra_attrs().clone(),
                     );
-                    #[cfg(feature = "cache_pygates")]
-                    new_inst.py_op.set(new_inst_obj).unwrap();
                     out_circ.push(py, new_inst)?;
                 } else {
                     out_circ.push(py, inst.clone())?;
diff --git a/crates/accelerate/src/unitary_synthesis.rs b/crates/accelerate/src/unitary_synthesis.rs
index 11e1f1f8eeb3..6870a40cc81f 100644
--- a/crates/accelerate/src/unitary_synthesis.rs
+++ b/crates/accelerate/src/unitary_synthesis.rs
@@ -12,8 +12,6 @@
 #![allow(clippy::too_many_arguments)]
 
 use std::f64::consts::PI;
-#[cfg(feature = "cache_pygates")]
-use std::sync::OnceLock;
 
 use approx::relative_eq;
 use hashbrown::{HashMap, HashSet};
@@ -320,8 +318,6 @@ fn py_run_main_loop(
                                 &[],
                                 Some(new_params),
                                 ExtraInstructionAttributes::default(),
-                                #[cfg(feature = "cache_pygates")]
-                                None,
                             )?;
                         }
                         out_dag.add_global_phase(py, &Param::Float(sequence.global_phase))?;
diff --git a/crates/circuit/src/circuit_data.rs b/crates/circuit/src/circuit_data.rs
index fb4e0f371bbd..629e073c8202 100644
--- a/crates/circuit/src/circuit_data.rs
+++ b/crates/circuit/src/circuit_data.rs
@@ -10,9 +10,6 @@
 // copyright notice, and modified files need to carry a notice indicating
 // that they have been altered from the originals.
 
-#[cfg(feature = "cache_pygates")]
-use std::sync::OnceLock;
-
 use crate::bit_data::BitData;
 use crate::circuit_instruction::{
     CircuitInstruction, ExtraInstructionAttributes, OperationFromPython,
@@ -405,7 +402,7 @@ impl CircuitData {
             *inst.extra_attrs_mut() = result.extra_attrs;
             #[cfg(feature = "cache_pygates")]
             {
-                inst.py_op = py_op.unbind().into();
+                *inst.py_op_mut() = py_op.unbind().into();
             }
         }
         Ok(())
@@ -515,7 +512,7 @@ impl CircuitData {
                     .collect(),
                 extra_attrs: inst.extra_attrs().clone(),
                 #[cfg(feature = "cache_pygates")]
-                py_op: inst.py_op.clone(),
+                py_op: inst.py_op().clone(),
             }
             .into_py(py)
         };
@@ -1374,7 +1371,7 @@ impl CircuitData {
                                 // Standard gates can all rebuild their definitions, so if the
                                 // cached py_op exists, update the `params` attribute and clear out
                                 // any existing cache.
-                                if let Some(borrowed) = previous.py_op.get() {
+                                if let Some(borrowed) = previous.py_op().get() {
                                     borrowed
                                         .bind(py)
                                         .getattr(params_attr)?
@@ -1447,10 +1444,6 @@ impl CircuitData {
                                 .params_view_mut()
                                 .swap_with_slice(&mut new_op.params);
                             *previous.extra_attrs_mut() = new_op.extra_attrs;
-                            #[cfg(feature = "cache_pygates")]
-                            {
-                                previous.py_op = op.into_py(py).into();
-                            }
                             for uuid in uuids.iter() {
                                 self.param_table.add_use(*uuid, usage)?
                             }
diff --git a/crates/circuit/src/converters.rs b/crates/circuit/src/converters.rs
index d0c676e572cb..1c61e6705f38 100644
--- a/crates/circuit/src/converters.rs
+++ b/crates/circuit/src/converters.rs
@@ -10,9 +10,6 @@
 // copyright notice, and modified files need to carry a notice indicating
 // that they have been altered from the originals.
 
-#[cfg(feature = "cache_pygates")]
-use std::sync::OnceLock;
-
 use ::pyo3::prelude::*;
 use hashbrown::HashMap;
 use pyo3::{
diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs
index 2d2d64afa56e..c11eb63d17c0 100644
--- a/crates/circuit/src/dag_circuit.rs
+++ b/crates/circuit/src/dag_circuit.rs
@@ -3253,7 +3253,7 @@ def _format(operand):
                             .into();
                             #[cfg(feature = "cache_pygates")]
                             {
-                                new_inst.py_op = new_op.unbind().into();
+                                *new_inst.py_op_mut() = new_op.unbind().into();
                             }
                         }
                     }
@@ -3271,10 +3271,6 @@ def _format(operand):
                             new_inst
                                 .extra_attrs_mut()
                                 .set_condition(new_condition.clone());
-                            #[cfg(feature = "cache_pygates")]
-                            {
-                                new_inst.py_op.take();
-                            }
                             match new_inst.op().view() {
                                 OperationRef::Instruction(py_inst) => {
                                     py_inst
@@ -3351,7 +3347,7 @@ def _format(operand):
             node.instruction.extra_attrs = new_weight.extra_attrs().clone();
             #[cfg(feature = "cache_pygates")]
             {
-                node.instruction.py_op = new_weight.py_op.clone();
+                node.instruction.py_op = new_weight.py_op().clone();
             }
             Ok(node.into_py(py))
         } else {
@@ -5065,19 +5061,8 @@ impl DAGCircuit {
         cargs: &[Clbit],
         params: Option<SmallVec<[Param; 3]>>,
         extra_attrs: ExtraInstructionAttributes,
-        #[cfg(feature = "cache_pygates")] py_op: Option<PyObject>,
     ) -> PyResult<NodeIndex> {
-        self.inner_apply_op(
-            py,
-            op,
-            qargs,
-            cargs,
-            params,
-            extra_attrs,
-            #[cfg(feature = "cache_pygates")]
-            py_op,
-            false,
-        )
+        self.inner_apply_op(py, op, qargs, cargs, params, extra_attrs, false)
     }
 
     /// Apply a [PackedOperation] to the front of the circuit.
@@ -5089,19 +5074,8 @@ impl DAGCircuit {
         cargs: &[Clbit],
         params: Option<SmallVec<[Param; 3]>>,
         extra_attrs: ExtraInstructionAttributes,
-        #[cfg(feature = "cache_pygates")] py_op: Option<PyObject>,
     ) -> PyResult<NodeIndex> {
-        self.inner_apply_op(
-            py,
-            op,
-            qargs,
-            cargs,
-            params,
-            extra_attrs,
-            #[cfg(feature = "cache_pygates")]
-            py_op,
-            true,
-        )
+        self.inner_apply_op(py, op, qargs, cargs, params, extra_attrs, true)
     }
 
     #[inline]
@@ -5114,7 +5088,6 @@ impl DAGCircuit {
         cargs: &[Clbit],
         params: Option<SmallVec<[Param; 3]>>,
         extra_attrs: ExtraInstructionAttributes,
-        #[cfg(feature = "cache_pygates")] py_op: Option<PyObject>,
         front: bool,
     ) -> PyResult<NodeIndex> {
         // Check that all qargs are within an acceptable range
@@ -5140,13 +5113,6 @@ impl DAGCircuit {
             }
             Ok(())
         })?;
-
-        #[cfg(feature = "cache_pygates")]
-        let py_op = if let Some(py_op) = py_op {
-            py_op.into()
-        } else {
-            OnceLock::new()
-        };
         let packed_instruction = PackedInstruction::new(
             op,
             self.qargs_interner.insert(qargs),
@@ -5642,7 +5608,7 @@ impl DAGCircuit {
                                     .collect(),
                                 extra_attrs: packed.extra_attrs().clone(),
                                 #[cfg(feature = "cache_pygates")]
-                                py_op: packed.py_op.clone(),
+                                py_op: packed.py_op().clone(),
                             },
                             sort_key: format!("{:?}", self.sort_key(id)).into_py(py),
                         },
@@ -6923,9 +6889,6 @@ impl DAGCircuit {
                 )));
         }
 
-        #[cfg(feature = "cache_pygates")]
-        let mut py_op_cache = Some(op.clone().unbind());
-
         let mut extra_attrs = new_op.extra_attrs.clone();
         // If either operation is a control-flow operation, propagate_condition is ignored
         if propagate_condition
@@ -6963,10 +6926,6 @@ impl DAGCircuit {
                         old_condition.downcast_bound::<PyTuple>(py)?,
                     )?;
                 }
-                #[cfg(feature = "cache_pygates")]
-                {
-                    py_op_cache = None;
-                }
             }
         };
         if new_wires != current_wires {
diff --git a/crates/circuit/src/packed_instruction.rs b/crates/circuit/src/packed_instruction.rs
index d52cbc1ac0c0..96d6d6fb250b 100644
--- a/crates/circuit/src/packed_instruction.rs
+++ b/crates/circuit/src/packed_instruction.rs
@@ -514,7 +514,7 @@ pub struct PackedInstruction {
     /// requires the GIL to even `get` (of course!), which makes implementing `Clone` hard for us.
     /// We can revisit once we're on PyO3 0.22+ and have been able to disable its `py-clone`
     /// feature.
-    pub py_op: OnceLock<Py<PyAny>>,
+    py_op: OnceLock<Py<PyAny>>,
 }
 
 impl PackedInstruction {
@@ -554,10 +554,17 @@ impl PackedInstruction {
         self.clbits
     }
 
+    /// Retrieves an immutable reference to the extra_attributes of this instruction.
     pub fn extra_attrs(&self) -> &ExtraInstructionAttributes {
         &self.extra_attrs
     }
 
+    /// Retrieves the cached py_gate immutably.
+    #[cfg(feature = "cache_pygates")]
+    pub fn py_op(&self) -> &OnceLock<PyObject> {
+        &self.py_op
+    }
+
     // Setters
 
     /// Retrieves an immutable reference to the instruction's underlying operation.
@@ -579,6 +586,7 @@ impl PackedInstruction {
         &mut self.clbits
     }
 
+    /// Retrieves a mutable reference to the extra_attributes of this instruction.
     pub fn extra_attrs_mut(&mut self) -> &mut ExtraInstructionAttributes {
         #[cfg(feature = "cache_pygates")]
         {
@@ -587,6 +595,12 @@ impl PackedInstruction {
         &mut self.extra_attrs
     }
 
+    /// Retrieves the cached py_gate immutably.
+    #[cfg(feature = "cache_pygates")]
+    pub fn py_op_mut(&mut self) -> &mut OnceLock<PyObject> {
+        &mut self.py_op
+    }
+
     /// Access the standard gate in this `PackedInstruction`, if it is one.  If the instruction
     /// refers to a Python-space object, `None` is returned.
     #[inline]