stupid half 2-3-4 tree

two_three_four_tree.rs

use std::cmp::max;
use std::collections::VecDeque;
use std::convert::TryInto;
use std::fmt::Debug;
use std::mem::MaybeUninit;
use std::ptr::slice_from_raw_parts;

use itertools::Itertools;
use itertools::MinMaxResult;

#[derive(Debug, Eq, PartialEq)]
struct Tree<V> {
    root: TreeNode<V>,
    size: usize,
}

#[derive(Debug, Eq, PartialEq)]
pub enum TreeNode<V> {
    Empty,
    Leaf(Vec<V>),
    Two(Branch<V, Box<TreeNode<V>>, 1, 2>),
    Three(Branch<V, Box<TreeNode<V>>, 2, 3>),
    Four(Branch<V, Box<TreeNode<V>>, 3, 4>),
}

//
// enum TreeNode<V> {
//     Empty,
//     Leaf(Vec<V>),
//     Two(V, Box<TreeNode<V>>, Box<TreeNode<V>>),
//     Three([V; 2], Box<TreeNode<V>>, Box<TreeNode<V>>,Box<TreeNode<V>>),
//     Four([V; 3], Box<TreeNode<V>>, Box<TreeNode<V>>,Box<TreeNode<V>>),
// }

#[derive(Debug, Eq, PartialEq)]
struct Branch<K, V, const NK: usize, const NV: usize> {
    keys: Vec<K>,
    values: VecDeque<V>,
}

impl<V> Tree<V>
where
    V: PartialEq + PartialOrd + Debug,
{
    fn empty() -> Self {
        Tree {
            root: TreeNode::Empty,
            size: 0,
        }
    }

    fn insert(&mut self, v: V) {
        let new_root = self.root.insert(v);
        if let Some(nr) = new_root {
            eprintln!("nr = {:?}", nr);
            let _ = std::mem::replace(&mut self.root, nr);
        }
    }
}

enum InsertResult {
    Ok(),
    Existed,
}

