Code shared from the Rust Playground

playground.rs

use core::mem::MaybeUninit;
use core::ptr::NonNull;
use core::marker::PhantomData;
use core::cmp::Ordering;

const MAX_N: usize = 16;

struct InternalNode<T> {
    leaf: LeafNode<T>,
    children: [BoxedNode<T>; MAX_N],
}

impl<T> InternalNode<T> {
    // do not drop this struct (uninitialized fields may be dropped)
    // this struct must be cleaned manually
    unsafe fn new() -> InternalNode<T> {
        InternalNode { 
            leaf: LeafNode::new(),
            children: MaybeUninit::uninit().assume_init()
        }
    }

    #[cfg(test)]
    unsafe fn set_value(&mut self, index: usize, value: T) {
        self.leaf.set_value(index, value)
    }

    #[cfg(test)]
    unsafe fn set_count(&mut self, count: usize) {
        self.leaf.set_count(count)
    }

    #[cfg(test)]
    unsafe fn set_child(&mut self, index: usize, child: BoxedNode<T>) {
        self.children[index] = child;
    }
}

struct LeafNode<T> {
    count: usize,
    values: [MaybeUninit<T>; MAX_N],
}

impl<T> LeafNode<T> {
    unsafe fn new() -> Self {
        LeafNode { 
            count: 0, 
            values: MaybeUninit::uninit().assume_init()
        }
    }

    unsafe fn get_value_unchecked(&self, index: usize) -> &T {
        core::mem::transmute(&self.values[index])
    }

    #[cfg(test)]
    unsafe fn set_value(&mut self, index: usize, value: T) {
        self.values[index] = MaybeUninit::new(value);
    }

    #[cfg(test)]
    unsafe fn set_count(&mut self, count: usize) {
        self.count = count;
    }
}

struct BoxedNode<T> {
    // actually this field could owns LeafNode or InternalNode
    owned_ptr: NonNull<LeafNode<T>>, 
}

impl<T> BoxedNode<T> {
    fn from_leaf(node: Box<LeafNode<T>>) -> Self {
        BoxedNode { owned_ptr: box_into_raw_non_null(node) }
    }

    fn from_internal(node: Box<InternalNode<T>>) -> Self {
        BoxedNode { owned_ptr: box_into_raw_non_null(node).cast() }
    }

    fn from_ptr(owned_ptr: NonNull<LeafNode<T>>) -> Self {
        BoxedNode { owned_ptr }
    }

    fn as_ptr(&self) -> NonNull<LeafNode<T>> {
        self.owned_ptr
    }
}

fn box_into_raw_non_null<T>(src: Box<T>) -> NonNull<T> {
    // note(unsafe): box pointers are guaranteed non-null
    unsafe { NonNull::new_unchecked(Box::into_raw(src)) }
}

// clean this up manually
pub struct Root<T> {
    node: BoxedNode<T>,
    height: usize,
}

impl<T> Root<T> {
    pub fn new_leaf() -> Self {
        Root { node: BoxedNode::from_leaf(Box::new(unsafe { LeafNode::new() } )), height: 0 }
    }

    pub fn as_ref(&self) -> NodeRef<marker::Immut<'_>, T, marker::LeafOrInternal> {
        NodeRef {
            height: self.height,
            ptr: self.node.as_ptr(),
            _marker: PhantomData,
        }
    }

    pub fn as_mut(&mut self) -> NodeRef<marker::Mut<'_>, T, marker::LeafOrInternal> {
        NodeRef {
            height: self.height,
            ptr: self.node.as_ptr(),
            _marker: PhantomData,
        }
    }

    pub fn into_ref(self) -> NodeRef<marker::Owned, T, marker::LeafOrInternal> {
        NodeRef {
            height: self.height,
            ptr: self.node.as_ptr(),
            _marker: PhantomData,
        }
    }
}

pub struct NodeRef<B, T, TR> {
    // leaf: 0, internal: >0, root: h-1
    height: usize,
    ptr: NonNull<LeafNode<T>>,
    _marker: PhantomData<(B, TR)>,
}

impl<B, T, TR> NodeRef<B, T, TR> {
    pub fn height(&self) -> usize {
        self.height
    }

    fn len(&self) -> usize {
        unsafe { self.ptr.as_ref() }.count
    }
    
    fn reborrow(&self) -> NodeRef<marker::Immut<'_>, T, TR> {
        NodeRef { height: self.height(), ptr: self.ptr, _marker: PhantomData }
    }

    fn values(&self) -> &[T] {
        self.reborrow().into_value_slice()
    } 
}

impl<'a, T, TR> NodeRef<marker::Immut<'a>, T, TR> {
    fn into_value_slice(&self) -> &'a [T] {
        let values = unsafe { &self.ptr.as_ref().values as *const _ as *const _ };
        unsafe { core::slice::from_raw_parts(values, self.len()) }
    }
}

