extended-collections 0.2.0

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

pub type Tree<T, U> = Option<Box<Node<T, U>>>;

pub fn merge<T, U>(l_tree: &mut Tree<T, U>, r_tree: Tree<T, U>)
where
    T: Ord,
{
    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, U>(tree: &mut Tree<T, U>, key: &T) -> (Tree<T, U>, Tree<T, U>)
where
    T: Ord,
{
    match tree.take() {
        Some(mut node) => {
            let mut ret;
            match key.cmp(&node.entry.key) {
                Ordering::Less => {
                    let mut res = split(&mut node.left, key);
                    *tree = node.left.take();
                    node.left = res.1;
                    node.update();
                    ret = (res.0, Some(node));
                },
                Ordering::Greater => {
                    ret = split(&mut node.right, key);
                    node.update();
                    *tree = Some(node);
                },
                Ordering::Equal => {
                    *tree = node.left.take();
                    let right = node.right.take();
                    node.update();
                    ret = (Some(node), right);
                },
            }
            ret
        },
        None => (None, None),
    }
}

pub fn insert<T, U>(tree: &mut Tree<T, U>, mut new_node: Node<T, U>) -> Option<Entry<T, U>>
where
    T: Ord,
{
    match tree {
        Some(ref mut node) => {
            if new_node.priority <= node.priority {
                match new_node.entry.key.cmp(&node.entry.key) {
                    Ordering::Less => {
                        let ret = insert(&mut node.left, new_node);
                        node.update();
                        return ret;
                    },
                    Ordering::Greater => {
                        let ret = insert(&mut node.right, new_node);
                        node.update();
                        return ret;
                    },
                    Ordering::Equal => {
                        let Node { ref mut entry, .. } = &mut **node;
                        return Some(mem::replace(entry, new_node.entry));
                    },
                }
            }
        },
        None => {
            *tree = Some(Box::new(new_node));
            return None;
        },
    }
    new_node.left = tree.take();
    let (dup_opt, right) = split(&mut new_node.left, &new_node.entry.key);
    new_node.right = right;
    new_node.update();
    *tree = Some(Box::new(new_node));
    dup_opt.map(|node| node.entry)
}

pub fn remove<T, U>(tree: &mut Tree<T, U>, key: &T) -> Option<Entry<T, U>>
where
    T: Ord,
{
    let mut new_tree;
    match tree {
        Some(ref mut node) => {
            match key.cmp(&node.entry.key) {
                Ordering::Less => {
                    let ret = remove(&mut node.left, key);
                    node.update();
                    return ret;
                },
                Ordering::Greater => {
                    let ret = remove(&mut node.right, key);
                    node.update();
                    return ret;
                },
                Ordering::Equal => {
                    new_tree = node.left.take();
                    merge(&mut new_tree, node.right.take());
                },
            }
        },
        None => return None,
    }
    mem::replace(tree, new_tree).map(|node| node.entry)
}

pub fn get<'a, T, U>(tree: &'a Tree<T, U>, key: &T) -> Option<&'a Entry<T, U>>
where
    T: Ord,
{
    tree.as_ref().and_then(|node| {
        match key.cmp(&node.entry.key) {
            Ordering::Less => get(&node.left, key),
            Ordering::Greater => get(&node.right, key),
            Ordering::Equal => Some(&node.entry),
        }
    })
}

pub fn get_mut<'a, T, U>(tree: &'a mut Tree<T, U>, key: &T) -> Option<&'a mut Entry<T, U>>
where
    T: Ord,
{
    tree.as_mut().and_then(|node| {
        match key.cmp(&node.entry.key) {
            Ordering::Less => get_mut(&mut node.left, key),
            Ordering::Greater => get_mut(&mut node.right, key),
            Ordering::Equal => Some(&mut node.entry),
        }
    })
}

pub fn ceil<'a, T, U>(tree: &'a Tree<T, U>, key: &T) -> Option<&'a Entry<T, U>>
where
    T: Ord,
{
    tree.as_ref().and_then(|node| {
        match key.cmp(&node.entry.key) {
            Ordering::Greater => ceil(&node.right, key),
            Ordering::Less => {
                match ceil(&node.left, key) {
                    None => Some(&node.entry),
                    res => res,
                }
            },
            Ordering::Equal => Some(&node.entry),
        }
    })
}

pub fn floor<'a, T, U>(tree: &'a Tree<T, U>, key: &T) -> Option<&'a Entry<T, U>>
where
    T: Ord,
{
    tree.as_ref().and_then(|node| {
        match key.cmp(&node.entry.key) {
            Ordering::Less => floor(&node.left, key),
            Ordering::Greater => {
                match floor(&node.right, key) {
                    None => Some(&node.entry),
                    res => res,
                }
            },
            Ordering::Equal => Some(&node.entry),
        }
    })
}

