Use IndexMap for performance boost in qubit allocation

backend/Cargo.toml

@@ -26,6 +26,10 @@ harness = false
 name = "grover"
 harness = false
 
+[[bench]]
+name = "qubit_management"
+harness = false
+
 [lib]
 crate-type = ["staticlib","rlib"]
 bench = false

backend/benches/qubit_management.rs

@@ -0,0 +1,21 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use criterion::{criterion_group, criterion_main, Criterion};
+use qir_backend::*;
+
+/// Benchmarks large number of qubit allocations and releases.
+pub fn allocate_release(c: &mut Criterion) {
+    c.bench_function("Allocate-Release 2k qubits", |b| {
+        b.iter(|| {
+            // Prepare a uniform superposition.
+            let qs = __quantum__rt__qubit_allocate_array(2_000);
+            unsafe {
+                __quantum__rt__qubit_release_array(qs);
+            }
+        })
+    });
+}
+
+criterion_group!(benches, allocate_release);
+criterion_main!(benches);

sparsesim/src/exp.rs

@@ -40,7 +40,7 @@ impl QuantumSim {
             .map(|c| {
                 *self
                     .id_map
-                    .get(c)
+                    .get(*c)
                     .unwrap_or_else(|| panic!("Unable to find qubit with id {c}"))
                     as u64
             })
@@ -51,7 +51,7 @@ impl QuantumSim {
             .map(|c| {
                 *self
                     .id_map
-                    .get(c)
+                    .get(*c)
                     .unwrap_or_else(|| panic!("Unable to find qubit with id {c}"))
                     as u64
             })

sparsesim/src/index_map.rs

@@ -0,0 +1,341 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use std::{
+    fmt::{self, Debug, Formatter},
+    iter::Enumerate,
+    marker::PhantomData,
+    ops::{Index, IndexMut},
+    option::Option,
+    slice, vec,
+};
+
+pub struct IndexMap<K, V> {
+    _keys: PhantomData<K>,
+    values: Vec<Option<V>>,
+}
+
+impl<K, V> IndexMap<K, V>
+where
+    K: Into<usize>,
+    V: Default,
+{
+    pub fn get_mut_or_default(&mut self, key: K) -> &mut V {
+        let index: usize = key.into();
+        if index >= self.values.len() {
+            self.values.resize_with(index + 1, Option::default);
+        }
+        self.values
+            .get_mut(index)
+            .expect("IndexMap::get_mut_or_default: index out of bounds")
+            .get_or_insert_with(Default::default)
+    }
+}
+
+impl<K, V> IndexMap<K, V> {
+    #[must_use]
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    #[must_use]
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            _keys: PhantomData,
+            values: Vec::with_capacity(capacity),
+        }
+    }
+
+    #[must_use]
+    pub fn is_empty(&self) -> bool {
+        self.values.is_empty()
+    }
+
+    // `Iter` does implement `Iterator`, but it has an additional bound on `K`.
+    #[allow(clippy::iter_not_returning_iterator)]
+    #[must_use]
+    pub fn iter(&self) -> Iter<K, V> {
+        Iter {
+            _keys: PhantomData,
+            base: self.values.iter().enumerate(),
+        }
+    }
+
+    // `Iter` does implement `Iterator`, but it has an additional bound on `K`.
+    #[allow(clippy::iter_not_returning_iterator)]
+    pub fn iter_mut(&mut self) -> IterMut<K, V> {
+        IterMut {
+            _keys: PhantomData,
+            base: self.values.iter_mut().enumerate(),
+        }
+    }
+
+    pub fn drain(&mut self) -> Drain<K, V> {
+        Drain {
+            _keys: PhantomData,
+            base: self.values.drain(..).enumerate(),
+        }
+    }
+
+    #[must_use]
+    pub fn values(&self) -> Values<V> {
+        Values {
+            base: self.values.iter(),
+        }
+    }
+
+    pub fn values_mut(&mut self) -> ValuesMut<V> {
+        ValuesMut {
+            base: self.values.iter_mut(),
+        }
+    }
+
+    pub fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(K, &V) -> bool,
+        K: From<usize>,
+    {
+        for (k, v) in self.values.iter_mut().enumerate() {
+            let remove = if let Some(value) = v {
+                !f(K::from(k), value)
+            } else {
+                false
+            };
+            if remove {
+                *v = None;
+            }
+        }
+    }
+
+    pub fn clear(&mut self) {
+        self.values.clear();
+    }
+}
+
+impl<K: Into<usize>, V> IndexMap<K, V> {
+    pub fn insert(&mut self, key: K, value: V) {
+        let index = key.into();
+        if index >= self.values.len() {
+            self.values.resize_with(index + 1, || None);
+        }
+        self.values[index] = Some(value);
+    }
+
+    pub fn contains_key(&self, key: K) -> bool {
+        let index: usize = key.into();
+        self.values.get(index).is_some_and(Option::is_some)
+    }
+
+    pub fn get(&self, key: K) -> Option<&V> {
+        let index: usize = key.into();
+        self.values.get(index).and_then(Option::as_ref)
+    }
+
+    pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
+        let index: usize = key.into();
+        self.values.get_mut(index).and_then(Option::as_mut)
+    }
+
+    pub fn remove(&mut self, key: K) {
+        let index: usize = key.into();
+        if index < self.values.len() {
+            self.values[index] = None;
+        }
+    }
+}
+
+impl<K, V: Clone> Clone for IndexMap<K, V> {
+    fn clone(&self) -> Self {
+        Self {
+            _keys: PhantomData,
+            values: self.values.clone(),
+        }
+    }
+}
+
+impl<K, V: Debug> Debug for IndexMap<K, V> {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        f.debug_struct("IndexMap")
+            .field(
+                "values",
+                &self
+                    .values
+                    .iter()
+                    .enumerate()
+                    .filter_map(|(k, v)| v.as_ref().map(|val| format!("{k:?}: {val:?}")))
+                    .collect::<Vec<_>>(),
+            )
+            .finish()
+    }
+}
+
+impl<K, V> Default for IndexMap<K, V> {
+    fn default() -> Self {
+        Self {
+            _keys: PhantomData,
+            values: Vec::default(),
+        }
+    }
+}
+
+impl<K: From<usize>, V> IntoIterator for IndexMap<K, V> {
+    type Item = (K, V);
+
+    type IntoIter = IntoIter<K, V>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        IntoIter {
+            _keys: PhantomData,
+            base: self.values.into_iter().enumerate(),
+        }
+    }
+}
+
+impl<'a, K: From<usize>, V> IntoIterator for &'a IndexMap<K, V> {
+    type Item = (K, &'a V);
+
+    type IntoIter = Iter<'a, K, V>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
+impl<K: Into<usize>, V> FromIterator<(K, V)> for IndexMap<K, V> {
+    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
+        let iter = iter.into_iter();
+        let mut map = Self::new();
+        let (lo, hi) = iter.size_hint();
+        map.values.reserve(hi.unwrap_or(lo));
+        for (key, value) in iter {
+            map.insert(key, value);
+        }
+        map
+    }
+}
+
+pub struct Iter<'a, K, V> {
+    _keys: PhantomData<K>,
+    base: Enumerate<slice::Iter<'a, Option<V>>>,
+}
+
+impl<'a, K: From<usize>, V> Iterator for Iter<'a, K, V> {
+    type Item = (K, &'a V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let (index, Some(value)) = self.base.next()? {
+                break Some((index.into(), value));
+            }
+        }
+    }
+}
+
+pub struct IterMut<'a, K, V> {
+    _keys: PhantomData<K>,
+    base: Enumerate<slice::IterMut<'a, Option<V>>>,
+}
+
+impl<K: From<usize>, V> DoubleEndedIterator for Iter<'_, K, V> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        loop {
+            if let (index, Some(value)) = self.base.next_back()? {
+                break Some((index.into(), value));
+            }
+        }
+    }
+}
+
+impl<'a, K: From<usize>, V> Iterator for IterMut<'a, K, V> {
+    type Item = (K, &'a mut V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let (index, Some(value)) = self.base.next()? {
+                break Some((index.into(), value));
+            }
+        }
+    }
+}
+
+pub struct IntoIter<K, V> {
+    _keys: PhantomData<K>,
+    base: Enumerate<vec::IntoIter<Option<V>>>,
+}
+
+impl<K: From<usize>, V> Iterator for IntoIter<K, V> {
+    type Item = (K, V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let (index, Some(value)) = self.base.next()? {
+                break Some((index.into(), value));
+            }
+        }
+    }
+}
+
+pub struct Drain<'a, K, V> {
+    _keys: PhantomData<K>,
+    base: Enumerate<vec::Drain<'a, Option<V>>>,
+}
+
+impl<K: From<usize>, V> Iterator for Drain<'_, K, V> {
+    type Item = (K, V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let (index, Some(value)) = self.base.next()? {
+                break Some((index.into(), value));
+            }
+        }
+    }
+}
+
+pub struct Values<'a, V> {
+    base: slice::Iter<'a, Option<V>>,
+}
+
+impl<'a, V> Iterator for Values<'a, V> {
+    type Item = &'a V;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let Some(value) = self.base.next()? {
+                break Some(value);
+            }
+        }
+    }
+}
+
+pub struct ValuesMut<'a, V> {
+    base: slice::IterMut<'a, Option<V>>,
+}
+
+impl<'a, V> Iterator for ValuesMut<'a, V> {
+    type Item = &'a mut V;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let Some(value) = self.base.next()? {
+                break Some(value);
+            }
+        }
+    }
+}
+
+impl Index<usize> for IndexMap<usize, usize> {
+    type Output = usize;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        self.get(index)
+            .expect("IndexMap::index: index out of bounds")
+    }
+}
+
+impl IndexMut<usize> for IndexMap<usize, usize> {
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+        self.get_mut(index)
+            .expect("IndexMap::index_mut: index out of bounds")
+    }
+}