// once handle is created operations on it is guaranteed to be valid
pub struct Handle<NR, TH> {
    node_ref: NR,
    index: usize,
    _marker: PhantomData<TH>,
}

impl<B, T, TR> Handle<NodeRef<B, T, TR>, marker::Slot> {
    pub fn new_slot(node_ref: NodeRef<B, T, TR>, index: usize) -> Self {
        debug_assert!(index < node_ref.len());
        Handle { node_ref, index, _marker: PhantomData }
    } 
}

impl<B, T, TR> Handle<NodeRef<B, T, TR>, marker::Edge> {
    pub fn new_edge(node_ref: NodeRef<B, T, TR>, index: usize) -> Self {
        debug_assert!(index <= node_ref.len());
        Handle { node_ref, index, _marker: PhantomData }
    } 
}

pub enum ForceResult<L, I> {
    Leaf(L),
    Internal(I),
}

impl<B, T> NodeRef<B, T, marker::LeafOrInternal> {
    pub fn force(self) -> ForceResult<
        NodeRef<B, T, marker::Leaf>,
        NodeRef<B, T, marker::Internal>,
    > {
        if self.height == 0 {
            ForceResult::Leaf(NodeRef {
                height: self.height,
                ptr: self.ptr,
                _marker: PhantomData
            })
        } else {
            ForceResult::Internal(NodeRef {
                height: self.height,
                ptr: self.ptr,
                _marker: PhantomData
            })
        }
    }
}

impl<B, T, TH> Handle<NodeRef<B, T, marker::LeafOrInternal>, TH> {
    pub fn force(self) -> ForceResult<
        Handle<NodeRef<B, T, marker::Leaf>, TH>,
        Handle<NodeRef<B, T, marker::Internal>, TH>
    > {
        match self.node_ref.force() {
            ForceResult::Leaf(node_ref) => {
                ForceResult::Leaf(Handle { node_ref, index: self.index, _marker: PhantomData })
            },
            ForceResult::Internal(node_ref) => {
                ForceResult::Internal(Handle { node_ref, index: self.index, _marker: PhantomData })
            },
        }
    }
}

impl<B, T> Handle<NodeRef<B, T, marker::Internal>, marker::Edge> {
    pub fn descend(self) -> NodeRef<B, T, marker::LeafOrInternal> {
        NodeRef {
            height: self.node_ref.height - 1,
            ptr: unsafe { 
                self.node_ref.ptr.cast::<InternalNode<T>>().as_ref().children[self.index].as_ptr()
            },
            _marker: PhantomData
        }
    }
}

impl<B, T> Handle<NodeRef<B, T, marker::LeafOrInternal>, marker::Slot> {
    pub fn get_value(&self) -> &T {
        unsafe { &self.node_ref.ptr.as_ref().get_value_unchecked(self.index) }
    }
}

pub enum SearchResult<B, T, FF, GD> {
    Found(Handle<NodeRef<B, T, FF>, marker::Slot>),
    GoDown(Handle<NodeRef<B, T, GD>, marker::Edge>),
}

impl<B, T, FF, GD> SearchResult<B, T, FF, GD> {
    pub fn is_found(&self) -> bool {
        if let SearchResult::Found(_) = self {
            true
        } else {
            false
        }
    }

    pub fn is_go_down(&self) -> bool {
        if let SearchResult::GoDown(_) = self {
            true
        } else {
            false
        }
    }
}

pub fn search_linear<B, T, TR, Q: ?Sized>(node_ref: &NodeRef<B, T, TR>, value: &Q) -> Result<usize, usize> 
where Q: Ord, T: Borrow<Q> {
    for (i, v) in node_ref.values().iter().enumerate() {
        match Q::cmp(value, v.borrow()) {
            Ordering::Greater => {},
            Ordering::Equal => return Ok(i),
            Ordering::Less => return Err(i),
        }
    }
    Err(node_ref.len())
}

pub fn search_node<B, T, TR, Q: ?Sized>(node_ref: NodeRef<B, T, TR>, value: &Q) -> SearchResult<B, T, TR, TR>
where Q: Ord, T: Borrow<Q> {
    match search_linear(&node_ref, value) {
        Ok(index) => SearchResult::Found(Handle::new_slot(node_ref, index)),
        Err(index) => SearchResult::GoDown(Handle::new_edge(node_ref, index))
    }
}

