fera-graph 0.2.0

// 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 http://mozilla.org/MPL/2.0/.

use std::ops::{Deref, DerefMut, IndexMut};
use std::collections::BinaryHeap;

pub trait PQTypes<'a, T> {
    type KeyEntry: DerefMut<Target = T>;
}

pub trait PriorityQueue<T, P>: for<'a> PQTypes<'a, P> {
    fn is_empty(&self) -> bool;

    fn insert(&mut self, item: T, prio: P);

    fn get_mut(&mut self, item: T) -> Option<<Self as PQTypes<P>>::KeyEntry>;

    fn pop(&mut self) -> Option<(T, P)>;
}

pub struct ArrayPq<T, P> {
    prop: P,
    table: Vec<T>,
}

impl<T, P> ArrayPq<T, P> {
    pub fn new(prop: P) -> Self {
        Self {
            prop,
            table: vec![],
        }
    }
}

impl<'a, T, P> PQTypes<'a, P::Output> for ArrayPq<T, P>
    where P: IndexMut<T>,
          P::Output: 'static + Sized
{
    type KeyEntry = &'a mut P::Output;
}

impl<T, P> PriorityQueue<T, P::Output> for ArrayPq<T, P>
    where T: Copy + PartialEq,
          P: IndexMut<T>,
          P::Output: 'static + Sized + Ord + Copy
{
    fn is_empty(&self) -> bool {
        self.table.is_empty()
    }

    fn insert(&mut self, item: T, prio: P::Output) {
        self.prop[item] = prio;
        self.table.push(item);
    }

    fn get_mut(&mut self, item: T) -> Option<<Self as PQTypes<P::Output>>::KeyEntry> {
        if self.table.contains(&item) {
            Some(&mut self.prop[item])
        } else {
            None
        }
    }

    fn pop(&mut self) -> Option<(T, P::Output)> {
        if let Some((i, _)) = self.table
               .iter()
               .enumerate()
               .max_by_key(|&(_, item)| self.prop[*item]) {
            let item = self.table.swap_remove(i);
            Some((item, self.prop[item]))
        } else {
            None
        }
    }
}

pub struct BinaryHeapPq<K, V, P> {
    count: usize,
    heap: BinaryHeap<(V, K)>,
    value: P,
}

pub struct BRPQKeyEntry<'a, K, V, P>
    where K: 'static + Ord + Copy,
          V: 'static + Ord + Copy,
          P: 'static + IndexMut<K, Output = Option<V>>
{
    pq: &'a mut BinaryHeapPq<K, V, P>,
    key: K,
    value: V,
}

impl<'a, K, V, P> Deref for BRPQKeyEntry<'a, K, V, P>
    where K: 'static + Ord + Copy,
          V: 'static + Ord + Copy,
          P: 'static + IndexMut<K, Output = Option<V>>
{
    type Target = V;

    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

impl<'a, K, V, P> DerefMut for BRPQKeyEntry<'a, K, V, P>
    where K: 'static + Ord + Copy,
          V: 'static + Ord + Copy,
          P: 'static + IndexMut<K, Output = Option<V>>
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.value
    }
}

impl<'a, K, V, P> PQTypes<'a, V> for BinaryHeapPq<K, V, P>
    where K: 'static + Ord + Copy,
          V: 'static + Ord + Copy,
          P: 'static + IndexMut<K, Output = Option<V>>
{
    type KeyEntry = BRPQKeyEntry<'a, K, V, P>;
}

impl<'a, K, V, P> Drop for BRPQKeyEntry<'a, K, V, P>
    where K: 'static + Ord + Copy,
          V: 'static + Ord + Copy,
          P: 'static + IndexMut<K, Output = Option<V>>
{
    fn drop(&mut self) {
        self.pq.insert(self.key, self.value);
    }
}

impl<K, V, P> PriorityQueue<K, V> for BinaryHeapPq<K, V, P>
    where K: 'static + Ord + Copy,
          V: 'static + Ord + Copy,
          P: 'static + IndexMut<K, Output = Option<V>>
{
    fn is_empty(&self) -> bool {
        self.count == 0
    }

    fn insert(&mut self, key: K, new: V) {
        if let Some(value) = self.value[key] {
            if value == new {
                return;
            }
        } else {
            self.count += 1;
        }
        self.value[key] = Some(new);
        self.heap.push((new, key));
    }

    fn get_mut(&mut self, key: K) -> Option<<Self as PQTypes<V>>::KeyEntry> {
        if let Some(value) = self.value[key] {
            Some(BRPQKeyEntry {
                     pq: self,
                     key,
                     value,
                 })
        } else {
            None
        }
    }

    fn pop(&mut self) -> Option<(K, V)> {
        if self.count == 0 {
            return None;
        }
        while let Some((value, key)) = self.heap.pop() {
            if Some(value) == self.value[key] {
                self.count -= 1;
                return Some((key, value));
            }
        }
        None
    }
}

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

    #[test]
    fn test_array() {
        let mut pq = ArrayPq::new(vec![0; 5]);
        pq.insert(0, 20);
        pq.insert(1, 30);
        pq.insert(2, 10);
        pq.insert(3, 40);

        assert_eq!(Some((3, 40)), pq.pop());
        assert_eq!(Some((1, 30)), pq.pop());
        assert_eq!(Some((0, 20)), pq.pop());
        assert_eq!(Some((2, 10)), pq.pop());
        assert_eq!(None, pq.pop());

        pq.insert(0, 20);
        pq.insert(1, 30);
        pq.insert(2, 10);
        pq.insert(3, 40);
        pq.insert(4, 35);

        *pq.get_mut(0).unwrap() = 5;
        *pq.get_mut(1).unwrap() = 4;
        *pq.get_mut(2).unwrap() = 3;
        *pq.get_mut(3).unwrap() = 2;
        *pq.get_mut(4).unwrap() = 1;
        assert_eq!(Some((0, 5)), pq.pop());
        assert_eq!(Some((1, 4)), pq.pop());
        assert_eq!(Some((2, 3)), pq.pop());
        assert_eq!(Some((3, 2)), pq.pop());
        assert_eq!(Some((4, 1)), pq.pop());
    }
}