pub fn min<T, U>(tree: &Tree<T, U>) -> Option<&Entry<T, U>>
where
    T: Ord,
{
    tree.as_ref().and_then(|node| {
        let mut curr = node;
        while let Some(ref left_node) = curr.left {
            curr = left_node;
        }
        Some(&curr.entry)
    })
}

pub fn max<T, U>(tree: &Tree<T, U>) -> Option<&Entry<T, U>>
where
    T: Ord,
{
    tree.as_ref().and_then(|node| {
        let mut curr = node;
        while let Some(ref right_node) = curr.right {
            curr = right_node;
        }
        Some(&curr.entry)
    })
}

pub fn union<T, U>(left_tree: Tree<T, U>, right_tree: Tree<T, U>, mut swapped: bool) -> Tree<T, U>
where
    T: Ord,
{
    match (left_tree, right_tree) {
        (Some(mut left_node), Some(mut right_node)) => {
            if left_node.priority < right_node.priority {
                mem::swap(&mut left_node, &mut right_node);
                swapped = !swapped;
            }
            {
                let Node {
                    left: ref mut left_subtree,
                    right: ref mut right_subtree,
                    ref mut entry,
                    ..
                } = &mut *left_node;
                let mut right_left_subtree = Some(right_node);
                let (duplicate_opt, right_right_subtree) = split(&mut right_left_subtree, &entry.key);
                let new_left_subtree = union(left_subtree.take(), right_left_subtree, swapped);
                let new_right_subtree = union(right_subtree.take(), right_right_subtree, swapped);
                *left_subtree = new_left_subtree;
                *right_subtree = new_right_subtree;
                if let Some(duplicate_node) = duplicate_opt {
                    if swapped {
                        *entry = duplicate_node.entry;
                    }
                }
            }
            left_node.update();
            Some(left_node)
        },
        (None, right_tree) => right_tree,
        (left_tree, None) => left_tree,
    }
}

pub fn intersection<T, U>(
    left_tree: Tree<T, U>,
    right_tree: Tree<T, U>,
    mut swapped: bool,
) -> Tree<T, U>
where
    T: Ord,
{
    match (left_tree, right_tree) {
        (Some(mut left_node), Some(mut right_node)) => {
            {
                if left_node.priority < right_node.priority {
                    mem::swap(&mut left_node, &mut right_node);
                    swapped = !swapped;
                }
                let Node {
                    left: ref mut left_subtree,
                    right: ref mut right_subtree,
                    ref mut entry,
                    ..
                } = &mut *left_node;
                let mut right_left_subtree = Some(right_node);
                let (duplicate_opt, right_right_subtree) = split(&mut right_left_subtree, &entry.key);
                let new_left_subtree = intersection(left_subtree.take(), right_left_subtree, swapped);
                let new_right_subtree = intersection(right_subtree.take(), right_right_subtree, swapped);
                *left_subtree = new_left_subtree;
                *right_subtree = new_right_subtree;
                match duplicate_opt {
                    Some(duplicate_node) => {
                        if swapped {
                            *entry = duplicate_node.entry;
                        }
                    },
                    None => {
                        merge(left_subtree, right_subtree.take());
                        return left_subtree.take();
                    },
                }
            }
            left_node.update();
            Some(left_node)
        },
        _ => None,
    }
}

pub fn difference<T, U>(
    left_tree: Tree<T, U>,
    right_tree: Tree<T, U>,
    mut swapped: bool,
    symmetric: bool,
) -> Tree<T, U>
where
    T: Ord,
{
    match (left_tree, right_tree) {
        (Some(mut left_node), Some(mut right_node)) => {
            {
                if left_node.priority < right_node.priority {
                    mem::swap(&mut left_node, &mut right_node);
                    swapped = !swapped;
                }
                let Node {
                    left: ref mut left_subtree,
                    right: ref mut right_subtree,
                    ref mut entry,
                    ..
                } = &mut *left_node;
                let mut right_left_subtree = Some(right_node);
                let (duplicate_opt, right_right_subtree) = split(&mut right_left_subtree, &entry.key);
                let new_left_subtree = difference(left_subtree.take(), right_left_subtree, swapped, symmetric);
                let new_right_subtree = difference(right_subtree.take(), right_right_subtree, swapped, symmetric);
                *left_subtree = new_left_subtree;
                *right_subtree = new_right_subtree;
                if duplicate_opt.is_some() || (swapped && !symmetric) {
                    merge(left_subtree, right_subtree.take());
                    return left_subtree.take();
                }
            }
            left_node.update();
            Some(left_node)
        },
        (mut left_tree, right_tree) => {
            if symmetric {
                merge(&mut left_tree, right_tree);
                left_tree
            } else if swapped {
                right_tree
            } else {
                left_tree
            }
        },
    }
}