copy from std::tree

src/lib.rs

@@ -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;

src/tree.rs

@@ -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);
+    }
+}