trie innit

Cargo.lock

@@ -18,6 +18,12 @@ dependencies = [
  "windows-sys",
 ]
 
+[[package]]
+name = "either"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
+
 [[package]]
 name = "equivalent"
 version = "1.0.2"
@@ -52,6 +58,7 @@ name = "werkzeug"
 version = "0.1.0"
 dependencies = [
  "atomic-wait",
+ "either",
  "hashbrown",
 ]
 

Cargo.toml

@@ -17,3 +17,4 @@ nightly = []
 # which does exactly and only a futex
 atomic-wait = "1.1.0"
 hashbrown = {version = "0.15", optional = true}
+either = "1.15.0"

src/lib.rs

@@ -2,7 +2,6 @@
 #![cfg_attr(
     feature = "nightly",
     feature(
-        strict_provenance_atomic_ptr,
         box_vec_non_null,
         maybe_uninit_slice,
         debug_closure_helpers,

src/option.rs

@@ -79,20 +79,20 @@ impl<T: From<U>, U> core::ops::FromResidual<U> for TransposedOption<T> {
     }
 }
 
-#[cfg(all(test, feature = "transposed-option"))]
+// #[cfg(all(test, feature = "transposed-option"))]
-mod tests {
+// mod tests {
-    use super::*;
+//     use super::*;
-    use TransposedOption::*;
+//     use TransposedOption::*;
 
-    #[test]
+//     #[test]
-    fn transposed_option_try() {
+//     fn transposed_option_try() {
-        let a: TransposedOption<i32> = try {
+//         let a: TransposedOption<i32> = try {
-            TransposedOption::Some(42)?;
+//             TransposedOption::Some(42)?;
-            None::<i32>?;
+//             None::<i32>?;
 
-            Some(3)
+//             Some(3)
-        };
+//         };
 
-        assert_eq!(a, TransposedOption::Some(42));
+//         assert_eq!(a, TransposedOption::Some(42));
-    }
+//     }
-}
+// }

src/tree.rs

@@ -15,6 +15,8 @@
 use alloc::boxed::Box;
 use core::{marker::PhantomData, mem::MaybeUninit, ptr::NonNull};
 
