use super::avl::{AvlMap, ReserveError, grown};
use super::identity::Vid;
use crate::engine::arena::ValueRef;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct RowMeta {
bit: u8,
len: u8,
start: u16,
}
#[derive(Clone, Debug, Default)]
struct SlotPage {
mask: u64,
meta: Vec<RowMeta>,
vals: Vec<ValueRef>,
}
impl SlotPage {
fn heap_bytes(&self) -> usize {
self.meta.capacity() * size_of::<RowMeta>() + self.vals.capacity() * size_of::<ValueRef>()
}
fn find(&self, bit: u8) -> Result<usize, usize> {
self.meta.binary_search_by_key(&bit, |m| m.bit)
}
}
pub const MAX_ARITY: usize = 255;
const NO_PAGE: u32 = u32::MAX;
#[derive(Debug)]
pub struct SlotStore {
pages: Vec<SlotPage>,
free: u32,
nfree: usize,
dir: AvlMap,
rows: usize,
payload: usize,
staged: Vec<(u32, SlotPage)>,
}
impl Clone for SlotStore {
fn clone(&self) -> Self {
let pages: Vec<SlotPage> = self.pages.clone();
let payload = pages.iter().map(SlotPage::heap_bytes).sum();
Self {
pages,
free: self.free,
nfree: self.nfree,
dir: self.dir.clone(),
rows: self.rows,
payload,
staged: self.staged.clone(),
}
}
}
impl Default for SlotStore {
fn default() -> Self {
Self {
pages: Vec::new(),
free: NO_PAGE,
nfree: 0,
dir: AvlMap::new(),
rows: 0,
payload: 0,
staged: Vec::new(),
}
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct SlotPlan {
pub touched: Vec<(u32, usize, usize)>,
pub new_pages: usize,
pub freed_pages: usize,
pub bytes_after: usize,
pub alloc_bytes: usize,
pub max_realloc_old: usize,
pub work: u64,
pub rows_delta: i64,
pub deltas_cap: usize,
slab_cap: usize,
dir_cap: usize,
}
impl SlotPlan {
pub fn scratch_bytes(&self) -> usize {
self.touched.capacity() * size_of::<(u32, usize, usize)>()
+ self.deltas_cap * size_of::<(u32, i64, i64)>()
}
}
fn tpush<T>(v: &mut Vec<T>, x: T) -> Result<(), ReserveError> {
if v.len() == v.capacity() {
v.try_reserve(1).map_err(|_| ReserveError)?;
}
v.push(x);
Ok(())
}
fn try_vec<T>(n: usize) -> Result<Vec<T>, ReserveError> {
let mut v = Vec::new();
v.try_reserve_exact(n).map_err(|_| ReserveError)?;
Ok(v)
}
impl SlotStore {
pub fn rows(&self) -> usize {
self.rows
}
pub fn heap_bytes(&self) -> usize {
self.pages.capacity() * size_of::<SlotPage>()
+ self.dir.heap_bytes()
+ self.payload
+ self.staged_bytes()
}
pub fn census_bytes(&self) -> usize {
self.pages.capacity() * size_of::<SlotPage>()
+ self.dir.heap_bytes()
+ self.pages.iter().map(SlotPage::heap_bytes).sum::<usize>()
+ self.staged_bytes()
}
pub fn pages(&self) -> usize {
self.dir.len()
}
fn staged_bytes(&self) -> usize {
self.staged.capacity() * size_of::<(u32, SlotPage)>()
+ self
.staged
.iter()
.map(|(_, p)| p.heap_bytes())
.sum::<usize>()
}
pub fn drop_staged(&mut self) {
self.staged = Vec::new();
}
fn slot_of_page(&self, page: u32) -> Option<u32> {
self.dir.get(u64::from(page))
}
pub fn get_cached(&self, id: Vid, cache: &mut Option<(u32, u32)>) -> Option<&[ValueRef]> {
let number = id >> 6;
let slot = match *cache {
Some((p, slot)) if p == number => slot,
_ => {
let slot = self.slot_of_page(number)?;
*cache = Some((number, slot));
slot
}
};
let page = &self.pages[slot as usize];
let bit = (id & 63) as u8;
if page.mask >> bit & 1 == 0 {
return None;
}
let m = page.meta[page.find(bit).ok()?];
Some(&page.vals[m.start as usize..m.start as usize + m.len as usize])
}
pub fn get(&self, id: Vid) -> Option<&[ValueRef]> {
let page = &self.pages[self.slot_of_page(id >> 6)? as usize];
let bit = (id & 63) as u8;
if page.mask >> bit & 1 == 0 {
return None;
}
let m = page.meta[page.find(bit).ok()?];
Some(&page.vals[m.start as usize..m.start as usize + m.len as usize])
}
pub fn plan(&self, removes: &[Vid], inserts: &[(Vid, usize)]) -> SlotPlan {
self.try_plan(removes, inserts)
.expect("slot plan allocation")
}
pub fn try_plan(
&self,
removes: &[Vid],
inserts: &[(Vid, usize)],
) -> Result<SlotPlan, ReserveError> {
let mut work = 0u64;
let mut deltas: Vec<(u32, i64, i64)> = Vec::new();
let mut cached: Option<(u32, Option<u32>)> = None;
for &id in removes {
work += 1;
let pg = id >> 6;
let slot = match cached {
Some((p, s)) if p == pg => s,
_ => {
let s = self.slot_of_page(pg);
cached = Some((pg, s));
s
}
};
let Some(slot) = slot else { continue };
let page = &self.pages[slot as usize];
let bit = (id & 63) as u8;
if page.mask >> bit & 1 == 0 {
continue;
}
let len = i64::from(page.meta[page.find(bit).expect("present row")].len);
match deltas.last_mut() {
Some(d) if d.0 == pg => {
d.1 -= 1;
d.2 -= len;
}
_ => tpush(&mut deltas, (pg, -1, -len))?,
}
}
for &(id, n) in inserts {
work += 1;
if n == 0 {
continue;
}
let pg = id >> 6;
match deltas.last_mut() {
Some(d) if d.0 == pg => {
d.1 += 1;
d.2 += n as i64;
}
_ => tpush(&mut deltas, (pg, 1, n as i64))?,
}
}
work += deltas.len() as u64;
deltas.sort_unstable_by_key(|d| d.0);
let deltas_cap = deltas.capacity();
let mut touched: Vec<(u32, usize, usize)> = Vec::new();
let (mut new_pages, mut freed_pages) = (0usize, 0usize);
let mut payload = self.payload;
let mut alloc = 0usize;
let mut max_old = 0usize;
let rows_delta: i64 = deltas.iter().map(|d| d.1).sum();
let mut i = 0;
while i < deltas.len() {
let pg = deltas[i].0;
let (mut dm, mut dv) = (0i64, 0i64);
while i < deltas.len() && deltas[i].0 == pg {
dm += deltas[i].1;
dv += deltas[i].2;
i += 1;
}
work += 1;
match self.slot_of_page(pg) {
Some(slot) => {
let p = &self.pages[slot as usize];
let m = (p.meta.len() as i64 + dm) as usize;
let v = (p.vals.len() as i64 + dv) as usize;
payload -= p.heap_bytes();
if m == 0 {
freed_pages += 1;
} else {
let (mc, vc) = (grown(p.meta.capacity(), m), grown(p.vals.capacity(), v));
let grown_bytes = mc * size_of::<RowMeta>() + vc * size_of::<ValueRef>();
payload += grown_bytes;
if mc > p.meta.capacity() {
alloc += (mc - p.meta.capacity()) * size_of::<RowMeta>();
max_old = max_old.max(p.meta.capacity() * size_of::<RowMeta>());
}
if vc > p.vals.capacity() {
alloc += (vc - p.vals.capacity()) * size_of::<ValueRef>();
max_old = max_old.max(p.vals.capacity() * size_of::<ValueRef>());
}
}
tpush(&mut touched, (pg, m, v))?;
}
None => {
debug_assert!(dm > 0, "rows removed from an absent page");
let (m, v) = (dm as usize, dv as usize);
new_pages += 1;
let b = m * size_of::<RowMeta>() + v * size_of::<ValueRef>();
payload += b;
alloc += b;
tpush(&mut touched, (pg, m, v))?;
}
}
}
drop(deltas);
let len = self.pages.len() + new_pages.saturating_sub(self.nfree);
let slab_cap = grown(self.pages.capacity(), len);
if slab_cap > self.pages.capacity() {
alloc += (slab_cap - self.pages.capacity()) * size_of::<SlotPage>();
max_old = max_old.max(self.pages.capacity() * size_of::<SlotPage>());
}
let mut dsh = self.dir.shadow();
for _ in 0..new_pages {
dsh.alloc();
}
for _ in 0..freed_pages {
dsh.release();
}
let dir_cap = dsh.cap;
if dir_cap > self.dir.shadow().cap {
alloc += (dir_cap - self.dir.shadow().cap) * AvlMap::NODE_BYTES;
max_old = max_old.max(self.dir.heap_bytes());
}
if new_pages > 0 {
alloc += new_pages * size_of::<(u32, SlotPage)>();
}
Ok(SlotPlan {
touched,
new_pages,
freed_pages,
bytes_after: slab_cap * size_of::<SlotPage>() + dir_cap * AvlMap::NODE_BYTES + payload,
alloc_bytes: alloc,
max_realloc_old: max_old,
work,
rows_delta,
deltas_cap,
slab_cap,
dir_cap,
})
}
pub fn planned_bytes(&self, plan: &SlotPlan) -> usize {
plan.bytes_after
}
pub fn try_reserve(&mut self, plan: &SlotPlan) -> Result<(), ReserveError> {
self.drop_staged();
if plan.slab_cap > self.pages.capacity() {
self.pages
.try_reserve_exact(plan.slab_cap - self.pages.len())
.map_err(|_| ReserveError)?;
}
self.dir.try_reserve_slots(plan.dir_cap)?;
let mut staged: Vec<(u32, SlotPage)> = try_vec(plan.new_pages)?;
for &(pg, m, v) in &plan.touched {
match self.slot_of_page(pg) {
Some(slot) if m > 0 => {
let p = &mut self.pages[slot as usize];
let before = p.heap_bytes();
let (mc, vc) = (grown(p.meta.capacity(), m), grown(p.vals.capacity(), v));
let mut r = Ok(());
if mc > p.meta.capacity() {
r = p.meta.try_reserve_exact(mc - p.meta.len());
}
if r.is_ok() && vc > p.vals.capacity() {
r = p.vals.try_reserve_exact(vc - p.vals.len());
}
self.payload += p.heap_bytes() - before;
r.map_err(|_| ReserveError)?;
}
Some(_) => {}
None => staged.push((
pg,
SlotPage {
mask: 0,
meta: try_vec(m)?,
vals: try_vec(v)?,
},
)),
}
}
self.staged = staged;
Ok(())
}
pub fn apply(&mut self, plan: &SlotPlan, removes: &[Vid], inserts: &[(Vid, &[ValueRef])]) {
for &id in removes {
self.remove(id);
}
for &(id, row) in inserts {
if !row.is_empty() {
self.insert(id, row);
}
}
for &(pg, m, _) in &plan.touched {
if m == 0
&& let Some(slot) = self.slot_of_page(pg)
{
self.free_page(pg, slot);
}
}
debug_assert!(
self.staged.iter().all(|(_, p)| p.heap_bytes() == 0),
"unused staged pages"
);
self.staged = Vec::new();
}
fn page_for_insert(&mut self, pg: u32) -> u32 {
if let Some(slot) = self.slot_of_page(pg) {
return slot;
}
let page = match self.staged.binary_search_by_key(&pg, |x| x.0) {
Ok(i) => std::mem::take(&mut self.staged[i].1),
Err(_) => {
debug_assert!(false, "unstaged slot page");
SlotPage::default()
}
};
self.payload += page.heap_bytes();
let slot = if self.free != NO_PAGE {
let slot = self.free;
self.free = self.pages[slot as usize].mask as u32;
self.nfree -= 1;
self.pages[slot as usize] = page;
slot
} else {
debug_assert!(
self.pages.len() < self.pages.capacity(),
"unreserved page slab"
);
self.pages.push(page);
(self.pages.len() - 1) as u32
};
self.dir.insert(u64::from(pg), slot);
slot
}
fn free_page(&mut self, pg: u32, slot: u32) {
let p = &mut self.pages[slot as usize];
debug_assert_eq!(p.mask, 0, "freeing a page with rows");
self.payload -= p.heap_bytes();
*p = SlotPage {
mask: u64::from(self.free),
meta: Vec::new(),
vals: Vec::new(),
};
self.free = slot;
self.nfree += 1;
self.dir.remove(u64::from(pg));
}
fn insert(&mut self, id: Vid, row: &[ValueRef]) {
assert!(row.len() <= MAX_ARITY, "slot row over {MAX_ARITY} literals");
let slot = self.page_for_insert(id >> 6);
let page = &mut self.pages[slot as usize];
let bit = (id & 63) as u8;
debug_assert_eq!(page.mask >> bit & 1, 0, "slot row exists");
debug_assert!(
page.meta.len() < page.meta.capacity()
&& page.vals.len() + row.len() <= page.vals.capacity(),
"unreserved slot row"
);
let at = page.find(bit).expect_err("absent row");
let start = page
.meta
.get(at)
.map_or(page.vals.len(), |m| m.start as usize);
let n = row.len();
for m in &mut page.meta[at..] {
m.start += n as u16;
}
page.meta.insert(
at,
RowMeta {
bit,
len: n as u8,
start: start as u16,
},
);
page.vals.extend_from_slice(row);
page.vals[start..].rotate_right(n);
page.mask |= 1 << bit;
self.rows += 1;
}
fn remove(&mut self, id: Vid) {
let Some(slot) = self.slot_of_page(id >> 6) else {
return;
};
let page = &mut self.pages[slot as usize];
let bit = (id & 63) as u8;
if page.mask >> bit & 1 == 0 {
return;
}
let at = page.find(bit).expect("present row");
let m = page.meta.remove(at);
let (s, n) = (m.start as usize, m.len as usize);
page.vals.drain(s..s + n);
for m in &mut page.meta[at..] {
m.start -= n as u16;
}
page.mask &= !(1 << bit);
self.rows -= 1;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rows_round_trip_and_plan_predicts_bytes() {
let mut s = SlotStore::default();
let v = ValueRef::from_raw;
let mut model: std::collections::BTreeMap<u32, Vec<ValueRef>> = Default::default();
let mut x: u32 = 5;
for step in 0..3000u32 {
x = x.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
let id = (x >> 8) % 300;
let (removes, inserts): (Vec<u32>, Vec<(u32, Vec<ValueRef>)>) =
if model.contains_key(&id) {
if step % 2 == 0 {
(vec![id], vec![])
} else {
let n = (x >> 24) as usize % 4;
(
vec![id],
vec![(id, (0..n as u32).map(|k| v(step + k)).collect())],
)
}
} else {
let n = (x >> 24) as usize % 4;
(
vec![],
vec![(id, (0..n as u32).map(|k| v(step * 7 + k)).collect())],
)
};
let ins: Vec<(u32, usize)> = inserts.iter().map(|(i, r)| (*i, r.len())).collect();
let plan = s.plan(&removes, &ins);
let predicted = s.planned_bytes(&plan);
s.try_reserve(&plan).unwrap();
let rows: Vec<(u32, &[ValueRef])> =
inserts.iter().map(|(i, r)| (*i, r.as_slice())).collect();
s.apply(&plan, &removes, &rows);
assert_eq!(s.heap_bytes(), predicted, "step {step}");
for id in &removes {
model.remove(id);
}
for (i, r) in &inserts {
if !r.is_empty() {
model.insert(*i, r.clone());
}
}
if step % 37 == 0 {
for id in 0..300u32 {
assert_eq!(s.get(id), model.get(&id).map(Vec::as_slice), "id {id}");
}
}
}
assert_eq!(s.rows(), model.len());
}
}