impl<V> TreeNode<V>
where
    V: PartialEq + PartialOrd + Debug,
{
    fn singleton(v: V) -> Self {
        let mut vec1 = Vec::with_capacity(3);
        vec1.push(v);
        TreeNode::Leaf(vec1)
    }

    // fn branch<const NK: usize, const NV: usize>(&self) -> &Branch<K, V, NK, NV> {
    //     match self {
    //         TreeNode::Tip(ref b) => b,
    //         TreeNode::Two(ref b) => b,
    //         TreeNode::Three(ref b) => b,
    //         TreeNode::Four(ref b) => b,
    //         TreeNode::Empty(ref b) => b,
    //     }
    // }

    fn keys(&self) -> &[V] {
        match self {
            TreeNode::Leaf(ref b) => &b[..],
            TreeNode::Two(ref b) => &b.keys,
            TreeNode::Three(ref b) => &b.keys,
            TreeNode::Four(ref b) => &b.keys,
            _ => &[],
        }
    }

    // fn values(&self) -> &[V] {
    //     match self {
    //         TreeNode::Two(ref b) => &b.values[..],
    //         TreeNode::Three(ref b) => &b.values[..],
    //         TreeNode::Four(ref b) => &b.values[..],
    //         leaf @ TreeNode::Leaf(_) => unsafe {
    //             let raw_parts = slice_from_raw_parts(leaf, 1);
    //             &*raw_parts
    //         },
    //     }
    // }

    fn split_four_nodes(&mut self) -> TreeNode<V> {
        match self {
            TreeNode::Leaf(ref mut values) if values.len() == 3 => {
                let right = values.pop();
                let new_key = values.pop();
                let left = values.pop();

                TreeNode::Two(Branch::new(
                    vec![new_key.unwrap()],
                    vec![
                        Box::new(TreeNode::singleton(left.unwrap())),
                        Box::new(TreeNode::singleton(right.unwrap())),
                    ],
                ))
            },
            TreeNode::Three(Branch {keys, values}) if keys.len() == 3 => {
                let right_key = keys.pop().unwrap();
                let new_key = keys.pop().unwrap();
                let left_key = keys.pop().unwrap();


                let ll = values.pop_front().unwrap();
                let lr = values.pop_front().unwrap();
                let rl = values.pop_front().unwrap();
                let rr = values.pop_front().unwrap();

                TreeNode::Two(Branch::new(vec![new_key],
                                          vec![
                                              Box::new(TreeNode::Two(Branch::new(vec![left_key], vec![ll, lr]))),
                                              Box::new(TreeNode::Two(Branch::new(vec![right_key], vec![rl, rr])))
                                          ]))
            }
            _ => todo!(),
        }
    }

    fn merge(&mut self, mut r: Self) -> Self {
        match (self, &mut r) {
            (
                TreeNode::Two(Branch {
                    keys: lk,
                    values: lv,
                    ..
                }),
                TreeNode::Two(Branch {
                    keys: rk,
                    values: rv,
                    ..
                }),
            ) => {
                let mut new_values: Vec<Box<TreeNode<V>>> = Vec::with_capacity(3);
                new_values.push(lv.pop_front().unwrap());
                let last = rv.pop_back().unwrap();
                new_values.push(rv.pop_back().unwrap());
                new_values.push(last);
                TreeNode::Three(Branch::new(
                    vec![lk.pop().unwrap(), rk.pop().unwrap()],
                    new_values,
                ))
            }
            (
                TreeNode::Three(Branch {
                    keys: lk,
                    values: lv,
                    ..
                }),
                TreeNode::Two(Branch {
                    keys: rk,
                    values: rv,
                    ..
                }),
            ) => {

                let mut new_keys = vec![];
                let mut new_values: VecDeque<_> = VecDeque::with_capacity(4);
                new_keys.append(lk);
                if let Some(idx)  = new_keys.iter().position(|k| k > &rk[0]) {
                    new_keys.insert(idx, rk.pop().unwrap());
                    todo!()
                } else {
                    new_keys.push(rk.pop().unwrap());
                    todo!()
                }


                TreeNode::Four(Branch {
                    keys: new_keys,
                    values: VecDeque::from(new_values)
                })
            }
            _ => todo!(),
        }
    }

    fn insert(&mut self, v: V) -> Option<Self> {
        return match self {
            TreeNode::Empty => Some(TreeNode::singleton(v)),
            TreeNode::Leaf(ref mut values) => {
                if values.len() == 3 {
                    let mut new_node = self.split_four_nodes();
                    new_node.insert(v);
                    return Some(new_node);
                }

                if let Some((idx, _)) = values.iter().find_position(|k| *k > &v) {
                    values.insert(idx, v)
                } else {
                    values.push(v)
                }

                None
            }
            TreeNode::Two(b) => {
                let mut target = if &b.keys[0] > &v {
                    &mut b.values[0]
                } else {
                    &mut b.values[1]
                };

                if let Some(new_root) = target.insert(v) {
                    return Some(TreeNode::merge(self, new_root));
                }

                None
            }
            TreeNode::Three(b) => {
                let mut target = if &b.keys[0] > &v {
                    &mut b.values[0]
                } else if &b.keys[1] > &v {
                    &mut b.values[1]
                } else if &b.keys[1] < &v {
                    &mut b.values[2]
                } else {
                    return None;
                };

                if let Some(new_root) = target.insert(v) {
                    return Some(TreeNode::merge(self, new_root));
                }

                None
            }

            _ => todo!(),
        };
    }
}

impl<T, V, const K: usize, const K1: usize> Branch<T, V, K, K1> {
    const NK: usize = K;
    const NV: usize = K1;

    fn empty() -> Self {
        Self {
            keys: unsafe { MaybeUninit::zeroed().assume_init() },
            values: unsafe { MaybeUninit::zeroed().assume_init() },
        }
    }

    fn new(ks: Vec<T>, vs: Vec<V>) -> Self {
        assert!(vs.len() == 0 || vs.len() == ks.len() + 1);
        Self {
            keys: ks,
            values: VecDeque::from(vs),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_model() {
        use TreeNode::{self, *};

        let mut tree = Tree::<i32> {
            root: TreeNode::singleton(1),
            size: 1,
        };

        tree.root = Two(Branch::new(
            vec![2],
            vec![
                Box::new(TreeNode::singleton(1)),
                Box::new(TreeNode::singleton(3)),
            ],
        ));

        eprintln!("tree.root.0 = {:?}", tree.root);
        // assert_eq!(
        //     tree,
        //     Tree::<i32> {
        //         root: Two(Branch::empty()),
        //     }
        // );

        let root_keys = tree.root.keys();
        eprintln!("root_keys = {:?}", root_keys);

        let mut tree1 = Tree::empty();

        tree1.insert(10);
        tree1.insert(30);
        tree1.insert(60);
        tree1.insert(20);
        tree1.insert(50);
        tree1.insert(40);
        eprintln!("tree1 = {:?}", tree1);
        tree1.insert(70);
        eprintln!("tree1 = {:?}", tree1);
        tree1.insert(80);
        tree1.insert(90);
        eprintln!("tree = {:?}", tree1);
    }
}