From 10e762e730aeeed8a659cddc432eeb21d63ce1c1 Mon Sep 17 00:00:00 2001
From: Janis <janis@nirgendwo.xyz>
Date: Fri, 8 Aug 2025 00:32:28 +0200
Subject: [PATCH] invested actual thought into this

---
 src/tree.rs | 109 ++++++++++++++++++++++++++++++++--------------------
 1 file changed, 67 insertions(+), 42 deletions(-)

diff --git a/src/tree.rs b/src/tree.rs
index 8eed009..49e24c9 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -1,5 +1,23 @@
+//! A Prefix-Tree (Trie) implementation in Rust.
+//! This Trie is structured with key-value pairs at the same level.
+// tree:
+//   root: { kvs: [('f', v0), ('g', v1), ('i', _)], edges: [child0, None, child1] }
+//     child0: { kvs: [('a', v2), ('b', v3)], edges: [] }
+//     child1: { kvs: [('a', v4)], edges: [] }
+// tree.find() -> ()
+// tree.find('f') -> v0
+// tree.find('g') -> v1
+// tree.find('fa') -> v2
+// tree.find('fb') -> v3
+// tree.find('ia') -> v4
+// tree.find('i') -> None
+
 use alloc::boxed::Box;
-use core::{marker::PhantomData, mem::MaybeUninit, ptr::NonNull};
+use core::{
+    marker::PhantomData,
+    mem::{ManuallyDrop, MaybeUninit},
+    ptr::NonNull,
+};
 
 mod marker {
     use core::marker::PhantomData;
@@ -38,9 +56,41 @@ const CAPACITY: usize = 16;
 struct LeafNode<K, V> {
     parent: Option<NonNull<InternalNode<K, V>>>,
     parent_idx: MaybeUninit<u16>,
+    value: ManuallyDrop<V>,
+    leaf: bool,
+}
+
+struct InternalNode<K, V> {
+    data: LeafNode<K, V>,
+    next: Option<NonNull<InternalNode<K, V>>>,
     len: u16,
     keys: [MaybeUninit<K>; CAPACITY],
-    values: [MaybeUninit<V>; CAPACITY],
+    edges: [MaybeUninit<BoxedNode<K, V>>; CAPACITY],
+}
+
+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)>,
+}
+
+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<K, V> LeafNode<K, V> {
@@ -53,41 +103,31 @@ impl<K, V> LeafNode<K, V> {
     }
 
     fn new() -> Box<Self> {
+        let mut this = Box::new_uninit();
         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 fn init(this: *mut 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
-
+            LeafNode::init(&raw mut (*this).data);
+            (&raw mut (*this).next).write(None);
+            (&raw mut (*this).edges).write([None; CAPACITY + 1]);
+        }
+    }
+    fn new() -> Box<Self> {
+        let mut this = Box::<Self>::new_uninit();
+        unsafe {
+            Self::init(this.as_mut_ptr());
             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();
@@ -101,8 +141,8 @@ impl<K, V> NodeRef<K, V, marker::Owned, marker::Leaf> {
 
 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);
+        let mut node = InternalNode::new();
+        node.edges[0] = Some(child.node);
 
         // SAFETY: `height` isn't zero
         unsafe { NodeRef::from_new_internal(node, child.height + 1) }
@@ -123,6 +163,8 @@ impl<K, V> NodeRef<K, V, marker::Owned, marker::Internal> {
     }
 }
 
+// here follows std::BTree stuff
+
 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
@@ -285,23 +327,6 @@ impl<'a, K, V> NodeRef<K, V, marker::Mut<'a>, marker::Internal> {
     }
 }
 
-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)) };