view src/lib.rs @ 1:04b6ed18abf3

Better splitting
author Lewin Bormann <lbo@spheniscida.de>
date Mon, 04 Jul 2022 18:14:18 -0700
parents e342ded1647e
children df21e58157f3
line wrap: on
line source

use ptree::TreeBuilder;
use std::cmp::Ordering;

const ORDER: usize = 5;

#[derive(Debug)]
struct Node<T> {
    entries: [Option<T>; ORDER - 1],
    children: [Option<usize>; ORDER],
    parent: Option<usize>,
}

impl<T: Ord> Node<T> {
    fn new(parent: Option<usize>) -> Node<T> {
        Node {
            entries: Default::default(),
            children: Default::default(),
            parent: parent,
        }
    }
}

#[derive(Debug)]
struct BTree<T> {
    v: Vec<Node<T>>,
    root: usize,
}

impl<T: std::fmt::Debug + Ord> BTree<T> {
    fn new(initial_capacity: usize) -> BTree<T> {
        BTree {
            v: Vec::with_capacity(initial_capacity),
            root: 0,
        }
    }

    fn free(&self, node: usize) -> usize {
        self.v[node]
            .entries
            .iter()
            .map(|e| if e.is_some() { 0 } else { 1 })
            .sum()
    }
    fn get(&self, node: usize, elem: usize) -> &Option<T> {
        &self.v[node].entries[elem]
    }
    fn get_mut(&mut self, node: usize, elem: usize) -> &mut Option<T> {
        &mut self.v[node].entries[elem]
    }
    fn get_child(&self, node: usize, child: usize) -> &Option<usize> {
        &self.v[node].children[child]
    }
    fn get_child_mut(&mut self, node: usize, child: usize) -> &mut Option<usize> {
        &mut self.v[node].children[child]
    }

    fn insert_inline(&mut self, node: usize, value: T) {
        assert!(self.free(node) >= 1);
        let node = &mut self.v[node];

        let mut loc = ORDER - 2;

        for i in 0..ORDER - 1 {
            if let Some(ref el) = node.entries[i] {
                if value.cmp(el) == Ordering::Less {
                    // Found right location. Shift remaining elements.
                    loc = i;
                    break;
                }
            } else {
                loc = i;
                break;
            }
        }

        for i in 0..(ORDER - 2 - loc) {
            node.entries[ORDER - 2 - i] = node.entries[ORDER - 3 - i].take();
            node.children[ORDER - 2 - i] = node.children[ORDER - 3 - i].take();
        }
        node.entries[loc] = Some(value);
    }
    fn insert(&mut self, value: T) -> bool {
        // 1. Find appropriate node
        // 2. Find appropriate order
        // 3. Insert or split.

        // Initialize empty tree
        if self.v.is_empty() {
            self.v.push(Node::new(None));
            self.v[0].entries[0] = Some(value);
            return true;
        }

        let mut node = self.root;
        'CHILDREN: loop {
            // 1. Descend as long as there are feasible children.
            // 2. Insert or split once no more children to be found.
            let mut found = false;
            let mut i = 0;
            'ELEMENTS: while i < ORDER - 1 {
                if let Some(el) = self.get(node, i) {
                    if value.cmp(el) == Ordering::Less {
                        found = true;
                        break 'ELEMENTS;
                    }
                } else {
                    // Past last filled element.
                    found = true;
                    break 'ELEMENTS;
                }
                i += 1
            }
            // There is an element greater than `value`?
            if let Some(child) = self.get_child(node, i) {
                // Descend into child node.
                node = *child;
                continue 'CHILDREN;
            } else if self.free(node) >= 1 {
                self.insert_inline(node, value);
                return true;
            } else {
                // We have to split.
                //self.split(node, Some(value), i);
                let new = self.split(node, None, 0);
                node = new;
                continue;
            }
        }
    }

    /// Split node `node` at `loc`, moving all values at `loc` and higher to the right node,
    /// inserting `value` into a new parent node.
    fn split(&mut self, node: usize, value: Option<T>, loc: usize) -> usize {
        let leftix = node;
        let rightix = self.v.len();
        let topix = self.v.len() + 1;

        let mut right = Node::new(Some(topix));
        let mut top = Node::new(self.v[node].parent);

        // TODO: We create a top node. In reality, the pivot value is to be inserted into the
        // existing parent node.
        if value.is_some() {
            // Put pivot entry into top node.
            top.entries[0] = value;

            // Copy greater elements to right node.
            for i in loc..ORDER - 1 {
                right.entries[i - loc] = self.get_mut(node, i).take();
            }

            // left
            top.children[0] = Some(leftix);
            // right
            top.children[1] = Some(rightix);
        } else {
            // Split without inserting value.
            let pivotix = (ORDER-1) / 2;
            let pivot = self.get_mut(node, pivotix).take();

            top.entries[0] = pivot;
            // Copy greater elements to right node.
            for i in pivotix..ORDER - 2 {
                right.entries[i - pivotix] = self.get_mut(node, i+1).take();
                println!("Inserting {:?} at {}", right.entries[i-pivotix], i-pivotix);
            }

            top.children[0] = Some(leftix);
            top.children[1] = Some(rightix);
        }
        // Update parent.
        {
            if let Some(parentix) = self.v[leftix].parent {
                let parent = &mut self.v[parentix];
                for i in 0..ORDER {
                    if parent.children[i] == Some(node) {
                        parent.children[i] = Some(topix);
                        break;
                    }
                }
            } else {
                self.root = topix;
            }
        }
        self.v[leftix].parent = Some(topix);
        right.parent = Some(topix);

        self.v.push(right);
        self.v.push(top);

        topix
    }
}
impl<T: std::fmt::Debug> BTree<T> {
    fn debug_print(&self) -> String {
        let mut t = TreeBuilder::new("btree".to_string());
        self.print_subtree(&mut t, self.root);
        let mut o: Vec<u8> = Vec::new();
        ptree::write_tree(&t.build(), &mut o);
        String::from_utf8(o).unwrap()
    }
    fn print_subtree(&self, tb: &mut TreeBuilder, node: usize) {
        let node = &self.v[node];
        for i in 0..ORDER - 1 {
            if let Some(ch) = node.children[i] {
                self.print_subtree(tb.begin_child(format!("<{}>", ch)), ch);
                tb.end_child();
            }
            tb.add_empty_child(format!("{:?}", node.entries[i]));
        }
    }
}

#[cfg(test)]
mod tests {
    use super::BTree;
    #[test]
    fn basic_test() {
        let mut bt = BTree::<i32>::new(16);

        let ns = &[2,4,6,8,10,12,14,1,16,18,20,7,22,24,26,28,30];

        for n in ns {
            bt.insert(*n);
            println!("Inserting {}", n);
            println!("root: {}", bt.root);
            println!("{}", bt.debug_print());
        }
            for (i, e) in bt.v.iter().enumerate() {
                println!("{} {:?}", i, e);
            }
    }
}