use core::{
alloc::Layout,
marker::PhantomData,
mem::{self, MaybeUninit},
ptr::NonNull,
slice,
};
use allocator_api2::alloc::{Allocator, handle_alloc_error};
use crate::BTreeInteger;
#[cfg(target_pointer_width = "64")]
pub(crate) const MAX_POOL_SIZE: usize = u32::MAX as usize;
#[cfg(target_pointer_width = "32")]
pub(crate) const MAX_POOL_SIZE: usize = i32::MAX as usize;
#[derive(Clone, Copy, Debug)]
pub(crate) struct NodePos<I: BTreeInteger> {
pos: u32,
marker: PhantomData<fn() -> I>,
}
macro_rules! pos {
($expr:expr) => {{
const { assert!($expr < K::Int::B) };
#[allow(unused_unsafe)]
unsafe {
$crate::node::NodePos::<K::Int>::new_unchecked($expr)
}
}};
}
impl<I: BTreeInteger> NodePos<I> {
#[inline]
pub(crate) const fn zero() -> Self {
Self {
pos: 0,
marker: PhantomData,
}
}
#[inline]
pub(crate) const unsafe fn new_unchecked(pos: usize) -> Self {
debug_assert!(pos < I::B);
Self {
pos: pos as u32,
marker: PhantomData,
}
}
#[inline]
pub(crate) fn index(self) -> usize {
self.pos as usize
}
#[inline]
pub(crate) unsafe fn next(self) -> Self {
debug_assert!(self.index() + 1 < I::B);
Self {
pos: self.pos + 1,
marker: PhantomData,
}
}
#[inline]
pub(crate) unsafe fn prev(self) -> Self {
debug_assert_ne!(self.pos, 0);
Self {
pos: self.pos - 1,
marker: PhantomData,
}
}
#[inline]
pub(crate) fn split_right_half(self) -> Option<Self> {
if self.index() >= I::B / 2 {
Some(Self {
pos: self.pos - I::B as u32 / 2,
marker: PhantomData,
})
} else {
None
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct NodeRef(u32);
impl NodeRef {
#[inline]
pub(crate) const fn zero() -> Self {
Self(0)
}
#[inline]
unsafe fn keys_ptr<I: BTreeInteger, V>(self, pool: &NodePool<I, V>) -> NonNull<I::Raw> {
pool.validate_noderef(self);
unsafe { pool.ptr.byte_add(self.0 as usize).cast() }
}
#[inline]
pub(crate) unsafe fn values_ptr<I: BTreeInteger, V>(
self,
pool: &NodePool<I, V>,
) -> NonNull<MaybeUninit<V>> {
pool.validate_noderef(self);
let values_offset = const { node_layout::<I, V>().1 };
unsafe {
let ptr = pool.ptr.byte_add(self.0 as usize);
ptr.byte_add(values_offset).cast::<MaybeUninit<V>>()
}
}
#[inline]
pub(crate) unsafe fn leaf_end<I: BTreeInteger, V>(self, pool: &NodePool<I, V>) -> NodePos<I> {
pool.validate_noderef(self);
unsafe { I::search(self.keys(pool), I::MAX) }
}
#[inline]
pub(crate) unsafe fn internal_end<I: BTreeInteger, V>(
self,
pool: &NodePool<I, V>,
) -> NodePos<I> {
pool.validate_noderef(self);
unsafe { I::search(self.keys(pool), I::MAX).next() }
}
#[inline]
pub(crate) unsafe fn keys<I: BTreeInteger, V>(self, pool: &NodePool<I, V>) -> &I::Keys {
unsafe { self.keys_ptr(pool).cast::<I::Keys>().as_ref() }
}
#[inline]
pub(crate) unsafe fn key<I: BTreeInteger, V>(
self,
pos: NodePos<I>,
pool: &NodePool<I, V>,
) -> I::Raw {
unsafe { self.keys_ptr(pool).add(pos.index()).read() }
}
#[inline]
pub(crate) unsafe fn set_key<I: BTreeInteger, V>(
self,
key: I::Raw,
pos: NodePos<I>,
pool: &mut NodePool<I, V>,
) {
unsafe { self.keys_ptr(pool).add(pos.index()).write(key) }
}
#[inline]
pub(crate) unsafe fn value<I: BTreeInteger, V>(
self,
pos: NodePos<I>,
pool: &NodePool<I, V>,
) -> &MaybeUninit<V> {
unsafe { self.values_ptr(pool).add(pos.index()).as_ref() }
}
#[inline]
pub(crate) unsafe fn value_mut<I: BTreeInteger, V>(
self,
pos: NodePos<I>,
pool: &mut NodePool<I, V>,
) -> &mut MaybeUninit<V> {
unsafe { self.values_ptr(pool).add(pos.index()).as_mut() }
}
#[inline]
pub(crate) unsafe fn next_leaf<I: BTreeInteger, V>(
self,
pool: &NodePool<I, V>,
) -> Option<NodeRef> {
pool.validate_noderef(self);
let next_leaf_offset = const { node_layout::<I, V>().2 };
let next_leaf = unsafe {
let ptr = pool.ptr.byte_add(self.0 as usize);
ptr.byte_add(next_leaf_offset).cast::<NodeRef>().read()
};
(next_leaf.0 != !0).then_some(next_leaf)
}
#[inline]
pub(crate) unsafe fn set_next_leaf<I: BTreeInteger, V>(
self,
next_leaf: Option<NodeRef>,
pool: &mut NodePool<I, V>,
) {
pool.validate_noderef(self);
let next_leaf_offset = const { node_layout::<I, V>().2 };
unsafe {
let ptr = pool.ptr.byte_add(self.0 as usize);
ptr.byte_add(next_leaf_offset)
.cast::<NodeRef>()
.write(next_leaf.unwrap_or(NodeRef(!0)));
}
}
#[inline]
pub(crate) unsafe fn insert_key<I: BTreeInteger, V>(
self,
key: I::Raw,
pos: NodePos<I>,
node_size: usize,
pool: &mut NodePool<I, V>,
) {
debug_assert!(node_size <= I::B);
debug_assert!(node_size > pos.index());
unsafe {
let ptr = self.keys_ptr(pool).add(pos.index());
let count = node_size - pos.index() - 1;
ptr.copy_to(ptr.add(1), count);
ptr.write(key);
}
}
#[inline]
pub(crate) unsafe fn insert_value<I: BTreeInteger, V>(
self,
value: V,
pos: NodePos<I>,
node_size: usize,
pool: &mut NodePool<I, V>,
) {
debug_assert!(node_size <= I::B);
debug_assert!(node_size > pos.index());
unsafe {
let ptr = self.values_ptr(pool).add(pos.index());
let count = node_size - pos.index() - 1;
ptr.copy_to(ptr.add(1), count);
ptr.write(MaybeUninit::new(value));
}
}
#[inline]
pub(crate) unsafe fn remove_key<I: BTreeInteger, V>(
self,
pos: NodePos<I>,
pool: &mut NodePool<I, V>,
) {
unsafe {
let ptr = self.keys_ptr(pool).add(pos.index());
let count = I::B - pos.index() - 1;
ptr.copy_from(ptr.add(1), count);
self.keys_ptr(pool).add(I::B - 1).write(I::MAX);
}
}
#[inline]
pub(crate) unsafe fn remove_value<I: BTreeInteger, V>(
self,
pos: NodePos<I>,
pool: &mut NodePool<I, V>,
) {
unsafe {
let ptr = self.values_ptr(pool).add(pos.index());
let count = I::B - pos.index() - 1;
ptr.copy_from(ptr.add(1), count);
}
}
#[allow(clippy::needless_pass_by_value)]
#[inline]
pub(crate) unsafe fn split_into<I: BTreeInteger, V>(
self,
dest: UninitNodeRef,
pool: &mut NodePool<I, V>,
) -> NodeRef {
unsafe {
self.keys_ptr(pool)
.add(I::B / 2)
.copy_to_nonoverlapping(dest.0.keys_ptr(pool), I::B / 2);
self.values_ptr(pool)
.add(I::B / 2)
.copy_to_nonoverlapping(dest.0.values_ptr(pool), I::B / 2);
slice::from_raw_parts_mut(self.keys_ptr(pool).add(I::B / 2).as_ptr(), I::B / 2)
.fill(I::MAX);
slice::from_raw_parts_mut(
dest.0
.keys_ptr(pool)
.add(I::B / 2)
.cast::<MaybeUninit<I::Raw>>()
.as_ptr(),
I::B / 2,
)
.fill(MaybeUninit::new(I::MAX));
}
dest.0
}
#[allow(clippy::needless_pass_by_value)]
#[inline]
pub(crate) unsafe fn merge_from<I: BTreeInteger, V>(
self,
src: NodeRef,
offset: NodePos<I>,
count: usize,
pool: &mut NodePool<I, V>,
) {
unsafe {
self.keys_ptr(pool)
.add(offset.index())
.copy_from_nonoverlapping(src.keys_ptr(pool), count);
self.values_ptr(pool)
.add(offset.index())
.copy_from_nonoverlapping(src.values_ptr(pool), count);
}
}
}
#[derive(Debug)]
pub(crate) struct UninitNodeRef(NodeRef);
impl UninitNodeRef {
#[inline]
pub(crate) unsafe fn init_keys<I: BTreeInteger, V>(self, pool: &mut NodePool<I, V>) -> NodeRef {
unsafe {
let ptr = self.0.keys_ptr(pool).cast::<MaybeUninit<I::Raw>>();
let slice = slice::from_raw_parts_mut(ptr.as_ptr(), I::B);
slice.fill(MaybeUninit::new(I::MAX));
}
self.0
}
}
#[inline]
pub(crate) const fn node_layout<I: BTreeInteger, V>() -> (Layout, usize, usize) {
const { assert!(I::B >= 4) };
const { assert!(I::B.is_multiple_of(2)) };
let keys = Layout::new::<I::Keys>();
let Ok(values) = Layout::array::<V>(I::B - 1) else {
panic!("Could not calculate node layout");
};
const fn max(a: usize, b: usize) -> usize {
if a > b { a } else { b }
}
let Ok(last_value) = Layout::from_size_align(
max(mem::size_of::<V>(), mem::size_of::<NodeRef>()),
max(mem::align_of::<V>(), mem::align_of::<NodeRef>()),
) else {
panic!("Could not calculate node layout");
};
let Ok((node, values_offset)) = keys.extend(values) else {
panic!("Could not calculate node layout");
};
let Ok((node, next_leaf_offset)) = node.extend(last_value) else {
panic!("Could not calculate node layout");
};
let Ok(layout) = node.align_to(4) else {
panic!("Could not calculate node layout");
};
(layout.pad_to_align(), values_offset, next_leaf_offset)
}
pub(crate) struct NodePool<I: BTreeInteger, V> {
ptr: NonNull<u8>,
capacity: u32,
len: u32,
free_list: u32,
marker: PhantomData<(I, V)>,
}
unsafe impl<I: BTreeInteger + Send, V: Send> Send for NodePool<I, V> {}
unsafe impl<I: BTreeInteger + Sync, V: Sync> Sync for NodePool<I, V> {}
impl<I: BTreeInteger, V> NodePool<I, V> {
#[inline]
pub(crate) fn new() -> Self {
Self {
ptr: NonNull::dangling(),
len: 0,
capacity: 0,
free_list: !0,
marker: PhantomData,
}
}
unsafe extern "C" fn grow(&mut self, alloc: &impl Allocator) {
let node_layout = const { node_layout::<I, V>().0 };
if self.capacity == 0 {
let new_layout = Layout::from_size_align(node_layout.size() * 2, node_layout.align())
.expect("exceeded BTree maximum allocation size");
assert!(
new_layout.size() <= MAX_POOL_SIZE,
"exceeded BTree maximum allocation size"
);
self.ptr = alloc
.allocate(new_layout)
.unwrap_or_else(|_| handle_alloc_error(new_layout))
.cast();
self.capacity = new_layout.size() as u32;
} else {
let old_layout = unsafe {
Layout::from_size_align_unchecked(self.capacity as usize, node_layout.align())
};
let new_layout =
Layout::from_size_align(self.capacity as usize * 2, node_layout.align())
.expect("exceeded BTree maximum allocation size");
assert!(
new_layout.size() <= MAX_POOL_SIZE,
"exceeded BTree maximum allocation size"
);
self.ptr = unsafe {
alloc
.grow(self.ptr, old_layout, new_layout)
.unwrap_or_else(|_| handle_alloc_error(new_layout))
.cast()
};
self.capacity = new_layout.size() as u32;
}
}
#[inline]
pub(crate) unsafe fn alloc_node(&mut self, alloc: &impl Allocator) -> UninitNodeRef {
let node_layout = const { node_layout::<I, V>().0 };
if self.free_list != !0 {
let node = UninitNodeRef(NodeRef(self.free_list));
self.free_list = unsafe { self.ptr.byte_add(self.free_list as usize).cast().read() };
return node;
}
if self.len == self.capacity {
unsafe { self.grow(alloc) };
}
let node = UninitNodeRef(NodeRef(self.len));
self.len += node_layout.size() as u32;
debug_assert!(self.len <= self.capacity);
node
}
#[inline]
pub(crate) unsafe fn free_node(&mut self, node: NodeRef) {
unsafe {
self.ptr
.byte_add(node.0 as usize)
.cast()
.write(self.free_list);
}
self.free_list = node.0;
}
pub(crate) fn clear(&mut self) {
self.len = 0;
self.free_list = !0;
}
pub(crate) fn clear_and_alloc_node(&mut self) -> UninitNodeRef {
let node_layout = const { node_layout::<I, V>().0 };
self.len = node_layout.size() as u32;
self.free_list = !0;
UninitNodeRef(NodeRef::zero())
}
#[inline]
pub(crate) unsafe fn clear_and_free(&mut self, alloc: &impl Allocator) {
self.clear();
let node_layout = const { node_layout::<I, V>().0 };
let layout = unsafe {
Layout::from_size_align_unchecked(self.capacity as usize, node_layout.align())
};
unsafe {
alloc.deallocate(self.ptr, layout);
}
self.capacity = 0;
}
#[inline]
pub(crate) fn validate_noderef(&self, node: NodeRef) {
let node_layout = const { node_layout::<I, V>().0 };
debug_assert_eq!(node.0 as usize % node_layout.size(), 0);
debug_assert!(node.0 < self.len);
}
}
#[cfg(test)]
mod tests {
use allocator_api2::alloc::Global;
use nonmax::NonMaxU32;
use super::NodePool;
#[test]
fn smoke() {
let mut pool = NodePool::<NonMaxU32, u32>::new();
let node = unsafe { pool.alloc_node(&Global).0 };
let node2 = unsafe { pool.alloc_node(&Global).0 };
unsafe {
pool.free_node(node);
}
let node3 = unsafe { pool.alloc_node(&Global).0 };
debug_assert_ne!(node, node2);
debug_assert_eq!(node, node3);
unsafe {
pool.clear_and_free(&Global);
}
}
}