extended-collections 0.2.0

An extension to the collections in the standard library with various data structures.
Documentation
use std::cmp::Ordering;
use std::mem;
use treap::node::ImplicitNode;

pub type Tree<T> = Option<Box<ImplicitNode<T>>>;

pub fn merge<T>(l_tree: &mut Tree<T>, r_tree: Tree<T>) {
    match (l_tree.take(), r_tree) {
        (Some(mut l_node), Some(mut r_node)) => {
            if l_node.priority > r_node.priority {
                merge(&mut l_node.right, Some(r_node));
                l_node.update();
                *l_tree = Some(l_node);
            } else {
                let mut new_tree = Some(l_node);
                merge(&mut new_tree, r_node.left.take());
                r_node.left = new_tree;
                r_node.update();
                *l_tree = Some(r_node);
            }
        },
        (new_tree, None) | (None, new_tree) => *l_tree = new_tree,
    }
}

pub fn split<T>(tree: &mut Tree<T>, index: usize, left_inclusive: bool) -> Tree<T> {
    match tree.take() {
        Some(mut node) => {
            let key = node.get_implicit_key();
            let cmp = index.cmp(&key);
            let ret;
            if cmp == Ordering::Less || (cmp == Ordering::Equal && left_inclusive) {
                let mut res = split(&mut node.left, index, left_inclusive);
                *tree = node.left.take();
                node.left = res;
                node.update();
                ret = Some(node);
            } else {
                ret = split(&mut node.right, index - key, left_inclusive);
                node.update();
                *tree = Some(node);
            }
            ret
        },
        None => None,
    }
}

pub fn insert<T>(tree: &mut Tree<T>, index: usize, new_node: ImplicitNode<T>) {
    assert!(1 <= index && index <= len(tree) + 1);
    let right = split(tree, index, true);
    merge(tree, Some(Box::new(new_node)));
    merge(tree, right);
}

pub fn remove<T>(tree: &mut Tree<T>, index: usize) -> T {
    assert!(1 <= index && index <= len(tree));
    let new_tree = {
        let node = tree.as_mut().expect("Expected non-empty tree.");
        let key = node.get_implicit_key();
        match index.cmp(&key) {
            Ordering::Less => {
                let ret = remove(&mut node.left, index);
                node.update();
                return ret;
            },
            Ordering::Greater => {
                let ret = remove(&mut node.right, index - key);
                node.update();
                return ret;
            },
            Ordering::Equal => {
                let ImplicitNode {
                    ref mut left,
                    ref mut right,
                    ..
                } = &mut **node;
                merge(left, right.take());
                left.take()
            },
        }
    };

    mem::replace(tree, new_tree)
        .expect("Expected non-empty tree.")
        .value
}

pub fn get<T>(tree: &Tree<T>, index: usize) -> Option<&T> {
    tree.as_ref().and_then(|node| {
        let key = node.get_implicit_key();
        match index.cmp(&key) {
            Ordering::Less => get(&node.left, index),
            Ordering::Greater => get(&node.right, index - key),
            Ordering::Equal => Some(&node.value),
        }
    })
}

pub fn get_mut<T>(tree: &mut Tree<T>, index: usize) -> Option<&mut T> {
    tree.as_mut().and_then(|node| {
        let key = node.get_implicit_key();
        match index.cmp(&key) {
            Ordering::Less => get_mut(&mut node.left, index),
            Ordering::Greater => get_mut(&mut node.right, index - key),
            Ordering::Equal => Some(&mut node.value),
        }
    })
}

pub fn len<T>(tree: &Tree<T>) -> usize {
    if let Some(ref node) = tree {
        node.len()
    } else {
        0
    }
}