Use qubit-index newtypes in Rust space

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))
     }
 }