use alloc::alloc::handle_alloc_error;
use core::alloc::Layout;
use core::fmt;
use core::marker::PhantomData;
use core::mem::offset_of;
#[cfg(feature = "file")]
use core::ops::Range;
use core::ptr::NonNull;
use core::sync::atomic::{AtomicBool, AtomicU16, Ordering};
use super::super::interleaving::Interleaving;
use super::internal::Internal;
use super::latch::NodeLatch;
use super::layout::{Header, InternalLayout, KeyArea, LeafLayout, NodeLayout};
use super::leaf::Leaf;
use super::node_plan::NodePlan;
use super::optimistic::RacingReads;
use super::root_page::RootPage;
use super::slots::{InlineKeySlots, KeyColumns, KeySlots, SeparatorSlots, Slots};
use crate::allocator::{AllocError, Allocator, PageKind};
use crate::inline_key::sealed::Sealed;
use crate::inline_key::{InlineKeyStrategy, InlineWidth};
pub(crate) unsafe trait PageType: Sized {
const KIND: PageKind;
fn page_layout() -> Layout;
}
pub(in crate::tree) unsafe trait NodeKind: PageType {
const IS_LEAF: bool;
fn has_room_for_hash_tags(inline_width: InlineWidth) -> bool;
}
pub(crate) struct FreshPage<P> {
page: NonNull<u8>,
_page_type: PhantomData<fn() -> P>,
}
pub(crate) enum PageByType<K, V, S> {
Leaf(FreshPage<Leaf<K, V, S>>),
Internal(FreshPage<Internal<K, V, S>>),
Root(FreshPage<RootPage<K, V, S>>),
}
impl<P: PageType> FreshPage<P> {
unsafe fn from_raw(page: NonNull<u8>) -> Self {
debug_assert!(
page.addr().get().is_multiple_of(P::page_layout().align()),
"node page is not aligned to its layout"
);
Self {
page,
_page_type: PhantomData,
}
}
}
impl<P> FreshPage<P> {
pub(in crate::tree) fn into_raw(self) -> NonNull<u8> {
self.page
}
}
pub(crate) unsafe trait PageAllocator {
fn allocate_page(&self, kind: PageKind, layout: Layout) -> NonNull<u8>;
unsafe fn deallocate_page(&self, kind: PageKind, page: NonNull<u8>, layout: Layout);
}
unsafe impl<A: Allocator + ?Sized> PageAllocator for A {
fn allocate_page(&self, kind: PageKind, layout: Layout) -> NonNull<u8> {
kind.allocate(self, layout)
.unwrap_or_else(|AllocError| handle_alloc_error(layout))
}
unsafe fn deallocate_page(&self, kind: PageKind, page: NonNull<u8>, layout: Layout) {
unsafe { kind.deallocate(self, page, layout) }
}
}
impl<P: PageType> FreshPage<P> {
pub(crate) fn allocate(alloc: &impl PageAllocator) -> Self {
let page = alloc.allocate_page(P::KIND, P::page_layout());
unsafe { Self::from_raw(page) }
}
pub(crate) fn try_allocate(alloc: &impl Allocator) -> Result<Self, Layout> {
let layout = P::page_layout();
let page = P::KIND
.allocate(alloc, layout)
.map_err(|AllocError| layout)?;
Ok(unsafe { Self::from_raw(page) })
}
pub(in crate::tree) fn deallocate(self, alloc: &impl PageAllocator) {
let page = self.into_raw();
unsafe { alloc.deallocate_page(P::KIND, page, P::page_layout()) }
}
}
impl<P> fmt::Debug for FreshPage<P> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("FreshPage").field(&self.page).finish()
}
}
impl<K, V, S: InlineKeyStrategy<K>> PageByType<K, V, S> {
pub(in crate::tree) fn deallocate(self, alloc: &impl PageAllocator) {
match self {
Self::Leaf(page) => page.deallocate(alloc),
Self::Internal(page) => page.deallocate(alloc),
Self::Root(page) => page.deallocate(alloc),
}
}
}
impl<K, V, S> From<FreshPage<Leaf<K, V, S>>> for PageByType<K, V, S> {
fn from(page: FreshPage<Leaf<K, V, S>>) -> Self {
Self::Leaf(page)
}
}
impl<K, V, S> From<FreshPage<Internal<K, V, S>>> for PageByType<K, V, S> {
fn from(page: FreshPage<Internal<K, V, S>>) -> Self {
Self::Internal(page)
}
}
impl<K, V, S> From<FreshPage<RootPage<K, V, S>>> for PageByType<K, V, S> {
fn from(page: FreshPage<RootPage<K, V, S>>) -> Self {
Self::Root(page)
}
}
#[repr(C)]
pub(in crate::tree) struct NodeHeader {
len: u16,
has_head: bool,
}
#[repr(C)]
pub(in crate::tree) struct InlineKeyNodeHeader {
header: NodeHeader,
prefix_len: u16,
}
#[cfg(feature = "file")]
impl NodeHeader {
pub(in crate::tree) const LEN_OFFSET: usize = offset_of!(Self, len);
pub(in crate::tree) const HAS_HEAD_OFFSET: usize = offset_of!(Self, has_head);
const PADDING: Range<usize> = Self::HAS_HEAD_OFFSET + size_of::<bool>()..size_of::<Self>();
}
impl InlineKeyNodeHeader {
pub(super) const MAX_PREFIX_LEN: usize = u16::MAX as usize;
}
#[repr(C)]
pub(in crate::tree) struct LatchedNodeHeader {
header: InlineKeyNodeHeader,
latch: NodeLatch,
}
impl LatchedNodeHeader {
#[inline(always)]
pub(in crate::tree) unsafe fn latch_of<'a>(page: NonNull<u8>) -> &'a NodeLatch {
unsafe {
page.byte_add(offset_of!(Self, latch))
.cast::<NodeLatch>()
.as_ref()
}
}
}
#[repr(transparent)]
pub(in crate::tree) struct NodePtr<K, V, S> {
tagged: NonNull<u8>,
_marker: PhantomData<*const (K, V)>,
_strategy: PhantomData<fn() -> S>,
}
impl<K, V, S> Clone for NodePtr<K, V, S> {
fn clone(&self) -> Self {
*self
}
}
impl<K, V, S> Copy for NodePtr<K, V, S> {}
impl<K, V, S> NodePtr<K, V, S> {
pub(in crate::tree) const LEAF_TAG: usize = 1;
#[inline(always)]
fn tagged_as<Kind: NodeKind>(page: NonNull<u8>) -> Self {
debug_assert_eq!(page.addr().get() & Self::LEAF_TAG, 0);
Self {
tagged: page.map_addr(|addr| addr | (usize::from(Kind::IS_LEAF) * Self::LEAF_TAG)),
_marker: PhantomData,
_strategy: PhantomData,
}
}
#[inline(always)]
pub(in crate::tree) fn tagged(self) -> NonNull<u8> {
self.tagged
}
#[inline(always)]
pub(in crate::tree) unsafe fn from_tagged(tagged: NonNull<u8>) -> Self {
Self {
tagged,
_marker: PhantomData,
_strategy: PhantomData,
}
}
#[inline(always)]
pub(in crate::tree) fn is_leaf(self) -> bool {
self.tagged.addr().get() & Self::LEAF_TAG != 0
}
#[inline(always)]
pub(in crate::tree) fn page(self) -> NonNull<u8> {
unsafe {
NonNull::new_unchecked(self.tagged.as_ptr().map_addr(|addr| addr & !Self::LEAF_TAG))
}
}
#[cfg(any(test, feature = "inspect"))]
pub(in crate::tree) fn addr(self) -> usize {
self.page().addr().get()
}
pub(in crate::tree) fn is(self, other: Self) -> bool {
self.tagged == other.tagged
}
pub(super) unsafe fn into_fresh_page<P: PageType>(self) -> FreshPage<P> {
FreshPage {
page: self.page(),
_page_type: PhantomData,
}
}
fn header(self) -> *mut NodeHeader {
self.page().cast::<NodeHeader>().as_ptr()
}
pub(super) unsafe fn len(self) -> usize {
usize::from(unsafe { (*self.header()).len })
}
pub(super) unsafe fn has_head(self) -> bool {
unsafe { (*self.header()).has_head }
}
pub(super) unsafe fn set_has_head(self, has_head: bool) {
unsafe { (*self.header()).has_head = has_head }
}
}
impl<K, V, S: InlineKeyStrategy<K>> NodePtr<K, V, S> {
#[inline(always)]
pub(in crate::tree) fn is_latched() -> bool {
<S as Sealed<K>>::LATCHED
}
#[inline(always)]
pub(in crate::tree) fn is_in_file() -> bool {
<S as Sealed<K>>::IN_FILE
}
#[inline(always)]
pub(in crate::tree) unsafe fn latch<'a>(self) -> &'a NodeLatch {
debug_assert!(Self::is_latched());
unsafe { LatchedNodeHeader::latch_of(self.page()) }
}
#[inline(always)]
pub(in crate::tree) unsafe fn load_racing_header(self) -> (usize, usize) {
let _reads = RacingReads::begin();
debug_assert!(Self::is_latched());
Interleaving::switch_point();
unsafe {
let header = self.page().cast::<InlineKeyNodeHeader>().as_ptr();
let len = AtomicU16::from_ptr(&raw mut (*header).header.len);
let prefix_len = AtomicU16::from_ptr(&raw mut (*header).prefix_len);
(
usize::from(len.load(Ordering::Relaxed)),
usize::from(prefix_len.load(Ordering::Relaxed)),
)
}
}
#[inline(always)]
pub(in crate::tree) unsafe fn load_racing_has_head(self) -> bool {
let _reads = RacingReads::begin();
debug_assert!(Self::is_latched());
Interleaving::switch_point();
unsafe { AtomicBool::from_ptr(&raw mut (*self.header()).has_head).load(Ordering::Relaxed) }
}
#[inline]
pub(in crate::tree) fn inline_width<Kind: NodeKind>() -> Option<InlineWidth> {
NodePlan::of::<S, Kind>().inline_width()
}
#[inline(always)]
pub(in crate::tree) fn has_inline_keys() -> bool {
NodePlan::of::<S, Leaf<K, V, S>>().inline_width().is_some()
}
#[inline(always)]
pub(super) fn leaf_has_room_for_hash_tags(inline_width: InlineWidth) -> bool {
NodeLayout::<K, V, S::Config, S::LeafAnnex>::leaf(
Some(inline_width),
true,
Self::is_latched(),
)
.hash_tags_offset
.is_some()
}
pub(in crate::tree) fn leaf_layout() -> LeafLayout {
const {
assert!(
size_of::<S::LeafAnnex>() == 0 || <S as Sealed<K>>::LATCHED,
"only a BBTreeMap's leaves keep an annex"
);
}
let plan = NodePlan::of::<S, Leaf<K, V, S>>();
NodeLayout::<K, V, S::Config, S::LeafAnnex>::leaf(
plan.inline_width(),
plan.has_hash_tags(),
Self::is_latched(),
)
}
pub(in crate::tree) fn internal_layout() -> InternalLayout {
NodeLayout::<K, V, S::Config>::internal(
Self::inline_width::<Internal<K, V, S>>(),
Self::is_latched(),
)
}
pub(in crate::tree) fn key_area<Kind: NodeKind>() -> KeyArea {
if Kind::IS_LEAF {
Self::leaf_layout().keys
} else {
Self::internal_layout().keys
}
}
pub(in crate::tree) fn capacity<Kind: NodeKind>() -> usize {
Self::key_area::<Kind>().capacity
}
pub(super) fn has_paired_leaf_entries() -> bool {
Self::leaf_layout().keys.key_stride != size_of::<K>()
}
#[inline(always)]
pub(super) fn leaf_annex(self) -> NonNull<S::LeafAnnex> {
let header = Header::leaf(Self::is_latched(), Layout::new::<S::LeafAnnex>());
unsafe { self.page().byte_add(header.annex_offset()).cast() }
}
pub(super) fn from_fresh_page<Kind: NodeKind>(page: FreshPage<Kind>) -> Self {
unsafe { Self::initialize::<Kind>(page.into_raw()) }
}
pub(super) unsafe fn initialize<Kind: NodeKind>(page: NonNull<u8>) -> Self {
let node = Self::tagged_as::<Kind>(page);
let header = NodeHeader {
len: 0,
has_head: false,
};
unsafe {
if Self::has_inline_keys() {
page.cast::<InlineKeyNodeHeader>()
.write(InlineKeyNodeHeader {
header,
prefix_len: 0,
});
} else {
page.cast::<NodeHeader>().write(header);
}
if Kind::IS_LEAF {
node.leaf_annex().write(S::LeafAnnex::default());
}
node.key_columns::<Kind>().initialize();
}
node
}
pub(super) unsafe fn move_node_to<Kind: NodeKind>(self, dest: NonNull<u8>) -> Self {
debug_assert_eq!(self.is_leaf(), Kind::IS_LEAF);
let size = Kind::page_layout().size();
unsafe {
let (source, dest) = (self.page().as_ptr(), dest.as_ptr());
if Self::is_latched() {
let latch = offset_of!(LatchedNodeHeader, latch);
let latch_end = latch + size_of::<NodeLatch>();
source.copy_to_nonoverlapping(dest, latch);
source
.add(latch_end)
.copy_to_nonoverlapping(dest.add(latch_end), size - latch_end);
} else {
source.copy_to_nonoverlapping(dest, size);
}
}
Self::tagged_as::<Kind>(dest)
}
#[cfg(feature = "file")]
pub(in crate::tree) unsafe fn copy_page_to(self, dest: &mut [u8], size: usize) {
let padding = NodeHeader::PADDING;
assert!(dest.len() >= size && size >= padding.end);
let source = self.page().as_ptr();
unsafe {
source.copy_to_nonoverlapping(dest.as_mut_ptr(), padding.start);
source
.add(padding.end)
.copy_to_nonoverlapping(dest.as_mut_ptr().add(padding.end), size - padding.end);
}
dest[padding].fill(0);
}
pub(super) unsafe fn prefix_len(self) -> usize {
if !Self::has_inline_keys() {
return 0;
}
let header = self.page().cast::<InlineKeyNodeHeader>().as_ptr();
usize::from(unsafe { (*header).prefix_len })
}
pub(super) unsafe fn set_prefix_len(self, prefix_len: usize) {
if Self::has_inline_keys() {
debug_assert!(prefix_len <= InlineKeyNodeHeader::MAX_PREFIX_LEN);
let prefix_len = prefix_len as u16;
let header = self.page().cast::<InlineKeyNodeHeader>().as_ptr();
unsafe { (*header).prefix_len = prefix_len };
}
}
pub(super) unsafe fn set_len(self, len: usize) {
let len = len as u16;
unsafe { (*self.header()).len = len }
}
#[inline(always)]
pub(super) unsafe fn set_len_refilling<Kind: NodeKind>(self, len: usize) {
debug_assert!(len <= Self::capacity::<Kind>());
unsafe {
let old_len = self.len();
self.key_columns::<Kind>().mark_vacant(len..old_len);
self.set_len(len);
}
}
pub(super) unsafe fn key_slots<Kind: NodeKind>(self) -> KeySlots<K, S, Kind> {
debug_assert!(Kind::IS_LEAF || !Self::has_inline_keys());
let area = Self::key_area::<Kind>();
let prefix_len = unsafe { self.prefix_len() };
unsafe {
let keys = Slots::with_stride(
self.page().byte_add(area.keys_offset).cast(),
area.capacity,
area.key_stride,
);
KeySlots::new(keys, self.key_columns::<Kind>(), prefix_len)
}
}
#[inline(always)]
pub(super) unsafe fn separator_slots(self) -> SeparatorSlots<K, V, S> {
unsafe {
if Self::has_inline_keys() {
SeparatorSlots::Windows {
columns: self.key_columns::<Internal<K, V, S>>(),
prefix_len: self.prefix_len(),
}
} else {
SeparatorSlots::Keys(self.key_slots::<Internal<K, V, S>>())
}
}
}
#[inline(always)]
pub(super) unsafe fn key_columns<Kind: NodeKind>(self) -> KeyColumns {
let plan = NodePlan::of::<S, Kind>();
let area = Self::key_area::<Kind>();
let hash_tags_offset = if plan.has_hash_tags() {
Self::leaf_layout().hash_tags_offset
} else {
None
};
unsafe {
let inline_keys =
Option::zip(plan.inline_width(), area.inline_keys_offset).map(|(width, offset)| {
InlineKeySlots::new(width, self.page().byte_add(offset), area.capacity)
});
let hash_tags = hash_tags_offset
.map(|offset| Slots::new(self.page().byte_add(offset).cast(), area.capacity));
KeyColumns::new(inline_keys, hash_tags, area.capacity)
}
}
pub(super) unsafe fn head_key_slot(self) -> NonNull<K> {
let offset = Self::leaf_layout()
.head_key_offset
.expect("a leaf with inline keys has a head key slot");
unsafe { self.page().byte_add(offset).cast() }
}
pub(super) unsafe fn value_slots(self) -> Slots<V> {
let layout = Self::leaf_layout();
unsafe {
Slots::with_stride(
self.page().byte_add(layout.values_offset()).cast(),
layout.keys.capacity,
layout.value_stride,
)
}
}
pub(super) unsafe fn head_and_value_slots(self) -> Slots<V> {
let layout = Self::leaf_layout();
unsafe {
Slots::with_stride(
self.page().byte_add(layout.head_offset).cast(),
layout.keys.capacity + 1,
layout.value_stride,
)
}
}
pub(super) unsafe fn edge_slots(self) -> Slots<NodePtr<K, V, S>> {
let layout = Self::internal_layout();
unsafe {
Slots::with_stride(
self.page().byte_add(layout.edges_offset).cast(),
layout.keys.capacity + 1,
layout.edge_stride,
)
}
}
#[inline(always)]
pub(super) unsafe fn file_edge_slots(self) -> Slots<usize> {
debug_assert!(Self::is_in_file());
let layout = Self::internal_layout();
unsafe {
Slots::with_stride(
self.page().byte_add(layout.edges_offset).cast(),
layout.keys.capacity + 1,
layout.edge_stride,
)
}
}
#[inline(always)]
pub(in crate::tree) unsafe fn at_offset(page: NonNull<u8>, offset: usize) -> Self {
unsafe { Self::from_tagged(page.byte_add(offset)) }
}
}
#[cfg(test)]
pub(crate) use test_pages::TestPages;
#[cfg(test)]
mod test_pages {
use core::alloc::Layout;
use core::cell::{Cell, RefCell};
use core::marker::PhantomData;
use core::ptr::NonNull;
use std::vec::Vec;
use super::{FreshPage, PageByType, PageType};
use crate::allocator::{AllocError, Allocator, Global, PageKind};
use crate::inline_key::{Auto, InlineKeyStrategy};
use crate::tree::node::internal::Internal;
use crate::tree::node::leaf::Leaf;
pub(crate) struct TestPages<K, V, S = Auto> {
allocated: RefCell<Vec<(NonNull<u8>, Layout)>>,
released: RefCell<Vec<(NonNull<u8>, Layout)>>,
budget: Cell<Option<usize>>,
kinds: RefCell<Vec<(NonNull<u8>, PageKind)>>,
kind_layouts: RefCell<Vec<(PageKind, Layout)>>,
_marker: PhantomData<fn() -> (K, V)>,
_strategy: PhantomData<fn() -> S>,
}
impl<K, V, S: InlineKeyStrategy<K>> TestPages<K, V, S> {
pub(crate) fn new() -> Self {
Self {
allocated: RefCell::new(Vec::new()),
released: RefCell::new(Vec::new()),
budget: Cell::new(None),
kinds: RefCell::new(Vec::new()),
kind_layouts: RefCell::new(Vec::new()),
_marker: PhantomData,
_strategy: PhantomData,
}
}
pub(crate) fn set_budget(&self, budget: Option<usize>) {
self.budget.set(budget);
}
pub(in crate::tree) fn fresh<P: PageType>(&self) -> FreshPage<P> {
FreshPage::allocate(self)
}
pub(in crate::tree) fn fresh_leaf(&self) -> FreshPage<Leaf<K, V, S>> {
self.fresh()
}
pub(in crate::tree) fn fresh_internal(&self) -> FreshPage<Internal<K, V, S>> {
self.fresh()
}
pub(in crate::tree) fn release(&self, page: impl Into<PageByType<K, V, S>>) {
page.into().deallocate(self);
}
pub(in crate::tree) fn allocated_count(&self) -> usize {
self.allocated.borrow().len()
}
pub(crate) fn in_use_count(&self) -> usize {
self.allocated.borrow().len() - self.released.borrow().len()
}
pub(crate) fn assert_all_released(&self) {
assert_eq!(
self.released.borrow().len(),
self.allocated.borrow().len(),
"pages were not released"
);
}
}
unsafe impl<K, V, S> Allocator for TestPages<K, V, S> {
fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
match self.budget.get() {
Some(0) => return Err(AllocError),
Some(left) => self.budget.set(Some(left - 1)),
None => {}
}
let mut released = self.released.borrow_mut();
let reusable = (released.iter()).rposition(|&(_, page_layout)| page_layout == layout);
if let Some(index) = reusable {
let block = released.remove(index).0;
unsafe { block.write_bytes(0, layout.size()) };
return Ok(block);
}
let block = Global.allocate_zeroed(layout)?;
self.allocated.borrow_mut().push((block, layout));
Ok(block)
}
unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
assert!(
self.allocated.borrow().contains(&(ptr, layout)),
"released a page from elsewhere, or with another layout"
);
assert!(
!(self.released.borrow().iter()).any(|&(released, _)| released == ptr),
"released a page twice"
);
unsafe { ptr.write_bytes(0xA5, layout.size()) };
self.released.borrow_mut().push((ptr, layout));
}
fn allocate_internal(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
self.allocate_page_of(PageKind::Internal, layout)
}
unsafe fn deallocate_internal(&self, page: NonNull<u8>, layout: Layout) {
unsafe { self.deallocate_page_of(PageKind::Internal, page, layout) }
}
fn allocate_leaf(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
self.allocate_page_of(PageKind::Leaf, layout)
}
unsafe fn deallocate_leaf(&self, page: NonNull<u8>, layout: Layout) {
unsafe { self.deallocate_page_of(PageKind::Leaf, page, layout) }
}
fn allocate_root(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
self.allocate_page_of(PageKind::Root, layout)
}
unsafe fn deallocate_root(&self, page: NonNull<u8>, layout: Layout) {
unsafe { self.deallocate_page_of(PageKind::Root, page, layout) }
}
}
impl<K, V, S> TestPages<K, V, S> {
fn allocate_page_of(
&self,
kind: PageKind,
layout: Layout,
) -> Result<NonNull<u8>, AllocError> {
let mut kind_layouts = self.kind_layouts.borrow_mut();
match kind_layouts.iter().find(|&&(asked, _)| asked == kind) {
Some(&(_, first)) => assert_eq!(first, layout, "a kind's pages have one layout"),
None => kind_layouts.push((kind, layout)),
}
let page = self.allocate_zeroed(layout)?;
self.kinds.borrow_mut().push((page, kind));
Ok(page)
}
unsafe fn deallocate_page_of(&self, kind: PageKind, page: NonNull<u8>, layout: Layout) {
let mut kinds = self.kinds.borrow_mut();
let index = (kinds.iter())
.position(|&(out, _)| out == page)
.expect("a page given back is out");
assert_eq!(
kinds.swap_remove(index).1,
kind,
"a page goes back as its kind"
);
drop(kinds);
unsafe { self.deallocate(page, layout) }
}
}
impl<K, V, S> Drop for TestPages<K, V, S> {
fn drop(&mut self) {
for (block, layout) in self.released.get_mut().drain(..) {
unsafe { Global.deallocate(block, layout) };
}
}
}
}
#[cfg(test)]
mod tests {
use core::mem::{align_of, size_of};
use core::ops::Range;
use std::string::String;
use std::vec::Vec;
use super::{FreshPage, NodeHeader, NodeKind, NodePtr, TestPages};
use crate::inline_key::{Auto, InlineKeyStrategy, InlineWidth, WithInlineKey};
use crate::node_config::DefaultNodeConfig;
use crate::tree::node::internal::Internal;
use crate::tree::node::layout::NodeLayout;
use crate::tree::node::leaf::Leaf;
use crate::tree::node::slots::{InlineKeySlots, Slots};
use crate::tree::test_support::{CacheLineAligned, PageAligned, Tracked};
struct SlotRange(Range<usize>, Vec<Range<usize>>);
impl SlotRange {
fn page_bytes<K, V, S: InlineKeyStrategy<K>, Kind: NodeKind>(
node: NodePtr<K, V, S>,
) -> Range<usize> {
let layout = Kind::page_layout();
let start = node.page().addr().get();
assert!(start.is_multiple_of(layout.align()));
start..start + layout.size()
}
fn of<T>(slots: Slots<T>, capacity: usize) -> Self {
let each: Vec<Range<usize>> = (0..capacity)
.map(|index| {
let start = unsafe { slots.slot(index) }.addr().get();
assert!(start.is_multiple_of(align_of::<T>()));
start..start + size_of::<T>()
})
.collect();
let start = each.iter().map(|slot| slot.start).min().unwrap();
let end = each.iter().map(|slot| slot.end).max().unwrap();
Self(start..end, each)
}
fn of_inline_keys(slots: InlineKeySlots, capacity: usize) -> Self {
match slots {
InlineKeySlots::Two(slots) => Self::of(slots, capacity),
InlineKeySlots::Four(slots) => Self::of(slots, capacity),
InlineKeySlots::Eight(slots) => Self::of(slots, capacity),
}
}
fn assert_within(&self, page: &Range<usize>) {
assert!(page.start <= self.0.start && self.0.end <= page.end);
}
fn assert_disjoint(&self, other: &Self) {
for a in &self.1 {
for b in &other.1 {
assert!(a.end <= b.start || b.end <= a.start);
}
}
}
}
fn assert_slots_are_aligned_and_disjoint<K, V, S: InlineKeyStrategy<K>>() {
let pages = TestPages::<K, V, S>::new();
let leaf = NodePtr::from_fresh_page(pages.fresh::<Leaf<K, V, S>>());
let internal = NodePtr::from_fresh_page(pages.fresh::<Internal<K, V, S>>());
let leaf_page = SlotRange::page_bytes::<K, V, S, Leaf<K, V, S>>(leaf);
let internal_page = SlotRange::page_bytes::<K, V, S, Internal<K, V, S>>(internal);
let leaf_capacity = NodePtr::<K, V, S>::capacity::<Leaf<K, V, S>>();
let internal_capacity = NodePtr::<K, V, S>::capacity::<Internal<K, V, S>>();
let keeps_keys = !NodePtr::<K, V, S>::has_inline_keys();
let (
leaf_keys,
leaf_columns,
values,
head_and_values,
internal_keys,
internal_columns,
edges,
) = unsafe {
(
leaf.key_slots::<Leaf<K, V, S>>(),
leaf.key_columns::<Leaf<K, V, S>>(),
SlotRange::of(leaf.value_slots(), leaf_capacity),
SlotRange::of(leaf.head_and_value_slots(), leaf_capacity + 1),
keeps_keys.then(|| internal.key_slots::<Internal<K, V, S>>()),
internal.key_columns::<Internal<K, V, S>>(),
SlotRange::of(internal.edge_slots(), internal_capacity + 1),
)
};
assert_eq!(
leaf_columns.inline_keys().map(InlineKeySlots::width),
NodePtr::<K, V, S>::inline_width::<Leaf<K, V, S>>()
);
assert_eq!(
internal_columns.inline_keys().map(InlineKeySlots::width),
NodePtr::<K, V, S>::inline_width::<Internal<K, V, S>>()
);
let leaf_inline_keys = (leaf_columns.inline_keys())
.map(|slots| SlotRange::of_inline_keys(slots, leaf_capacity));
let internal_inline_keys = (internal_columns.inline_keys())
.map(|slots| SlotRange::of_inline_keys(slots, internal_capacity));
let head_key = NodePtr::<K, V, S>::has_inline_keys().then(|| unsafe {
let slot = leaf.head_key_slot();
assert!(slot.addr().get().is_multiple_of(align_of::<K>()));
let start = slot.addr().get();
SlotRange(
start..start + size_of::<K>(),
std::vec![start..start + size_of::<K>()],
)
});
let leaf_keys = SlotRange::of(leaf_keys.keys(), leaf_capacity);
let internal_keys = internal_keys.map(|keys| SlotRange::of(keys.keys(), internal_capacity));
unsafe {
assert!(
internal
.key_columns::<Internal<K, V, S>>()
.hash_tags()
.is_none()
);
if let Some(tags) = leaf.key_columns::<Leaf<K, V, S>>().hash_tags() {
assert!(
tags.slice_ptr(leaf_capacity)
.as_ref()
.iter()
.all(|&tag| tag == 0)
);
let tags = SlotRange::of(tags, leaf_capacity);
tags.assert_within(&leaf_page);
tags.assert_disjoint(&leaf_keys);
tags.assert_disjoint(&head_and_values);
tags.assert_disjoint(leaf_inline_keys.as_ref().unwrap());
if let Some(head_key) = &head_key {
tags.assert_disjoint(head_key);
}
}
}
for slots in [&leaf_keys, &values, &head_and_values] {
slots.assert_within(&leaf_page);
}
if let Some(head_key) = &head_key {
head_key.assert_within(&leaf_page);
head_key.assert_disjoint(&leaf_keys);
head_key.assert_disjoint(&head_and_values);
head_key.assert_disjoint(leaf_inline_keys.as_ref().unwrap());
}
edges.assert_within(&internal_page);
if let Some(internal_keys) = &internal_keys {
internal_keys.assert_within(&internal_page);
internal_keys.assert_disjoint(&edges);
}
let keys_offset = leaf_keys.0.start - leaf_page.start;
assert!(size_of::<NodeHeader>() <= keys_offset);
assert_eq!(values.1, head_and_values.1[1..]);
leaf_keys.assert_disjoint(&head_and_values);
match (leaf_inline_keys, internal_inline_keys) {
(None, None) => {
let internal_keys = internal_keys.expect("a node without inline keys keeps keys");
assert_eq!(internal_keys.0.start - internal_page.start, keys_offset);
}
(Some(leaf_inline_keys), Some(internal_inline_keys)) => {
leaf_inline_keys.assert_within(&leaf_page);
internal_inline_keys.assert_within(&internal_page);
assert_eq!(leaf_inline_keys.0.start - leaf_page.start, 8);
assert_eq!(internal_inline_keys.0.start - internal_page.start, 8);
assert!(leaf_inline_keys.0.end <= leaf_keys.0.start);
assert!(
internal_keys.is_none(),
"a node with inline keys keeps no keys"
);
leaf_inline_keys.assert_disjoint(&head_and_values);
internal_inline_keys.assert_disjoint(&edges);
}
_ => unreachable!("both kinds have inline keys or neither"),
}
unsafe {
pages.release(leaf.into_fresh_page::<Leaf<K, V, S>>());
pages.release(internal.into_fresh_page::<Internal<K, V, S>>());
}
pages.assert_all_released();
}
#[test]
fn slots_are_aligned_and_disjoint() {
assert_slots_are_aligned_and_disjoint::<u64, u64, Auto>();
assert_slots_are_aligned_and_disjoint::<u8, (), Auto>();
assert_slots_are_aligned_and_disjoint::<(), (), Auto>();
assert_slots_are_aligned_and_disjoint::<u64, [u8; 1000], Auto>();
assert_slots_are_aligned_and_disjoint::<[u8; 77], (), Auto>();
assert_slots_are_aligned_and_disjoint::<CacheLineAligned, u64, Auto>();
assert_slots_are_aligned_and_disjoint::<u8, CacheLineAligned, Auto>();
assert_slots_are_aligned_and_disjoint::<u64, PageAligned, Auto>();
assert_slots_are_aligned_and_disjoint::<String, Vec<u8>, Auto>();
assert_slots_are_aligned_and_disjoint::<String, u64, Auto>();
assert_slots_are_aligned_and_disjoint::<String, (), Auto>();
assert_slots_are_aligned_and_disjoint::<String, [u64; 5], Auto>();
assert_slots_are_aligned_and_disjoint::<Vec<u8>, u64, Auto>();
assert_slots_are_aligned_and_disjoint::<String, u64, WithInlineKey>();
assert_slots_are_aligned_and_disjoint::<&str, PageAligned, Auto>();
assert_slots_are_aligned_and_disjoint::<[u8; 20], (), Auto>();
assert_slots_are_aligned_and_disjoint::<[u8; 16], u64, Auto>();
assert_slots_are_aligned_and_disjoint::<[u8; 32], u64, Auto>();
assert_slots_are_aligned_and_disjoint::<Tracked<'_>, Tracked<'_>, WithInlineKey>();
assert_slots_are_aligned_and_disjoint::<Tracked<'_>, (), WithInlineKey>();
assert_slots_are_aligned_and_disjoint::<u64, u64, WithInlineKey>();
assert_slots_are_aligned_and_disjoint::<i64, (), WithInlineKey>();
assert_slots_are_aligned_and_disjoint::<u128, [u8; 1000], WithInlineKey>();
assert_slots_are_aligned_and_disjoint::<i128, CacheLineAligned, WithInlineKey>();
}
#[test]
#[cfg_attr(not(target_pointer_width = "64"), ignore = "assumes 64-bit pointers")]
fn only_inline_keys_change_the_layout() {
assert!(!NodePtr::<[u64; 1], u64, Auto>::has_inline_keys());
assert_eq!(
NodePtr::<[u64; 1], u64, Auto>::leaf_layout(),
NodeLayout::<[u64; 1], u64, DefaultNodeConfig>::PLAIN_LEAF
);
assert_eq!(
NodePtr::<[u64; 1], u64, Auto>::internal_layout(),
NodeLayout::<[u64; 1], u64, DefaultNodeConfig>::PLAIN_INTERNAL
);
let leaf = NodePtr::<[u64; 1], u64, Auto>::leaf_layout();
let internal = NodePtr::<[u64; 1], u64, Auto>::internal_layout();
assert_eq!((leaf.page.size(), leaf.page.align()), (1024, 1024));
assert_eq!((internal.page.size(), internal.page.align()), (1024, 1024));
assert_eq!(leaf.keys, internal.keys);
assert_eq!(
(
leaf.keys.capacity,
leaf.keys.keys_offset,
leaf.keys.inline_keys_offset
),
(63, 8, None)
);
assert_eq!((leaf.head_offset, internal.edges_offset), (512, 512));
assert_eq!(size_of::<NodeHeader>(), 4);
assert!(NodePtr::<String, u64, Auto>::has_inline_keys());
assert_eq!(
NodePtr::<String, u64, Auto>::leaf_layout(),
NodeLayout::<String, u64, DefaultNodeConfig>::EIGHT_BYTE_INLINE_HASH_TAGGED_LEAF
);
assert_eq!(
NodePtr::<Vec<u8>, u64, Auto>::internal_layout(),
NodeLayout::<Vec<u8>, u64, DefaultNodeConfig>::FOUR_BYTE_INLINE_INTERNAL
);
assert_eq!(
NodePtr::<String, u64, Auto>::inline_width::<Leaf<String, u64, Auto>>(),
Some(InlineWidth::Eight)
);
assert_eq!(
NodePtr::<String, u64, Auto>::inline_width::<Internal<String, u64, Auto>>(),
Some(InlineWidth::Four)
);
assert_eq!(
NodePtr::<[u8; 32], u64, WithInlineKey>::internal_layout(),
NodeLayout::<[u8; 32], u64, DefaultNodeConfig>::EIGHT_BYTE_INLINE_INTERNAL
);
assert!(NodePtr::<Tracked<'_>, u64, WithInlineKey>::has_inline_keys());
assert!(!NodePtr::<Tracked<'_>, u64, Auto>::has_inline_keys());
assert!(NodePtr::<i64, u64, Auto>::has_inline_keys());
assert_eq!(
NodePtr::<u64, u64, Auto>::leaf_layout(),
NodeLayout::<u64, u64, DefaultNodeConfig>::TWO_BYTE_INLINE_LEAF
);
assert_eq!(
NodePtr::<u64, u64, WithInlineKey>::leaf_layout(),
NodeLayout::<u64, u64, DefaultNodeConfig>::TWO_BYTE_INLINE_LEAF
);
assert_eq!(
NodePtr::<i128, u64, WithInlineKey>::internal_layout(),
NodeLayout::<i128, u64, DefaultNodeConfig>::TWO_BYTE_INLINE_INTERNAL
);
}
#[test]
fn a_new_node_is_empty() {
let pages = TestPages::<u32, u64>::new();
let node = NodePtr::<u32, u64, Auto>::from_fresh_page(pages.fresh_leaf());
assert_eq!(unsafe { (node.len(), node.has_head()) }, (0, false));
assert_eq!(unsafe { node.prefix_len() }, 0);
pages.release(unsafe { node.into_fresh_page::<Leaf<u32, u64, Auto>>() });
pages.assert_all_released();
}
#[test]
fn a_new_node_with_inline_keys_has_no_prefix() {
let pages = TestPages::<String, u64>::new();
pages.release(pages.fresh_internal());
let node = NodePtr::<String, u64, Auto>::from_fresh_page(pages.fresh_internal());
assert_eq!(unsafe { (node.len(), node.prefix_len()) }, (0, 0));
pages.release(unsafe { node.into_fresh_page::<Internal<String, u64, Auto>>() });
pages.assert_all_released();
}
#[test]
fn released_pages_are_handed_out_again() {
let pages = TestPages::<u64, u64>::new();
let first = pages.fresh_leaf().into_raw();
let second = pages.fresh_leaf().into_raw();
assert_ne!(first, second);
unsafe {
pages.release(FreshPage::<Leaf<u64, u64, Auto>>::from_raw(first));
pages.release(FreshPage::<Leaf<u64, u64, Auto>>::from_raw(second));
}
assert_eq!(pages.in_use_count(), 0);
let reused = pages.fresh_leaf();
assert_eq!(pages.in_use_count(), 1);
assert_eq!(reused.into_raw(), second);
pages.release(unsafe { FreshPage::<Leaf<u64, u64, Auto>>::from_raw(second) });
pages.assert_all_released();
}
#[test]
fn released_pages_are_handed_out_for_their_layout_only() {
let pages = TestPages::<u32, u64>::new();
pages.release(pages.fresh_leaf());
pages.release(pages.fresh_internal());
assert_eq!(pages.allocated_count(), 1);
pages.assert_all_released();
let pages = TestPages::<u64, [u8; 1000]>::new();
pages.release(pages.fresh_leaf());
let internal = pages.fresh_internal();
assert_eq!((pages.allocated_count(), pages.in_use_count()), (2, 1));
pages.release(internal);
pages.release(pages.fresh_leaf());
assert_eq!(pages.allocated_count(), 2);
pages.assert_all_released();
}
}