AllocRef support for BTreeMap

src/btree/mod.rs

@@ -0,0 +1,37 @@
+#![allow(
+    explicit_outlives_requirements,
+    single_use_lifetimes,
+    trivial_casts,
+)]
+#![allow(
+    clippy::pedantic,
+    clippy::nursery,
+    clippy::type_complexity
+)]
+pub mod map;
+mod navigate;
+mod node;
+mod search;
+pub mod set;
+#[cfg(all(test, std))]
+mod tests;
+
+#[doc(hidden)]
+trait Recover<Q: ?Sized> {
+    type Key;
+
+    fn get(&self, key: &Q) -> Option<&Self::Key>;
+    fn take(&mut self, key: &Q) -> Option<Self::Key>;
+    fn replace(&mut self, key: Self::Key) -> Option<Self::Key>;
+}
+
+#[inline(always)]
+pub unsafe fn unwrap_unchecked<T>(val: Option<T>) -> T {
+    val.unwrap_or_else(|| {
+        if cfg!(debug_assertions) {
+            panic!("'unchecked' unwrap on None in BTreeMap");
+        } else {
+            core::intrinsics::unreachable();
+        }
+    })
+}

src/btree/navigate.rs

@@ -0,0 +1,248 @@
+use core::ptr;
+
+use super::{
+    node::{marker, ForceResult::*, Handle, NodeRef},
+    unwrap_unchecked,
+};
+
+use crate::alloc::AllocRef;
+
+impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
+    /// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
+    /// on the right side, which is either in the same leaf node or in an ancestor node.
+    /// If the leaf edge is the last one in the tree, returns [`Result::Err`] with the root node.
+    pub fn next_kv(
+        self,
+    ) -> Result<
+        Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>,
+        NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
+    > {
+        let mut edge = self.forget_node_type();
+        loop {
+            edge = match edge.right_kv() {
+                Ok(internal_kv) => return Ok(internal_kv),
+                Err(last_edge) => match last_edge.into_node().ascend() {
+                    Ok(parent_edge) => parent_edge.forget_node_type(),
+                    Err(root) => return Err(root.forget_type()),
+                },
+            }
+        }
+    }
+
+    /// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
+    /// on the left side, which is either in the same leaf node or in an ancestor node.
+    /// If the leaf edge is the first one in the tree, returns [`Result::Err`] with the root node.
+    pub fn next_back_kv(
+        self,
+    ) -> Result<
+        Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>,
+        NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
+    > {
+        let mut edge = self.forget_node_type();
+        loop {
+            edge = match edge.left_kv() {
+                Ok(internal_kv) => return Ok(internal_kv),
+                Err(last_edge) => match last_edge.into_node().ascend() {
+                    Ok(parent_edge) => parent_edge.forget_node_type(),
+                    Err(root) => return Err(root.forget_type()),
+                },
+            }
+        }
+    }
+}
+
+macro_rules! def_next_kv_uncheched_dealloc {
+    { unsafe fn $name:ident : $adjacent_kv:ident } => {
+        /// Given a leaf edge handle into an owned tree, returns a handle to the next KV,
+        /// while deallocating any node left behind.
+        /// Unsafe for two reasons:
+        /// - The caller must ensure that the leaf edge is not the last one in the tree.
+        /// - The node pointed at by the given handle, and its ancestors, may be deallocated,
+        ///   while the reference to those nodes in the surviving ancestors is left dangling;
+        ///   thus using the returned handle to navigate further is dangerous.
+        unsafe fn $name <K, V, A: AllocRef>(
+            leaf_edge: Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
+            mut alloc: A
+        ) -> Handle<NodeRef<marker::Owned, K, V, marker::LeafOrInternal>, marker::KV> {
+            let mut edge = leaf_edge.forget_node_type();
+            loop {
+                edge = match edge.$adjacent_kv() {
+                    Ok(internal_kv) => return internal_kv,
+                    Err(last_edge) => {
+                        let parent_edge = last_edge.into_node().deallocate_and_ascend(&mut alloc);
+                        unwrap_unchecked(parent_edge).forget_node_type()
+                    }
+                }
+            }
+        }
+    };
+}
+
+def_next_kv_uncheched_dealloc! {unsafe fn next_kv_unchecked_dealloc: right_kv}
+def_next_kv_uncheched_dealloc! {unsafe fn next_back_kv_unchecked_dealloc: left_kv}
+
+/// This replaces the value behind the `v` unique reference by calling the
+/// relevant function.
+///
+/// Safety: The change closure must not panic.
+#[inline]
+unsafe fn replace<T, R>(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R {
+    let value = ptr::read(v);
+    let (new_value, ret) = change(value);
+    ptr::write(v, new_value);
+    ret
+}
+
+impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge> {
+    /// Moves the leaf edge handle to the next leaf edge and returns references to the
+    /// key and value in between.
+    /// Unsafe because the caller must ensure that the leaf edge is not the last one in the tree.
+    pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
+        replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_kv();
+            let kv = unwrap_unchecked(kv.ok());
+            (kv.next_leaf_edge(), kv.into_kv())
+        })
+    }
+
+    /// Moves the leaf edge handle to the previous leaf edge and returns references to the
+    /// key and value in between.
+    /// Unsafe because the caller must ensure that the leaf edge is not the first one in the tree.
+    pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
+        replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_back_kv();
+            let kv = unwrap_unchecked(kv.ok());
+            (kv.next_back_leaf_edge(), kv.into_kv())
+        })
+    }
+}
+
+impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
+    /// Moves the leaf edge handle to the next leaf edge and returns references to the
+    /// key and value in between.
+    /// Unsafe for two reasons:
+    /// - The caller must ensure that the leaf edge is not the last one in the tree.
+    /// - Using the updated handle may well invalidate the returned references.
+    pub unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
+        let kv = replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_kv();
+            let kv = unwrap_unchecked(kv.ok());
+            (ptr::read(&kv).next_leaf_edge(), kv)
+        });
+        // Doing the descend (and perhaps another move) invalidates the references
+        // returned by `into_kv_mut`, so we have to do this last.
+        kv.into_kv_mut()
+    }
+
+    /// Moves the leaf edge handle to the previous leaf and returns references to the
+    /// key and value in between.
+    /// Unsafe for two reasons:
+    /// - The caller must ensure that the leaf edge is not the first one in the tree.
+    /// - Using the updated handle may well invalidate the returned references.
+    pub unsafe fn next_back_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
+        let kv = replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_back_kv();
+            let kv = unwrap_unchecked(kv.ok());
+            (ptr::read(&kv).next_back_leaf_edge(), kv)
+        });
+        // Doing the descend (and perhaps another move) invalidates the references
+        // returned by `into_kv_mut`, so we have to do this last.
+        kv.into_kv_mut()
+    }
+}
+
+impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
+    /// Moves the leaf edge handle to the next leaf edge and returns the key and value
+    /// in between, while deallocating any node left behind.
+    /// Unsafe for three reasons:
+    /// - The caller must ensure that the leaf edge is not the last one in the tree
+    ///   and is not a handle previously resulting from counterpart `next_back_unchecked`.
+    /// - If the leaf edge is the last edge of a node, that node and possibly ancestors
+    ///   will be deallocated, while the reference to those nodes in the surviving ancestor
+    ///   is left dangling; thus further use of the leaf edge handle is dangerous.
+    ///   It is, however, safe to call this method again on the updated handle.
+    ///   if the two preconditions above hold.
+    /// - Using the updated handle may well invalidate the returned references.
+    pub unsafe fn next_unchecked<A: AllocRef>(&mut self, alloc: A) -> (K, V) {
+        replace(self, |leaf_edge| {
+            let kv = next_kv_unchecked_dealloc(leaf_edge, alloc);
+            let k = ptr::read(kv.reborrow().into_kv().0);
+            let v = ptr::read(kv.reborrow().into_kv().1);
+            (kv.next_leaf_edge(), (k, v))
+        })
+    }
+
+    /// Moves the leaf edge handle to the previous leaf edge and returns the key
+    /// and value in between, while deallocating any node left behind.
+    /// Unsafe for three reasons:
+    /// - The caller must ensure that the leaf edge is not the first one in the tree
+    ///   and is not a handle previously resulting from counterpart `next_unchecked`.
+    /// - If the lead edge is the first edge of a node, that node and possibly ancestors
+    ///   will be deallocated, while the reference to those nodes in the surviving ancestor
+    ///   is left dangling; thus further use of the leaf edge handle is dangerous.
+    ///   It is, however, safe to call this method again on the updated handle.
+    ///   if the two preconditions above hold.
+    /// - Using the updated handle may well invalidate the returned references.
+    pub unsafe fn next_back_unchecked<A: AllocRef>(&mut self, alloc: A) -> (K, V) {
+        replace(self, |leaf_edge| {
+            let kv = next_back_kv_unchecked_dealloc(leaf_edge, alloc);
+            let k = ptr::read(kv.reborrow().into_kv().0);
+            let v = ptr::read(kv.reborrow().into_kv().1);
+            (kv.next_back_leaf_edge(), (k, v))
+        })
+    }
+}
+
+impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
+    /// Returns the leftmost leaf edge in or underneath a node - in other words, the edge
+    /// you need first when navigating forward (or last when navigating backward).
+    #[inline]
+    pub fn first_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
+        let mut node = self;
+        loop {
+            match node.force() {
+                Leaf(leaf) => return leaf.first_edge(),
+                Internal(internal) => node = internal.first_edge().descend(),
+            }
+        }
+    }
+
+    /// Returns the rightmost leaf edge in or underneath a node - in other words, the edge
+    /// you need last when navigating forward (or first when navigating backward).
+    #[inline]
+    pub fn last_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
+        let mut node = self;
+        loop {
+            match node.force() {
+                Leaf(leaf) => return leaf.last_edge(),
+                Internal(internal) => node = internal.last_edge().descend(),
+            }
+        }
+    }
+}
+
+impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
+    /// Returns the leaf edge closest to a KV for forward navigation.
+    pub fn next_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
+        match self.force() {
+            Leaf(leaf_kv) => leaf_kv.right_edge(),
+            Internal(internal_kv) => {
+                let next_internal_edge = internal_kv.right_edge();
+                next_internal_edge.descend().first_leaf_edge()
+            }
+        }
+    }
+
+    /// Returns the leaf edge closest to a KV for backward navigation.
+    pub fn next_back_leaf_edge(
+        self,
+    ) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
+        match self.force() {
+            Leaf(leaf_kv) => leaf_kv.left_edge(),
+            Internal(internal_kv) => {
+                let next_internal_edge = internal_kv.left_edge();
+                next_internal_edge.descend().last_leaf_edge()
+            }
+        }
+    }
+}