copy from std::tree
This commit is contained in:
parent
9db415a7c9
commit
7939d6df47
|
|
@ -1,5 +1,8 @@
|
|||
#![cfg_attr(not(feature = "std"), no_std)]
|
||||
#![cfg_attr(feature = "nightly", feature(strict_provenance_atomic_ptr))]
|
||||
#![cfg_attr(
|
||||
feature = "nightly",
|
||||
feature(strict_provenance_atomic_ptr, box_vec_non_null)
|
||||
)]
|
||||
#![cfg_attr(feature = "transposed-option", feature(try_trait_v2))]
|
||||
|
||||
#[cfg(any(test, feature = "std", feature = "alloc"))]
|
||||
|
|
@ -23,5 +26,8 @@ pub mod smallbox;
|
|||
pub mod sync;
|
||||
pub mod util;
|
||||
|
||||
#[cfg(feature = "alloc")]
|
||||
pub mod tree;
|
||||
|
||||
pub use cachepadded::CachePadded;
|
||||
pub use mem::can_transmute;
|
||||
|
|
|
|||
400
src/tree.rs
Normal file
400
src/tree.rs
Normal file
|
|
@ -0,0 +1,400 @@
|
|||
use alloc::boxed::Box;
|
||||
use core::{marker::PhantomData, mem::MaybeUninit, ptr::NonNull};
|
||||
|
||||
mod marker {
|
||||
use core::marker::PhantomData;
|
||||
|
||||
pub struct LeafOrInternal;
|
||||
pub struct Leaf;
|
||||
pub struct Internal;
|
||||
|
||||
pub struct Edge;
|
||||
pub struct KV;
|
||||
|
||||
pub struct Owned;
|
||||
pub struct Mut<'a>(PhantomData<&'a mut ()>);
|
||||
pub struct ValMut<'a>(PhantomData<&'a mut ()>);
|
||||
pub struct Immutable<'a>(PhantomData<&'a ()>);
|
||||
pub struct Dying;
|
||||
pub struct DormantMut;
|
||||
|
||||
pub trait BorrowType {
|
||||
const TRAVERSAL_PERMIT: bool = true;
|
||||
}
|
||||
|
||||
impl BorrowType for Owned {
|
||||
const TRAVERSAL_PERMIT: bool = false;
|
||||
}
|
||||
|
||||
impl<'a> BorrowType for Mut<'a> {}
|
||||
impl<'a> BorrowType for ValMut<'a> {}
|
||||
impl<'a> BorrowType for Immutable<'a> {}
|
||||
impl BorrowType for Dying {}
|
||||
impl BorrowType for DormantMut {}
|
||||
}
|
||||
|
||||
const CAPACITY: usize = 16;
|
||||
|
||||
struct LeafNode<K, V> {
|
||||
parent: Option<NonNull<InternalNode<K, V>>>,
|
||||
parent_idx: MaybeUninit<u16>,
|
||||
len: u16,
|
||||
keys: [MaybeUninit<K>; CAPACITY],
|
||||
values: [MaybeUninit<V>; CAPACITY],
|
||||
}
|
||||
|
||||
impl<K, V> LeafNode<K, V> {
|
||||
unsafe fn init(this: *mut Self) {
|
||||
unsafe {
|
||||
(&raw mut (*this).parent).write(None);
|
||||
(&raw mut (*this).len).write(0);
|
||||
// parent_idx and keys/values may be left uninitialized
|
||||
}
|
||||
}
|
||||
|
||||
fn new() -> Box<Self> {
|
||||
unsafe {
|
||||
let mut this = Box::new_uninit();
|
||||
Self::init(this.as_mut_ptr());
|
||||
this.assume_init()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct InternalNode<K, V> {
|
||||
data: LeafNode<K, V>,
|
||||
next: Option<NonNull<InternalNode<K, V>>>,
|
||||
edges: [MaybeUninit<BoxedNode<K, V>>; CAPACITY * 2 - 1],
|
||||
}
|
||||
|
||||
impl<K, V> InternalNode<K, V> {
|
||||
unsafe fn new() -> Box<Self> {
|
||||
unsafe {
|
||||
let mut this = Box::<Self>::new_uninit();
|
||||
LeafNode::init(&raw mut (*this.as_mut_ptr()).data);
|
||||
(&raw mut (*this.as_mut_ptr()).next).write(None);
|
||||
// edges may be left uninitialized
|
||||
|
||||
this.assume_init()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type BoxedNode<K, V> = NonNull<LeafNode<K, V>>;
|
||||
|
||||
struct NodeRef<K, V, BorrowType, Type> {
|
||||
height: usize,
|
||||
node: NonNull<LeafNode<K, V>>,
|
||||
_marker: PhantomData<(BorrowType, Type)>,
|
||||
}
|
||||
|
||||
impl<K, V> NodeRef<K, V, marker::Owned, marker::Leaf> {
|
||||
fn new_leaf() -> Self {
|
||||
let node = LeafNode::new();
|
||||
NodeRef {
|
||||
height: 0,
|
||||
node: Box::into_non_null(node),
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V> NodeRef<K, V, marker::Owned, marker::Internal> {
|
||||
fn new_internal(child: Root<K, V>) -> Self {
|
||||
let mut node = unsafe { InternalNode::new() };
|
||||
node.edges[0].write(child.node);
|
||||
|
||||
// SAFETY: `height` isn't zero
|
||||
unsafe { NodeRef::from_new_internal(node, child.height + 1) }
|
||||
}
|
||||
|
||||
/// # Safety
|
||||
/// `height` must not be zero.
|
||||
unsafe fn from_new_internal(internal: Box<InternalNode<K, V>>, height: usize) -> Self {
|
||||
debug_assert!(height > 0);
|
||||
let node = NonNull::from(Box::leak(internal)).cast();
|
||||
let mut this = NodeRef {
|
||||
height,
|
||||
node,
|
||||
_marker: PhantomData,
|
||||
};
|
||||
this.borrow_mut().correct_all_childrens_parent_links();
|
||||
this
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V, Type> NodeRef<K, V, marker::Owned, Type> {
|
||||
/// Mutably borrows the owned root node. Unlike `reborrow_mut`, this is safe
|
||||
/// because the return value cannot be used to destroy the root, and there
|
||||
/// cannot be other references to the tree.
|
||||
pub(super) fn borrow_mut(&mut self) -> NodeRef<K, V, marker::Mut<'_>, Type> {
|
||||
NodeRef {
|
||||
height: self.height,
|
||||
node: self.node,
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
/// Irreversibly transitions to a reference that permits traversal and offers
|
||||
/// destructive methods and little else.
|
||||
pub(super) fn into_dying(self) -> NodeRef<K, V, marker::Dying, Type> {
|
||||
NodeRef {
|
||||
height: self.height,
|
||||
node: self.node,
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V, BorrowType, Type> NodeRef<K, V, BorrowType, Type> {
|
||||
fn len(&self) -> usize {
|
||||
unsafe { (*self.node.as_ptr()).len as usize }
|
||||
}
|
||||
|
||||
fn height(&self) -> usize {
|
||||
self.height
|
||||
}
|
||||
|
||||
fn reborrow(&self) -> NodeRef<K, V, marker::Immutable<'_>, Type> {
|
||||
NodeRef {
|
||||
height: self.height,
|
||||
node: self.node,
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
fn as_leaf_ptr(this: &Self) -> *mut LeafNode<K, V> {
|
||||
this.node.as_ptr()
|
||||
}
|
||||
}
|
||||
|
||||
impl<BorrowType, K, V> NodeRef<K, V, BorrowType, marker::Internal> {
|
||||
/// Exposes the data of an internal node.
|
||||
///
|
||||
/// Returns a raw ptr to avoid invalidating other references to this node.
|
||||
fn as_internal_ptr(this: &Self) -> *mut InternalNode<K, V> {
|
||||
// SAFETY: the static node type is `Internal`.
|
||||
this.node.as_ptr() as *mut InternalNode<K, V>
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K, V> NodeRef<K, V, marker::Mut<'a>, marker::Internal> {
|
||||
/// Borrows exclusive access to the data of an internal node.
|
||||
fn as_internal_mut(&mut self) -> &mut InternalNode<K, V> {
|
||||
let ptr = Self::as_internal_ptr(self);
|
||||
unsafe { &mut *ptr }
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K: 'a, V: 'a> NodeRef<K, V, marker::Mut<'a>, marker::LeafOrInternal> {
|
||||
fn set_parent_link(&mut self, parent: NonNull<InternalNode<K, V>>, parent_idx: usize) {
|
||||
unsafe {
|
||||
let leaf_ptr = Self::as_leaf_ptr(self);
|
||||
(*leaf_ptr).parent = Some(parent);
|
||||
(*leaf_ptr).parent_idx.write(parent_idx as u16);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K: 'a, V: 'a, Type> NodeRef<K, V, marker::Mut<'a>, Type> {
|
||||
unsafe fn key_area_mut<I, Output>(&mut self, index: I) -> &mut Output
|
||||
where
|
||||
I: core::slice::SliceIndex<[MaybeUninit<K>], Output = Output>,
|
||||
Output: ?Sized,
|
||||
{
|
||||
unsafe {
|
||||
self.as_leaf_mut()
|
||||
.keys
|
||||
.as_mut_slice()
|
||||
.get_unchecked_mut(index)
|
||||
}
|
||||
}
|
||||
unsafe fn val_area_mut<I, Output>(&mut self, index: I) -> &mut Output
|
||||
where
|
||||
I: core::slice::SliceIndex<[MaybeUninit<V>], Output = Output>,
|
||||
Output: ?Sized,
|
||||
{
|
||||
unsafe {
|
||||
self.as_leaf_mut()
|
||||
.values
|
||||
.as_mut_slice()
|
||||
.get_unchecked_mut(index)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K: 'a, V: 'a> NodeRef<K, V, marker::Mut<'a>, marker::Internal> {
|
||||
unsafe fn edge_area_mut<I, Output>(&mut self, index: I) -> &mut Output
|
||||
where
|
||||
I: core::slice::SliceIndex<[MaybeUninit<BoxedNode<K, V>>], Output = Output>,
|
||||
Output: ?Sized,
|
||||
{
|
||||
unsafe {
|
||||
self.as_internal_mut()
|
||||
.edges
|
||||
.as_mut_slice()
|
||||
.get_unchecked_mut(index)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K: 'a, V: 'a, Type> NodeRef<K, V, marker::Mut<'a>, Type> {
|
||||
/// Borrows exclusive access to the length of the node.
|
||||
pub(super) fn len_mut(&mut self) -> &mut u16 {
|
||||
&mut self.as_leaf_mut().len
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K, V, Type> NodeRef<K, V, marker::Mut<'a>, Type> {
|
||||
unsafe fn reborrow_mut(&mut self) -> NodeRef<K, V, marker::Mut<'a>, Type> {
|
||||
NodeRef {
|
||||
height: self.height,
|
||||
node: self.node,
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
fn as_leaf_mut(&mut self) -> &mut LeafNode<K, V> {
|
||||
unsafe { &mut *Self::as_leaf_ptr(self) }
|
||||
}
|
||||
|
||||
fn into_leaf_mut(mut self) -> &'a mut LeafNode<K, V> {
|
||||
// SAFETY: `self` is a mutable reference to a `NodeRef` with a mutable borrow type.
|
||||
unsafe { &mut *Self::as_leaf_ptr(&mut self) }
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K, V> NodeRef<K, V, marker::Mut<'a>, marker::Internal> {
|
||||
unsafe fn correct_childrens_parent_links<R>(&mut self, range: R)
|
||||
where
|
||||
R: Iterator<Item = usize>,
|
||||
{
|
||||
for i in range {
|
||||
assert!(i < self.len());
|
||||
unsafe {
|
||||
Handle::new_edge(self.reborrow_mut(), i).correct_parent_link();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn correct_all_childrens_parent_links(&mut self) {
|
||||
let len = self.len();
|
||||
unsafe {
|
||||
self.correct_childrens_parent_links(0..=len);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct Handle<Node, Type> {
|
||||
node: Node,
|
||||
idx: usize,
|
||||
_marker: PhantomData<Type>,
|
||||
}
|
||||
|
||||
impl<BorrowType, NodeType, K, V> Handle<NodeRef<K, V, BorrowType, NodeType>, marker::Edge> {
|
||||
unsafe fn new_edge(node: NodeRef<K, V, BorrowType, NodeType>, idx: usize) -> Self {
|
||||
assert!(idx < node.len());
|
||||
Handle {
|
||||
node,
|
||||
idx,
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K, V> Handle<NodeRef<K, V, marker::Mut<'a>, marker::Internal>, marker::Edge> {
|
||||
fn correct_parent_link(self) {
|
||||
let ptr = unsafe { NonNull::new_unchecked(NodeRef::as_internal_ptr(&self.node)) };
|
||||
let idx = self.idx;
|
||||
let mut child = self.descend();
|
||||
child.set_parent_link(ptr, idx);
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K, V, NodeType, HandleType> Handle<NodeRef<K, V, marker::Mut<'a>, NodeType>, HandleType> {
|
||||
unsafe fn reborrow_mut(
|
||||
&mut self,
|
||||
) -> Handle<NodeRef<K, V, marker::Mut<'a>, NodeType>, HandleType> {
|
||||
Handle {
|
||||
node: unsafe { self.node.reborrow_mut() },
|
||||
idx: self.idx,
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V, BorrowType> Handle<NodeRef<K, V, BorrowType, marker::Internal>, marker::Edge> {
|
||||
fn descend(self) -> NodeRef<K, V, BorrowType, marker::LeafOrInternal> {
|
||||
let parent_ptr = NodeRef::as_internal_ptr(&self.node);
|
||||
let node = unsafe {
|
||||
(*parent_ptr)
|
||||
.edges
|
||||
.get_unchecked(self.idx)
|
||||
.assume_init_read()
|
||||
};
|
||||
NodeRef {
|
||||
height: self.node.height - 1,
|
||||
node,
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K: 'a, V: 'a> Handle<NodeRef<K, V, marker::Mut<'a>, marker::Internal>, marker::Edge> {
|
||||
fn insert_fit(&mut self, key: K, val: V, edge: Root<K, V>) {
|
||||
debug_assert!(self.node.len() < CAPACITY);
|
||||
debug_assert!(edge.height == self.node.height - 1);
|
||||
let new_len = self.node.len() + 1;
|
||||
|
||||
unsafe {
|
||||
slice_insert(self.node.key_area_mut(..new_len), self.idx, key);
|
||||
slice_insert(self.node.val_area_mut(..new_len), self.idx, val);
|
||||
slice_insert(
|
||||
self.node.edge_area_mut(..new_len + 1),
|
||||
self.idx + 1,
|
||||
edge.node,
|
||||
);
|
||||
*self.node.len_mut() = new_len as u16;
|
||||
|
||||
self.node
|
||||
.correct_childrens_parent_links(self.idx + 1..new_len + 1);
|
||||
}
|
||||
}
|
||||
|
||||
fn insert(mut self, key: K, val: V, edge: Root<K, V>) {
|
||||
if self.node.len() < CAPACITY {
|
||||
self.insert_fit(key, val, edge);
|
||||
} else {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, K: 'a, V: 'a> Handle<NodeRef<K, V, marker::Mut<'a>, marker::Leaf>, marker::Edge> {}
|
||||
|
||||
type Root<K, V> = NodeRef<K, V, marker::Owned, marker::LeafOrInternal>;
|
||||
|
||||
struct Tree<K, V> {
|
||||
root: Option<Root<K, V>>,
|
||||
_marker: PhantomData<alloc::boxed::Box<(K, V)>>,
|
||||
}
|
||||
|
||||
impl<K, V> Tree<K, V> {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
root: None,
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn slice_insert<T>(slice: &mut [MaybeUninit<T>], idx: usize, value: T) {
|
||||
unsafe {
|
||||
let len = slice.len();
|
||||
debug_assert!(len > idx);
|
||||
let slice_ptr = slice.as_mut_ptr();
|
||||
if len > idx + 1 {
|
||||
core::ptr::copy(slice_ptr.add(idx), slice_ptr.add(idx + 1), len - idx - 1);
|
||||
}
|
||||
(*slice_ptr.add(idx)).write(value);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue