use core::{borrow::Borrow, fmt, iter::FromIterator, mem, num::NonZeroU32, ops, slice};
use hash32::{BuildHasher, BuildHasherDefault, FnvHasher, Hash, Hasher};
use crate::Vec;
pub type FnvIndexMap<K, V, const N: usize> = IndexMap<K, V, BuildHasherDefault<FnvHasher>, N>;
#[derive(Clone, Copy, Eq, PartialEq)]
struct HashValue(u16);
impl HashValue {
fn desired_pos(&self, mask: usize) -> usize {
usize::from(self.0) & mask
}
fn probe_distance(&self, mask: usize, current: usize) -> usize {
current.wrapping_sub(self.desired_pos(mask) as usize) & mask
}
}
#[doc(hidden)]
#[derive(Clone)]
pub struct Bucket<K, V> {
hash: HashValue,
key: K,
value: V,
}
#[doc(hidden)]
#[derive(Clone, Copy, PartialEq)]
pub struct Pos {
nz: NonZeroU32,
}
impl Pos {
fn new(index: usize, hash: HashValue) -> Self {
Pos {
nz: unsafe {
NonZeroU32::new_unchecked(
((u32::from(hash.0) << 16) + index as u32).wrapping_add(1),
)
},
}
}
fn hash(&self) -> HashValue {
HashValue((self.nz.get().wrapping_sub(1) >> 16) as u16)
}
fn index(&self) -> usize {
self.nz.get().wrapping_sub(1) as u16 as usize
}
}
pub enum Inserted<V> {
Done,
Swapped { prev_value: V },
RobinHood { probe: usize, old_pos: Pos },
}
macro_rules! probe_loop {
($probe_var: ident < $len: expr, $body: expr) => {
loop {
if $probe_var < $len {
$body
$probe_var += 1;
} else {
$probe_var = 0;
}
}
}
}
struct CoreMap<K, V, const N: usize> {
entries: Vec<Bucket<K, V>, N>,
indices: [Option<Pos>; N],
}
impl<K, V, const N: usize> CoreMap<K, V, N> {
const fn new() -> Self {
const INIT: Option<Pos> = None;
CoreMap {
entries: Vec::new(),
indices: [INIT; N],
}
}
}
impl<K, V, const N: usize> CoreMap<K, V, N>
where
K: Eq + Hash,
{
fn capacity() -> usize {
N
}
fn mask() -> usize {
Self::capacity() - 1
}
fn find<Q>(&self, hash: HashValue, query: &Q) -> Option<(usize, usize)>
where
K: Borrow<Q>,
Q: ?Sized + Eq,
{
let mut probe = hash.desired_pos(Self::mask());
let mut dist = 0;
probe_loop!(probe < self.indices.len(), {
if let Some(pos) = self.indices[probe] {
let entry_hash = pos.hash();
let i = pos.index();
debug_assert!(i < self.entries.len());
if dist > entry_hash.probe_distance(Self::mask(), probe) {
return None;
} else if entry_hash == hash
&& unsafe { self.entries.get_unchecked(i).key.borrow() == query }
{
return Some((probe, i));
}
} else {
return None;
}
dist += 1;
});
}
fn insert_phase_1(&mut self, hash: HashValue, key: K, value: V) -> Inserted<V> {
let mut probe = hash.desired_pos(Self::mask());
let mut dist = 0;
let inserted;
probe_loop!(probe < self.indices.len(), {
let pos = &mut self.indices[probe];
if let Some(pos) = *pos {
let entry_hash = pos.hash();
let i = pos.index();
debug_assert!(i < self.entries.len());
let their_dist = entry_hash.probe_distance(Self::mask(), probe);
if their_dist < dist {
let index = self.entries.len();
inserted = Inserted::RobinHood {
probe: probe,
old_pos: Pos::new(index, hash),
};
break;
} else if entry_hash == hash && unsafe { self.entries.get_unchecked(i).key == key }
{
return Inserted::Swapped {
prev_value: mem::replace(
unsafe { &mut self.entries.get_unchecked_mut(i).value },
value,
),
};
}
} else {
let index = self.entries.len();
*pos = Some(Pos::new(index, hash));
inserted = Inserted::Done;
break;
}
dist += 1;
});
unsafe { self.entries.push_unchecked(Bucket { hash, key, value }) }
inserted
}
fn insert_phase_2(&mut self, mut probe: usize, mut old_pos: Pos) {
probe_loop!(probe < self.indices.len(), {
let pos = unsafe { self.indices.get_unchecked_mut(probe) };
let mut is_none = true; if let Some(pos) = pos.as_mut() {
old_pos = mem::replace(pos, old_pos);
is_none = false;
}
if is_none {
*pos = Some(old_pos);
break;
}
});
}
fn remove_found(&mut self, probe: usize, found: usize) -> (K, V) {
self.indices[probe] = None;
let entry = unsafe { self.entries.swap_remove_unchecked(found) };
if let Some(entry) = self.entries.get(found) {
let mut probe = entry.hash.desired_pos(Self::mask());
probe_loop!(probe < self.indices.len(), {
if let Some(pos) = self.indices[probe] {
if pos.index() >= self.entries.len() {
self.indices[probe] = Some(Pos::new(found, entry.hash));
break;
}
}
});
}
self.backward_shift_after_removal(probe);
(entry.key, entry.value)
}
fn backward_shift_after_removal(&mut self, probe_at_remove: usize) {
let mut last_probe = probe_at_remove;
let mut probe = probe_at_remove + 1;
probe_loop!(probe < self.indices.len(), {
if let Some(pos) = self.indices[probe] {
let entry_hash = pos.hash();
if entry_hash.probe_distance(Self::mask(), probe) > 0 {
unsafe { *self.indices.get_unchecked_mut(last_probe) = self.indices[probe] }
self.indices[probe] = None;
} else {
break;
}
} else {
break;
}
last_probe = probe;
});
}
}
impl<K, V, const N: usize> Clone for CoreMap<K, V, N>
where
K: Eq + Hash + Clone,
V: Clone,
{
fn clone(&self) -> Self {
Self {
entries: self.entries.clone(),
indices: self.indices.clone(),
}
}
}
pub struct IndexMap<K, V, S, const N: usize> {
core: CoreMap<K, V, N>,
build_hasher: S,
}
impl<K, V, S, const N: usize> IndexMap<K, V, BuildHasherDefault<S>, N> {
pub const fn new() -> Self {
crate::sealed::greater_than_1::<N>();
crate::sealed::power_of_two::<N>();
IndexMap {
build_hasher: BuildHasherDefault::new(),
core: CoreMap::new(),
}
}
}
impl<K, V, S, const N: usize> IndexMap<K, V, S, N>
where
K: Eq + Hash,
S: BuildHasher,
{
pub fn capacity(&self) -> usize {
N
}
pub fn keys(&self) -> impl Iterator<Item = &K> {
self.core.entries.iter().map(|bucket| &bucket.key)
}
pub fn values(&self) -> impl Iterator<Item = &V> {
self.core.entries.iter().map(|bucket| &bucket.value)
}
pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
self.core.entries.iter_mut().map(|bucket| &mut bucket.value)
}
pub fn iter(&self) -> Iter<'_, K, V> {
Iter {
iter: self.core.entries.iter(),
}
}
pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
IterMut {
iter: self.core.entries.iter_mut(),
}
}
pub fn len(&self) -> usize {
self.core.entries.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn clear(&mut self) {
self.core.entries.clear();
for pos in self.core.indices.iter_mut() {
*pos = None;
}
}
pub fn get<Q>(&self, key: &Q) -> Option<&V>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
self.find(key)
.map(|(_, found)| unsafe { &self.core.entries.get_unchecked(found).value })
}
pub fn contains_key<Q>(&self, key: &Q) -> bool
where
K: Borrow<Q>,
Q: ?Sized + Eq + Hash,
{
self.find(key).is_some()
}
pub fn get_mut<'v, Q>(&'v mut self, key: &Q) -> Option<&'v mut V>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
if let Some((_, found)) = self.find(key) {
Some(unsafe { &mut self.core.entries.get_unchecked_mut(found).value })
} else {
None
}
}
pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>, (K, V)> {
if self.core.entries.is_full() {
Err((key, value))
} else {
Ok(match self.insert_phase_1(key, value) {
Inserted::Swapped { prev_value } => Some(prev_value),
Inserted::Done => None,
Inserted::RobinHood { probe, old_pos } => {
self.core.insert_phase_2(probe, old_pos);
None
}
})
}
}
pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
self.swap_remove(key)
}
pub fn swap_remove<Q>(&mut self, key: &Q) -> Option<V>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
self.find(key)
.map(|(probe, found)| self.core.remove_found(probe, found).1)
}
fn find<Q>(&self, key: &Q) -> Option<(usize, usize)>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
if self.len() == 0 {
return None;
}
let h = hash_with(key, &self.build_hasher);
self.core.find(h, key)
}
fn insert_phase_1(&mut self, key: K, value: V) -> Inserted<V> {
let hash = hash_with(&key, &self.build_hasher);
self.core.insert_phase_1(hash, key, value)
}
}
impl<'a, K, Q, V, S, const N: usize> ops::Index<&'a Q> for IndexMap<K, V, S, N>
where
K: Eq + Hash + Borrow<Q>,
Q: ?Sized + Eq + Hash,
S: BuildHasher,
{
type Output = V;
fn index(&self, key: &Q) -> &V {
self.get(key).expect("key not found")
}
}
impl<'a, K, Q, V, S, const N: usize> ops::IndexMut<&'a Q> for IndexMap<K, V, S, N>
where
K: Eq + Hash + Borrow<Q>,
Q: ?Sized + Eq + Hash,
S: BuildHasher,
{
fn index_mut(&mut self, key: &Q) -> &mut V {
self.get_mut(key).expect("key not found")
}
}
impl<K, V, S, const N: usize> Clone for IndexMap<K, V, S, N>
where
K: Eq + Hash + Clone,
V: Clone,
S: Clone,
{
fn clone(&self) -> Self {
Self {
core: self.core.clone(),
build_hasher: self.build_hasher.clone(),
}
}
}
impl<K, V, S, const N: usize> fmt::Debug for IndexMap<K, V, S, N>
where
K: Eq + Hash + fmt::Debug,
V: fmt::Debug,
S: BuildHasher,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_map().entries(self.iter()).finish()
}
}
impl<K, V, S, const N: usize> Default for IndexMap<K, V, S, N>
where
K: Eq + Hash,
S: BuildHasher + Default,
{
fn default() -> Self {
IndexMap {
build_hasher: <_>::default(),
core: CoreMap::new(),
}
}
}
impl<K, V, S, S2, const N: usize, const N2: usize> PartialEq<IndexMap<K, V, S2, N2>>
for IndexMap<K, V, S, N>
where
K: Eq + Hash,
V: Eq,
S: BuildHasher,
S2: BuildHasher,
{
fn eq(&self, other: &IndexMap<K, V, S2, N2>) -> bool {
self.len() == other.len()
&& self
.iter()
.all(|(key, value)| other.get(key).map_or(false, |v| *value == *v))
}
}
impl<K, V, S, const N: usize> Eq for IndexMap<K, V, S, N>
where
K: Eq + Hash,
V: Eq,
S: BuildHasher,
{
}
impl<K, V, S, const N: usize> Extend<(K, V)> for IndexMap<K, V, S, N>
where
K: Eq + Hash,
S: BuildHasher,
{
fn extend<I>(&mut self, iterable: I)
where
I: IntoIterator<Item = (K, V)>,
{
for (k, v) in iterable {
self.insert(k, v).ok().unwrap();
}
}
}
impl<'a, K, V, S, const N: usize> Extend<(&'a K, &'a V)> for IndexMap<K, V, S, N>
where
K: Eq + Hash + Copy,
V: Copy,
S: BuildHasher,
{
fn extend<I>(&mut self, iterable: I)
where
I: IntoIterator<Item = (&'a K, &'a V)>,
{
self.extend(iterable.into_iter().map(|(&key, &value)| (key, value)))
}
}
impl<K, V, S, const N: usize> FromIterator<(K, V)> for IndexMap<K, V, S, N>
where
K: Eq + Hash,
S: BuildHasher + Default,
{
fn from_iter<I>(iterable: I) -> Self
where
I: IntoIterator<Item = (K, V)>,
{
let mut map = IndexMap::default();
map.extend(iterable);
map
}
}
impl<'a, K, V, S, const N: usize> IntoIterator for &'a IndexMap<K, V, S, N>
where
K: Eq + Hash,
S: BuildHasher,
{
type Item = (&'a K, &'a V);
type IntoIter = Iter<'a, K, V>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'a, K, V, S, const N: usize> IntoIterator for &'a mut IndexMap<K, V, S, N>
where
K: Eq + Hash,
S: BuildHasher,
{
type Item = (&'a K, &'a mut V);
type IntoIter = IterMut<'a, K, V>;
fn into_iter(self) -> Self::IntoIter {
self.iter_mut()
}
}
pub struct Iter<'a, K, V> {
iter: slice::Iter<'a, Bucket<K, V>>,
}
impl<'a, K, V> Iterator for Iter<'a, K, V> {
type Item = (&'a K, &'a V);
fn next(&mut self) -> Option<Self::Item> {
self.iter.next().map(|bucket| (&bucket.key, &bucket.value))
}
}
impl<'a, K, V> Clone for Iter<'a, K, V> {
fn clone(&self) -> Self {
Self {
iter: self.iter.clone(),
}
}
}
pub struct IterMut<'a, K, V> {
iter: slice::IterMut<'a, Bucket<K, V>>,
}
impl<'a, K, V> Iterator for IterMut<'a, K, V> {
type Item = (&'a K, &'a mut V);
fn next(&mut self) -> Option<Self::Item> {
self.iter
.next()
.map(|bucket| (&bucket.key, &mut bucket.value))
}
}
fn hash_with<K, S>(key: &K, build_hasher: &S) -> HashValue
where
K: ?Sized + Hash,
S: BuildHasher,
{
let mut h = build_hasher.build_hasher();
key.hash(&mut h);
HashValue(h.finish() as u16)
}
#[cfg(test)]
mod tests {
use crate::FnvIndexMap;
use core::mem;
#[test]
fn size() {
const CAP: usize = 4;
assert_eq!(
mem::size_of::<FnvIndexMap<i16, u16, CAP>>(),
CAP * mem::size_of::<u32>() + CAP * (mem::size_of::<i16>() + mem::size_of::<u16>() + mem::size_of::<u16>() ) + mem::size_of::<usize>() )
}
#[test]
fn partial_eq() {
{
let mut a: FnvIndexMap<_, _, 4> = FnvIndexMap::new();
a.insert("k1", "v1").unwrap();
let mut b: FnvIndexMap<_, _, 4> = FnvIndexMap::new();
b.insert("k1", "v1").unwrap();
assert!(a == b);
b.insert("k2", "v2").unwrap();
assert!(a != b);
}
{
let mut a: FnvIndexMap<_, _, 4> = FnvIndexMap::new();
a.insert("k1", "v1").unwrap();
a.insert("k2", "v2").unwrap();
let mut b: FnvIndexMap<_, _, 4> = FnvIndexMap::new();
b.insert("k2", "v2").unwrap();
b.insert("k1", "v1").unwrap();
assert!(a == b);
}
}
}