use crate::prefetch;
#[cfg(not(feature = "std"))]
use alloc::{vec, vec::Vec};
use core::iter;
const FINGERPRINT_SHIFT: u32 = 56;
const EMPTY: u32 = u32::MAX;
const MAX_DISPLACEMENT: u32 = (1 << 24) - 2;
const DISPLACEMENT_UNIT: u32 = 1 << 8;
const MIN_SLOTS: usize = 16;
#[inline]
const fn slot_meta(hash: u64, displacement: u32) -> u32 {
(displacement << 8) | (hash >> FINGERPRINT_SHIFT) as u32
}
#[inline]
pub const fn partition(hash: u64, partitions: usize) -> usize {
assert!(
partitions.is_power_of_two() && partitions <= 1 << 16,
"partitions must be a power of two no larger than 2^16"
);
(hash >> (FINGERPRINT_SHIFT - partitions.trailing_zeros())) as usize & (partitions - 1)
}
#[repr(C, packed)]
struct Slot<H> {
meta: u32,
handle: H,
}
impl<H: Copy> Clone for Slot<H> {
fn clone(&self) -> Self {
*self
}
}
impl<H: Copy> Copy for Slot<H> {}
pub struct Table<H> {
slots: Vec<Slot<H>>,
len: usize,
}
impl<H> core::fmt::Debug for Table<H> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Table")
.field("len", &self.len)
.field("slots", &self.slots.len())
.finish()
}
}
impl<H: Copy + Eq + Default> Table<H> {
pub const fn new() -> Self {
Self {
slots: Vec::new(),
len: 0,
}
}
pub fn with_capacity(entries: usize) -> Self {
if entries == 0 {
return Self::new();
}
Self {
slots: Self::empty_slots(Self::slots_for(entries)),
len: 0,
}
}
fn slots_for(entries: usize) -> usize {
entries
.checked_mul(4)
.and_then(|scaled| scaled.div_ceil(3).checked_next_power_of_two())
.expect("capacity overflow")
.max(MIN_SLOTS)
}
fn empty_slots(count: usize) -> Vec<Slot<H>> {
vec![
Slot {
meta: EMPTY,
handle: H::default(),
};
count
]
}
pub const fn len(&self) -> usize {
self.len
}
pub const fn is_empty(&self) -> bool {
self.len == 0
}
pub const fn capacity(&self) -> usize {
self.slots.len() * 3 / 4
}
pub fn reserve(&mut self, additional: usize, mut rehash: impl FnMut(H) -> u64) {
let entries = self.len.checked_add(additional).expect("capacity overflow");
if entries > self.capacity() {
self.rebuild(Self::slots_for(entries), &mut rehash);
}
}
#[inline]
pub fn get(&self, hash: u64) -> impl Iterator<Item = H> + '_ {
self.probe(hash).map(|slot| self.slots[slot].handle)
}
#[inline]
pub fn prefetch(&self, hash: u64) {
if !prefetch::ENABLED {
return;
}
let mask = self.slots.len().wrapping_sub(1);
if let Some(slot) = self.slots.get(hash as usize & mask) {
prefetch::prefetch(slot);
}
}
#[inline]
pub fn insert(&mut self, hash: u64, handle: H, mut rehash: impl FnMut(H) -> u64) {
if self.len == self.capacity() {
self.grow(&mut rehash);
}
self.place(hash, handle);
self.len += 1;
}
#[inline]
pub fn find_or_insert(
&mut self,
hash: u64,
matches: impl FnMut(H) -> bool,
handle: H,
mut rehash: impl FnMut(H) -> u64,
) -> (H, bool) {
if self.len == self.capacity() {
self.grow(&mut rehash);
}
let (target, displacement) = match self.walk(hash, matches) {
Ok(found) => return (found, false),
Err(stop) => stop,
};
self.place_at(target, displacement, hash, handle);
self.len += 1;
(handle, true)
}
#[inline]
pub fn remove(&mut self, hash: u64, handle: H) -> bool {
let Some(mut hole) = self.probe(hash).find(|&slot| {
let candidate = self.slots[slot].handle;
candidate == handle
}) else {
return false;
};
let mask = self.slots.len() - 1;
loop {
let next = (hole + 1) & mask;
let Slot { meta, handle } = self.slots[next];
if meta == EMPTY || meta < DISPLACEMENT_UNIT {
break;
}
self.slots[hole] = Slot {
meta: meta - DISPLACEMENT_UNIT,
handle,
};
hole = next;
}
self.slots[hole].meta = EMPTY;
self.len -= 1;
true
}
#[inline]
fn probe(&self, hash: u64) -> impl Iterator<Item = usize> + '_ {
let mask = self.slots.len().wrapping_sub(1);
let mut slot = hash as usize & mask;
let mut displacement = 0;
iter::from_fn(move || {
loop {
let meta = self.slots.get(slot)?.meta;
if meta == EMPTY || meta >> 8 < displacement {
return None;
}
let current = slot;
let matched = meta == slot_meta(hash, displacement);
slot = (slot + 1) & mask;
displacement += 1;
if matched {
return Some(current);
}
}
})
}
#[inline]
fn walk(&self, hash: u64, mut matches: impl FnMut(H) -> bool) -> Result<H, (usize, u32)> {
let mask = self.slots.len() - 1;
let mut slot = hash as usize & mask;
let mut displacement = 0;
loop {
let Slot { meta, handle } = self.slots[slot];
if meta == EMPTY || meta >> 8 < displacement {
return Err((slot, displacement));
}
if meta == slot_meta(hash, displacement) && matches(handle) {
return Ok(handle);
}
assert!(
displacement < MAX_DISPLACEMENT,
"the run from the home slot of hash {hash:#x} must end within the maximum displacement"
);
slot = (slot + 1) & mask;
displacement += 1;
}
}
#[inline]
fn place(&mut self, hash: u64, handle: H) {
let (target, displacement) = match self.walk(hash, |_| false) {
Ok(_) => unreachable!("a walk that accepts no candidate must stop at a slot"),
Err(stop) => stop,
};
self.place_at(target, displacement, hash, handle);
}
#[inline]
fn place_at(&mut self, target: usize, displacement: u32, hash: u64, handle: H) {
let mask = self.slots.len() - 1;
let mut slot = target;
loop {
let meta = self.slots[slot].meta;
if meta == EMPTY {
break;
}
assert!(
meta >> 8 < MAX_DISPLACEMENT,
"shifting the run from slot {target} must not move an entry past the maximum displacement"
);
slot = (slot + 1) & mask;
}
while slot != target {
let previous = slot.wrapping_sub(1) & mask;
let Slot { meta, handle } = self.slots[previous];
self.slots[slot] = Slot {
meta: meta + DISPLACEMENT_UNIT,
handle,
};
slot = previous;
}
self.slots[target] = Slot {
meta: slot_meta(hash, displacement),
handle,
};
}
fn grow(&mut self, rehash: &mut impl FnMut(H) -> u64) {
self.rebuild((self.slots.len() * 2).max(MIN_SLOTS), rehash);
}
fn rebuild(&mut self, slots: usize, rehash: &mut impl FnMut(H) -> u64) {
let mut table = Self {
slots: Self::empty_slots(slots),
len: self.len,
};
for Slot { meta, handle } in self.slots.iter().copied() {
if meta != EMPTY {
table.place(rehash(handle), handle);
}
}
*self = table;
}
}
impl<H: Copy + Eq + Default> Default for Table<H> {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{TestRng, test_rng};
use core::cell::Cell;
use rand::{Rng as _, RngExt as _};
use std::collections::BTreeMap;
impl<H: Copy + Eq + Default> Table<H> {
fn metas(&self) -> Vec<u32> {
self.slots.iter().map(|slot| slot.meta).collect()
}
fn check_invariants(&self, hash_of: impl Fn(H) -> u64) {
let slots = self.slots.len();
assert_eq!(self.capacity(), slots * 3 / 4);
assert!(self.len <= self.capacity());
if slots == 0 {
return;
}
assert!(slots.is_power_of_two() && slots >= MIN_SLOTS);
let mask = slots - 1;
let mut occupied = 0;
for (slot, Slot { meta, handle }) in self.slots.iter().copied().enumerate() {
if meta == EMPTY {
continue;
}
occupied += 1;
let hash = hash_of(handle);
let displacement = slot.wrapping_sub(hash as usize) & mask;
assert!(displacement <= MAX_DISPLACEMENT as usize);
assert_eq!(meta, slot_meta(hash, displacement as u32));
let next = self.slots[(slot + 1) & mask].meta;
assert!(next == EMPTY || next >> 8 <= (meta >> 8) + 1);
}
assert_eq!(occupied, self.len);
}
}
const fn spread(i: u64) -> u64 {
i.wrapping_mul(0x9E37_79B9_7F4A_7C15)
}
fn sorted<H: Ord>(handles: impl Iterator<Item = H>) -> Vec<H> {
let mut handles: Vec<H> = handles.collect();
handles.sort();
handles
}
#[test]
fn test_find_or_insert_finds_or_places() {
let hashes = [spread(1), spread(1), spread(2)];
let rehash = |handle: usize| hashes[handle];
let mut table = Table::new();
assert_eq!(
table.find_or_insert(hashes[0], |_| false, 0, rehash),
(0, true)
);
assert_eq!(
table.find_or_insert(hashes[0], |handle| handle == 0, 7, rehash),
(0, false)
);
assert_eq!(
table.find_or_insert(hashes[1], |handle| handle == 1, 1, rehash),
(1, true)
);
assert_eq!(table.get(hashes[0]).collect::<Vec<_>>(), vec![0, 1]);
let mut offered = Vec::new();
let accept = |handle| {
offered.push(handle);
true
};
assert_eq!(
table.find_or_insert(hashes[0], accept, 7, rehash),
(0, false)
);
assert_eq!(offered, vec![0]);
assert_eq!(
table.find_or_insert(hashes[2], |handle| handle == 2, 2, rehash),
(2, true)
);
assert_eq!(table.len(), 3);
table.check_invariants(rehash);
}
#[test]
fn test_find_or_insert_grows_and_prefetch_is_harmless() {
let mut rng = test_rng();
let hashes: Vec<u64> = (0..1000).map(|_| rng.random()).collect();
let rehash = |handle: usize| hashes[handle];
Table::<usize>::new().prefetch(hashes[0]);
let mut table = Table::new();
for (handle, &hash) in hashes.iter().enumerate() {
table.prefetch(hash);
assert_eq!(
table.find_or_insert(hash, |candidate| candidate == handle, handle, rehash),
(handle, true)
);
assert_eq!(
table.find_or_insert(hash, |candidate| candidate == handle, usize::MAX, rehash),
(handle, false)
);
}
assert_eq!(table.len(), hashes.len());
table.check_invariants(rehash);
for (handle, &hash) in hashes.iter().enumerate() {
assert!(table.get(hash).any(|candidate| candidate == handle));
}
}
#[test]
fn test_find_or_insert_grows_a_full_table_before_walking() {
let hashes: Vec<u64> = (0..=12)
.map(|handle| 0xAB00_0000_0000_000F | ((handle % 2) << 4))
.collect();
let rehash = |handle: usize| hashes[handle];
let full = || {
let mut table = Table::new();
for (handle, &hash) in hashes.iter().enumerate().take(12) {
table.insert(hash, handle, rehash);
}
assert_eq!(table.len(), table.capacity());
table
};
for (accept, want) in [(1, (12, true)), (2, (2, false))] {
let mut table = full();
let mut offered = Vec::new();
let found = table.find_or_insert(
hashes[12],
|handle| {
offered.push(handle);
handle == accept
},
12,
rehash,
);
assert_eq!(found, want);
assert_eq!(table.capacity(), 24);
assert!(!offered.is_empty() && offered.iter().all(|handle| handle % 2 == 0));
table.check_invariants(rehash);
}
}
#[test]
fn test_find_or_insert_offers_the_candidates_of_get() {
let mut rng = test_rng();
let mut hashes = vec![
0xAB00_0000_0000_0005,
0xAB00_0100_0000_0005,
0xAB00_0000_0000_0006,
0xCD00_0000_0000_0005,
];
hashes.extend((0..200).map(|_| rng.next_u64()));
let seeded = hashes.len();
hashes.extend_from_within(..seeded);
hashes.extend((0..50).map(|_| rng.next_u64()));
let rehash = |handle: usize| hashes[handle];
let mut table = Table::new();
for (handle, &hash) in hashes.iter().enumerate().take(seeded) {
table.insert(hash, handle, rehash);
}
table.reserve(hashes.len() - seeded, rehash);
for (handle, &hash) in hashes.iter().enumerate().skip(seeded) {
let before: Vec<usize> = table.get(hash).collect();
let mut offered = Vec::new();
let matches = |candidate: usize| {
offered.push(candidate);
false
};
assert_eq!(
table.find_or_insert(hash, matches, handle, rehash),
(handle, true)
);
assert_eq!(
offered, before,
"the walk for hash {hash:#x} must offer the candidates get yields, in order"
);
let after: Vec<usize> = table.get(hash).collect();
assert_eq!(
after,
[before, vec![handle]].concat(),
"the handle inserted under hash {hash:#x} must follow its existing candidates"
);
}
table.check_invariants(rehash);
}
#[test]
fn test_basic_insert_get_remove() {
let hashes = [spread(1), spread(1), spread(2), spread(3), spread(4)];
let rehash = |handle: usize| hashes[handle];
let mut table = Table::new();
assert!(table.is_empty());
assert_eq!(table.slots.capacity(), 0);
assert_eq!(table.get(hashes[0]).next(), None);
assert!(!table.remove(hashes[0], 0));
for (handle, &hash) in hashes.iter().enumerate().take(4) {
table.insert(hash, handle, rehash);
}
assert_eq!(table.len(), 4);
assert_eq!(sorted(table.get(hashes[0])), [0, 1]);
assert_eq!(sorted(table.get(hashes[2])), [2]);
assert_eq!(sorted(table.get(hashes[3])), [3]);
assert_eq!(table.get(hashes[4]).next(), None);
assert!(table.remove(hashes[0], 0));
assert_eq!(sorted(table.get(hashes[0])), [1]);
assert_eq!(table.len(), 3);
let metas = table.metas();
assert!(!table.remove(hashes[0], 0));
assert!(!table.remove(hashes[4], 4));
assert_eq!(table.metas(), metas);
assert_eq!(table.len(), 3);
table.insert(hashes[1], 1, rehash);
assert_eq!(sorted(table.get(hashes[1])), [1, 1]);
assert!(table.remove(hashes[1], 1));
assert_eq!(sorted(table.get(hashes[1])), [1]);
table.check_invariants(rehash);
}
#[test]
fn test_fingerprint_false_candidates() {
let keys = ["a", "b", "c", "d"];
let hashes = [
0xAB00_0000_0000_0005,
0xAB00_0100_0000_0005,
0xAB00_0000_0000_0006,
0xCD00_0000_0000_0005,
];
let rehash = |handle: usize| hashes[handle];
let mut table = Table::new();
for (handle, &hash) in hashes.iter().enumerate() {
table.insert(hash, handle, rehash);
}
assert_eq!(sorted(table.get(hashes[0])), [0, 1]);
let confirmed: Vec<_> = table.get(hashes[0]).filter(|&h| keys[h] == "a").collect();
assert_eq!(confirmed, [0]);
assert_eq!(sorted(table.get(hashes[2])), [2]);
assert_eq!(sorted(table.get(hashes[3])), [3]);
table.check_invariants(rehash);
}
#[test]
fn test_growth_rehashes_every_entry() {
let mut rng = test_rng();
let mut hashes = Vec::new();
let calls = Cell::new(0);
let mut table = Table::new();
let mut growths = 0;
for handle in 0..5_000 {
hashes.push(rng.next_u64());
let (slots, len, before) = (table.slots.len(), table.len(), calls.get());
table.insert(hashes[handle], handle, |handle| {
calls.set(calls.get() + 1);
hashes[handle]
});
assert_eq!(table.len(), handle + 1);
if table.slots.len() == slots {
assert_eq!(calls.get(), before);
continue;
}
growths += 1;
assert_eq!(table.slots.len(), (slots * 2).max(MIN_SLOTS));
assert_eq!(len, slots * 3 / 4);
assert_eq!(calls.get() - before, len);
table.check_invariants(|handle| hashes[handle]);
for (handle, &hash) in hashes.iter().enumerate() {
assert!(table.get(hash).any(|h| h == handle));
}
}
assert_eq!(growths, 10);
assert_eq!(table.slots.len(), 8_192);
}
#[test]
fn test_shared_hash_run() {
const ENTRIES: usize = 2_000;
let mut table = Table::new();
for handle in 0..ENTRIES {
table.insert(spread(1), handle, |_| spread(1));
}
assert_eq!(table.slots.len(), 4_096);
assert_eq!(table.get(spread(1)).count(), ENTRIES);
table.check_invariants(|_| spread(1));
for handle in (0..ENTRIES).step_by(2) {
assert!(table.remove(spread(1), handle));
}
assert_eq!(
sorted(table.get(spread(1))),
sorted((1..ENTRIES).step_by(2))
);
table.check_invariants(|_| spread(1));
}
#[test]
fn test_matches_multimap_model() {
const HANDLES: usize = 1_024;
const HASHES: usize = 512;
const OPERATIONS: usize = 50_000;
let mut rng = TestRng::new(5);
let mut pool: Vec<u64> = Vec::with_capacity(HASHES);
for i in 0..HASHES {
let fingerprint = rng.random_range(0..4u64) << 56;
let middle = rng.next_u64() & 0x00FF_FFFF_FF00_0000;
let low = if i % 2 == 1 {
pool[i - 1]
} else {
rng.next_u64()
} & 0xFF_FFFF;
pool.push(fingerprint | middle | low);
}
let hash_of: Vec<u64> = (0..HANDLES)
.map(|_| pool[rng.random_range(0..pool.len())])
.collect();
let rehash_calls = Cell::new(0);
let rehash = |handle: usize| {
rehash_calls.set(rehash_calls.get() + 1);
hash_of[handle]
};
let mut model: BTreeMap<u64, Vec<usize>> = BTreeMap::new();
let mut table = Table::new();
let (mut removed, mut false_candidates, mut max_len) = (0, 0, 0);
for op in 0..OPERATIONS {
let inserts = if (op / 5_000) % 2 == 0 { 6 } else { 3 };
let roll = rng.random_range(0..10);
let (slots, calls) = (table.slots.len(), rehash_calls.get());
let hash = if roll < inserts {
let handle = rng.random_range(0..HANDLES);
table.insert(hash_of[handle], handle, rehash);
model.entry(hash_of[handle]).or_default().push(handle);
hash_of[handle]
} else if roll < 9 {
let handle = rng.random_range(0..HANDLES);
let hash = hash_of[handle];
let handles = model.get_mut(&hash);
let position = handles
.as_ref()
.and_then(|handles| handles.iter().position(|&h| h == handle));
assert_eq!(table.remove(hash, handle), position.is_some());
if let (Some(handles), Some(position)) = (handles, position) {
handles.swap_remove(position);
if handles.is_empty() {
model.remove(&hash);
}
removed += 1;
}
hash
} else if rng.random_bool(0.5) {
pool[rng.random_range(0..pool.len())]
} else {
rng.next_u64()
};
if table.slots.len() == slots {
assert_eq!(rehash_calls.get(), calls);
}
assert_eq!(table.len(), model.values().map(Vec::len).sum::<usize>());
max_len = max_len.max(table.len());
let mask = table.slots.len().wrapping_sub(1);
let expected = sorted(
model
.iter()
.filter(|&(&other, _)| {
other >> 56 == hash >> 56 && (other ^ hash) & mask as u64 == 0
})
.flat_map(|(_, handles)| handles.iter().copied()),
);
let actual = sorted(table.get(hash));
assert_eq!(actual, expected);
false_candidates += actual.iter().filter(|&&h| hash_of[h] != hash).count();
if op % 1_000 == 0 {
table.check_invariants(|handle| hash_of[handle]);
}
}
table.check_invariants(|handle| hash_of[handle]);
assert!(removed > 10_000 && false_candidates > 1_000 && max_len > 2_000);
}
#[test]
fn test_with_capacity_never_rehashes() {
let table = Table::<u64>::with_capacity(0);
assert_eq!(table.slots.capacity(), 0);
let mut rng = test_rng();
for entries in [1, 12, 13, 100, 384, 385, 1_000] {
let mut table = Table::with_capacity(entries);
let slots = table.slots.len();
assert!(table.capacity() >= entries);
assert!(slots == MIN_SLOTS || slots / 2 * 3 / 4 < entries);
for handle in 0..entries {
table.insert(rng.next_u64(), handle, |_| {
unreachable!("rehash without growth")
});
}
assert_eq!(table.slots.len(), slots);
assert_eq!(table.len(), entries);
}
}
#[test]
fn test_reserve_grows_once_and_keeps_entries() {
let mut rng = test_rng();
let hashes: Vec<u64> = (0..1_000).map(|_| rng.next_u64()).collect();
let rehash = |handle: usize| hashes[handle];
let calls = Cell::new(0);
let counted = |handle: usize| {
calls.set(calls.get() + 1);
hashes[handle]
};
let mut table = Table::new();
table.reserve(0, rehash);
assert_eq!(table.slots.capacity(), 0);
for (handle, &hash) in hashes.iter().enumerate().take(600) {
table.insert(hash, handle, rehash);
}
let slots = table.slots.len();
table.reserve(1, counted);
assert_eq!(table.slots.len(), slots);
assert_eq!(calls.get(), 0);
table.reserve(hashes.len() - 600, counted);
assert!(table.capacity() >= hashes.len());
assert_eq!(calls.get(), 600);
let slots = table.slots.len();
for (handle, &hash) in hashes.iter().enumerate().skip(600) {
table.insert(hash, handle, |_| {
unreachable!("reserved inserts must not rehash")
});
}
assert_eq!(table.slots.len(), slots);
assert_eq!(table.len(), hashes.len());
table.check_invariants(rehash);
for (handle, &hash) in hashes.iter().enumerate() {
assert!(table.get(hash).any(|candidate| candidate == handle));
}
}
#[test]
fn test_partition_reads_neither_home_nor_fingerprint_bits() {
assert_eq!(partition(0xAB << (FINGERPRINT_SHIFT - 8), 256), 0xAB);
let mut rng = test_rng();
for _ in 0..1_000 {
let hash: u64 = rng.random();
for partitions in [1, 2, 256, 1 << 16] {
let group = partition(hash, partitions);
assert!(
group < partitions,
"group must be below the partition count"
);
assert_eq!(
partition(hash ^ (0xFF << FINGERPRINT_SHIFT), partitions),
group
);
assert_eq!(partition(hash ^ ((1 << 40) - 1), partitions), group);
}
}
}
#[test]
#[should_panic(expected = "partitions must be a power of two")]
fn test_partition_rejects_a_count_that_is_not_a_power_of_two() {
partition(0, 3);
}
}