I'm mainly looking for tips in how make my code more idiomatic and elegant, but any general reviews are welcome.
pub mod set {
use ::std::cmp;
use ::std::cmp::Ordering::{Less, Greater, Equal};
use ::std::mem;
pub struct AvlTree<T : Ord> {
root : Option<Box<Node<T>>>
}
struct Node<T : Ord> {
key : T,
height : i32,
left : AvlTree<T>,
right : AvlTree<T>,
}
impl<T> Node<T> where T : Ord{
fn new(key: T) -> Option<Box<Node<T>>>{
Some(Box::new(
Node {
key : key,
height : 0,
left : AvlTree::new(),
right : AvlTree::new()
}
))
}
}
impl<T> AvlTree<T> where T : Ord{
pub fn new() -> AvlTree<T> {
AvlTree { root: None }
}
fn rotate<CA, CB>(&mut self, ch_a: &CA, ch_b: &CB)
where CA: Fn(&mut Box<Node<T>>) -> &mut AvlTree<T>,
CB: Fn(&mut Box<Node<T>>) -> &mut AvlTree<T> {
let mut root = self.node_take();
let mut child = ch_a(&mut root).node_take();
ch_a(&mut root).root = ch_b(&mut child).root.take();
ch_b(&mut child).root = Some(root);
ch_b(&mut child).update_height();
self.root = Some(child);
self.update_height();
}
fn double_rotate<CA, CB>(&mut self, ch_a: &CA, ch_b: &CB)
where CA: Fn(&mut Box<Node<T>>) -> &mut AvlTree<T>,
CB: Fn(&mut Box<Node<T>>) -> &mut AvlTree<T> {
{
let child = ch_b(self.m_box());
if ch_b(child.m_box()).h() < ch_a(child.m_box()).h() {
child.rotate(ch_a, ch_b);
}
}
self.rotate(ch_b, ch_a);
}
fn restructure(& mut self) {
self.update_height();
let ord = self.root.as_ref()
.map(|_| self.degree())
.map(|d| (d, d.cmp(&0)));
match ord {
Some((d, Less)) if -1 > d => {
self.double_rotate(
&|x| &mut x.left,
&|x| &mut x.right)
},
Some((d, Greater)) if d > 1 => {
self.double_rotate(
&|x| &mut x.right,
&|x| &mut x.left)
},
_ => { },
}
}
fn update_height(&mut self) {
self.root.as_mut()
.map(|r| r.height = cmp::max(r.right.h(), r.left.h()) + 1);
}
fn m_box(&mut self) -> &mut Box<Node<T>> {
self.root.as_mut().expect("root can't be None")
}
fn r_box(&self) -> & Box<Node<T>> {
self.root.as_ref().expect("root can't be None")
}
fn node_take(&mut self) -> Box<Node<T>> {
self.root.take().expect("root can't be None")
}
fn h(&self) -> i32 {
self.root.as_ref().map_or(-1, |x| x.height)
}
fn degree(&self) -> i32 {
self.root
.as_ref()
.map(|x| x.left.h() - x.right.h())
.unwrap_or(0)
}
pub fn insert(&mut self, key : T) {
match self.map_to_ord(&key) {
Some(Less) => self.m_box().left.insert(key),
Some(Greater) => self.m_box().right.insert(key),
None => self.root = Node::new(key),
_ => return,
}
self.restructure();
}
pub fn remove(&mut self, key: &T) {
match self.map_to_ord(key) {
Some(Less) => self.m_box().left.remove(key),
Some(Greater) => self.m_box().right.remove(key),
Some(Equal) => self.remove_self(key),
None => return,
};
self.restructure();
}
fn min(&mut self) -> &mut Box<Node<T>> {
let n = self.m_box();
if n.left.root.is_none(){
return n;
}
n.left.min()
}
fn map_to_ord(&self, key: &T) -> Option<cmp::Ordering> {
self.root.as_ref()
.map(|p| &p.key)
.map(|k| key.cmp(k))
}
pub fn contains(& self, key: &T) -> bool {
match self.map_to_ord(key) {
Some(Less) => self.r_box().left.contains(key),
Some(Greater) => self.r_box().right.contains(key),
Some(Equal) => true,
None => false,
}
}
fn remove_self(&mut self, key: &T) {
let mut c = self.node_take();
self.root = match (c.left.root.take(),
c.right.root.take()) {
(None, right) => right,
(left, None) => left,
(l@ Some(_), r@ Some(_)) => {
let mut right = AvlTree{root: r};
mem::swap(
&mut right.min().key,
&mut c.key);
c.right = right;
c.left.root = l;
c.right.remove(key);
Some(c)
}
};
}
pub fn iter<'a>(&'a self) -> TreeIterator<'a, T> {
TreeIterator {
current : self.root.as_ref(),
path: Vec::new(),
}
}
}
pub struct TreeIterator<'a, T: 'a> where T: Ord {
current : Option<&'a Box<Node<T>>>,
path: Vec<&'a Box<Node<T>>>,
}
impl<'a, T> IntoIterator for & 'a AvlTree<T> where T: Ord {
type Item = &'a T;
type IntoIter = TreeIterator<'a, T>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'a, T> Iterator for TreeIterator<'a, T> where T: Ord {
type Item = &'a T;
fn next<'b>(&'b mut self) -> Option<Self::Item> {
let mut c = self.current;
while let Some(n) = c {
self.path.push(n);
c = n.left.root.as_ref();
}
self.path.pop().map(|n|{
self.current = n.right.root.as_ref();
&n.key
})
}
}
}