Реализация Rust Persistent Red Black Tree

Ниже представлена ​​моя реализация постоянного красно-черного дерева в Rust.

У меня есть несколько вопросов о возможных улучшениях. В настоящее время данные и узлы хранятся в указателях со счетчиком, на которые имеются ссылки. Это лучший способ сделать это?

Аналогичным образом, операторы сопоставления с образцом для функций баланса довольно подробны из-за использования rc. Есть ли способ написать это более кратко?

Кроме того, поскольку я использую указатели rc, следует ли мне реализовать трейт drop?

#[allow(clippy::module_inception)]
pub mod red_black_tree {
    use std::cmp::Ordering;
    use std::rc::Rc;

    pub enum RedBlackTree<T>
    where
        T: Ord,
    {
        Node {
            color: Color,
            data: Rc<T>,
            left: Rc<RedBlackTree<T>>,
            right: Rc<RedBlackTree<T>>,
        },
        Leaf,
    }

    #[derive(Clone)]
    pub enum Color {
        Red,
        Black,
    }

    impl<T> RedBlackTree<T>
    where
        T: Ord,
    {
        pub fn new() -> Self {
            RedBlackTree::Leaf
        }

        pub fn contains(&self, item: T) -> bool {
            match self {
                RedBlackTree::Node {
                    color: _,
                    data,
                    left,
                    right,
                } => match item.cmp(&data) {
                    Ordering::Less => left.contains(item),
                    Ordering::Equal => true,
                    Ordering::Greater => right.contains(item),
                },
                RedBlackTree::Leaf => false,
            }
        }

        pub fn insert(&self, item: T) -> Self {
            match self {
                RedBlackTree::Node {
                    color,
                    data,
                    left,
                    right,
                } => match item.cmp(&data) {
                    Ordering::Less => RedBlackTree::Node {
                        color: color.clone(),
                        data: Rc::clone(data),
                        left: Rc::new(left.insert(item)),
                        right: Rc::clone(right),
                    }
                    .balance()
                    .make_black(),
                    Ordering::Equal => RedBlackTree::Node {
                        color: color.clone(),
                        data: Rc::clone(data),
                        left: Rc::clone(left),
                        right: Rc::clone(right),
                    }
                    .make_black(),
                    Ordering::Greater => RedBlackTree::Node {
                        color: color.clone(),
                        data: Rc::clone(data),
                        left: Rc::clone(left),
                        right: Rc::new(right.insert(item)),
                    }
                    .balance()
                    .make_black(),
                },
                RedBlackTree::Leaf => RedBlackTree::Node {
                    color: Color::Black,
                    data: Rc::new(item),
                    left: Rc::new(RedBlackTree::new()),
                    right: Rc::new(RedBlackTree::new()),
                },
            }
        }

        pub fn get(&self, item: &T) -> Option<Rc<T>> {
            match self {
                RedBlackTree::Node {
                    color: _,
                    data,
                    left,
                    right,
                } => match item.cmp(&data) {
                    Ordering::Less => left.get(item),
                    Ordering::Equal => Option::from(Rc::clone(data)),
                    Ordering::Greater => right.get(item),
                },
                RedBlackTree::Leaf => Option::None,
            }
        }

        fn make_black(&self) -> Self {
            match self {
                RedBlackTree::Node {
                    color: _,
                    data,
                    left,
                    right,
                } => RedBlackTree::Node {
                    color: Color::Black,
                    data: Rc::clone(data),
                    left: Rc::clone(left),
                    right: Rc::clone(right),
                },
                RedBlackTree::Leaf => RedBlackTree::new(),
            }
        }

        fn balance(self) -> Self {
            if let RedBlackTree::Node {
                color: Color::Black,
                data: parent_data,
                left: parent_left,
                right: parent_right,
            } = &self
            {
                if let RedBlackTree::Node {
                    color: Color::Red,
                    data: child_data,
                    left: child_left,
                    right: child_right,
                } = Rc::as_ref(&parent_left)
                {
                    if let RedBlackTree::Node {
                        color: Color::Red,
                        data: grandchild_data,
                        left: grandchild_left,
                        right: grandchild_right,
                    } = Rc::as_ref(&child_left)
                    {
                        return RedBlackTree::from(
                            grandchild_left,
                            grandchild_right,
                            child_right,
                            parent_right,
                            grandchild_data,
                            child_data,
                            parent_data,
                        );
                    } else if let RedBlackTree::Node {
                        color: Color::Red,
                        data: grandchild_data,
                        left: grandchild_left,
                        right: grandchild_right,
                    } = Rc::as_ref(&child_right)
                    {
                        return RedBlackTree::from(
                            child_left,
                            grandchild_left,
                            grandchild_right,
                            parent_right,
                            child_data,
                            grandchild_data,
                            parent_data,
                        );
                    }
                } else if let RedBlackTree::Node {
                    color: Color::Red,
                    data: child_data,
                    left: child_left,
                    right: child_right,
                } = Rc::as_ref(&parent_right)
                {
                    if let RedBlackTree::Node {
                        color: Color::Red,
                        data: grandchild_data,
                        left: grandchild_left,
                        right: grandchild_right,
                    } = Rc::as_ref(&child_left)
                    {
                        return RedBlackTree::from(
                            parent_left,
                            grandchild_left,
                            grandchild_right,
                            child_right,
                            parent_data,
                            grandchild_data,
                            child_data,
                        );
                    } else if let RedBlackTree::Node {
                        color: Color::Red,
                        data: grandchild_data,
                        left: grandchild_left,
                        right: grandchild_right,
                    } = Rc::as_ref(&child_right)
                    {
                        return RedBlackTree::from(
                            parent_left,
                            child_left,
                            grandchild_left,
                            grandchild_right,
                            parent_data,
                            child_data,
                            grandchild_data,
                        );
                    }
                }
            }
            self.clone()
        }

        #[allow(clippy::many_single_char_names)]
        fn from(
            a: &Rc<RedBlackTree<T>>,
            b: &Rc<RedBlackTree<T>>,
            c: &Rc<RedBlackTree<T>>,
            d: &Rc<RedBlackTree<T>>,
            x: &Rc<T>,
            y: &Rc<T>,
            z: &Rc<T>,
        ) -> RedBlackTree<T> {
            RedBlackTree::Node {
                color: Color::Red,
                data: Rc::clone(y),
                left: Rc::new(RedBlackTree::Node {
                    color: Color::Black,
                    data: Rc::clone(x),
                    left: Rc::clone(a),
                    right: Rc::clone(b),
                }),
                right: Rc::new(RedBlackTree::Node {
                    color: Color::Black,
                    data: Rc::clone(z),
                    left: Rc::clone(c),
                    right: Rc::clone(d),
                }),
            }
        }
    }

