1
\$\begingroup\$

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
            })                             
        }
    }
}
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.