pub fn search_tree<B, T, Q: ?Sized>(mut node_ref: NodeRef<B, T, marker::LeafOrInternal>, value: &Q) 
-> SearchResult<B, T, marker::LeafOrInternal, marker::Leaf> 
where Q: Ord, T: Borrow<Q> {
    loop {
        match search_node(node_ref, value) {
            SearchResult::Found(handle) => return SearchResult::Found(handle),
            SearchResult::GoDown(handle) => match handle.force() {
                ForceResult::Leaf(leaf) => return SearchResult::GoDown(leaf),
                ForceResult::Internal(internal) => {
                    node_ref = internal.descend();
                }
            }
        }
    }
}

mod marker {
    use core::marker::PhantomData;

    // Borrow or ownership type (B)
    pub struct Immut<'a>(PhantomData<&'a ()>);
    pub struct Mut<'a>(PhantomData<&'a ()>);
    pub enum Owned {}

    // Handle type (TH)
    pub enum Edge {}
    pub enum Slot {}

    pub enum Leaf {}
    pub enum Internal {}
    pub enum LeafOrInternal {}
}

use core::borrow::Borrow;

pub struct BTreeSet<T> {
    root: Root<T>,
    length: usize,
}

impl<T: Ord> BTreeSet<T> {
    pub fn new() -> Self {
        let leaf = Box::new(unsafe { LeafNode::new() });
        let root = Root { node: BoxedNode::from_leaf(leaf), height: 0 };
        BTreeSet { root, length: 0 }
    }

    pub fn contains<Q: ?Sized>(&self, key: &Q) -> bool
    where 
        T: Borrow<Q>,
        Q: Ord 
    {
        search_tree(self.root.as_ref(), key).is_found()
    }

    pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&T>
    where
        T: Borrow<Q>,
        Q: Ord
    {
        match search_tree(self.root.as_ref(), key) {
            SearchResult::Found(handle) => {
                Some(handle.get_value()) // ??
            },
            SearchResult::GoDown(_) => None,
        }
    }

    // true: success
    // false: value presents and is not inserted
    pub fn insert(&mut self, value: T) -> bool {
        match search_tree(self.root.as_ref(), &value) {
            SearchResult::Found(handle) => {
                todo!()
            },
            SearchResult::GoDown(handle) => {
                todo!()
            },
        }
    }
}

impl<T> Drop for BTreeSet<T> {
    fn drop(&mut self) {
        todo!()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_search_linear_node() {
        let mut leaf = Box::new(unsafe { LeafNode::new() });
        unsafe { 
            leaf.set_value(0, 100);
            leaf.set_value(1, 200);
            leaf.set_value(2, 300);
            leaf.set_count(3);
        }
        let root = Root { node: BoxedNode::from_leaf(leaf), height: 0 };
        assert_eq!(Err(0), search_linear(&root.as_ref(), &50));
        assert_eq!(Ok(0), search_linear(&root.as_ref(), &100));
        assert_eq!(Err(1), search_linear(&root.as_ref(), &150));
        assert_eq!(Err(3), search_linear(&root.as_ref(), &350));
        core::mem::forget(root); 
    }
    
    #[test]
    fn test_search_tree() {
        let leaf_1 = unsafe {
            let mut leaf = Box::new(LeafNode::new());
            leaf.set_value(0, 100);
            leaf.set_value(1, 200);
            leaf.set_value(2, 300);
            leaf.set_count(3);
            BoxedNode::from_leaf(leaf)
        };
        let leaf_2 = unsafe {
            let mut leaf = Box::new(LeafNode::new());
            leaf.set_value(0, 1100);
            leaf.set_value(1, 1200);
            leaf.set_value(2, 1300);
            leaf.set_count(3);
            BoxedNode::from_leaf(leaf)
        };
        let mut internal = Box::new(unsafe { InternalNode::new() });
        unsafe { 
            internal.set_value(0, 1000);
            internal.set_count(1);
            internal.set_child(0, leaf_1);
            internal.set_child(1, leaf_2);
        }
        let root = Root { node: BoxedNode::from_internal(internal), height: 1 };
        // Found
        assert!(search_tree(root.as_ref(), &100).is_found());  // H=0, i=0
        assert!(search_tree(root.as_ref(), &200).is_found());  // H=0, i=1
        assert!(search_tree(root.as_ref(), &1000).is_found()); // H=1, i=0
        assert!(search_tree(root.as_ref(), &1300).is_found()); // H=0, i=2
        // GoDown
        assert!(search_tree(root.as_ref(), &50).is_go_down());   // H=0, i=0
        assert!(search_tree(root.as_ref(), &150).is_go_down());  // H=0, i=1
        assert!(search_tree(root.as_ref(), &350).is_go_down());  // H=0, i=3
        assert!(search_tree(root.as_ref(), &1150).is_go_down()); // H=0, i=1
        assert!(search_tree(root.as_ref(), &2000).is_go_down()); // H=0, i=3
        core::mem::forget(root); 
    }
}