Fix pickling of NLayout (backport #12298) (#12300)

* Fix pickling of `NLayout` (#12298)

Previously the module of this was set incorrectly (stemming from its
move in gh-9064), at which point the `__getstate__`/`__setstate__`
pickling wouldn't work correctly any more.  Also, however, there was no
`__getnewargs__` and `new` didn't have a zero-argument form, so this
wouldn't have worked either.

(cherry picked from commit 136548c)

* Switch back to PyO3 0.20

---------

Co-authored-by: Jake Lishman <jake.lishman@ibm.com>

crates/accelerate/src/nlayout.rs

@@ -91,7 +91,7 @@ impl VirtualQubit {
 ///         physical qubit index on the coupling graph.
 ///     logical_qubits (int): The number of logical qubits in the layout
 ///     physical_qubits (int): The number of physical qubits in the layout
-#[pyclass(module = "qiskit._accelerate.stochastic_swap")]
+#[pyclass(module = "qiskit._accelerate.nlayout")]
 #[derive(Clone, Debug)]
 pub struct NLayout {
     virt_to_phys: Vec<PhysicalQubit>,
@@ -117,13 +117,12 @@ impl NLayout {
         res
     }
 
-    fn __getstate__(&self) -> (Vec<PhysicalQubit>, Vec<VirtualQubit>) {
-        (self.virt_to_phys.clone(), self.phys_to_virt.clone())
-    }
-
-    fn __setstate__(&mut self, state: (Vec<PhysicalQubit>, Vec<VirtualQubit>)) {
-        self.virt_to_phys = state.0;
-        self.phys_to_virt = state.1;
+    fn __reduce__(&self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok((
+            py.get_type::<Self>().getattr("from_virtual_to_physical")?,
+            (self.virt_to_phys.to_object(py),),
+        )
+            .into_py(py))
     }
 
     /// Return the layout mapping.

test/python/transpiler/test_layout.py

@@ -13,12 +13,14 @@
 """Tests the layout object"""
 
 import copy
+import pickle
 import unittest
 import numpy
 
 from qiskit.circuit import QuantumRegister, Qubit
 from qiskit.transpiler.layout import Layout
 from qiskit.transpiler.exceptions import LayoutError
+from qiskit._accelerate.nlayout import NLayout
 from test import QiskitTestCase  # pylint: disable=wrong-import-order
 
 
@@ -429,5 +431,54 @@ def test_layout_contains(self):
         self.assertNotIn(1, layout)
 
 
+class TestNLayout(QiskitTestCase):
+    """This is a private class, so mostly doesn't need direct tests."""
+
+    def test_pickle(self):
+        """Test that the layout roundtrips through pickle."""
+        v2p = [3, 5, 1, 2, 0, 4]
+        size = len(v2p)
+        layout = NLayout.from_virtual_to_physical(v2p)
+        self.assertEqual([layout.virtual_to_physical(x) for x in range(size)], v2p)
+        roundtripped = pickle.loads(pickle.dumps(layout))
+        self.assertEqual([roundtripped.virtual_to_physical(x) for x in range(size)], v2p)
+
+        # No changes to `layout`.
+        roundtripped.swap_virtual(0, 1)
+        expected = [5, 3, 1, 2, 0, 4]
+        self.assertEqual([layout.virtual_to_physical(x) for x in range(size)], v2p)
+        self.assertEqual([roundtripped.virtual_to_physical(x) for x in range(size)], expected)
+
+    def test_copy(self):
+        """Test that the layout roundtrips through copy."""
+        v2p = [3, 5, 1, 2, 0, 4]
+        size = len(v2p)
+        layout = NLayout.from_virtual_to_physical(v2p)
+        self.assertEqual([layout.virtual_to_physical(x) for x in range(size)], v2p)
+        roundtripped = copy.copy(layout)
+        self.assertEqual([roundtripped.virtual_to_physical(x) for x in range(size)], v2p)
+
+        # No changes to `layout`.
+        roundtripped.swap_virtual(0, 1)
+        expected = [5, 3, 1, 2, 0, 4]
+        self.assertEqual([layout.virtual_to_physical(x) for x in range(size)], v2p)
+        self.assertEqual([roundtripped.virtual_to_physical(x) for x in range(size)], expected)
+
+    def test_deepcopy(self):
+        """Test that the layout roundtrips through deepcopy."""
+        v2p = [3, 5, 1, 2, 0, 4]
+        size = len(v2p)
+        layout = NLayout.from_virtual_to_physical(v2p)
+        self.assertEqual([layout.virtual_to_physical(x) for x in range(size)], v2p)
+        roundtripped = copy.deepcopy(layout)
+        self.assertEqual([roundtripped.virtual_to_physical(x) for x in range(size)], v2p)
+
+        # No changes to `layout`.
+        roundtripped.swap_virtual(0, 1)
+        expected = [5, 3, 1, 2, 0, 4]
+        self.assertEqual([layout.virtual_to_physical(x) for x in range(size)], v2p)
+        self.assertEqual([roundtripped.virtual_to_physical(x) for x in range(size)], expected)
+
+
 if __name__ == "__main__":
     unittest.main()