Preserve bit locations through pickle (#13981)

* Preserve bit locations through pickle

Previously when we were pickling a DAGCircuit the bit locations fields
were forgotten about. So when we loaded a DAGCircuit from a pickle the
bit locations were empty. This would cause any call to find_bit() to
raise an error because there was no entry for any of the bits in the
dag. This commit fixes this by reconstructing by including the bit
locations fields in the object state returned by __getstate__ and
populating the fields from the provided state in __setstate__/

Fixes #13976

* Support pickle for BitLocations

* Adjust pickle impl

* Make BitLocations importable from Python so pickle can use it

Author: mtreinish

crates/circuit/src/dag_circuit.rs

@@ -346,13 +346,25 @@ fn reject_new_register(reg: &Bound<PyAny>) -> PyResult<()> {
 
 #[pyclass(module = "qiskit._accelerate.circuit")]
 #[derive(Clone, Debug)]
-struct BitLocations {
+pub(crate) struct BitLocations {
     #[pyo3(get)]
     index: usize,
     #[pyo3(get)]
     registers: Py<PyList>,
 }
 
+#[pymethods]
+impl BitLocations {
+    #[new]
+    fn new(index: usize, registers: Py<PyList>) -> Self {
+        Self { index, registers }
+    }
+
+    fn __getnewargs__(&self, py: Python) -> (usize, Py<PyList>) {
+        (self.index, self.registers.clone_ref(py))
+    }
+}
+
 #[derive(Copy, Clone, Debug)]
 enum DAGVarType {
     Input = 0,
@@ -528,6 +540,8 @@ impl DAGCircuit {
         out_dict.set_item("qregs", self.qregs.clone_ref(py))?;
         out_dict.set_item("cregs", self.cregs.clone_ref(py))?;
         out_dict.set_item("global_phase", self.global_phase.clone())?;
+        out_dict.set_item("qubit_locations", self.qubit_locations.clone_ref(py))?;
+        out_dict.set_item("clbit_locations", self.clbit_locations.clone_ref(py))?;
         out_dict.set_item(
             "qubit_io_map",
             self.qubit_io_map
@@ -617,6 +631,8 @@ impl DAGCircuit {
         self.qregs = dict_state.get_item("qregs")?.unwrap().extract()?;
         self.cregs = dict_state.get_item("cregs")?.unwrap().extract()?;
         self.global_phase = dict_state.get_item("global_phase")?.unwrap().extract()?;
+        self.qubit_locations = dict_state.get_item("qubit_locations")?.unwrap().extract()?;
+        self.clbit_locations = dict_state.get_item("clbit_locations")?.unwrap().extract()?;
         self.op_names = dict_state.get_item("op_name")?.unwrap().extract()?;
         self.vars_by_type = dict_state.get_item("vars_by_type")?.unwrap().extract()?;
         let binding = dict_state.get_item("vars_info")?.unwrap();

crates/circuit/src/lib.rs

@@ -125,6 +125,7 @@ impl From<Clbit> for BitType {
 pub fn circuit(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_class::<circuit_data::CircuitData>()?;
     m.add_class::<circuit_instruction::CircuitInstruction>()?;
+    m.add_class::<dag_circuit::BitLocations>()?;
     m.add_class::<dag_circuit::DAGCircuit>()?;
     m.add_class::<dag_node::DAGNode>()?;
     m.add_class::<dag_node::DAGInNode>()?;

test/python/dagcircuit/test_dagcircuit.py

@@ -15,6 +15,9 @@
 from __future__ import annotations
 
 from collections import Counter
+import io
+import copy
+import pickle
 import unittest
 
 from ddt import ddt, data
@@ -161,6 +164,80 @@ def test_dag_get_qubits(self):
             ],
         )
 
+    def test_pickle_bit_locations_with_reg(self):
+        """Test bit locations preserved through pickle."""
+        dag = DAGCircuit()
+        qr = QuantumRegister(2, "qr")
+        cr = ClassicalRegister(1, "cr")
+        dag.add_qreg(qr)
+        dag.add_creg(cr)
+        self.assertEqual(dag.find_bit(dag.qubits[1]).index, 1)
+        self.assertEqual(dag.find_bit(dag.qubits[1]).registers, [(qr, 1)])
+        self.assertEqual(dag.find_bit(dag.clbits[0]).index, 0)
+        self.assertEqual(dag.find_bit(dag.clbits[0]).registers, [(cr, 0)])
+        with io.BytesIO() as buf:
+            pickle.dump(dag, buf)
+            buf.seek(0)
+            output = pickle.load(buf)
+        self.assertEqual(output.find_bit(output.qubits[1]).index, 1)
+        self.assertEqual(output.find_bit(output.qubits[1]).registers, [(qr, 1)])
+        self.assertEqual(output.find_bit(output.clbits[0]).index, 0)
+        self.assertEqual(output.find_bit(output.clbits[0]).registers, [(cr, 0)])
+
+    def test_deepcopy_bit_locations_with_reg(self):
+        """Test bit locations preserved through pickle."""
+        dag = DAGCircuit()
+        qr = QuantumRegister(2, "qr")
+        cr = ClassicalRegister(1, "cr")
+        dag.add_qreg(qr)
+        dag.add_creg(cr)
+        self.assertEqual(dag.find_bit(dag.qubits[1]).index, 1)
+        self.assertEqual(dag.find_bit(dag.qubits[1]).registers, [(qr, 1)])
+        self.assertEqual(dag.find_bit(dag.clbits[0]).index, 0)
+        self.assertEqual(dag.find_bit(dag.clbits[0]).registers, [(cr, 0)])
+        output = copy.deepcopy(dag)
+        self.assertEqual(output.find_bit(output.qubits[1]).index, 1)
+        self.assertEqual(output.find_bit(output.qubits[1]).registers, [(qr, 1)])
+        self.assertEqual(output.find_bit(output.clbits[0]).index, 0)
+        self.assertEqual(output.find_bit(output.clbits[0]).registers, [(cr, 0)])
+
+    def test_pickle_bit_locations_with_no_reg(self):
+        """Test bit locations preserved through pickle."""
+        dag = DAGCircuit()
+        qubits = [Qubit(), Qubit()]
+        clbits = [Clbit()]
+        dag.add_qubits(qubits)
+        dag.add_clbits(clbits)
+        self.assertEqual(dag.find_bit(dag.qubits[1]).index, 1)
+        self.assertEqual(dag.find_bit(dag.qubits[1]).registers, [])
+        self.assertEqual(dag.find_bit(dag.clbits[0]).index, 0)
+        self.assertEqual(dag.find_bit(dag.clbits[0]).registers, [])
+        with io.BytesIO() as buf:
+            pickle.dump(dag, buf)
+            buf.seek(0)
+            output = pickle.load(buf)
+        self.assertEqual(output.find_bit(output.qubits[1]).index, 1)
+        self.assertEqual(output.find_bit(output.qubits[1]).registers, [])
+        self.assertEqual(output.find_bit(output.clbits[0]).index, 0)
+        self.assertEqual(output.find_bit(output.clbits[0]).registers, [])
+
+    def test_deepcopy_bit_locations_with_no_reg(self):
+        """Test bit locations preserved through pickle."""
+        dag = DAGCircuit()
+        qubits = [Qubit(), Qubit()]
+        clbits = [Clbit()]
+        dag.add_qubits(qubits)
+        dag.add_clbits(clbits)
+        self.assertEqual(dag.find_bit(dag.qubits[1]).index, 1)
+        self.assertEqual(dag.find_bit(dag.qubits[1]).registers, [])
+        self.assertEqual(dag.find_bit(dag.clbits[0]).index, 0)
+        self.assertEqual(dag.find_bit(dag.clbits[0]).registers, [])
+        output = copy.deepcopy(dag)
+        self.assertEqual(output.find_bit(output.qubits[1]).index, 1)
+        self.assertEqual(output.find_bit(output.qubits[1]).registers, [])
+        self.assertEqual(output.find_bit(output.clbits[0]).index, 0)
+        self.assertEqual(output.find_bit(output.clbits[0]).registers, [])
+
     def test_add_reg_duplicate(self):
         """add_qreg with the same register twice is not allowed."""
         dag = DAGCircuit()