[DAGCircuit Oxidation] Refactor bit management in CircuitData

crates/circuit/src/bit_data.rs

@@ -0,0 +1,163 @@
+// This code is part of Qiskit.
+//
+// (C) Copyright IBM 2024
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+use crate::BitType;
+use hashbrown::HashMap;
+use pyo3::exceptions::{PyRuntimeError, PyValueError};
+use pyo3::prelude::*;
+use pyo3::types::PyList;
+use std::hash::{Hash, Hasher};
+
+/// Private wrapper for Python-side Bit instances that implements
+/// [Hash] and [Eq], allowing them to be used in Rust hash-based
+/// sets and maps.
+///
+/// Python's `hash()` is called on the wrapped Bit instance during
+/// construction and returned from Rust's [Hash] trait impl.
+/// The impl of [PartialEq] first compares the native Py pointers
+/// to determine equality. If these are not equal, only then does
+/// it call `repr()` on both sides, which has a significant
+/// performance advantage.
+#[derive(Clone, Debug)]
+struct BitAsKey {
+    /// Python's `hash()` of the wrapped instance.
+    hash: isize,
+    /// The wrapped instance.
+    bit: PyObject,
+}
+
+impl BitAsKey {
+    pub fn new(bit: &Bound<PyAny>) -> Self {
+        BitAsKey {
+            // This really shouldn't fail, but if it does,
+            // we'll just use 0.
+            hash: bit.hash().unwrap_or(0),
+            bit: bit.into_py(bit.py()),
+        }
+    }
+}
+
+impl Hash for BitAsKey {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        state.write_isize(self.hash);
+    }
+}
+
+impl PartialEq for BitAsKey {
+    fn eq(&self, other: &Self) -> bool {
+        self.bit.is(&other.bit)
+            || Python::with_gil(|py| {
+                self.bit
+                    .bind(py)
+                    .repr()
+                    .unwrap()
+                    .eq(other.bit.bind(py).repr().unwrap())
+                    .unwrap()
+            })
+    }
+}
+
+impl Eq for BitAsKey {}
+
+#[derive(Clone, Debug)]
+pub(crate) struct BitData<T> {
+    /// The public field name (i.e. `qubits` or `clbits`).
+    description: String,
+    /// Registered Python bits.
+    bits: Vec<PyObject>,
+    /// Maps Python bits to native type.
+    indices: HashMap<BitAsKey, T>,
+    /// The bits registered, cached as a PyList.
+    cached: Py<PyList>,
+}
+
+impl<T> BitData<T>
+where
+    T: From<BitType>,
+    BitType: From<T>,
+{
+    pub fn new(py: Python<'_>, description: String) -> Self {
+        BitData {
+            description,
+            bits: Vec::new(),
+            indices: HashMap::new(),
+            cached: PyList::empty_bound(py).unbind(),
+        }
+    }
+
+    /// Gets the number of bits.
+    pub fn len(&self) -> usize {
+        self.bits.len()
+    }
+
+    /// Gets a reference to the underlying vector of Python bits.
+    #[inline]
+    pub fn bits(&self) -> &Vec<PyObject> {
+        &self.bits
+    }
+
+    /// Gets a reference to the cached Python list, maintained by
+    /// this instance.
+    #[inline]
+    pub fn cached(&self) -> &Py<PyList> {
+        &self.cached
+    }
+
+    /// Finds the native bit index of the given Python bit.
+    #[inline]
+    pub fn find(&self, bit: &Bound<PyAny>) -> Option<&T> {
+        self.indices.get(&BitAsKey::new(bit))
+    }
+
+    /// Gets the Python bit corresponding to the given native
+    /// bit index.
+    #[inline]
+    pub fn get(&self, index: T) -> Option<&PyObject> {
+        self.bits.get(<BitType as From<T>>::from(index) as usize)
+    }
+
+    /// Adds a new Python bit.
+    pub fn add(&mut self, py: Python, bit: &Bound<PyAny>, strict: bool) -> PyResult<()> {
+        if self.bits.len() != self.cached.bind(bit.py()).len() {
+            return Err(PyRuntimeError::new_err(
+            format!("This circuit's {} list has become out of sync with the circuit data. Did something modify it?", self.description)
+            ));
+        }
+        let idx: BitType = self.bits.len().try_into().map_err(|_| {
+            PyRuntimeError::new_err(format!(
+                "The number of {} in the circuit has exceeded the maximum capacity",
+                self.description
+            ))
+        })?;
+        if self
+            .indices
+            .try_insert(BitAsKey::new(bit), idx.into())
+            .is_ok()
+        {
+            self.bits.push(bit.into_py(py));
+            self.cached.bind(py).append(bit)?;
+        } else if strict {
+            return Err(PyValueError::new_err(format!(
+                "Existing bit {:?} cannot be re-added in strict mode.",
+                bit
+            )));
+        }
+        Ok(())
+    }
+
+    /// Called during Python garbage collection, only!.
+    /// Note: INVALIDATES THIS INSTANCE.
+    pub fn dispose(&mut self) {
+        self.indices.clear();
+        self.bits.clear();
+    }
+}