use alloc::collections::TryReserveError;
use core::mem::MaybeUninit;
use std::collections::hash_map::{self, HashMap};
use std::hash;
use std::iter::{Extend, FromIterator, FusedIterator};
use std::marker::PhantomData;
use std::ops::{Index, IndexMut};
use super::{Key, KeyData};
use crate::util::is_older_version;
#[derive(Debug, Clone)]
struct Slot<T> {
version: u32,
value: T,
}
#[derive(Debug, Clone)]
pub struct SparseSecondaryMap<K: Key, V, S: hash::BuildHasher = hash_map::RandomState> {
slots: HashMap<u32, Slot<V>, S>,
_k: PhantomData<fn(K) -> K>,
}
impl<K: Key, V> SparseSecondaryMap<K, V, hash_map::RandomState> {
pub fn new() -> Self {
Self::with_capacity(0)
}
pub fn with_capacity(capacity: usize) -> Self {
Self {
slots: HashMap::with_capacity(capacity),
_k: PhantomData,
}
}
}
impl<K: Key, V, S: hash::BuildHasher> SparseSecondaryMap<K, V, S> {
pub fn with_hasher(hash_builder: S) -> Self {
Self {
slots: HashMap::with_hasher(hash_builder),
_k: PhantomData,
}
}
pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> Self {
Self {
slots: HashMap::with_capacity_and_hasher(capacity, hash_builder),
_k: PhantomData,
}
}
pub fn len(&self) -> usize {
self.slots.len()
}
pub fn is_empty(&self) -> bool {
self.slots.is_empty()
}
pub fn capacity(&self) -> usize {
self.slots.capacity()
}
pub fn reserve(&mut self, additional: usize) {
self.slots.reserve(additional);
}
pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
self.slots.try_reserve(additional)
}
pub fn contains_key(&self, key: K) -> bool {
let kd = key.data();
self.slots
.get(&kd.idx)
.map_or(false, |slot| slot.version == kd.version.get())
}
pub fn insert(&mut self, key: K, value: V) -> Option<V> {
if key.is_null() {
return None;
}
let kd = key.data();
if let Some(slot) = self.slots.get_mut(&kd.idx) {
if slot.version == kd.version.get() {
return Some(std::mem::replace(&mut slot.value, value));
}
if is_older_version(kd.version.get(), slot.version) {
return None;
}
*slot = Slot {
version: kd.version.get(),
value,
};
return None;
}
self.slots.insert(
kd.idx,
Slot {
version: kd.version.get(),
value,
},
);
None
}
pub fn remove(&mut self, key: K) -> Option<V> {
let kd = key.data();
if let hash_map::Entry::Occupied(entry) = self.slots.entry(kd.idx) {
if entry.get().version == kd.version.get() {
return Some(entry.remove_entry().1.value);
}
}
None
}
pub fn retain<F>(&mut self, mut f: F)
where
F: FnMut(K, &mut V) -> bool,
{
self.slots.retain(|&idx, slot| {
let key = KeyData::new(idx, slot.version).into();
f(key, &mut slot.value)
})
}
pub fn clear(&mut self) {
self.slots.clear();
}
pub fn drain(&mut self) -> Drain<'_, K, V> {
Drain {
inner: self.slots.drain(),
_k: PhantomData,
}
}
pub fn get(&self, key: K) -> Option<&V> {
let kd = key.data();
self.slots
.get(&kd.idx)
.filter(|slot| slot.version == kd.version.get())
.map(|slot| &slot.value)
}
pub unsafe fn get_unchecked(&self, key: K) -> &V {
debug_assert!(self.contains_key(key));
self.get(key).unwrap_unchecked()
}
pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
let kd = key.data();
self.slots
.get_mut(&kd.idx)
.filter(|slot| slot.version == kd.version.get())
.map(|slot| &mut slot.value)
}
pub unsafe fn get_unchecked_mut(&mut self, key: K) -> &mut V {
debug_assert!(self.contains_key(key));
self.get_mut(key).unwrap_unchecked()
}
pub fn get_disjoint_mut<const N: usize>(&mut self, keys: [K; N]) -> Option<[&mut V; N]> {
let mut ptrs: [MaybeUninit<*mut V>; N] = [(); N].map(|_| MaybeUninit::uninit());
let mut i = 0;
while i < N {
let kd = keys[i].data();
match self.slots.get_mut(&kd.idx) {
Some(Slot { version, value }) if *version == kd.version.get() => {
ptrs[i] = MaybeUninit::new(&mut *value);
*version ^= 1;
},
_ => break,
}
i += 1;
}
for k in &keys[0..i] {
match self.slots.get_mut(&k.data().idx) {
Some(Slot { version, .. }) => {
*version ^= 1;
},
_ => unsafe { core::hint::unreachable_unchecked() },
}
}
if i == N {
Some(unsafe { core::mem::transmute_copy::<_, [&mut V; N]>(&ptrs) })
} else {
None
}
}
pub unsafe fn get_disjoint_unchecked_mut<const N: usize>(
&mut self,
keys: [K; N],
) -> [&mut V; N] {
let mut ptrs: [MaybeUninit<*mut V>; N] = MaybeUninit::uninit().assume_init();
for i in 0..N {
ptrs[i] = MaybeUninit::new(self.get_unchecked_mut(keys[i]));
}
core::mem::transmute_copy::<_, [&mut V; N]>(&ptrs)
}
pub fn iter(&self) -> Iter<'_, K, V> {
Iter {
inner: self.slots.iter(),
_k: PhantomData,
}
}
pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
IterMut {
inner: self.slots.iter_mut(),
_k: PhantomData,
}
}
pub fn keys(&self) -> Keys<'_, K, V> {
Keys { inner: self.iter() }
}
pub fn values(&self) -> Values<'_, K, V> {
Values { inner: self.iter() }
}
pub fn values_mut(&mut self) -> ValuesMut<'_, K, V> {
ValuesMut {
inner: self.iter_mut(),
}
}
pub fn entry(&mut self, key: K) -> Option<Entry<'_, K, V>> {
if key.is_null() {
return None;
}
let kd = key.data();
if let hash_map::Entry::Occupied(o) = self.slots.entry(kd.idx) {
if o.get().version != kd.version.get() {
if is_older_version(o.get().version, kd.version.get()) {
o.remove();
} else {
return None;
}
}
}
Some(match self.slots.entry(kd.idx) {
hash_map::Entry::Occupied(inner) => {
Entry::Occupied(OccupiedEntry {
inner,
kd,
_k: PhantomData,
})
},
hash_map::Entry::Vacant(inner) => Entry::Vacant(VacantEntry {
inner,
kd,
_k: PhantomData,
}),
})
}
}
impl<K, V, S> Default for SparseSecondaryMap<K, V, S>
where
K: Key,
S: hash::BuildHasher + Default,
{
fn default() -> Self {
Self::with_hasher(Default::default())
}
}
impl<K, V, S> Index<K> for SparseSecondaryMap<K, V, S>
where
K: Key,
S: hash::BuildHasher,
{
type Output = V;
fn index(&self, key: K) -> &V {
match self.get(key) {
Some(r) => r,
None => panic!("invalid SparseSecondaryMap key used"),
}
}
}
impl<K, V, S> IndexMut<K> for SparseSecondaryMap<K, V, S>
where
K: Key,
S: hash::BuildHasher,
{
fn index_mut(&mut self, key: K) -> &mut V {
match self.get_mut(key) {
Some(r) => r,
None => panic!("invalid SparseSecondaryMap key used"),
}
}
}
impl<K, V, S> PartialEq for SparseSecondaryMap<K, V, S>
where
K: Key,
V: PartialEq,
S: hash::BuildHasher,
{
fn eq(&self, other: &Self) -> bool {
if self.len() != other.len() {
return false;
}
self.iter().all(|(key, value)| {
other
.get(key)
.map_or(false, |other_value| *value == *other_value)
})
}
}
impl<K, V, S> Eq for SparseSecondaryMap<K, V, S>
where
K: Key,
V: Eq,
S: hash::BuildHasher,
{
}
impl<K, V, S> FromIterator<(K, V)> for SparseSecondaryMap<K, V, S>
where
K: Key,
S: hash::BuildHasher + Default,
{
fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
let mut sec = Self::default();
sec.extend(iter);
sec
}
}
impl<K, V, S> Extend<(K, V)> for SparseSecondaryMap<K, V, S>
where
K: Key,
S: hash::BuildHasher,
{
fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
let iter = iter.into_iter();
for (k, v) in iter {
self.insert(k, v);
}
}
}
impl<'a, K, V, S> Extend<(K, &'a V)> for SparseSecondaryMap<K, V, S>
where
K: Key,
V: 'a + Copy,
S: hash::BuildHasher,
{
fn extend<I: IntoIterator<Item = (K, &'a V)>>(&mut self, iter: I) {
let iter = iter.into_iter();
for (k, v) in iter {
self.insert(k, *v);
}
}
}
#[derive(Debug)]
pub struct OccupiedEntry<'a, K: Key, V> {
inner: hash_map::OccupiedEntry<'a, u32, Slot<V>>,
kd: KeyData,
_k: PhantomData<fn(K) -> K>,
}
#[derive(Debug)]
pub struct VacantEntry<'a, K: Key, V> {
inner: hash_map::VacantEntry<'a, u32, Slot<V>>,
kd: KeyData,
_k: PhantomData<fn(K) -> K>,
}
#[derive(Debug)]
pub enum Entry<'a, K: Key, V> {
Occupied(OccupiedEntry<'a, K, V>),
Vacant(VacantEntry<'a, K, V>),
}
impl<'a, K: Key, V> Entry<'a, K, V> {
pub fn or_insert(self, default: V) -> &'a mut V {
self.or_insert_with(|| default)
}
pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> &'a mut V {
match self {
Entry::Occupied(x) => x.into_mut(),
Entry::Vacant(x) => x.insert(default()),
}
}
pub fn key(&self) -> K {
match self {
Entry::Occupied(entry) => entry.kd.into(),
Entry::Vacant(entry) => entry.kd.into(),
}
}
pub fn and_modify<F>(self, f: F) -> Self
where
F: FnOnce(&mut V),
{
match self {
Entry::Occupied(mut entry) => {
f(entry.get_mut());
Entry::Occupied(entry)
},
Entry::Vacant(entry) => Entry::Vacant(entry),
}
}
}
impl<'a, K: Key, V: Default> Entry<'a, K, V> {
pub fn or_default(self) -> &'a mut V {
self.or_insert_with(Default::default)
}
}
impl<'a, K: Key, V> OccupiedEntry<'a, K, V> {
pub fn key(&self) -> K {
self.kd.into()
}
pub fn remove_entry(self) -> (K, V) {
(self.kd.into(), self.remove())
}
pub fn get(&self) -> &V {
&self.inner.get().value
}
pub fn get_mut(&mut self) -> &mut V {
&mut self.inner.get_mut().value
}
pub fn into_mut(self) -> &'a mut V {
&mut self.inner.into_mut().value
}
pub fn insert(&mut self, value: V) -> V {
std::mem::replace(self.get_mut(), value)
}
pub fn remove(self) -> V {
self.inner.remove().value
}
}
impl<'a, K: Key, V> VacantEntry<'a, K, V> {
pub fn key(&self) -> K {
self.kd.into()
}
pub fn insert(self, value: V) -> &'a mut V {
&mut self
.inner
.insert(Slot {
version: self.kd.version.get(),
value,
})
.value
}
}
#[derive(Debug)]
pub struct Drain<'a, K: Key + 'a, V: 'a> {
inner: hash_map::Drain<'a, u32, Slot<V>>,
_k: PhantomData<fn(K) -> K>,
}
#[derive(Debug)]
pub struct IntoIter<K: Key, V> {
inner: hash_map::IntoIter<u32, Slot<V>>,
_k: PhantomData<fn(K) -> K>,
}
#[derive(Debug)]
pub struct Iter<'a, K: Key + 'a, V: 'a> {
inner: hash_map::Iter<'a, u32, Slot<V>>,
_k: PhantomData<fn(K) -> K>,
}
impl<'a, K: 'a + Key, V: 'a> Clone for Iter<'a, K, V> {
fn clone(&self) -> Self {
Iter {
inner: self.inner.clone(),
_k: self._k,
}
}
}
#[derive(Debug)]
pub struct IterMut<'a, K: Key + 'a, V: 'a> {
inner: hash_map::IterMut<'a, u32, Slot<V>>,
_k: PhantomData<fn(K) -> K>,
}
#[derive(Debug)]
pub struct Keys<'a, K: Key + 'a, V: 'a> {
inner: Iter<'a, K, V>,
}
impl<'a, K: 'a + Key, V: 'a> Clone for Keys<'a, K, V> {
fn clone(&self) -> Self {
Keys {
inner: self.inner.clone(),
}
}
}
#[derive(Debug)]
pub struct Values<'a, K: Key + 'a, V: 'a> {
inner: Iter<'a, K, V>,
}
impl<'a, K: 'a + Key, V: 'a> Clone for Values<'a, K, V> {
fn clone(&self) -> Self {
Values {
inner: self.inner.clone(),
}
}
}
#[derive(Debug)]
pub struct ValuesMut<'a, K: Key + 'a, V: 'a> {
inner: IterMut<'a, K, V>,
}
impl<'a, K: Key, V> Iterator for Drain<'a, K, V> {
type Item = (K, V);
fn next(&mut self) -> Option<(K, V)> {
self.inner.next().map(|(idx, slot)| {
let key = KeyData::new(idx, slot.version).into();
(key, slot.value)
})
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<'a, K: Key, V> Drop for Drain<'a, K, V> {
fn drop(&mut self) {
self.for_each(|_drop| {});
}
}
impl<K: Key, V> Iterator for IntoIter<K, V> {
type Item = (K, V);
fn next(&mut self) -> Option<(K, V)> {
self.inner.next().map(|(idx, slot)| {
let key = KeyData::new(idx, slot.version).into();
(key, slot.value)
})
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<'a, K: Key, V> Iterator for Iter<'a, K, V> {
type Item = (K, &'a V);
fn next(&mut self) -> Option<(K, &'a V)> {
self.inner.next().map(|(&idx, slot)| {
let key = KeyData::new(idx, slot.version).into();
(key, &slot.value)
})
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<'a, K: Key, V> Iterator for IterMut<'a, K, V> {
type Item = (K, &'a mut V);
fn next(&mut self) -> Option<(K, &'a mut V)> {
self.inner.next().map(|(&idx, slot)| {
let key = KeyData::new(idx, slot.version).into();
(key, &mut slot.value)
})
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<'a, K: Key, V> Iterator for Keys<'a, K, V> {
type Item = K;
fn next(&mut self) -> Option<K> {
self.inner.next().map(|(key, _)| key)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<'a, K: Key, V> Iterator for Values<'a, K, V> {
type Item = &'a V;
fn next(&mut self) -> Option<&'a V> {
self.inner.next().map(|(_, value)| value)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<'a, K: Key, V> Iterator for ValuesMut<'a, K, V> {
type Item = &'a mut V;
fn next(&mut self) -> Option<&'a mut V> {
self.inner.next().map(|(_, value)| value)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<'a, K, V, S> IntoIterator for &'a SparseSecondaryMap<K, V, S>
where
K: Key,
S: hash::BuildHasher,
{
type Item = (K, &'a V);
type IntoIter = Iter<'a, K, V>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'a, K, V, S> IntoIterator for &'a mut SparseSecondaryMap<K, V, S>
where
K: Key,
S: hash::BuildHasher,
{
type Item = (K, &'a mut V);
type IntoIter = IterMut<'a, K, V>;
fn into_iter(self) -> Self::IntoIter {
self.iter_mut()
}
}
impl<K, V, S> IntoIterator for SparseSecondaryMap<K, V, S>
where
K: Key,
S: hash::BuildHasher,
{
type Item = (K, V);
type IntoIter = IntoIter<K, V>;
fn into_iter(self) -> Self::IntoIter {
IntoIter {
inner: self.slots.into_iter(),
_k: PhantomData,
}
}
}
impl<'a, K: Key, V> FusedIterator for Iter<'a, K, V> {}
impl<'a, K: Key, V> FusedIterator for IterMut<'a, K, V> {}
impl<'a, K: Key, V> FusedIterator for Keys<'a, K, V> {}
impl<'a, K: Key, V> FusedIterator for Values<'a, K, V> {}
impl<'a, K: Key, V> FusedIterator for ValuesMut<'a, K, V> {}
impl<'a, K: Key, V> FusedIterator for Drain<'a, K, V> {}
impl<K: Key, V> FusedIterator for IntoIter<K, V> {}
impl<'a, K: Key, V> ExactSizeIterator for Iter<'a, K, V> {}
impl<'a, K: Key, V> ExactSizeIterator for IterMut<'a, K, V> {}
impl<'a, K: Key, V> ExactSizeIterator for Keys<'a, K, V> {}
impl<'a, K: Key, V> ExactSizeIterator for Values<'a, K, V> {}
impl<'a, K: Key, V> ExactSizeIterator for ValuesMut<'a, K, V> {}
impl<'a, K: Key, V> ExactSizeIterator for Drain<'a, K, V> {}
impl<K: Key, V> ExactSizeIterator for IntoIter<K, V> {}
#[cfg(feature = "serde")]
mod serialize {
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use super::*;
use crate::SecondaryMap;
impl<K, V, H> Serialize for SparseSecondaryMap<K, V, H>
where
K: Key,
V: Serialize,
H: hash::BuildHasher,
{
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut serde_sec = SecondaryMap::new();
for (k, v) in self {
serde_sec.insert(k, v);
}
serde_sec.serialize(serializer)
}
}
impl<'de, K, V, S> Deserialize<'de> for SparseSecondaryMap<K, V, S>
where
K: Key,
V: Deserialize<'de>,
S: hash::BuildHasher + Default,
{
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let serde_sec: SecondaryMap<K, V> = Deserialize::deserialize(deserializer)?;
let mut sec = Self::default();
for (k, v) in serde_sec {
sec.insert(k, v);
}
Ok(sec)
}
}
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use quickcheck::quickcheck;
use crate::*;
#[test]
fn custom_hasher() {
type FastSparseSecondaryMap<K, V> = SparseSecondaryMap<K, V, fxhash::FxBuildHasher>;
let mut sm = SlotMap::new();
let mut sec = FastSparseSecondaryMap::default();
let key1 = sm.insert(42);
sec.insert(key1, 1234);
assert_eq!(sec[key1], 1234);
assert_eq!(sec.len(), 1);
let sec2 = sec
.iter()
.map(|(k, &v)| (k, v))
.collect::<FastSparseSecondaryMap<_, _>>();
assert_eq!(sec, sec2);
}
#[test]
fn disjoint() {
let mut sm = SlotMap::new();
let mut sec = SparseSecondaryMap::new();
for i in 0..20usize {
sm.insert(i);
}
sm.retain(|_, i| *i % 2 == 0);
for (i, k) in sm.keys().enumerate() {
sec.insert(k, i);
}
let keys: Vec<_> = sm.keys().collect();
for i in 0..keys.len() {
for j in 0..keys.len() {
if let Some([r0, r1]) = sec.get_disjoint_mut([keys[i], keys[j]]) {
*r0 ^= *r1;
*r1 = r1.wrapping_add(*r0);
} else {
assert!(i == j);
}
}
}
for i in 0..keys.len() {
for j in 0..keys.len() {
for k in 0..keys.len() {
if let Some([r0, r1, r2]) = sec.get_disjoint_mut([keys[i], keys[j], keys[k]]) {
*r0 ^= *r1;
*r0 = r0.wrapping_add(*r2);
*r1 ^= *r0;
*r1 = r1.wrapping_add(*r2);
*r2 ^= *r0;
*r2 = r2.wrapping_add(*r1);
} else {
assert!(i == j || j == k || i == k);
}
}
}
}
}
quickcheck! {
fn qc_secmap_equiv_hashmap(operations: Vec<(u8, u32)>) -> bool {
let mut hm = HashMap::new();
let mut hm_keys = Vec::new();
let mut unique_key = 0u32;
let mut sm = SlotMap::new();
let mut sec = SparseSecondaryMap::new();
let mut sm_keys = Vec::new();
#[cfg(not(feature = "serde"))]
let num_ops = 4;
#[cfg(feature = "serde")]
let num_ops = 5;
for (op, val) in operations {
match op % num_ops {
0 => {
hm.insert(unique_key, val);
hm_keys.push(unique_key);
unique_key += 1;
let k = sm.insert(val);
sec.insert(k, val);
sm_keys.push(k);
}
1 => {
if hm_keys.is_empty() { continue; }
let idx = val as usize % hm_keys.len();
sm.remove(sm_keys[idx]);
if hm.remove(&hm_keys[idx]) != sec.remove(sm_keys[idx]) {
return false;
}
}
2 => {
if hm_keys.is_empty() { continue; }
let idx = val as usize % hm_keys.len();
let (hm_key, sm_key) = (&hm_keys[idx], sm_keys[idx]);
if hm.contains_key(hm_key) != sec.contains_key(sm_key) ||
hm.get(hm_key) != sec.get(sm_key) {
return false;
}
}
3 => {
sec = sec.clone();
}
#[cfg(feature = "serde")]
4 => {
let ser = serde_json::to_string(&sec).unwrap();
sec = serde_json::from_str(&ser).unwrap();
}
_ => unreachable!(),
}
}
let mut secv: Vec<_> = sec.values().collect();
let mut hmv: Vec<_> = hm.values().collect();
secv.sort();
hmv.sort();
secv == hmv
}
}
}