splay.rs

use std::ptr;
use std::mem;

struct Node<T> {
    index: usize,
    parent: *mut Node<T>,
    children: [*mut Node<T>; 2],
    value: T,
}

impl<T> Node<T> {
    fn new(value: T) -> Node<T> {
        Node {
            index: !0,
            parent: ptr::null_mut(),
            children: [ptr::null_mut(), ptr::null_mut()],
            value: value,
        }
    }
}

struct Tree<T> {
    node: *mut Node<T>,
}

impl<T> Tree<T> {
    fn new(value: T) -> Tree<T> {
        Tree { node: Box::into_raw(Box::new(Node::new(value))) }
    }

    fn parent(&mut self) -> Option<usize> {
        let parent = unsafe { (*self.node).parent };
        if parent == ptr::null_mut() {
            None
        } else {
            let index = unsafe { (*self.node).index };
            self.node = parent;
            Some(index)
        }
    }

    fn child(&mut self, index: usize) -> bool {
        let child = unsafe { *(*self.node).children.get_unchecked(index) };
        if child == ptr::null_mut() {
            false
        } else {
            self.node = child;
            true
        }
    }

    fn link(&mut self, index: usize, mut tree: Tree<T>) {
        if unsafe { *(*self.node).children.get_unchecked(index) } != ptr::null_mut() {
            //panic!("Cannot link a tree because another node is linked at index {}", index);
            return;
        }
        while tree.parent().is_some() {}
        unsafe {
            (*tree.node).index = index;
            (*tree.node).parent = self.node;
            *(*self.node).children.get_unchecked_mut(index) = tree.node;
        }
        mem::forget(tree);
    }

    fn take(&mut self, index: usize) -> Option<Tree<T>> {
        let child = unsafe { *(*self.node).children.get_unchecked(index) };
        if child == ptr::null_mut() {
            None
        } else {
            unsafe {
                (*child).parent = ptr::null_mut();
                *(*self.node).children.get_unchecked_mut(index) = ptr::null_mut();
            }
            Some(Tree { node: child })
        }
    }

    fn value(&self) -> &mut T {
        unsafe { &mut (*self.node).value }
    }
}

impl<T> Drop for Tree<T> {
    fn drop(&mut self) {
        while self.parent().is_some() {}
        unsafe {
            for &mut child in (*self.node).children.iter_mut() {
                if child != ptr::null_mut() {
                    (*child).parent = ptr::null_mut();
                    Tree { node: child };
                }
            }
            Box::from_raw(self.node);
        }
    }
}

struct Splay<T> {
    tree: Option<Tree<T>>,
}

impl<T> Splay<T> where T: PartialOrd {
    fn new() -> Splay<T> {
        Splay { tree: None }
    }

    fn rotate(&self, tree: &mut Tree<T>, index: usize) {
        let mut child = match tree.take(index) {
            None => return,
            Some(x) => x,
        };
        if let Some(grandchild) = child.take(index ^ 1) {
            tree.link(index, grandchild);
        }

        if let Some(i) = tree.parent() {
            let t = match tree.take(i) {
                Some(x) => x,
                None => return,
            };
            tree.link(i, child);
            tree.child(i);
            tree.link(index ^ 1, t);
        } else {
            mem::swap(tree, &mut child);
            tree.link(index ^ 1, child);
        }
    }

    fn splay(&self, mut tree: &mut Tree<T>) {
        loop {
            let i1 = match tree.parent() {
                None => break,
                Some(index) => index,
            };
            let i2 = match tree.parent() {
                None => {
                    self.rotate(tree, i1);
                    break;
                },
                Some(index) => index,
            };

            if i1 == i2 {
                self.rotate(tree, i2);
                self.rotate(tree, i1);
            } else {
                tree.child(i2);
                self.rotate(tree, i1);
                tree.parent();
                self.rotate(tree, i2);
            }
        }
    }

    fn insert(&mut self, value: T) {
        match self.tree.take() {
            None => self.tree = Some(Tree::new(value)),
            Some(mut tree) => {
                loop {
                    let index = if value < *tree.value() { 0 } else { 1 };

                    if !tree.child(index) {
                        tree.link(index, Tree::new(value));
                        tree.child(index);
                        self.splay(&mut tree);
                        self.tree = Some(tree);
                        break;
                    }
                }
            }
        }
    }
}

fn print<T>(tree: &mut Tree<T>, depth: usize) where T: std::fmt::Display {
    if tree.child(0) {
        print(tree, depth + 1);
        tree.parent();
    }
    println!("{:width$}{}", "", tree.value(), width = depth * 3);
    if tree.child(1) {
        print(tree, depth + 1);
        tree.parent();
    }
}

fn main() {
    let mut splay = Splay::new();

    let mut num = 1u32;
    for _ in 0..1000000 {
        num = num.wrapping_mul(17).wrapping_add(255);
        splay.insert(num);
    }

    // print(&mut splay.tree.unwrap(), 0);
}