trie innit

This commit is contained in:
janis 2025-09-17 00:57:15 +02:00
parent 3e3b096174
commit 5e1a4a245a
Signed by: janis
SSH key fingerprint: SHA256:bB1qbbqmDXZNT0KKD5c2Dfjg53JGhj7B3CFcLIzSqq8
5 changed files with 167 additions and 95 deletions

7
Cargo.lock generated
View file

@ -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",
]

View file

@ -16,4 +16,5 @@ nightly = []
# While I could use libc / windows for this, why not just use this tiny crate
# which does exactly and only a futex
atomic-wait = "1.1.0"
hashbrown = {version = "0.15", optional = true}
hashbrown = {version = "0.15", optional = true}
either = "1.15.0"

View file

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

View file

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

View file

@ -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>> {
super::subtree::Subtree {
root: self.handle.node,
}
pub fn into_subtree(self) -> super::subtree::Subtree<'a, K, V, marker::Mut<'a>> {
Subtree::new_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(super) root: NodeRef<BorrowType, K, V>,
pub struct Subtree<'tree, K, V, BorrowType> {
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?
let root = unsafe { self.root.reborrow_mut() };
match root.search_tree(&mut key_seq) {
search::SearchResult::Found(handle) => Occupied(OccupiedEntry {
handle,
_marker: PhantomData,
}),
search::SearchResult::GoDown(handle) => Vacant(VacantEntry {
key: key_seq.into(),
handle: TreeOrHandle::Other(handle),
_marker: PhantomData,
}),
search::SearchResult::Insert(key, handle) => Vacant(VacantEntry {
key: OnceAndIter::once(key, key_seq),
handle: TreeOrHandle::Other(handle),
_marker: PhantomData,
}),
match self.root_mut() {
Either::Right(tree) => {
return Vacant(VacantEntry {
key: key_seq.into(),
handle: TreeOrHandle::Tree(tree),
_marker: PhantomData,
});
}
Either::Left(root) => match root.search_tree(&mut key_seq) {
search::SearchResult::Found(handle) => Occupied(OccupiedEntry {
handle,
_marker: PhantomData,
}),
search::SearchResult::GoDown(handle) => Vacant(VacantEntry {
key: key_seq.into(),
handle: TreeOrHandle::Other(handle),
_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>>> {
let root = self.root.as_mut()?.borrow_mut().dormant();
Some(subtree::Subtree {
root: unsafe { root.awaken() },
})
pub fn as_subtree_mut<'a>(&'a mut self) -> subtree::Subtree<'a, K, V, marker::Mut<'a>> {
match self.root.as_mut() {
Some(node) => {
let dormant = node.borrow_mut().dormant();
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;
use entry::{Entry::*, VacantEntry};
let (tree, dormant) = DormantMutRef::new(self);
let entry = match tree.as_subtree_mut() {
Some(mut subtree) => {
let entry = subtree.entry(key_seq);
// SAFETY: extending the lifetime is fine because we borrow the tree for 'a,
// and no references to the subtree are live after this.
// The same could be achieved using `dormant.reborrow()` a bunch
// 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
let mut subtree = self.as_subtree_mut();
let entry = subtree.entry(key_seq);
// SAFETY: extending the lifetime is fine because we borrow the tree for 'a,
// and no references to the subtree are live after this.
// The same could be achieved using `dormant.reborrow()` a bunch
// 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)
}
}
}