+use crate::tree::subtree::Subtree;
+
 mod marker {
     use core::marker::PhantomData;
 
@@ -52,7 +54,7 @@ mod marker {
 }
 
 #[derive(Debug)]
-struct LeafNode<K, V> {
+pub(crate) struct LeafNode<K, V> {
     parent: Option<BoxedNode<K, V>>,
     parent_idx: MaybeUninit<u16>,
     value: Option<V>,
@@ -65,13 +67,13 @@ struct LeafNode<K, V> {
 type BoxedNode<K, V> = NonNull<LeafNode<K, V>>;
 
 #[derive(Debug)]
-struct NodeRef<BorrowType, K, V> {
+pub(crate) struct NodeRef<BorrowType, K, V> {
     node: NonNull<LeafNode<K, V>>,
     _marker: PhantomData<BorrowType>,
 }
 
 #[derive(Debug)]
-struct Handle<Node, Type> {
+pub(crate) struct Handle<Node, Type> {
     node: Node,
     idx: usize,
     _marker: PhantomData<Type>,
@@ -207,6 +209,7 @@ impl<BorrowType: marker::BorrowType, K, V> Handle<NodeRef<BorrowType, K, V>, mar
 }
 
 impl<K, V> NodeRef<marker::Dying, K, V> {
+    #[allow(dead_code)]
     pub(super) fn deallocate_and_ascend(
         self,
     ) -> Option<Handle<NodeRef<marker::Dying, K, V>, marker::Edge>> {
@@ -544,6 +547,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Immut<'a>, K, V>, marker::Value> {
 }
 
 impl<K, V> NodeRef<marker::Dying, K, V> {
+    #[allow(dead_code)]
     /// Borrows exclusive access to the leaf portion of a dying leaf or internal node.
     fn as_leaf_dying(&mut self) -> &mut LeafNode<K, V> {
         let ptr = Self::as_leaf_ptr(self);
@@ -553,6 +557,7 @@ impl<K, V> NodeRef<marker::Dying, K, V> {
 }
 
 impl<K, V> NodeRef<marker::Dying, K, V> {
+    #[allow(dead_code)]
     pub(super) unsafe fn into_value(mut self) -> Option<V> {
         let leaf = self.as_leaf_dying();
         leaf.value.take()
@@ -864,23 +869,36 @@ mod search {
     }
 }
 
-enum TreeOr<'a, K, V, T> {
+pub(crate) enum TreeOr<'a, K, V, T> {
     Tree(borrow::DormantMutRef<'a, Tree<K, V>>),
     Other(T),
 }
 
+impl<'a, K, V, T> TreeOr<'a, K, V, T> {
+    #[allow(dead_code)]
+    fn as_tree_mut(&mut self) -> &mut borrow::DormantMutRef<'a, Tree<K, V>> {
+        match self {
+            TreeOr::Tree(tree) => tree,
+            TreeOr::Other(_) => panic!("no tree present"),
+        }
+    }
+
+    #[allow(dead_code)]
+    fn as_other_mut(&mut self) -> &mut T {
+        match self {
+            TreeOr::Tree(_) => panic!("no other present"),
+            TreeOr::Other(other) => other,
+        }
+    }
+}
+
 type TreeOrHandle<'a, BorrowType, K, V, HandleType> =
     TreeOr<'a, K, V, Handle<NodeRef<BorrowType, K, V>, HandleType>>;
 
-enum HandleOrTree<'a, BorrowType, K, V, HandleType> {
-    Handle(Handle<NodeRef<BorrowType, K, V>, HandleType>),
-    Tree(borrow::DormantMutRef<'a, Tree<K, V>>),
-}
-
 mod entry {
     use core::marker::PhantomData;
 
-    use crate::tree::TreeOrHandle;
+    use crate::tree::{TreeOrHandle, subtree::Subtree};
 
     use super::{Handle, NodeRef, marker};
 
@@ -949,10 +967,8 @@ mod entry {
             unsafe { self.handle.value_mut() }
         }
 
-        pub fn into_subtree(self) -> super::subtree::Subtree<K, V, marker::Mut<'a>> {
+        pub fn into_subtree(self) -> super::subtree::Subtree<'a, K, V, marker::Mut<'a>> {
-            super::subtree::Subtree {
+            Subtree::new_root(self.handle.node)
-                root: self.handle.node,
-            }
         }
     }
 
@@ -998,7 +1014,9 @@ mod entry {
 mod subtree {
     use core::marker::PhantomData;
 
-    use crate::tree::{TreeOrHandle, search};
+    use either::Either;
+
+    use crate::tree::{TreeOr, TreeOrHandle, borrow::DormantMutRef, search};
 
     use super::{
         NodeRef, OnceAndIter,
@@ -1007,11 +1025,59 @@ mod subtree {
     };
 
     // BorrowType may be one of `Immut`, `Mut`.
-    pub struct Subtree<K, V, BorrowType> {
+    pub struct Subtree<'tree, K, V, BorrowType> {
-        pub(super) root: NodeRef<BorrowType, K, V>,
+        pub(super) root: TreeOr<'tree, K, V, NodeRef<BorrowType, K, V>>,
     }
 
-    impl<K, V, BorrowType> Subtree<K, V, BorrowType>
+    impl<'tree, K, V, BorrowType> Subtree<'tree, K, V, BorrowType> {
+        fn root(&self) -> Option<&NodeRef<BorrowType, K, V>> {
+            match &self.root {
+                TreeOr::Tree(_) => None,
+                TreeOr::Other(node) => Some(node),
+            }
+        }
+
+        pub(crate) fn new_root(node: NodeRef<BorrowType, K, V>) -> Self {
+            Self {
+                root: TreeOr::Other(node),
+            }
+        }
+
+        pub(crate) fn new_empty(tree: DormantMutRef<'tree, super::Tree<K, V>>) -> Self {
+            Self {
+                root: TreeOr::Tree(tree),
+            }
+        }
+    }
+
+    impl<'tree, K, V> Subtree<'tree, K, V, marker::Mut<'tree>> {
+        /// Returns a mutable reference to the root node of this subtree.
+        /// This function rechecks whether the tree is empty.
+        fn root_mut<'a>(
+            &'a mut self,
+        ) -> Either<NodeRef<marker::Mut<'a>, K, V>, DormantMutRef<'a, super::Tree<K, V>>> {
+            let node = match &mut self.root {
+                TreeOr::Tree(tree) => {
+                    let (tree, borrow) = DormantMutRef::new(unsafe { tree.reborrow() });
+                    match tree.root.as_mut() {
+                        Some(node) => node.borrow_mut().dormant(),
+                        None => return Either::Right(borrow),
+                    }
+                }
+                TreeOr::Other(node) => node.dormant(),
+            };
+
+            self.root = TreeOr::Other(unsafe { node.awaken() });
+
+            let TreeOr::Other(node) = &mut self.root else {
+                unreachable!()
+            };
+
+            return Either::Left(unsafe { node.reborrow_mut() });
+        }
+    }
+
+    impl<'tree, K, V, BorrowType> Subtree<'tree, K, V, BorrowType>
     where
         K: Ord,
     {
@@ -1019,7 +1085,7 @@ mod subtree {
         where
             Q: Iterator<Item = K>,
         {
-            let root = self.root.reborrow();
+            let root = self.root()?.reborrow();
             match root.search_tree(&mut key_seq) {
                 search::SearchResult::Found(handle) => Some(unsafe { handle.into_value() }),
                 _ => {
@@ -1029,13 +1095,16 @@ mod subtree {
             }
         }
 
-        pub fn get_subtree<Q>(&self, mut key_seq: Q) -> Option<Subtree<K, V, marker::Immut<'_>>>
+        pub fn get_subtree<Q>(
+            &'_ self,
+            mut key_seq: Q,
+        ) -> Option<Subtree<'_, K, V, marker::Immut<'_>>>
         where
             Q: Iterator<Item = K>,
         {
-            let root = self.root.reborrow();
+            let root = self.root()?.reborrow();
             match root.search_tree(&mut key_seq) {
-                search::SearchResult::Found(handle) => Some(Subtree { root: handle.node }),
+                search::SearchResult::Found(handle) => Some(Subtree::new_root(handle.node)),
                 _ => {
                     // key not found
                     None
@@ -1044,7 +1113,7 @@ mod subtree {
         }
     }
 
-    impl<'a, K, V> Subtree<K, V, marker::Mut<'a>>
+    impl<'tree, K, V> Subtree<'tree, K, V, marker::Mut<'tree>>
     where
         K: Ord,
     {
@@ -1052,7 +1121,7 @@ mod subtree {
         where
             Q: Iterator<Item = K>,
         {
-            let root = unsafe { self.root.reborrow_mut() };
+            let root = self.root_mut().left()?;
             match root.search_tree(&mut key_seq) {
                 search::SearchResult::Found(handle) => Some(handle.into_value_mut()),
                 _ => {
@@ -1065,13 +1134,13 @@ mod subtree {
         pub fn get_subtree_mut<Q>(
             &'_ mut self,
             mut key_seq: Q,
-        ) -> Option<Subtree<K, V, marker::Mut<'_>>>
+        ) -> Option<Subtree<'_, K, V, marker::Mut<'_>>>
         where
             Q: Iterator<Item = K>,
         {
-            let root = unsafe { self.root.reborrow_mut() };
+            let root = self.root_mut().left()?;
             match root.search_tree(&mut key_seq) {
-                search::SearchResult::Found(handle) => Some(Subtree { root: handle.node }),
+                search::SearchResult::Found(handle) => Some(Subtree::new_root(handle.node)),
                 _ => {
                     // key not found
                     None
@@ -1085,24 +1154,30 @@ mod subtree {
         {
             use Entry::*;
 
-            // SAFETY: this is actually our borrow?
+            match self.root_mut() {
-            let root = unsafe { self.root.reborrow_mut() };
+                Either::Right(tree) => {
-
+                    return Vacant(VacantEntry {
-            match root.search_tree(&mut key_seq) {
+                        key: key_seq.into(),
-                search::SearchResult::Found(handle) => Occupied(OccupiedEntry {
+                        handle: TreeOrHandle::Tree(tree),
-                    handle,
+                        _marker: PhantomData,
-                    _marker: PhantomData,
+                    });
-                }),
+                }
-                search::SearchResult::GoDown(handle) => Vacant(VacantEntry {
+                Either::Left(root) => match root.search_tree(&mut key_seq) {
-                    key: key_seq.into(),
+                    search::SearchResult::Found(handle) => Occupied(OccupiedEntry {
-                    handle: TreeOrHandle::Other(handle),
+                        handle,
-                    _marker: PhantomData,
+                        _marker: PhantomData,
-                }),
+                    }),
-                search::SearchResult::Insert(key, handle) => Vacant(VacantEntry {
+                    search::SearchResult::GoDown(handle) => Vacant(VacantEntry {
-                    key: OnceAndIter::once(key, key_seq),
+                        key: key_seq.into(),
-                    handle: TreeOrHandle::Other(handle),
+                        handle: TreeOrHandle::Other(handle),
-                    _marker: PhantomData,
+                        _marker: PhantomData,
-                }),
+                    }),
+                    search::SearchResult::Insert(key, handle) => Vacant(VacantEntry {
+                        key: OnceAndIter::once(key, key_seq),
+                        handle: TreeOrHandle::Other(handle),
+                        _marker: PhantomData,
+                    }),
+                },
             }
         }
     }
@@ -1123,7 +1198,7 @@ mod borrow {
     /// the compiler to follow. A `DormantMutRef` allows you to check borrowing
     /// yourself, while still expressing its stacked nature, and encapsulating
     /// the raw pointer code needed to do this without undefined behavior.
-    pub(super) struct DormantMutRef<'a, T> {
+    pub(crate) struct DormantMutRef<'a, T> {
         ptr: NonNull<T>,
         _marker: PhantomData<&'a mut T>,
     }
@@ -1307,45 +1382,35 @@ impl<K, V> Tree<K, V>
 where
     K: Ord,
 {
-    pub fn as_subtree_mut<'a>(&'a mut self) -> Option<subtree::Subtree<K, V, marker::Mut<'a>>> {
+    pub fn as_subtree_mut<'a>(&'a mut self) -> subtree::Subtree<'a, K, V, marker::Mut<'a>> {
-        let root = self.root.as_mut()?.borrow_mut().dormant();
+        match self.root.as_mut() {
-
+            Some(node) => {
-        Some(subtree::Subtree {
+                let dormant = node.borrow_mut().dormant();
-            root: unsafe { root.awaken() },
+                Subtree::new_root(unsafe { dormant.awaken() })
-        })
+            }
+            None => {
+                let (_, dormant) = borrow::DormantMutRef::new(self);
+                Subtree::new_empty(dormant)
+            }
+        }
     }
 
     pub fn entry<'a, Q>(&'a mut self, key_seq: Q) -> entry::Entry<'a, OnceAndIter<Q, K>, K, V>
     where
         Q: Iterator<Item = K>,
     {
-        use borrow::DormantMutRef;
+        let mut subtree = self.as_subtree_mut();
-        use entry::{Entry::*, VacantEntry};
+        let entry = subtree.entry(key_seq);
-
+        // SAFETY: extending the lifetime is fine because we borrow the tree for 'a,
-        let (tree, dormant) = DormantMutRef::new(self);
+        // and no references to the subtree are live after this.
-
+        // The same could be achieved using `dormant.reborrow()` a bunch
-        let entry = match tree.as_subtree_mut() {
+        // of times while destructuring the entry.
-            Some(mut subtree) => {
+        unsafe {
-                let entry = subtree.entry(key_seq);
+            core::mem::transmute::<
-                // SAFETY: extending the lifetime is fine because we borrow the tree for 'a,
+                entry::Entry<'_, OnceAndIter<Q, K>, K, V>,
-                // and no references to the subtree are live after this.
+                entry::Entry<'a, OnceAndIter<Q, K>, K, V>,
-                // The same could be achieved using `dormant.reborrow()` a bunch
+            >(entry)
-                // of times while destructuring the entry.
+        }
-                unsafe {
-                    core::mem::transmute::<
-                        entry::Entry<'_, OnceAndIter<Q, K>, K, V>,
-                        entry::Entry<'a, OnceAndIter<Q, K>, K, V>,
-                    >(entry)
-                }
-            }
-            None => Vacant(VacantEntry {
-                key: key_seq.into(),
-                handle: TreeOrHandle::Tree(dormant),
-                _marker: PhantomData,
-            }),
-        };
-
-        entry
     }
 }