use super::ARC_HEADER;
use std::sync::Arc;
use super::chunked_vec::{CHUNK, FANOUT};
use crate::varint;
const SHIFT: u32 = CHUNK.trailing_zeros();
const MASK: usize = CHUNK - 1;
const FAN_SHIFT: u32 = FANOUT.trailing_zeros();
const FAN_MASK: usize = FANOUT - 1;
const SPAN_SHIFT: u32 = SHIFT + FAN_SHIFT;
const SPAN_MASK: usize = (CHUNK * FANOUT) - 1;
const EMPTY: u32 = u32::MAX;
const LOOSE: u32 = 1 << 31;
const LARGE_MIN: usize = 1 << 16;
const LOOSE_MAX: u32 = 32;
const LOOSE_SLACK: u32 = 8;
const SEG_SLOTS: usize = 64;
const SEG_SHIFT: u32 = SEG_SLOTS.trailing_zeros();
const SEGMENTS: usize = CHUNK / SEG_SLOTS;
const SEG_MAX: usize = 2 * (SEG_SLOTS + 1) * (LARGE_MIN + 3);
const _: () = assert!(SEG_MAX < LOOSE as usize);
const _: () = assert!(CHUNK.is_multiple_of(SEG_SLOTS) && SEG_SLOTS.is_power_of_two());
const _: () = assert!(2 * CHUNK + LOOSE_SLACK as usize + 1 < LOOSE as usize);
#[derive(Clone, Copy)]
enum Place {
Empty,
Segment(usize),
Loose(usize),
}
impl Place {
#[inline]
fn of(offset: u32) -> Place {
if offset == EMPTY {
Place::Empty
} else if offset & LOOSE != 0 {
Place::Loose((offset & !LOOSE) as usize)
} else {
Place::Segment(offset as usize)
}
}
}
#[inline]
fn segment_of(slot: usize) -> usize {
slot >> SEG_SHIFT
}
#[derive(Clone, Default)]
struct Segment {
buf: Option<Arc<[u8]>>,
len: u32,
garbage: u32,
}
impl Segment {
#[inline]
fn bytes(&self) -> &[u8] {
self.buf.as_deref().unwrap_or(&[])
}
#[inline]
fn live_bytes(&self) -> usize {
(self.len - self.garbage) as usize
}
#[inline]
fn record(&self, offset: usize) -> &[u8] {
let bytes = self.bytes();
let mut pos = offset;
let len = varint::read_len(bytes, &mut pos);
&bytes[pos..pos + len]
}
#[inline]
fn stored_len(&self, offset: usize) -> usize {
let mut pos = offset;
let len = varint::read_len(self.bytes(), &mut pos);
pos - offset + len
}
}
#[derive(Clone)]
struct Leaf {
offsets: [u32; CHUNK],
segments: [Segment; SEGMENTS],
loose: Vec<Option<Arc<[u8]>>>,
loose_garbage: u32,
loose_small: u32,
live: u32,
}
impl Leaf {
fn new() -> Self {
Leaf {
offsets: [EMPTY; CHUNK],
segments: Default::default(),
loose: Vec::new(),
loose_garbage: 0,
loose_small: 0,
live: 0,
}
}
#[inline]
fn get(&self, slot: usize) -> Option<&[u8]> {
match Place::of(self.offsets[slot]) {
Place::Empty => None,
Place::Segment(offset) => Some(self.segments[segment_of(slot)].record(offset)),
Place::Loose(index) => self.loose[index].as_deref(),
}
}
#[inline]
fn contains(&self, slot: usize) -> bool {
self.offsets[slot] != EMPTY
}
fn release(&mut self, slot: usize) {
match Place::of(self.offsets[slot]) {
Place::Empty => return,
Place::Segment(offset) => {
let segment = &mut self.segments[segment_of(slot)];
segment.garbage += segment.stored_len(offset) as u32;
}
Place::Loose(index) => {
let record = self.loose[index].take();
if record.is_some_and(|record| record.len() < LARGE_MIN) {
self.loose_small -= 1;
}
self.loose_garbage += 1;
}
}
self.offsets[slot] = EMPTY;
self.live -= 1;
}
fn set(&mut self, slot: usize, bytes: &[u8]) {
let update = self.contains(slot);
self.release(slot);
self.live += 1;
if bytes.len() >= LARGE_MIN {
self.push_loose(slot, bytes);
return;
}
let need = varint::len(bytes.len() as u64) + bytes.len();
if self.append(slot, bytes, need) {
return;
}
if update && self.loose_small < LOOSE_MAX {
self.loose_small += 1;
self.push_loose(slot, bytes);
return;
}
self.rebuild(segment_of(slot), need);
let appended = self.append(slot, bytes, need);
assert!(appended, "rebuild leaves the segment unshared, with room");
}
#[inline]
fn append(&mut self, slot: usize, bytes: &[u8], need: usize) -> bool {
let segment = &mut self.segments[segment_of(slot)];
let start = segment.len as usize;
let Some(buf) = segment
.buf
.as_mut()
.and_then(Arc::get_mut)
.filter(|buf| buf.len() - start >= need)
else {
return false;
};
let mut pos = start;
varint::write(buf, &mut pos, bytes.len() as u64);
buf[pos..pos + bytes.len()].copy_from_slice(bytes);
segment.len = (pos + bytes.len()) as u32;
self.offsets[slot] = start as u32;
true
}
fn push_loose(&mut self, slot: usize, bytes: &[u8]) {
if self.loose_garbage >= LOOSE_SLACK && self.loose_garbage as usize > self.loose.len() / 2 {
self.compact_loose();
}
self.offsets[slot] = LOOSE | self.loose.len() as u32;
self.loose.push(Some(bytes.into()));
}
#[inline]
fn foldable(&self, index: usize) -> Option<&[u8]> {
self.loose[index]
.as_deref()
.filter(|record| record.len() < LARGE_MIN)
}
fn remove(&mut self, slot: usize) -> bool {
let place = Place::of(self.offsets[slot]);
if matches!(place, Place::Empty) {
return false;
}
self.release(slot);
if self.live == 0 {
*self = Leaf::new();
} else if let Place::Segment(_) = place {
let seg = segment_of(slot);
let segment = &self.segments[seg];
if segment.garbage == segment.len {
self.segments[seg] = Segment::default();
} else if segment.garbage > 4096 && segment.garbage > segment.len / 2 {
self.rebuild(seg, 0);
}
}
true
}
#[inline]
fn slots(seg: usize) -> std::ops::Range<usize> {
seg * SEG_SLOTS..(seg + 1) * SEG_SLOTS
}
fn spare(&self, seg: usize, live: usize, live_bytes: usize, need: usize) -> usize {
if need == 0 {
return 0;
}
let empty = SEG_SLOTS.saturating_sub(live + 1);
let others = (self.live - 1) as usize;
if live == 0 && others > 0 {
let leaf_bytes: usize = self.segments.iter().map(Segment::live_bytes).sum();
debug_assert_eq!(self.segments[seg].live_bytes(), 0);
empty * (leaf_bytes / others).max(1)
} else {
let avg = live_bytes.checked_div(live).unwrap_or(need);
(empty.min(live.max(1)) * avg).max(live_bytes / 8)
}
}
fn rebuild(&mut self, seg: usize, need: usize) {
let mut live_bytes = self.segments[seg].live_bytes();
let mut live = 0;
for &offset in &self.offsets[Self::slots(seg)] {
match Place::of(offset) {
Place::Empty => {}
Place::Segment(_) => live += 1,
Place::Loose(index) => {
if let Some(record) = self.foldable(index) {
live += 1;
live_bytes += varint::len(record.len() as u64) + record.len();
}
}
}
}
let size = live_bytes + need + self.spare(seg, live, live_bytes, need);
debug_assert!(size <= SEG_MAX, "segment of {size} bytes");
let mut buf: Arc<[u8]> = std::iter::repeat_n(0, size).collect();
let out = Arc::get_mut(&mut buf).expect("a new buffer is unshared");
let old = &self.segments[seg];
let mut end = 0;
for slot in Self::slots(seg) {
let (record, folded) = match Place::of(self.offsets[slot]) {
Place::Empty => continue,
Place::Segment(offset) => (old.record(offset), None),
Place::Loose(index) => match self.loose[index].as_deref() {
Some(record) if record.len() < LARGE_MIN => (record, Some(index)),
_ => continue,
},
};
self.offsets[slot] = end as u32;
varint::write(out, &mut end, record.len() as u64);
out[end..end + record.len()].copy_from_slice(record);
end += record.len();
if let Some(index) = folded {
self.loose[index] = None;
self.loose_small -= 1;
self.loose_garbage += 1;
}
}
self.segments[seg] = Segment {
buf: Some(buf),
len: end as u32,
garbage: 0,
};
}
fn compact_loose(&mut self) {
let mut loose = Vec::with_capacity(self.loose.len() - self.loose_garbage as usize);
for offset in &mut self.offsets {
if let Place::Loose(index) = Place::of(*offset) {
*offset = LOOSE | loose.len() as u32;
loose.push(self.loose[index].take());
}
}
self.loose = loose;
self.loose_garbage = 0;
}
fn heap_bytes(&self) -> usize {
self.segments
.iter()
.filter_map(|segment| segment.buf.as_ref())
.map(|buf| ARC_HEADER + buf.len())
.sum::<usize>()
+ self.loose.capacity() * size_of::<Option<Arc<[u8]>>>()
+ self
.loose
.iter()
.flatten()
.map(|record| ARC_HEADER + record.len())
.sum::<usize>()
}
#[cfg(test)]
fn check_invariants(&self) {
let mut live = 0;
let mut loose_small = 0;
let mut referenced = vec![false; self.loose.len()];
let mut used = [0usize; SEGMENTS];
for (slot, &offset) in self.offsets.iter().enumerate() {
match Place::of(offset) {
Place::Empty => continue,
Place::Segment(offset) => {
let segment = &self.segments[segment_of(slot)];
let stored = segment.stored_len(offset);
assert!(
offset + stored <= segment.len as usize,
"slot {slot} past len"
);
used[segment_of(slot)] += stored;
}
Place::Loose(index) => {
let record = self.loose[index].as_deref().expect("live loose entry");
assert!(
!referenced[index],
"loose entry {index} shared by two slots"
);
referenced[index] = true;
loose_small += usize::from(record.len() < LARGE_MIN);
}
}
live += 1;
}
assert_eq!(live, self.live as usize, "live");
assert_eq!(loose_small, self.loose_small as usize, "loose_small");
assert!(self.loose_small <= LOOSE_MAX);
let released = self.loose.iter().filter(|r| r.is_none()).count();
assert_eq!(released, self.loose_garbage as usize, "loose_garbage");
assert!(
self.loose
.iter()
.zip(&referenced)
.all(|(r, &seen)| r.is_none() || seen),
"an unreferenced live loose entry"
);
for (seg, segment) in self.segments.iter().enumerate() {
assert!(segment.len as usize <= segment.bytes().len());
assert_eq!(used[seg], segment.live_bytes(), "segment {seg} live bytes");
}
}
}
type Interior = Arc<[Arc<Leaf>]>;
pub(super) struct RecordSlab {
root: Arc<[Interior]>,
len: usize,
}
impl Default for RecordSlab {
fn default() -> Self {
Self {
root: Arc::new([]),
len: 0,
}
}
}
impl Clone for RecordSlab {
fn clone(&self) -> Self {
Self {
root: Arc::clone(&self.root),
len: self.len,
}
}
}
impl std::fmt::Debug for RecordSlab {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_map()
.entries(self.iter().map(|(i, bytes)| (i, bytes.len())))
.finish()
}
}
impl RecordSlab {
#[inline]
pub(super) fn len(&self) -> usize {
self.len
}
#[inline]
fn leaf(&self, i: usize) -> &Leaf {
debug_assert!(i < self.len);
&self.root[i >> SPAN_SHIFT][(i >> SHIFT) & FAN_MASK]
}
#[inline]
pub(super) fn get(&self, i: usize) -> Option<&[u8]> {
if i >= self.len {
return None;
}
self.leaf(i).get(i & MASK)
}
#[inline]
pub(super) fn contains(&self, i: usize) -> bool {
i < self.len && self.leaf(i).contains(i & MASK)
}
fn leaf_mut(&mut self, i: usize) -> &mut Leaf {
assert!(i < self.len, "record slot {i} out of bounds");
let node = &mut Arc::make_mut(&mut self.root)[i >> SPAN_SHIFT];
Arc::make_mut(&mut Arc::make_mut(node)[(i >> SHIFT) & FAN_MASK])
}
pub(super) fn set(&mut self, i: usize, bytes: &[u8]) -> bool {
let leaf = self.leaf_mut(i);
let was_empty = !leaf.contains(i & MASK);
leaf.set(i & MASK, bytes);
was_empty
}
pub(super) fn remove(&mut self, i: usize) -> bool {
if !self.contains(i) {
return false;
}
self.leaf_mut(i).remove(i & MASK)
}
pub(super) fn try_reserve(
&self,
additional: usize,
) -> Result<(), std::collections::TryReserveError> {
let leaves = self
.len
.checked_add(additional)
.map(|total| total.div_ceil(CHUNK))
.and_then(|leaves| leaves.checked_mul(size_of::<Leaf>()));
if leaves.is_none_or(|bytes| bytes > isize::MAX as usize) {
return Err(Vec::<u8>::new()
.try_reserve_exact(usize::MAX)
.expect_err("usize::MAX bytes exceeds isize::MAX"));
}
Ok(())
}
pub(super) fn grow_to(&mut self, new_len: usize) {
while self.len.div_ceil(CHUNK) < new_len.div_ceil(CHUNK) {
let first = self.len.div_ceil(CHUNK) * CHUNK;
let leaf = Arc::new(Leaf::new());
if first & SPAN_MASK == 0 {
self.root = self
.root
.iter()
.cloned()
.chain(std::iter::once(Arc::from([leaf])))
.collect();
} else {
let root = Arc::make_mut(&mut self.root);
let node = root.last_mut().expect("first slot is not a span start");
*node = node.iter().cloned().chain(std::iter::once(leaf)).collect();
}
self.len = first + 1;
}
self.len = self.len.max(new_len);
}
fn leaves_from(
&self,
from: usize,
) -> impl Iterator<Item = (usize, &Leaf, std::ops::Range<usize>)> + '_ {
let first = from / CHUNK;
(first..self.len.div_ceil(CHUNK)).map(move |l| {
let base = l * CHUNK;
let start = if l == first { from & MASK } else { 0 };
let end = (self.len - base).min(CHUNK);
(base, self.leaf(base), start..end)
})
}
pub(super) fn iter_from(&self, from: usize) -> impl Iterator<Item = (usize, &[u8])> + '_ {
self.leaves_from(from).flat_map(|(base, leaf, slots)| {
slots.filter_map(move |s| leaf.get(s).map(|bytes| (base + s, bytes)))
})
}
pub(super) fn iter(&self) -> impl Iterator<Item = (usize, &[u8])> + '_ {
self.iter_from(0)
}
pub(super) fn slots_from(&self, from: usize) -> impl Iterator<Item = usize> + '_ {
self.leaves_from(from).flat_map(|(base, leaf, slots)| {
slots
.filter(move |&s| leaf.contains(s))
.map(move |s| base + s)
})
}
pub(super) fn heap_bytes(&self) -> usize {
let tree = ARC_HEADER
+ self
.root
.iter()
.map(|node| {
size_of::<Interior>() + ARC_HEADER + node.len() * size_of::<Arc<Leaf>>()
})
.sum::<usize>();
let leaves: usize = self
.root
.iter()
.flat_map(|node| node.iter())
.map(|leaf| ARC_HEADER + size_of::<Leaf>() + leaf.heap_bytes())
.sum();
tree + leaves
}
#[cfg(test)]
fn check_invariants(&self) {
let leaves = self.len.div_ceil(CHUNK);
assert_eq!(self.root.len(), leaves.div_ceil(FANOUT), "interior nodes");
for (n, node) in self.root.iter().enumerate() {
let expected = if n + 1 < self.root.len() {
FANOUT
} else {
leaves - n * FANOUT
};
assert_eq!(node.len(), expected, "leaves of node {n}");
}
for (base, leaf, slots) in self.leaves_from(0) {
leaf.check_invariants();
assert!(
(slots.end..CHUNK).all(|s| !leaf.contains(s)),
"a record past len in the leaf at {base}"
);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::memory::test_util::splitmix;
use std::collections::BTreeMap;
fn record(seed: u64) -> Vec<u8> {
let len = match seed % 50 {
0 => LARGE_MIN + (seed % 7) as usize,
1..=5 => 200 + (seed % 300) as usize,
_ => (seed % 40) as usize,
};
(0..len).map(|i| (seed as usize + i) as u8).collect()
}
fn check(slab: &RecordSlab, model: &BTreeMap<usize, Vec<u8>>, len: usize) {
slab.check_invariants();
assert_eq!(slab.len(), len);
for i in 0..len + 3 {
assert_eq!(slab.get(i), model.get(&i).map(|v| v.as_slice()), "slot {i}");
assert_eq!(slab.contains(i), model.contains_key(&i));
}
let all: Vec<(usize, Vec<u8>)> = slab.iter().map(|(i, b)| (i, b.to_vec())).collect();
let expected: Vec<(usize, Vec<u8>)> = model.iter().map(|(&i, v)| (i, v.clone())).collect();
assert_eq!(all, expected);
for from in [0, 1, CHUNK - 1, CHUNK, CHUNK + 1, len / 2, len, len + 1] {
assert!(slab
.iter_from(from)
.map(|(i, _)| i)
.eq(model.range(from..).map(|(&i, _)| i)));
assert!(slab
.slots_from(from)
.eq(model.range(from..).map(|(&i, _)| i)));
}
}
fn step(
slab: &mut RecordSlab,
model: &mut BTreeMap<usize, Vec<u8>>,
len: usize,
x: u64,
removes: u64,
) {
let i = (splitmix(x) % len as u64) as usize;
if splitmix(x ^ 0xA5A5) % 100 < removes {
assert_eq!(slab.remove(i), model.remove(&i).is_some());
} else {
let r = record(splitmix(x ^ 0xFEED));
assert_eq!(slab.set(i, &r), model.insert(i, r).is_none());
}
slab.check_invariants();
}
fn run_model(seed: u64, steps: u64, removes: u64) {
let mut slab = RecordSlab::default();
let mut model = BTreeMap::new();
let len = CHUNK * 3 + 17;
for i in 0..len {
slab.grow_to(i + 1);
if splitmix(seed ^ i as u64) % 11 != 3 {
let r = record(seed ^ i as u64);
assert!(slab.set(i, &r));
model.insert(i, r);
}
}
check(&slab, &model, len);
let snapshot = (slab.clone(), model.clone());
let mut fork: Option<(RecordSlab, BTreeMap<usize, Vec<u8>>)> = None;
for n in 0..steps {
let x = seed.wrapping_mul(0x9E37) ^ n;
step(&mut slab, &mut model, len, x, removes);
if n % 997 == 0 {
if let Some((f, fm)) = &fork {
check(f, fm, len);
}
fork = Some((slab.clone(), model.clone()));
}
if let Some((f, fm)) = &mut fork {
step(f, fm, len, !x, removes);
}
}
check(&slab, &model, len);
check(&snapshot.0, &snapshot.1, len);
if let Some((f, fm)) = &fork {
check(f, fm, len);
}
for i in 0..len {
assert_eq!(slab.remove(i), model.remove(&i).is_some());
}
slab.check_invariants();
assert!(slab.iter().next().is_none());
}
#[test]
fn matches_a_map_through_sets_removes_and_clones() {
for seed in 0..4 {
run_model(seed, 6_000, 25);
}
}
#[test]
fn matches_a_map_when_removes_dominate() {
run_model(17, 8_000, 70);
}
#[test]
fn grows_across_interior_boundaries() {
let mut slab = RecordSlab::default();
let span = CHUNK * FANOUT;
slab.grow_to(span + 5);
assert!(slab.set(span + 4, b"far"));
assert!(slab.set(0, b"near"));
assert_eq!(slab.get(span + 4), Some(&b"far"[..]));
assert_eq!(slab.get(0), Some(&b"near"[..]));
assert_eq!(slab.get(span + 5), None);
assert_eq!(slab.iter().count(), 2);
slab.grow_to(3);
assert_eq!(slab.len(), span + 5);
}
fn shares_segment(a: &RecordSlab, b: &RecordSlab, seg: usize) -> bool {
match (
&a.root[0][0].segments[seg].buf,
&b.root[0][0].segments[seg].buf,
) {
(Some(x), Some(y)) => Arc::ptr_eq(x, y),
_ => false,
}
}
fn full_leaf() -> RecordSlab {
let mut slab = RecordSlab::default();
slab.grow_to(CHUNK);
for i in 0..CHUNK {
slab.set(i, &[i as u8; 40]);
}
slab
}
#[test]
fn a_write_after_a_clone_copies_one_segment_at_most() {
let mut slab = full_leaf();
let snapshot = slab.clone();
slab.set(SEG_SLOTS + 3, b"new");
assert!(!Arc::ptr_eq(&slab.root[0][0], &snapshot.root[0][0]));
assert!((0..SEGMENTS).all(|seg| shares_segment(&slab, &snapshot, seg)));
assert_eq!(slab.get(SEG_SLOTS + 3), Some(&b"new"[..]));
assert_eq!(
snapshot.get(SEG_SLOTS + 3),
Some(&[(SEG_SLOTS + 3) as u8; 40][..])
);
slab.remove(SEG_SLOTS + 5);
slab.set(SEG_SLOTS + 5, b"again");
for seg in 0..SEGMENTS {
assert_eq!(
shares_segment(&slab, &snapshot, seg),
seg != 1,
"segment {seg}"
);
}
assert_eq!(slab.root[0][0].loose_small, 0);
assert_eq!(slab.get(SEG_SLOTS + 3), Some(&b"new"[..]));
assert_eq!(slab.get(SEG_SLOTS + 5), Some(&b"again"[..]));
}
#[test]
fn updates_past_the_loose_cap_rebuild_their_segment() {
let mut slab = full_leaf();
for i in 0..LOOSE_MAX as usize {
slab.set(i * 7, &[0xAB; 41]);
}
assert_eq!(slab.root[0][0].loose_small, LOOSE_MAX);
let folded = (0..LOOSE_MAX as usize)
.filter(|i| i * 7 < SEG_SLOTS)
.count() as u32;
slab.set(1, &[0xCD; 41]);
assert_eq!(slab.root[0][0].loose_small, LOOSE_MAX - folded);
for i in 0..CHUNK {
let expected = if i == 1 {
vec![0xCD; 41]
} else if i % 7 == 0 && i / 7 < LOOSE_MAX as usize {
vec![0xAB; 41]
} else {
vec![i as u8; 40]
};
assert_eq!(slab.get(i), Some(expected.as_slice()), "slot {i}");
}
}
#[test]
fn a_full_leaf_of_small_records_costs_little_beyond_its_bytes() {
let mut slab = RecordSlab::default();
let n = CHUNK * 8;
for i in 0..n {
slab.grow_to(i + 1);
slab.set(i, &[7u8; 11]);
}
let per_record = slab.heap_bytes() as f64 / n as f64;
assert!(per_record < 17.5, "{per_record}");
}
}