use yo_common::hash_key;
use crate::elem::{Elements, Full};
use crate::scan::{Cursor, MAX_PARTS};
pub const PARTITION_AT: usize = 262_144;
pub const PART_MIN: u32 = 4;
pub const PART_TARGET: usize = PARTITION_AT / PART_MIN as usize;
pub const PART_BIT: u32 = 32;
pub const BLOCK: usize = 64;
#[must_use]
pub fn parts_for(n: usize) -> u32 {
layout(u32::try_from(n.div_ceil(PART_TARGET)).unwrap_or(MAX_PARTS))
}
const fn layout(parts: u32) -> u32 {
if parts >= MAX_PARTS {
return MAX_PARTS;
}
let want = parts.next_power_of_two();
if want < PART_MIN { PART_MIN } else { want }
}
#[derive(Debug, Clone)]
pub struct Parts<V> {
tables: Vec<Elements<V>>,
lens: Vec<u32>,
blocks: Vec<u32>,
total: usize,
mask: u32,
}
impl<V: Copy> Parts<V> {
#[must_use]
pub fn with_parts(parts: u32) -> Parts<V> {
let parts = layout(parts);
let n = parts as usize;
Parts {
tables: (0..n).map(|_| Elements::new()).collect(),
lens: vec![0; n],
blocks: vec![0; n.div_ceil(BLOCK)],
total: 0,
mask: parts - 1,
}
}
#[inline]
fn gained(&mut self, at: usize) {
self.lens[at] += 1;
self.blocks[at / BLOCK] += 1;
self.total += 1;
}
#[inline]
fn lost(&mut self, at: usize) {
self.lens[at] -= 1;
self.blocks[at / BLOCK] -= 1;
self.total -= 1;
}
#[must_use]
pub fn from_table(table: &Elements<V>, parts: u32) -> Parts<V> {
let mut p = Parts::with_parts(parts);
for (name, value) in table.iter() {
let h = hash_key(name);
let at = p.part_of(h) as usize;
let _ = p.tables[at].insert_hashed(h, name, *value);
p.gained(at);
}
p
}
#[inline]
#[must_use]
pub const fn parts(&self) -> u32 {
self.mask + 1
}
#[inline]
const fn part_of(&self, h: u64) -> u32 {
((h >> PART_BIT) as u32) & self.mask
}
#[inline]
#[must_use]
pub const fn len(&self) -> usize {
self.total
}
#[inline]
#[must_use]
pub const fn is_empty(&self) -> bool {
self.total == 0
}
#[inline]
#[must_use]
pub fn hash_of(name: &[u8]) -> u64 {
hash_key(name)
}
#[inline]
#[must_use]
pub fn get(&self, name: &[u8]) -> Option<&V> {
self.get_hashed(hash_key(name), name)
}
#[inline]
#[must_use]
pub fn get_hashed(&self, h: u64, name: &[u8]) -> Option<&V> {
self.tables[self.part_of(h) as usize].get_hashed(h, name)
}
#[inline]
pub fn get_mut(&mut self, name: &[u8]) -> Option<&mut V> {
let h = hash_key(name);
let at = self.part_of(h) as usize;
self.tables[at].get_hashed_mut(h, name)
}
#[inline]
#[must_use]
pub fn contains(&self, name: &[u8]) -> bool {
self.get(name).is_some()
}
#[inline]
#[must_use]
pub fn contains_hashed(&self, h: u64, name: &[u8]) -> bool {
self.get_hashed(h, name).is_some()
}
pub fn insert(&mut self, name: &[u8], value: V) -> Result<Option<V>, Full> {
let h = hash_key(name);
let at = self.part_of(h) as usize;
let was = self.tables[at].insert_hashed(h, name, value)?;
if was.is_none() {
self.gained(at);
}
Ok(was)
}
pub fn remove(&mut self, name: &[u8]) -> Option<V> {
let h = hash_key(name);
let at = self.part_of(h) as usize;
let was = self.tables[at].remove_hashed(h, name)?;
self.lost(at);
Some(was)
}
#[inline]
#[must_use]
pub fn locate(&self, idx: usize) -> Option<(usize, usize)> {
if idx >= self.total {
return None;
}
let mut seen = 0usize;
let mut at = 0usize;
for &n in &self.blocks {
let n = n as usize;
if idx < seen + n {
break;
}
seen += n;
at += BLOCK;
}
for &n in &self.lens[at..] {
let n = n as usize;
if idx < seen + n {
return Some((at, idx - seen));
}
seen += n;
at += 1;
}
None
}
#[inline]
#[must_use]
pub fn at(&self, idx: usize) -> Option<(&[u8], &V)> {
let (at, within) = self.locate(idx)?;
self.tables[at].at(within)
}
#[inline]
pub fn at_mut(&mut self, idx: usize) -> Option<&mut V> {
let (at, within) = self.locate(idx)?;
self.tables[at].at_mut(within)
}
pub fn remove_at(&mut self, idx: usize) -> Option<V> {
let (at, within) = self.locate(idx)?;
let was = self.tables[at].remove_at(within)?;
self.lost(at);
Some(was)
}
pub fn take_at(&mut self, idx: usize) -> Option<(Vec<u8>, V)> {
let (at, within) = self.locate(idx)?;
let got = self.tables[at].take_at(within)?;
self.lost(at);
Some(got)
}
pub fn iter(&self) -> impl Iterator<Item = (&[u8], &V)> {
self.tables.iter().flat_map(Elements::iter)
}
pub fn payloads_mut(&mut self) -> impl Iterator<Item = &mut V> {
self.tables.iter_mut().flat_map(Elements::payloads_mut)
}
pub fn scan<F>(&self, cursor: Cursor, count: usize, mut f: F) -> Cursor
where
F: FnMut(&[u8], &V),
{
if self.total == 0 {
return Cursor::END;
}
let parts = self.parts();
let (mut part, mut idx) = if cursor.is_end() {
(parts - 1, None)
} else {
let here = cursor.rebase(parts);
(here.part().min(parts - 1), here.idx())
};
let mut left = count.max(1);
loop {
let table = &self.tables[part as usize];
let mut at = if table.is_empty() {
None
} else {
let top = table.len() - 1;
Some(match idx {
Some(i) => (i as usize).min(top),
None => top,
})
};
while let Some(row) = at {
if left == 0 {
return Cursor::at(parts, part, row as u64);
}
let (name, value) = table.at(row).expect("the row is inside the table");
f(name, value);
left -= 1;
at = row.checked_sub(1);
}
if part == 0 {
return Cursor::END;
}
part -= 1;
idx = None;
if left == 0 {
return Cursor::top(parts, part);
}
}
}
pub fn grow_to(&mut self, parts: u32) -> bool {
let parts = layout(parts);
if parts <= self.parts() {
return false;
}
let mut grown = Parts::with_parts(parts);
for table in &self.tables {
for (name, value) in table.iter() {
let h = hash_key(name);
let at = grown.part_of(h) as usize;
let _ = grown.tables[at].insert_hashed(h, name, *value);
grown.gained(at);
}
}
*self = grown;
true
}
#[must_use]
pub fn wants_parts(&self) -> Option<u32> {
let want = parts_for(self.total);
(want > self.parts()).then_some(want)
}
pub fn clear(&mut self) {
for table in &mut self.tables {
table.clear();
}
self.lens.fill(0);
self.blocks.fill(0);
self.total = 0;
}
#[must_use]
pub fn memory_bytes(&self) -> usize {
self.tables
.iter()
.map(Elements::memory_bytes)
.sum::<usize>()
+ (self.lens.len() + self.blocks.len()) * size_of::<u32>()
}
#[must_use]
pub fn slot_bytes(&self) -> usize {
self.tables.iter().map(Elements::slot_bytes).sum()
}
#[must_use]
pub fn row_bytes(&self) -> usize {
self.tables.iter().map(Elements::row_bytes).sum()
}
#[must_use]
pub fn name_bytes(&self) -> usize {
self.tables.iter().map(Elements::name_bytes).sum()
}
#[must_use]
pub fn dead_name_bytes(&self) -> usize {
self.tables.iter().map(Elements::dead_name_bytes).sum()
}
#[must_use]
pub fn lengths(&self) -> &[u32] {
&self.lens
}
#[must_use]
pub fn block_lengths(&self) -> &[u32] {
&self.blocks
}
}
#[cfg(test)]
mod tests {
use super::*;
fn filled(n: usize, parts: u32) -> Parts<u32> {
let mut p = Parts::with_parts(parts);
for i in 0..n {
p.insert(format!("member:{i}").as_bytes(), i as u32)
.expect("room");
}
p
}
#[test]
fn the_layout_is_a_power_of_two_and_never_two() {
assert_eq!(Parts::<()>::with_parts(0).parts(), PART_MIN);
assert_eq!(Parts::<()>::with_parts(1).parts(), PART_MIN);
assert_eq!(Parts::<()>::with_parts(2).parts(), PART_MIN);
assert_eq!(Parts::<()>::with_parts(3).parts(), PART_MIN);
assert_eq!(Parts::<()>::with_parts(4).parts(), 4);
assert_eq!(Parts::<()>::with_parts(5).parts(), 8);
assert_eq!(Parts::<()>::with_parts(u32::MAX).parts(), MAX_PARTS);
for n in [0, 1, 100, PARTITION_AT - 1, PARTITION_AT] {
assert_ne!(parts_for(n), 2, "there is no layout with two partitions");
}
assert_eq!(parts_for(PARTITION_AT), PART_MIN);
assert_eq!(parts_for(PARTITION_AT + 1), 8);
assert_eq!(parts_for(1_000_000), 16);
}
#[test]
fn what_goes_in_comes_out_of_the_partition_it_went_into() {
let p = filled(2_000, 8);
assert_eq!(p.len(), 2_000);
for i in 0..2_000 {
let name = format!("member:{i}");
assert_eq!(p.get(name.as_bytes()), Some(&(i as u32)));
}
assert!(!p.contains(b"member:2000"));
assert_eq!(p.get(b"nothing"), None);
}
#[test]
fn the_lengths_add_up_to_the_total_and_the_spread_is_even() {
let p = filled(8_000, 8);
assert_eq!(p.lengths().len(), 8);
assert_eq!(
p.lengths().iter().map(|&n| n as usize).sum::<usize>(),
8_000
);
for &n in p.lengths() {
assert!((700..1_300).contains(&n), "one partition holds {n}");
}
}
#[test]
fn writing_a_member_again_replaces_it_and_moves_no_length() {
let mut p = filled(100, 4);
let before = p.lengths().to_vec();
assert_eq!(p.insert(b"member:7", 700), Ok(Some(7)));
assert_eq!(p.len(), 100);
assert_eq!(p.lengths(), before.as_slice());
assert_eq!(p.get(b"member:7"), Some(&700));
}
#[test]
fn removing_takes_it_out_of_one_partition_only() {
let mut p = filled(1_000, 4);
let before = p.lengths().to_vec();
assert_eq!(p.remove(b"member:500"), Some(500));
assert_eq!(p.len(), 999);
assert!(!p.contains(b"member:500"));
assert_eq!(p.remove(b"member:500"), None);
let after = p.lengths();
let moved: Vec<usize> = (0..4).filter(|&i| after[i] != before[i]).collect();
assert_eq!(moved.len(), 1, "one partition changed length");
}
#[test]
fn a_draw_reaches_every_member_exactly_once() {
let p = filled(1_000, 8);
let mut seen = vec![false; 1_000];
for i in 0..p.len() {
let (_, &v) = p.at(i).expect("inside the collection");
assert!(!seen[v as usize], "position {i} handed back a repeat");
seen[v as usize] = true;
}
assert!(seen.into_iter().all(|s| s));
assert!(p.at(1_000).is_none());
}
#[test]
fn a_collection_drained_one_draw_at_a_time_stays_correct() {
let mut p = filled(500, 4);
let mut seen = Vec::new();
while !p.is_empty() {
let n = p.len();
let (name, _) = p.at(n - 1).expect("inside");
seen.push(name.to_vec());
assert!(p.remove_at(n - 1).is_some());
assert_eq!(p.len(), n - 1);
}
seen.sort();
seen.dedup();
assert_eq!(seen.len(), 500, "every member came out once");
assert_eq!(p.lengths().iter().sum::<u32>(), 0);
}
#[test]
fn both_levels_of_the_cache_agree_with_the_partitions() {
fn agrees(p: &Parts<u32>) {
for (at, table) in p.tables.iter().enumerate() {
assert_eq!(p.lengths()[at] as usize, table.len(), "partition {at}");
}
for (b, &sum) in p.block_lengths().iter().enumerate() {
let run = &p.lengths()[b * BLOCK..((b + 1) * BLOCK).min(p.lengths().len())];
assert_eq!(sum, run.iter().sum::<u32>(), "block {b}");
}
assert_eq!(
p.len(),
p.block_lengths().iter().map(|&n| n as usize).sum::<usize>()
);
}
let mut p = filled(300, 128);
assert_eq!(p.block_lengths().len(), 2, "128 partitions is two blocks");
agrees(&p);
p.insert(b"fresh", 1).expect("room");
agrees(&p);
p.insert(b"fresh", 2).expect("room");
agrees(&p);
p.remove(b"fresh");
agrees(&p);
p.remove_at(0);
agrees(&p);
p.take_at(p.len() - 1);
agrees(&p);
p.grow_to(256);
agrees(&p);
p.clear();
agrees(&p);
}
#[test]
fn a_small_layout_has_exactly_one_block() {
for parts in [PART_MIN, 8, 32, BLOCK as u32] {
let p: Parts<()> = Parts::with_parts(parts);
assert_eq!(p.block_lengths().len(), 1, "{parts} partitions");
}
let p: Parts<()> = Parts::with_parts(MAX_PARTS);
assert_eq!(p.block_lengths().len(), MAX_PARTS as usize / BLOCK);
}
#[test]
fn the_two_level_resolve_agrees_with_the_flat_one() {
let p = filled(5_000, 128);
for idx in 0..p.len() {
let mut seen = 0usize;
let flat = p
.lengths()
.iter()
.enumerate()
.find_map(|(at, &n)| {
let n = n as usize;
if idx < seen + n {
Some((at, idx - seen))
} else {
seen += n;
None
}
})
.expect("inside the collection");
assert_eq!(p.locate(idx), Some(flat), "position {idx}");
}
assert_eq!(p.locate(p.len()), None);
assert_eq!(p.locate(usize::MAX), None);
}
#[test]
fn taking_by_position_hands_back_the_name() {
let mut p = filled(50, 4);
let (name, value) = p.take_at(0).expect("inside");
assert_eq!(p.len(), 49);
assert!(!p.contains(&name));
assert_eq!(p.get(&name), None);
assert!(value < 50);
}
#[test]
fn a_walk_sees_everything_once() {
let p = filled(3_000, 8);
let mut seen: Vec<u32> = p.iter().map(|(_, &v)| v).collect();
assert_eq!(seen.len(), 3_000);
seen.sort_unstable();
seen.dedup();
assert_eq!(seen.len(), 3_000);
}
#[test]
fn payloads_can_be_rewritten_in_place() {
let mut p = filled(200, 4);
for v in p.payloads_mut() {
*v += 1;
}
assert_eq!(p.get(b"member:0"), Some(&1));
assert_eq!(p.get(b"member:199"), Some(&200));
*p.at_mut(0).expect("inside") = 12345;
assert_eq!(p.iter().filter(|&(_, &v)| v == 12345).count(), 1);
}
#[test]
fn a_scan_sees_everything_once_at_every_page_size() {
let p = filled(1_000, 8);
for page in [1, 2, 7, 64, 999, 1_000, 5_000] {
let mut seen = Vec::new();
let mut cursor = Cursor::START;
let mut rounds = 0;
loop {
cursor = p.scan(cursor, page, |_, &v| seen.push(v));
rounds += 1;
assert!(rounds < 10_000, "the scan is not finishing");
if cursor.is_end() {
break;
}
}
seen.sort_unstable();
let before = seen.len();
seen.dedup();
assert_eq!(before, 1_000, "page {page} returned {before}");
assert_eq!(seen.len(), 1_000, "page {page} repeated a member");
}
}
#[test]
fn a_scan_finishes_at_the_largest_layout() {
let p = filled(8_000, MAX_PARTS);
assert_eq!(p.parts(), MAX_PARTS);
assert_eq!(
Cursor::at(p.parts(), p.parts() - 1, 0).parts(),
p.parts(),
"a cursor has to be able to name the layout it was issued under"
);
let mut seen = Vec::new();
let mut cursor = Cursor::START;
let mut rounds = 0;
loop {
cursor = p.scan(cursor, 10, |_, &v| seen.push(v));
rounds += 1;
assert!(rounds < 20_000, "the scan is not finishing");
if cursor.is_end() {
break;
}
}
seen.sort_unstable();
let before = seen.len();
seen.dedup();
assert_eq!(before, 8_000);
assert_eq!(seen.len(), 8_000);
}
#[test]
fn a_scan_of_an_empty_collection_is_over_at_once() {
let p: Parts<u32> = Parts::with_parts(8);
let mut hits = 0;
assert!(p.scan(Cursor::START, 10, |_, _| hits += 1).is_end());
assert_eq!(hits, 0);
}
#[test]
fn a_page_of_zero_still_makes_progress() {
let p = filled(4, 4);
let mut hits = 0;
let cursor = p.scan(Cursor::START, 0, |_, _| hits += 1);
assert_eq!(hits, 1, "a zero page is read as one, not as none");
assert!(!cursor.is_end() || p.len() == 1);
}
#[test]
fn a_scan_survives_the_collection_growing_underneath_it() {
let mut p = filled(4_000, 4);
let mut seen = Vec::new();
let mut cursor = p.scan(Cursor::START, 700, |_, &v| seen.push(v));
assert!(!cursor.is_end());
assert!(p.grow_to(16));
assert_eq!(p.parts(), 16);
assert_eq!(p.len(), 4_000);
let mut rounds = 0;
loop {
cursor = p.scan(cursor, 700, |_, &v| seen.push(v));
rounds += 1;
assert!(rounds < 1_000, "the scan is not finishing");
if cursor.is_end() {
break;
}
}
seen.sort_unstable();
seen.dedup();
assert_eq!(seen.len(), 4_000, "the growth lost a member");
}
#[test]
fn a_scan_under_removals_returns_everything_that_stayed() {
let mut p = filled(2_000, 8);
let mut seen = Vec::new();
let mut cursor = p.scan(Cursor::START, 300, |_, &v| seen.push(v));
for i in 0..500 {
p.remove(format!("member:{i}").as_bytes());
}
let mut rounds = 0;
loop {
cursor = p.scan(cursor, 300, |_, &v| seen.push(v));
rounds += 1;
assert!(rounds < 1_000, "the scan is not finishing");
if cursor.is_end() {
break;
}
}
seen.sort_unstable();
seen.dedup();
for v in 500..2_000u32 {
assert!(
seen.contains(&v),
"member:{v} stayed and was never returned"
);
}
}
#[test]
fn growing_splits_a_partition_and_moves_nothing_else() {
let mut p = filled(4_000, 4);
let before: Vec<(Vec<u8>, usize)> = (0..4)
.flat_map(|at| {
p.tables[at]
.iter()
.map(move |(n, _)| (n.to_vec(), at))
.collect::<Vec<_>>()
})
.collect();
assert!(p.grow_to(8));
for (name, was) in before {
let h = hash_key(&name);
let now = p.part_of(h) as usize;
assert!(
now == was || now == was + 4,
"{name:?} moved from {was} to {now}, which is not a split"
);
}
assert_eq!(p.len(), 4_000);
assert_eq!(
p.lengths().iter().map(|&n| n as usize).sum::<usize>(),
4_000
);
}
#[test]
fn growing_stops_at_the_ceiling_and_never_shrinks() {
let mut p: Parts<()> = Parts::with_parts(8);
assert!(!p.grow_to(4), "a smaller layout is not a growth");
assert!(!p.grow_to(8), "the same layout is not a growth");
assert!(p.grow_to(MAX_PARTS));
assert_eq!(p.parts(), MAX_PARTS);
assert!(!p.grow_to(MAX_PARTS), "there is nowhere above the ceiling");
}
#[test]
fn the_layout_is_asked_for_rather_than_decided() {
let mut p: Parts<u32> = Parts::with_parts(PART_MIN);
assert_eq!(p.wants_parts(), None);
for i in 0..PART_TARGET * 5 {
p.insert(format!("m{i}").as_bytes(), i as u32)
.expect("room");
}
assert_eq!(p.wants_parts(), Some(8));
assert!(p.grow_to(8));
assert_eq!(p.wants_parts(), None);
}
#[test]
fn promotion_carries_every_element_across() {
let mut one = Elements::<u32>::new();
for i in 0..5_000u32 {
one.insert(format!("member:{i}").as_bytes(), i)
.expect("room");
}
let p = Parts::from_table(&one, 8);
assert_eq!(p.len(), 5_000);
assert_eq!(p.parts(), 8);
for i in 0..5_000u32 {
assert_eq!(p.get(format!("member:{i}").as_bytes()), Some(&i));
}
assert_eq!(
p.lengths().iter().map(|&n| n as usize).sum::<usize>(),
5_000
);
}
#[test]
fn clearing_keeps_the_layout_and_forgets_the_elements() {
let mut p = filled(1_000, 8);
let bytes = p.memory_bytes();
p.clear();
assert!(p.is_empty());
assert_eq!(p.len(), 0);
assert_eq!(p.parts(), 8);
assert!(p.lengths().iter().all(|&n| n == 0));
assert!(p.at(0).is_none());
assert!(p.memory_bytes() <= bytes, "clearing does not allocate");
p.insert(b"back", 1).expect("room");
assert_eq!(p.len(), 1);
}
#[test]
fn the_memory_accounting_adds_the_partitions_up() {
let p = filled(2_000, 8);
assert_eq!(
p.memory_bytes(),
p.slot_bytes() + p.row_bytes() + p.name_bytes() + 9 * size_of::<u32>()
);
assert_eq!(p.dead_name_bytes(), 0);
assert!(p.name_bytes() > 0);
}
}