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::mem;

/// An ordered map by integer key, where the implementation is based on Patricia tree.
/// See article "Fast Mergeable Integer Maps" by Chris Okasaki and Andrew Gill.
#[derive(Debug, PartialEq, Eq)]
pub enum IntMap<K, V> {

    /// The branch item.
    Br(K, K, Rc<IntMap<K, V>>, Rc<IntMap<K, V>>),

    /// The leaf item.
    Lf(K, V),

    /// An empty map.
    Nil
}

impl<K, V> Clone for IntMap<K, V>
    where K: Clone + Copy,
          V: Clone
{
    fn clone(&self) -> Self {
        match self {
            &IntMap::Br(prefix, bit, ref left, ref right) => {
                IntMap::Br(prefix, bit, left.clone(), right.clone())
            },
            &IntMap::Lf(key, ref val) => {
                IntMap::Lf(key, val.clone())
            },
            &IntMap::Nil => {
                IntMap::Nil
            }
        }
    }
}

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

    /// Create an empty map.
    pub fn empty() -> Self {
        IntMap::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 {
                &IntMap::Br(prefix, bit, ref left, ref right) => {
                    if !K::match_prefix(*key, prefix, bit) {
                        return None;
                    } else if K::zero_bit(*key, bit) {
                        p = left.deref();
                    } else {
                        p = right.deref();
                    }
                },
                &IntMap::Lf(ref key0, ref val0) if *key == *key0 => {
                    return Some(val0);
                },
                &IntMap::Lf(_, _) => {
                    return None;
                },
                &IntMap::Nil => {
                    return None;
                }
            }
        }
    }

    /// Get the first key-value pair.
    pub fn first_key_value(&self) -> Option<(&K, &V)> {
        let mut p = self;
        let mut first_br = true;
        loop {
            match p {
                &IntMap::Br(_, bit, ref left, ref right) => {
                    if first_br && K::negative_bit(bit) {
                        first_br = false;
                        p = right.deref();
                    } else {
                        p = left.deref();
                    }
                },
                &IntMap::Lf(ref key, ref val) => {
                    return Some((key, val));
                },
                &IntMap::Nil => {
                    return None;
                }
            }
        }
    }

    /// Get the last key-value pair.
    pub fn last_key_value(&self) -> Option<(&K, &V)> {
        let mut p = self;
        let mut first_br = true;
        loop {
            match p {
                &IntMap::Br(_, bit, ref left, ref right) => {
                    if first_br && K::negative_bit(bit) {
                        first_br = false;
                        p = left.deref();
                    } else {
                        p = right.deref();
                    }
                },
                &IntMap::Lf(ref key, ref val) => {
                    return Some((key, val));
                },
                &IntMap::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 {
            &IntMap::Br(_, bit, ref left, ref right) => {
                if K::negative_bit(bit) {
                    match right.find(pred) {
                        z @ Some(_) => z,
                        None => left.find(pred)
                    }
                } else {
                    match left.find(pred) {
                        z @ Some(_) => z,
                        None => right.find(pred)
                    }
                }
            },
            &IntMap::Lf(ref key, ref val) => {
                pred(key, val)
            },
            &IntMap::Nil => {
                None
            }
        }
    }

    /// Insert a new key-value pair.
    pub fn insert(&self, key: K, val: V) -> IntMap<K, V> {
        match self {
            &IntMap::Br(prefix, bit, ref left, ref right) => {
                if K::match_prefix(key, prefix, bit) {
                    if K::zero_bit(key, bit) {
                        IntMap::Br(prefix, bit, Rc::new(left.insert(key, val)), right.clone())
                    } else {
                        IntMap::Br(prefix, bit, left.clone(), Rc::new(right.insert(key, val)))
                    }
                } else {
                    Self::join(key, IntMap::Lf(key, val), prefix, &self)
                }
            },
            &IntMap::Lf(key0, _) if key == key0 => {
                IntMap::Lf(key, val)
            },
            &IntMap::Lf(key0, _) => {
                Self::join(key, IntMap::Lf(key, val), key0, &self)
            },
            &IntMap::Nil => {
                IntMap::Lf(key, val)
            }
        }
    }

    /// Remove the key-value pair.
    pub fn remove(&self, key: &K) -> IntMap<K, V> {
        match self {
            &IntMap::Br(prefix, bit, ref left, ref right) => {
                if K::match_prefix(*key, prefix, bit) {
                    if K::zero_bit(*key, bit) {
                        Self::br_left(prefix, bit, left.remove(key), right.clone())
                    } else {
                        Self::br_right(prefix, bit, left.clone(), right.remove(key))
                    }
                } else {
                    self.clone()
                }
            },
            &IntMap::Lf(key0, _) if *key == key0 => {
                IntMap::Nil
            },
            &IntMap::Lf(_, _) => {
                self.clone()
            },
            &IntMap::Nil => {
                IntMap::Nil
            }
        }
    }

    /// Join the trees.
    fn join(prefix0: K, tree0: Self, prefix1: K, tree1: &Self) -> Self {
        let m = K::branching_bit(prefix0, prefix1);
        if K::zero_bit(prefix0, m) {
            IntMap::Br(K::mask(prefix0, m), m, Rc::new(tree0), Rc::new(tree1.clone()))
        } else {
            IntMap::Br(K::mask(prefix0, m), m, Rc::new(tree1.clone()), Rc::new(tree0))
        }
    }

    /// Create by the updated left branch.
    fn br_left(prefix: K, bit: K, left: Self, right: Rc<Self>) -> Self {
        if left.is_empty() {
            right.deref().clone()
        } else {
            IntMap::Br(prefix, bit, Rc::new(left), right)
        }
    }

    /// Create by the updated right branch.
    fn br_right(prefix: K, bit: K, left: Rc<Self>, right: Self) -> Self {
        if right.is_empty() {
            left.deref().clone()
        } else {
            IntMap::Br(prefix, bit, left, Rc::new(right))
        }
    }
}

/// A trait for integer map keys.
pub trait IntMapKey {

    /// Whether the prefix matches.
    fn match_prefix(key: Self, prefix: Self, bit: Self) -> bool;

    /// Is the bit zero?
    fn zero_bit(prefix: Self, bit: Self) -> bool;

    /// Is the bit negative?
    fn negative_bit(bit: Self) -> bool;

    /// Get the branching bit.
    fn branching_bit(prefix0: Self, prefix: Self) -> Self;

    /// Get the mask prefix.
    fn mask(prefix: Self, bit: Self) -> Self;
}

impl IntMapKey for i64 {

    /// Whether the prefix matches.
    fn match_prefix(key: Self, prefix: Self, bit: Self) -> bool {
        Self::mask(key, bit) == prefix
    }

    /// Is the bit zero?
    fn zero_bit(prefix: Self, bit: Self) -> bool {
        (prefix & bit) == 0
    }

    /// Is the bit negative?
    fn negative_bit(bit: Self) -> bool {
        bit < 0
    }

    /// Get the branching bit.
    fn branching_bit(prefix0: Self, prefix: Self) -> Self {
        let p0  = prefix0 as u64;
        let p1  = prefix as u64;
        let x   = p0 ^ p1;
        let bit = (1 as u64) << (8 * mem::size_of_val(&x) - 1 - (x.leading_zeros() as usize)) as u64;

        bit as i64
    }

    /// Get the mask prefix.
    fn mask(prefix: Self, bit: Self) -> Self {
        let prefix = prefix as u64;
        let bit    = bit as u64;
        let mask   = prefix & (((bit as i64).wrapping_neg() as u64) ^ bit);

        mask as i64
    }
}

impl IntMapKey for u64 {

    /// Whether the prefix matches.
    fn match_prefix(key: Self, prefix: Self, bit: Self) -> bool {
        Self::mask(key, bit) == prefix
    }

    /// Is the bit zero?
    fn zero_bit(prefix: Self, bit: Self) -> bool {
        (prefix & bit) == 0
    }

    /// Is the bit negative?
    fn negative_bit(_bit: Self) -> bool {
        false
    }

    /// Get the branching bit.
    fn branching_bit(prefix0: Self, prefix: Self) -> Self {
        let p0  = prefix0;
        let p1  = prefix;
        let x   = p0 ^ p1;
        let bit = (1 as u64) << (8 * mem::size_of_val(&x) - 1 - (x.leading_zeros() as usize));

        bit
    }

    /// Get the mask prefix.
    fn mask(prefix: Self, bit: Self) -> Self {
        let mask = prefix & (((bit as i64).wrapping_neg() as u64) ^ bit);
        mask
    }
}

impl IntMapKey for isize {

    /// Whether the prefix matches.
    fn match_prefix(key: Self, prefix: Self, bit: Self) -> bool {
        Self::mask(key, bit) == prefix
    }

    /// Is the bit zero?
    fn zero_bit(prefix: Self, bit: Self) -> bool {
        (prefix & bit) == 0
    }

    /// Is the bit negative?
    fn negative_bit(bit: Self) -> bool {
        bit < 0
    }

    /// Get the branching bit.
    fn branching_bit(prefix0: Self, prefix: Self) -> Self {
        let p0  = prefix0 as usize;
        let p1  = prefix as usize;
        let x   = p0 ^ p1;
        let bit = (1 as usize) << (8 * mem::size_of_val(&x) - 1 - (x.leading_zeros() as usize)) as usize;

        bit as isize
    }

    /// Get the mask prefix.
    fn mask(prefix: Self, bit: Self) -> Self {
        let prefix = prefix as usize;
        let bit    = bit as usize;
        let mask   = prefix & (((bit as isize).wrapping_neg() as usize) ^ bit);

        mask as isize
    }
}