    impl<T> Clone for RedBlackTree<T>
    where
        T: Ord,
    {
        fn clone(&self) -> Self {
            match self {
                RedBlackTree::Node {
                    color,
                    data,
                    left,
                    right,
                } => RedBlackTree::Node {
                    color: color.clone(),
                    data: Rc::clone(data),
                    left: Rc::clone(left),
                    right: Rc::clone(right),
                },
                RedBlackTree::Leaf => RedBlackTree::new(),
            }
        }
    }

    impl<T> Default for RedBlackTree<T>
    where
        T: Ord,
    {
        fn default() -> Self {
            RedBlackTree::<T>::new()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::red_black_tree::*;
    use std::cmp::*;

    struct Point {
        x: i64,
        y: i64,
    }

    impl Point {
        fn new(x: i64, y: i64) -> Point {
            Point { x, y }
        }
        fn magnitude_squared(&self) -> u64 {
            (self.x as u64).pow(2) + (self.y as u64).pow(2)
        }
    }

    impl PartialEq for Point {
        fn eq(&self, other: &Self) -> bool {
            self.magnitude_squared() == other.magnitude_squared()
        }
    }

    impl Eq for Point {}

    impl PartialOrd for Point {
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
            Some(self.cmp(other))
        }
    }

    impl Ord for Point {
        fn cmp(&self, other: &Self) -> Ordering {
            self.magnitude_squared().cmp(&other.magnitude_squared())
        }
    }

    #[test]
    fn test_empty() {
        let tree = RedBlackTree::<i64>::new();
        assert!(!tree.contains(0));
        assert!(!tree.contains(5));
        assert!(!tree.contains(-20));
    }

    #[test]
    fn test_insert() {
        let mut tree: RedBlackTree<char> = RedBlackTree::new();

        assert!(!tree.contains('a'));
        assert!(!tree.contains('b'));
        assert!(!tree.contains('c'));

        tree = tree.insert('a');

        assert!(tree.contains('a'));
        assert!(!tree.contains('b'));
        assert!(!tree.contains('c'));

        tree = tree.insert('b');

        assert!(tree.contains('a'));
        assert!(tree.contains('b'));
        assert!(!tree.contains('c'));
    }

    #[test]
    fn test_get() {
        let mut tree: RedBlackTree<Point> = RedBlackTree::new();

        tree = tree.insert(Point::new(0, 0));
        tree = tree.insert(Point::new(1, 1));
        tree = tree.insert(Point::new(2, 2));
        tree = tree.insert(Point::new(3, 4));

        assert_eq!(tree.get(&Point::new(0, 0)).unwrap().x, 0);
        assert_eq!(tree.get(&Point::new(0, 0)).unwrap().y, 0);

        assert_eq!(tree.get(&Point::new(0, 5)).unwrap().x, 3);
        assert_eq!(tree.get(&Point::new(0, 5)).unwrap().y, 4);
    }
}


Изменить: функция ребалансировки основана на этой диаграмме (источник):

введите описание изображения здесь

0

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *