dvcompute_utils 2.0.0

// Copyright (c) 2020-2022  David Sorokin <davsor@mail.ru>, based in Yoshkar-Ola, Russia
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

use std::rc::Rc;
use std::ops::Deref;
use std::cmp::Ordering;

/// An ordered map based on red-black trees, where the very basic idea is described in article
/// "Functional Pearls: Red-Black Trees in a Functional Setting" by Chris Okasaki.
#[derive(Debug, PartialEq, Eq)]
pub enum OrdMap<K, V> {

    /// The tree item.
    Tree(OrdMapColor, K, V, Rc<OrdMap<K, V>>, Rc<OrdMap<K, V>>),

    /// An empty map.
    Nil
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OrdMapColor {

    /// The red color.
    Red,

    /// The black color.
    Black
}

impl<K, V> Clone for OrdMap<K, V>
    where K: Clone,
          V: Clone
{
    fn clone(&self) -> Self {
        match self {
            &OrdMap::Tree(color, ref key, ref val, ref left, ref right) => {
                OrdMap::Tree(color, key.clone(), val.clone(), left.clone(), right.clone())
            },
            &OrdMap::Nil => {
                OrdMap::Nil
            }
        }
    }
}

impl<K, V> OrdMap<K, V>
    where K: Clone + Eq + Ord,
          V: Clone
{
    /// Test whether the map is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        match &self {
            &OrdMap::Tree(_, _, _, _, _) => false,
            &OrdMap::Nil => true
        }
    }

    /// Create an empty map.
    pub fn empty() -> Self {
        OrdMap::Nil
    }

    /// Collect items into vector.
    pub fn collect(&self) -> Vec<(K, V)> {
        let mut vs = Vec::new();
        let mut f = |key: &K, val: &V| {
            vs.push((key.clone(), val.clone()));
            None
        };
        let _: Option<isize> = self.find(&mut f);
        vs
    }

    // Test whether there is an item with the specified key.
    pub fn contains_key(&self, key: &K) -> bool {
        self.get(key).is_some()
    }

    // Get the item by its key.
    pub fn get(&self, key: &K) -> Option<&V> {
        let mut p = self;
        loop {
            match p {
                &OrdMap::Tree(_, ref key0, ref val0, ref left, ref right) => {
                    match key.cmp(key0) {
                        Ordering::Less => {
                            p = left.deref();
                        },
                        Ordering::Greater => {
                            p = right.deref();
                        },
                        Ordering::Equal => {
                            return Some(val0);
                        }
                    }
                },
                &OrdMap::Nil => {
                    return None;
                }
            }
        }
    }

    /// Get the first key-value pair.
    pub fn first_key_value(&self) -> Option<(&K, &V)> {
        let mut p = self;
        loop {
            match p {
                &OrdMap::Tree(_, ref key, ref val, ref left, _) if left.is_empty() => {
                    return Some((key, val));
                },
                &OrdMap::Tree(_, _, _, ref left, _) => {
                    p = left.deref();
                }
                &OrdMap::Nil => {
                    return None;
                }
            }
        }
    }

    /// Get the last key-value pair.
    pub fn last_key_value(&self) -> Option<(&K, &V)> {
        let mut p = self;
        loop {
            match p {
                &OrdMap::Tree(_, ref key, ref val, _, ref right) if right.is_empty() => {
                    return Some((key, val));
                },
                &OrdMap::Tree(_, _, _, _, ref right) => {
                    p = right.deref();
                }
                &OrdMap::Nil => {
                    return None;
                }
            }
        }
    }

    /// Try to find a new key-value pair satistying to the specified predicate.
    pub fn find<F, T>(&self, pred: &mut F) -> Option<T>
        where F: FnMut(&K, &V) -> Option<T>
    {
        match self {
            &OrdMap::Tree(_, ref key, ref val, ref left, ref right) => {
                match left.find(pred) {
                    z @ Some(_) => z,
                    None => {
                        match pred(key, val) {
                            z @ Some(_) => z,
                            None => right.find(pred)
                        }
                    }
                }
            },
            &OrdMap::Nil => {
                None
            }
        }
    }

    /// Insert a new key-value pair.
    pub fn insert(&self, key: K, val: V) -> OrdMap<K, V> {
        self.ins(key, val).make_black()
    }

    /// Insert a new key-value pair without post-processing.
    fn ins(&self, key: K, val: V) -> OrdMap<K, V> {
        match self {
            &OrdMap::Tree(color0, ref key0, ref val0, ref left, ref right) => {
                match key.cmp(key0) {
                    Ordering::Less => {
                        Self::balance(color0, key0, val0, &Rc::new(left.ins(key, val)), right)
                    },
                    Ordering::Greater => {
                        Self::balance(color0, key0, val0, left, &Rc::new(right.ins(key, val)))
                    },
                    Ordering::Equal => {
                        OrdMap::Tree(color0, key, val, left.clone(), right.clone())
                    }
                }
            },
            &OrdMap::Nil => {
                OrdMap::Tree(OrdMapColor::Red, key, val, Rc::new(OrdMap::Nil), Rc::new(OrdMap::Nil))
            }
        }
    }

    /// Make the root node black.
    fn make_black(&self) -> OrdMap<K, V> {
        match self {
            &OrdMap::Tree(_, ref key, ref val, ref left, ref right) => {
                OrdMap::Tree(OrdMapColor::Black, key.clone(), val.clone(), left.clone(), right.clone())
            },
            &OrdMap::Nil => {
                OrdMap::Nil
            }
        }

    }

    /// Balance a tree node.
    fn balance(color: OrdMapColor, key: &K, val: &V, left: &Rc<OrdMap<K, V>>, right: &Rc<OrdMap<K, V>>) -> OrdMap<K, V> {
        if color == OrdMapColor::Black {
            match left.deref() {
                &OrdMap::Tree(OrdMapColor::Red, ref key2, ref val2, ref left2, ref right2) => {
                    match left2.deref() {
                        &OrdMap::Tree(OrdMapColor::Red, ref key3, ref val3, ref left3, ref right3) => {
                            return OrdMap::Tree(OrdMapColor::Red, key2.clone(), val2.clone(),
                                Rc::new(OrdMap::Tree(OrdMapColor::Black, key3.clone(), val3.clone(), left3.clone(), right3.clone())),
                                Rc::new(OrdMap::Tree(OrdMapColor::Black, key.clone(), val.clone(), right2.clone(), right.clone())))
                        },
                        _ => {}
                    };
                    match right2.deref() {
                        &OrdMap::Tree(OrdMapColor::Red, ref key3, ref val3, ref left3, ref right3) => {
                            return OrdMap::Tree(OrdMapColor::Red, key3.clone(), val3.clone(),
                                Rc::new(OrdMap::Tree(OrdMapColor::Black, key2.clone(), val2.clone(), left2.clone(), left3.clone())),
                                Rc::new(OrdMap::Tree(OrdMapColor::Black, key.clone(), val.clone(), right3.clone(), right.clone())))
                        },
                        _ => {}
                    };
                },
                _ => {}
            }
            match right.deref() {
                &OrdMap::Tree(OrdMapColor::Red, ref key2, ref val2, ref left2, ref right2) => {
                    match left2.deref() {
                        &OrdMap::Tree(OrdMapColor::Red, ref key3, ref val3, ref left3, ref right3) => {
                            return OrdMap::Tree(OrdMapColor::Red, key3.clone(), val3.clone(),
                                Rc::new(OrdMap::Tree(OrdMapColor::Black, key.clone(), val.clone(), left.clone(), left3.clone())),
                                Rc::new(OrdMap::Tree(OrdMapColor::Black, key2.clone(), val2.clone(), right3.clone(), right2.clone())))
                        },
                        _ => {}
                    };
                    match right2.deref() {
                        &OrdMap::Tree(OrdMapColor::Red, ref key3, ref val3, ref left3, ref right3) => {
                            return OrdMap::Tree(OrdMapColor::Red, key2.clone(), val2.clone(),
                                Rc::new(OrdMap::Tree(OrdMapColor::Black, key.clone(), val.clone(), left.clone(), left2.clone())),
                                Rc::new(OrdMap::Tree(OrdMapColor::Black, key3.clone(), val3.clone(), left3.clone(), right3.clone())))
                        },
                        _ => {}
                    };
                },
                _ => {}
            }
        }

        OrdMap::Tree(color, key.clone(), val.clone(), left.clone(), right.clone())
    }

    /// Get the depth of the tree.
    /// Note that it has complexity O(n) unlike other methods.
    /// This method is destined for testing only!
    pub fn depth(&self) -> usize {
        match self {
            &OrdMap::Tree(_, _, _, ref left, ref right) => {
                1 + left.depth().max(right.depth())
            },
            &OrdMap::Nil => {
                0
            }
        }
    }

    /// Get the size of the tree.
    /// Note that it has complexity O(n) unlike other methods.
    /// This method is destined for testing only!
    pub fn size(&self) -> usize {
        match self {
            &OrdMap::Tree(_, _, _, ref left, ref right) => {
                1 + left.size() + right.size()
            },
            &OrdMap::Nil => {
                0
            }
        }
    }
}