Use qubit-index newtypes in Rust space

This converts all of our Rust-space components that are concerned with
virtual and physical qubits (via the `NLayout` class) to represent those
integers with newtypes wrapping them as half-size indices, and changes
the `NLayout` API to only allow access via these types and not by raw
integers.

The way this is done should have no overhead in time usage from Rust,
and is actually a reduction in memory usage because all the qubits are
stored at half the width they were previously (for most systems).  This
is done to add type safety to all our components that were concerned
with the mixing of these two qubits.  The implementation of this commit
already turned up the logical bug fixed by Qiskitgh-10756.

crates/accelerate/src/edge_collections.rs

@@ -12,14 +12,15 @@
 
 use numpy::IntoPyArray;
 use pyo3::prelude::*;
-use pyo3::Python;
+
+use crate::nlayout::PhysicalQubit;
 
 /// A simple container that contains a vector representing edges in the
 /// coupling map that are found to be optimal by the swap mapper.
 #[pyclass(module = "qiskit._accelerate.stochastic_swap")]
 #[derive(Clone, Debug)]
 pub struct EdgeCollection {
-    pub edges: Vec<usize>,
+    pub edges: Vec<PhysicalQubit>,
 }
 
 impl Default for EdgeCollection {
@@ -42,7 +43,7 @@ impl EdgeCollection {
     ///     edge_start (int): The beginning edge.
     ///     edge_end (int): The end of the edge.
     #[pyo3(text_signature = "(self, edge_start, edge_end, /)")]
-    pub fn add(&mut self, edge_start: usize, edge_end: usize) {
+    pub fn add(&mut self, edge_start: PhysicalQubit, edge_end: PhysicalQubit) {
         self.edges.push(edge_start);
         self.edges.push(edge_end);
     }
@@ -57,11 +58,11 @@ impl EdgeCollection {
         self.edges.clone().into_pyarray(py).into()
     }
 
-    fn __getstate__(&self) -> Vec<usize> {
+    fn __getstate__(&self) -> Vec<PhysicalQubit> {
         self.edges.clone()
     }
 
-    fn __setstate__(&mut self, state: Vec<usize>) {
+    fn __setstate__(&mut self, state: Vec<PhysicalQubit>) {
         self.edges = state
     }
 }

crates/accelerate/src/error_map.rs

@@ -13,6 +13,8 @@
 use pyo3::exceptions::PyIndexError;
 use pyo3::prelude::*;
 
+use crate::nlayout::PhysicalQubit;
+
 use hashbrown::HashMap;
 
 /// A mapping that represents the avg error rate for a particular edge in
@@ -34,7 +36,7 @@ use hashbrown::HashMap;
 #[pyclass(mapping, module = "qiskit._accelerate.error_map")]
 #[derive(Clone, Debug)]
 pub struct ErrorMap {
-    pub error_map: HashMap<[usize; 2], f64>,
+    pub error_map: HashMap<[PhysicalQubit; 2], f64>,
 }
 
 #[pymethods]
@@ -60,45 +62,45 @@ impl ErrorMap {
     /// construct the error map iteratively with :meth:`.add_error` instead of
     /// constructing an intermediate dict and using this constructor.
     #[staticmethod]
-    fn from_dict(error_map: HashMap<[usize; 2], f64>) -> Self {
+    fn from_dict(error_map: HashMap<[PhysicalQubit; 2], f64>) -> Self {
         ErrorMap { error_map }
     }
 
-    fn add_error(&mut self, index: [usize; 2], error_rate: f64) {
+    fn add_error(&mut self, index: [PhysicalQubit; 2], error_rate: f64) {
         self.error_map.insert(index, error_rate);
     }
 
-    // The pickle protocol methods can't return `HashMap<[usize; 2], f64>` to Python, because by
-    // PyO3's natural conversion as of 0.17.3 it will attempt to construct a `dict[list[int],
-    // float]`, where `list[int]` is unhashable in Python.
+    // The pickle protocol methods can't return `HashMap<[T; 2], f64>` to Python, because by PyO3's
+    // natural conversion as of 0.17.3 it will attempt to construct a `dict[list[T], float]`, where
+    // `list[T]` is unhashable in Python.
 
-    fn __getstate__(&self) -> HashMap<(usize, usize), f64> {
+    fn __getstate__(&self) -> HashMap<(PhysicalQubit, PhysicalQubit), f64> {
         self.error_map
             .iter()
             .map(|([a, b], value)| ((*a, *b), *value))
             .collect()
     }
 
-    fn __setstate__(&mut self, state: HashMap<[usize; 2], f64>) {
+    fn __setstate__(&mut self, state: HashMap<[PhysicalQubit; 2], f64>) {
         self.error_map = state;
     }
 
     fn __len__(&self) -> PyResult<usize> {
         Ok(self.error_map.len())
     }
 
-    fn __getitem__(&self, key: [usize; 2]) -> PyResult<f64> {
+    fn __getitem__(&self, key: [PhysicalQubit; 2]) -> PyResult<f64> {
         match self.error_map.get(&key) {
             Some(data) => Ok(*data),
             None => Err(PyIndexError::new_err("No node found for index")),
         }
     }
 
-    fn __contains__(&self, key: [usize; 2]) -> PyResult<bool> {
+    fn __contains__(&self, key: [PhysicalQubit; 2]) -> PyResult<bool> {
         Ok(self.error_map.contains_key(&key))
     }
 
-    fn get(&self, py: Python, key: [usize; 2], default: Option<PyObject>) -> PyObject {
+    fn get(&self, py: Python, key: [PhysicalQubit; 2], default: Option<PyObject>) -> PyObject {
         match self.error_map.get(&key).copied() {
             Some(val) => val.to_object(py),
             None => match default {

crates/accelerate/src/nlayout.rs

@@ -11,9 +11,76 @@
 // that they have been altered from the originals.
 
 use pyo3::prelude::*;
+use pyo3::types::PyList;
 
 use hashbrown::HashMap;
 
+/// A newtype for the different categories of qubits used within layouts.  This is to enforce
+/// significantly more type safety when dealing with mixtures of physical and virtual qubits, as we
+/// typically are when dealing with layouts.  In Rust space, `NLayout` only works in terms of the
+/// correct newtype, meaning that it's not possible to accidentally pass the wrong type of qubit to
+/// a lookup.  We can't enforce the same rules on integers in Python space without runtime
+/// overhead, so we just allow conversion to and from any valid `PyLong`.
+macro_rules! qubit_newtype {
+    ($id: ident) => {
+        #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+        pub struct $id(u32);
+
+        impl $id {
+            #[inline]
+            pub fn new(val: u32) -> Self {
+                Self(val)
+            }
+            #[inline]
+            pub fn index(&self) -> usize {
+                self.0 as usize
+            }
+        }
+
+        impl pyo3::IntoPy<PyObject> for $id {
+            fn into_py(self, py: Python<'_>) -> PyObject {
+                self.0.into_py(py)
+            }
+        }
+        impl pyo3::ToPyObject for $id {
+            fn to_object(&self, py: Python<'_>) -> PyObject {
+                self.0.to_object(py)
+            }
+        }
+
+        impl pyo3::FromPyObject<'_> for $id {
+            fn extract(ob: &PyAny) -> PyResult<Self> {
+                Ok(Self(ob.extract()?))
+            }
+        }
+
+        unsafe impl numpy::Element for $id {
+            const IS_COPY: bool = true;
+
+            fn get_dtype(py: Python<'_>) -> &numpy::PyArrayDescr {
+                u32::get_dtype(py)
+            }
+        }
+    };
+}
+
+qubit_newtype!(PhysicalQubit);
+impl PhysicalQubit {
+    /// Get the virtual qubit that currently corresponds to this index of physical qubit in the
+    /// given layout.
+    pub fn to_virt(self, layout: &NLayout) -> VirtualQubit {
+        layout.phys_to_virt[self.index()]
+    }
+}
+qubit_newtype!(VirtualQubit);
+impl VirtualQubit {
+    /// Get the physical qubit that currently corresponds to this index of virtual qubit in the
+    /// given layout.
+    pub fn to_phys(self, layout: &NLayout) -> PhysicalQubit {
+        layout.virt_to_phys[self.index()]
+    }
+}
+
 /// An unsigned integer Vector based layout class
 ///
 /// This class tracks the layout (or mapping between virtual qubits in the the
@@ -27,114 +94,125 @@ use hashbrown::HashMap;
 #[pyclass(module = "qiskit._accelerate.stochastic_swap")]
 #[derive(Clone, Debug)]
 pub struct NLayout {
-    pub logic_to_phys: Vec<usize>,
-    pub phys_to_logic: Vec<usize>,
-}
-
-impl NLayout {
-    pub fn swap(&mut self, idx1: usize, idx2: usize) {
-        self.phys_to_logic.swap(idx1, idx2);
-        self.logic_to_phys[self.phys_to_logic[idx1]] = idx1;
-        self.logic_to_phys[self.phys_to_logic[idx2]] = idx2;
-    }
+    virt_to_phys: Vec<PhysicalQubit>,
+    phys_to_virt: Vec<VirtualQubit>,
 }
 
 #[pymethods]
 impl NLayout {
     #[new]
-    #[pyo3(text_signature = "(qubit_indices, logical_qubits, physical_qubits, /)")]
     fn new(
-        qubit_indices: HashMap<usize, usize>,
-        logical_qubits: usize,
+        qubit_indices: HashMap<VirtualQubit, PhysicalQubit>,
+        virtual_qubits: usize,
         physical_qubits: usize,
     ) -> Self {
         let mut res = NLayout {
-            logic_to_phys: vec![std::usize::MAX; logical_qubits],
-            phys_to_logic: vec![std::usize::MAX; physical_qubits],
+            virt_to_phys: vec![PhysicalQubit(std::u32::MAX); virtual_qubits],
+            phys_to_virt: vec![VirtualQubit(std::u32::MAX); physical_qubits],
         };
-        for (key, value) in qubit_indices {
-            res.logic_to_phys[key] = value;
-            res.phys_to_logic[value] = key;
+        for (virt, phys) in qubit_indices {
+            res.virt_to_phys[virt.index()] = phys;
+            res.phys_to_virt[phys.index()] = virt;
         }
         res
     }
 
-    fn __getstate__(&self) -> [Vec<usize>; 2] {
-        [self.logic_to_phys.clone(), self.phys_to_logic.clone()]
+    fn __getstate__(&self) -> (Vec<PhysicalQubit>, Vec<VirtualQubit>) {
+        (self.virt_to_phys.clone(), self.phys_to_virt.clone())
     }
 
-    fn __setstate__(&mut self, state: [Vec<usize>; 2]) {
-        self.logic_to_phys = state[0].clone();
-        self.phys_to_logic = state[1].clone();
+    fn __setstate__(&mut self, state: (Vec<PhysicalQubit>, Vec<VirtualQubit>)) {
+        self.virt_to_phys = state.0;
+        self.phys_to_virt = state.1;
     }
 
-    /// Return the layout mapping
+    /// Return the layout mapping.
     ///
     /// .. note::
     ///
-    ///     this copies the data from Rust to Python and has linear
-    ///     overhead based on the number of qubits.
+    ///     This copies the data from Rust to Python and has linear overhead based on the number of
+    ///     qubits.
     ///
     /// Returns:
-    ///     list: A list of 2 element lists in the form:
-    ///     ``[[logical_qubit, physical_qubit], ...]``. Where the logical qubit
-    ///     is the index in the qubit index in the circuit.
+    ///     list: A list of 2 element lists in the form ``[(virtual_qubit, physical_qubit), ...]``,
+    ///     where the virtual qubit is the index in the qubit index in the circuit.
     ///
     #[pyo3(text_signature = "(self, /)")]
-    fn layout_mapping(&self) -> Vec<[usize; 2]> {
-        (0..self.logic_to_phys.len())
-            .map(|i| [i, self.logic_to_phys[i]])
-            .collect()
+    fn layout_mapping(&self, py: Python<'_>) -> Py<PyList> {
+        PyList::new(py, self.iter_virtual()).into()
     }
 
-    /// Get physical bit from logical bit
-    #[pyo3(text_signature = "(self, logical_bit, /)")]
-    fn logical_to_physical(&self, logical_bit: usize) -> usize {
-        self.logic_to_phys[logical_bit]
+    /// Get physical bit from virtual bit
+    #[pyo3(text_signature = "(self, virtual, /)")]
+    pub fn virtual_to_physical(&self, r#virtual: VirtualQubit) -> PhysicalQubit {
+        self.virt_to_phys[r#virtual.index()]
     }
 
-    /// Get logical bit from physical bit
-    #[pyo3(text_signature = "(self, physical_bit, /)")]
-    pub fn physical_to_logical(&self, physical_bit: usize) -> usize {
-        self.phys_to_logic[physical_bit]
+    /// Get virtual bit from physical bit
+    #[pyo3(text_signature = "(self, physical, /)")]
+    pub fn physical_to_virtual(&self, physical: PhysicalQubit) -> VirtualQubit {
+        self.phys_to_virt[physical.index()]
     }
 
     /// Swap the specified virtual qubits
     #[pyo3(text_signature = "(self, bit_a, bit_b, /)")]
-    pub fn swap_logical(&mut self, bit_a: usize, bit_b: usize) {
-        self.logic_to_phys.swap(bit_a, bit_b);
-        self.phys_to_logic[self.logic_to_phys[bit_a]] = bit_a;
-        self.phys_to_logic[self.logic_to_phys[bit_b]] = bit_b;
+    pub fn swap_virtual(&mut self, bit_a: VirtualQubit, bit_b: VirtualQubit) {
+        self.virt_to_phys.swap(bit_a.index(), bit_b.index());
+        self.phys_to_virt[self.virt_to_phys[bit_a.index()].index()] = bit_a;
+        self.phys_to_virt[self.virt_to_phys[bit_b.index()].index()] = bit_b;
     }
 
     /// Swap the specified physical qubits
     #[pyo3(text_signature = "(self, bit_a, bit_b, /)")]
-    pub fn swap_physical(&mut self, bit_a: usize, bit_b: usize) {
-        self.swap(bit_a, bit_b)
+    pub fn swap_physical(&mut self, bit_a: PhysicalQubit, bit_b: PhysicalQubit) {
+        self.phys_to_virt.swap(bit_a.index(), bit_b.index());
+        self.virt_to_phys[self.phys_to_virt[bit_a.index()].index()] = bit_a;
+        self.virt_to_phys[self.phys_to_virt[bit_b.index()].index()] = bit_b;
     }
 
     pub fn copy(&self) -> NLayout {
         self.clone()
     }
 
     #[staticmethod]
-    pub fn generate_trivial_layout(num_qubits: usize) -> Self {
+    pub fn generate_trivial_layout(num_qubits: u32) -> Self {
         NLayout {
-            logic_to_phys: (0..num_qubits).collect(),
-            phys_to_logic: (0..num_qubits).collect(),
+            virt_to_phys: (0..num_qubits).map(PhysicalQubit).collect(),
+            phys_to_virt: (0..num_qubits).map(VirtualQubit).collect(),
         }
     }
 
     #[staticmethod]
-    pub fn from_logical_to_physical(logic_to_phys: Vec<usize>) -> Self {
-        let mut phys_to_logic = vec![std::usize::MAX; logic_to_phys.len()];
-        for (logic, phys) in logic_to_phys.iter().enumerate() {
-            phys_to_logic[*phys] = logic;
-        }
-        NLayout {
-            logic_to_phys,
-            phys_to_logic,
+    pub fn from_virtual_to_physical(virt_to_phys: Vec<PhysicalQubit>) -> PyResult<Self> {
+        let mut phys_to_virt = vec![VirtualQubit(std::u32::MAX); virt_to_phys.len()];
+        for (virt, phys) in virt_to_phys.iter().enumerate() {
+            phys_to_virt[phys.index()] = VirtualQubit(virt.try_into()?);
         }
+        Ok(NLayout {
+            virt_to_phys,
+            phys_to_virt,
+        })
+    }
+}
+
+impl NLayout {
+    /// Iterator of `(VirtualQubit, PhysicalQubit)` pairs, in order of the `VirtualQubit` indices.
+    pub fn iter_virtual(
+        &'_ self,
+    ) -> impl ExactSizeIterator<Item = (VirtualQubit, PhysicalQubit)> + '_ {
+        self.virt_to_phys
+            .iter()
+            .enumerate()
+            .map(|(v, p)| (VirtualQubit::new(v as u32), *p))
+    }
+    /// Iterator of `(PhysicalQubit, VirtualQubit)` pairs, in order of the `PhysicalQubit` indices.
+    pub fn iter_physical(
+        &'_ self,
+    ) -> impl ExactSizeIterator<Item = (PhysicalQubit, VirtualQubit)> + '_ {
+        self.phys_to_virt
+            .iter()
+            .enumerate()
+            .map(|(p, v)| (PhysicalQubit::new(p as u32), *v))
     }
 }