use std::borrow::Borrow;
use std::collections::hash_map::RandomState;
use std::fmt::{Debug, Display};
use std::hash::{BuildHasher, Hash};
use std::iter::FromIterator;
use crate::editable_binary_heap::{BinaryHeap, BinaryHeapIterator};
use crate::mediator::{
Mediator, MediatorEntry, MediatorIndex, OccupiedEntry as MediatorOccupiedEntry,
VacantEntry as MediatorVacantEntry,
};
#[derive(Clone)]
pub struct KeyedPriorityQueue<TKey, TPriority, S = RandomState>
where
TKey: Hash + Eq,
TPriority: Ord,
S: BuildHasher,
{
heap: BinaryHeap<TPriority>,
key_to_pos: Mediator<TKey, S>,
}
impl<TKey: Hash + Eq, TPriority: Ord> KeyedPriorityQueue<TKey, TPriority, RandomState> {
#[inline]
pub fn new() -> Self {
Self::with_capacity_and_hasher(0, RandomState::default())
}
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
Self::with_capacity_and_hasher(capacity, RandomState::default())
}
}
impl<TKey: Hash + Eq, TPriority: Ord, S: BuildHasher> KeyedPriorityQueue<TKey, TPriority, S> {
#[inline]
pub fn with_hasher(hasher: S) -> Self {
Self::with_capacity_and_hasher(0, hasher)
}
#[inline]
pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self {
Self {
heap: BinaryHeap::with_capacity(capacity),
key_to_pos: Mediator::with_capacity_and_hasher(capacity, hasher),
}
}
#[inline]
pub fn reserve(&mut self, additional: usize) {
self.heap.reserve(additional);
self.key_to_pos.reserve(additional);
}
pub fn push(&mut self, key: TKey, priority: TPriority) -> Option<TPriority> {
match self.entry(key) {
Entry::Vacant(entry) => {
entry.set_priority(priority);
None
}
Entry::Occupied(entry) => Some(entry.set_priority(priority)),
}
}
pub fn pop(&mut self) -> Option<(TKey, TPriority)> {
let (to_remove, _) = self.heap.most_prioritized_idx()?;
Some(self.remove_internal(to_remove))
}
pub fn peek(&self) -> Option<(&TKey, &TPriority)> {
let (first_idx, heap_idx) = self.heap.most_prioritized_idx()?;
let (key, _) = self.key_to_pos.get_index(first_idx);
let (_, priority) = self
.heap
.look_into(heap_idx)
.expect("Checked using key_to_pos");
Some((key, priority))
}
pub fn entry(&mut self, key: TKey) -> Entry<TKey, TPriority, S> {
let key_to_pos = &mut self.key_to_pos;
let heap = &mut self.heap;
match key_to_pos.entry(key) {
MediatorEntry::Vacant(internal_entry) => Entry::Vacant(VacantEntry {
internal_entry,
heap,
}),
MediatorEntry::Occupied(internal_entry) => Entry::Occupied(OccupiedEntry {
internal_entry,
heap,
}),
}
}
pub fn get_priority<Q>(&self, key: &Q) -> Option<&TPriority>
where
TKey: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
let heap_idx = self.key_to_pos.get(key)?;
Some(
self.heap
.look_into(heap_idx)
.expect("Must contain if key_to_pos contain")
.1,
)
}
#[inline]
pub fn set_priority<Q>(
&mut self,
key: &Q,
priority: TPriority,
) -> Result<TPriority, SetPriorityNotFoundError>
where
TKey: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
let map_pos = match self.key_to_pos.get_full(key) {
None => return Err(SetPriorityNotFoundError {}),
Some((idx, _, _)) => idx,
};
Ok(self.set_priority_internal(map_pos, priority))
}
#[inline]
pub fn remove<Q>(&mut self, key: &Q) -> Option<TPriority>
where
TKey: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
let (_, priority) = self.remove_entry(key)?;
Some(priority)
}
#[inline]
pub fn remove_entry<Q>(&mut self, key: &Q) -> Option<(TKey, TPriority)>
where
TKey: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
let (index, _, _) = self.key_to_pos.get_full(key)?;
Some(self.remove_internal(index))
}
#[inline]
pub fn len(&self) -> usize {
debug_assert_eq!(self.key_to_pos.len(), self.heap.usize_len());
self.key_to_pos.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
debug_assert_eq!(self.heap.is_empty(), self.key_to_pos.is_empty());
self.key_to_pos.is_empty()
}
#[inline]
pub fn clear(&mut self) {
self.heap.clear();
self.key_to_pos.clear();
}
pub fn iter(&self) -> KeyedPriorityQueueBorrowIter<TKey, TPriority, S> {
KeyedPriorityQueueBorrowIter {
key_to_pos: &self.key_to_pos,
heap_iterator: self.heap.iter(),
}
}
fn remove_internal(&mut self, position: MediatorIndex) -> (TKey, TPriority) {
let key_to_pos = &mut self.key_to_pos;
let heap = &mut self.heap;
let (_, heap_to_rem) = key_to_pos.get_index(position);
let (removed_idx, priority) = heap
.remove(heap_to_rem, |index, heap_idx| {
*key_to_pos.get_index_mut(index) = heap_idx
})
.expect("Checked by key_to_pos");
debug_assert_eq!(position, removed_idx);
let (removed_key, _) = key_to_pos.swap_remove_index(position);
if MediatorIndex(key_to_pos.len()) != removed_idx {
let (_, heap_idx_of_moved) = key_to_pos.get_index(removed_idx);
heap.change_outer_pos(removed_idx, heap_idx_of_moved);
}
(removed_key, priority)
}
fn set_priority_internal(&mut self, position: MediatorIndex, priority: TPriority) -> TPriority {
let heap = &mut self.heap;
let key_to_pos = &mut self.key_to_pos;
let (_, heap_idx) = key_to_pos.get_index(position);
heap.change_priority(heap_idx, priority, |index, heap_idx| {
*key_to_pos.get_index_mut(index) = heap_idx
})
}
}
pub enum Entry<'a, TKey: Eq + Hash, TPriority: Ord, S: BuildHasher> {
Occupied(OccupiedEntry<'a, TKey, TPriority, S>),
Vacant(VacantEntry<'a, TKey, TPriority, S>),
}
pub struct OccupiedEntry<'a, TKey, TPriority, S = RandomState>
where
TKey: 'a + Eq + Hash,
TPriority: 'a + Ord,
S: BuildHasher,
{
internal_entry: MediatorOccupiedEntry<'a, TKey, S>,
heap: &'a mut BinaryHeap<TPriority>,
}
impl<'a, TKey, TPriority, S> OccupiedEntry<'a, TKey, TPriority, S>
where
TKey: 'a + Eq + Hash,
TPriority: 'a + Ord,
S: BuildHasher,
{
#[inline]
pub fn get_priority(&self) -> &TPriority {
let heap_idx = self.internal_entry.get_heap_idx();
self.heap.look_into(heap_idx).expect("Must be in queue").1
}
#[inline]
pub fn set_priority(mut self, priority: TPriority) -> TPriority {
let heap_idx = self.internal_entry.get_heap_idx();
let heap = &mut self.heap;
let key_to_pos = unsafe {
self.internal_entry.transform_to_map()
};
heap.change_priority(heap_idx, priority, |index, heap_idx| {
*key_to_pos.get_index_mut(index) = heap_idx;
})
}
#[inline]
pub fn get_key(&self) -> &TKey {
self.internal_entry.get_key()
}
pub fn remove(mut self) -> (TKey, TPriority) {
let heap_idx = self.internal_entry.get_heap_idx();
let key_to_pos = unsafe {
self.internal_entry.transform_to_map()
};
let heap = &mut self.heap;
let (removed_idx, priority) = heap
.remove(heap_idx, |index, heap_idx| {
*key_to_pos.get_index_mut(index) = heap_idx
})
.expect("Checked by key_to_pos");
let (removed_key, _) = key_to_pos.swap_remove_index(removed_idx);
if MediatorIndex(key_to_pos.len()) != removed_idx {
let (_, heap_idx_of_moved) = key_to_pos.get_index(removed_idx);
heap.change_outer_pos(removed_idx, heap_idx_of_moved);
}
(removed_key, priority)
}
}
pub struct VacantEntry<'a, TKey, TPriority, S = RandomState>
where
TKey: 'a + Eq + Hash,
TPriority: 'a + Ord,
S: BuildHasher,
{
internal_entry: MediatorVacantEntry<'a, TKey, S>,
heap: &'a mut BinaryHeap<TPriority>,
}
impl<'a, TKey, TPriority, S> VacantEntry<'a, TKey, TPriority, S>
where
TKey: 'a + Eq + Hash,
TPriority: 'a + Ord,
S: BuildHasher,
{
#[inline]
pub fn set_priority(self, priority: TPriority) {
let heap = self.heap;
let internal_entry = self.internal_entry;
let (key_to_pos, mediator_index) = unsafe {
internal_entry.insert(heap.len())
};
heap.push(mediator_index, priority, |index, val| {
*key_to_pos.get_index_mut(index) = val
});
}
#[inline]
pub fn get_key(&self) -> &TKey {
self.internal_entry.get_key()
}
}
impl<TKey: Hash + Eq + Debug, TPriority: Ord + Debug, S: BuildHasher> Debug
for KeyedPriorityQueue<TKey, TPriority, S>
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
write!(f, "[")?;
for entry in self.iter() {
write!(f, "{:?}", entry)?;
}
write!(f, "]")
}
}
impl<TKey: Hash + Eq, TPriority: Ord, S: BuildHasher + Default> Default
for KeyedPriorityQueue<TKey, TPriority, S>
{
#[inline]
fn default() -> Self {
Self::with_capacity_and_hasher(0, S::default())
}
}
impl<TKey: Hash + Eq, TPriority: Ord, S: BuildHasher + Default> FromIterator<(TKey, TPriority)>
for KeyedPriorityQueue<TKey, TPriority, S>
{
fn from_iter<T: IntoIterator<Item = (TKey, TPriority)>>(iter: T) -> Self {
let (heap, key_to_pos) = BinaryHeap::produce_from_iter_hash(iter);
Self { heap, key_to_pos }
}
}
impl<TKey: Hash + Eq, TPriority: Ord> IntoIterator for KeyedPriorityQueue<TKey, TPriority> {
type Item = (TKey, TPriority);
type IntoIter = KeyedPriorityQueueIterator<TKey, TPriority>;
fn into_iter(self) -> Self::IntoIter {
Self::IntoIter { queue: self }
}
}
pub struct KeyedPriorityQueueIterator<TKey, TPriority, S = RandomState>
where
TKey: Hash + Eq,
TPriority: Ord,
S: BuildHasher,
{
queue: KeyedPriorityQueue<TKey, TPriority, S>,
}
impl<TKey: Hash + Eq, TPriority: Ord, S: BuildHasher> Iterator
for KeyedPriorityQueueIterator<TKey, TPriority, S>
{
type Item = (TKey, TPriority);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.queue.pop()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.queue.len(), Some(self.queue.len()))
}
#[inline]
fn count(self) -> usize
where
Self: Sized,
{
self.queue.len()
}
}
pub struct KeyedPriorityQueueBorrowIter<'a, TKey, TPriority, S = RandomState>
where
TKey: 'a + Hash + Eq,
TPriority: 'a,
S: BuildHasher,
{
heap_iterator: BinaryHeapIterator<'a, TPriority>,
key_to_pos: &'a Mediator<TKey, S>,
}
impl<'a, TKey: 'a + Hash + Eq, TPriority: 'a, S: BuildHasher> Iterator
for KeyedPriorityQueueBorrowIter<'a, TKey, TPriority, S>
{
type Item = (&'a TKey, &'a TPriority);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
let heap_iterator = &mut self.heap_iterator;
let key_to_pos = &self.key_to_pos;
heap_iterator.next().map(|(index, priority)| {
let (key, _) = key_to_pos.get_index(index);
(key, priority)
})
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.heap_iterator.size_hint()
}
#[inline]
fn count(self) -> usize
where
Self: Sized,
{
self.heap_iterator.count()
}
}
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Copy, Clone, Default)]
pub struct SetPriorityNotFoundError;
impl Display for SetPriorityNotFoundError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
write!(f, "Key not found in KeyedPriorityQueue during set_priority")
}
}
impl std::error::Error for SetPriorityNotFoundError {}
#[cfg(test)]
mod tests {
use super::KeyedPriorityQueue;
#[test]
fn test_priority() {
let mut items = [1, 4, 5, 2, 3];
let mut queue = KeyedPriorityQueue::<i32, i32>::with_capacity(items.len());
for (i, &x) in items.iter().enumerate() {
queue.push(x, x);
assert_eq!(queue.len(), i + 1);
}
assert_eq!(queue.len(), items.len());
items.sort_unstable_by_key(|&x| -x);
for &x in items.iter() {
assert_eq!(queue.pop(), Some((x, x)));
}
assert_eq!(queue.pop(), None);
}
#[test]
fn test_peek() {
let items = [
("first", 5),
("second", 4),
("third", 3),
("fourth", 2),
("fifth", 1),
];
let mut queue: KeyedPriorityQueue<&str, i32> = items.iter().cloned().collect();
while queue.len() > 0 {
let (&key, &priority) = queue.peek().unwrap();
let (key1, priority1) = queue.pop().unwrap();
assert_eq!(key, key1);
assert_eq!(priority, priority1);
}
assert_eq!(queue.peek(), None);
}
#[test]
fn test_get_priority() {
let items = [
("first", 5),
("second", 4),
("third", 3),
("fourth", 2),
("fifth", 1),
];
let queue: KeyedPriorityQueue<&str, i32> = items.iter().cloned().collect();
for &(key, priority) in items.iter() {
let &real = queue.get_priority(&key).unwrap();
assert_eq!(real, priority);
}
let mut queue = queue;
while let Some(_) = queue.pop() {}
for &(key, _) in items.iter() {
assert_eq!(queue.get_priority(&key), None);
}
}
#[test]
fn test_change_priority() {
let items = [
("first", 5),
("second", 4),
("third", 3),
("fourth", 2),
("fifth", 1),
];
let mut queue: KeyedPriorityQueue<&str, i32> = items.iter().cloned().collect();
assert_eq!(
queue.set_priority(&"HELLO", 64),
Err(super::SetPriorityNotFoundError::default())
);
let old_priority = *queue.get_priority(&"fifth").unwrap();
assert_eq!(queue.set_priority(&"fifth", old_priority + 10), Ok(1));
assert_eq!(queue.get_priority(&"fifth"), Some(&11));
assert_eq!(queue.pop(), Some(("fifth", 11)));
let old_priority = *queue.get_priority(&"first").unwrap();
assert_eq!(queue.set_priority(&"first", old_priority - 10), Ok(5));
assert_eq!(queue.get_priority(&"first"), Some(&-5));
queue.pop();
queue.pop();
queue.pop();
assert_eq!(queue.pop(), Some(("first", -5)));
}
#[test]
fn test_remove_items() {
let mut items = [1, 4, 5, 2, 3];
let mut queue: KeyedPriorityQueue<i32, i32> = items.iter().map(|&x| (x, x)).collect();
assert_eq!(queue.remove_entry(&3), Some((3, 3)));
assert_eq!(queue.remove_entry(&20), None);
assert_eq!(queue.len(), items.len() - 1);
assert_eq!(queue.get_priority(&3), None);
items.sort_unstable_by_key(|&x| -x);
for x in items.iter().cloned().filter(|&x| x != 3) {
assert_eq!(queue.pop(), Some((x, x)));
}
assert_eq!(queue.pop(), None);
}
#[test]
fn test_iteration() {
let items = [
("first", 5),
("second", 4),
("third", 3),
("fourth", 2),
("fifth", 1),
];
let queue: KeyedPriorityQueue<&str, i32> = items.iter().rev().cloned().collect();
let mut iter = queue.into_iter();
assert_eq!(iter.next(), Some(("first", 5)));
assert_eq!(iter.next(), Some(("second", 4)));
assert_eq!(iter.next(), Some(("third", 3)));
assert_eq!(iter.next(), Some(("fourth", 2)));
assert_eq!(iter.next(), Some(("fifth", 1)));
assert_eq!(iter.next(), None);
}
#[test]
fn test_multiple_push() {
let mut queue = KeyedPriorityQueue::new();
queue.push(0, 1);
assert_eq!(queue.peek(), Some((&0, &1)));
queue.push(0, 5);
assert_eq!(queue.peek(), Some((&0, &5)));
queue.push(0, 7);
assert_eq!(queue.peek(), Some((&0, &7)));
queue.push(0, 9);
assert_eq!(queue.peek(), Some((&0, &9)));
}
#[test]
fn test_borrow_keys() {
let mut queue: KeyedPriorityQueue<String, i32> = KeyedPriorityQueue::new();
queue.push("Hello".to_string(), 5);
let string = "Hello".to_string();
let string_ref: &String = &string;
let str_ref: &str = &string;
assert_eq!(queue.get_priority(string_ref), Some(&5));
assert_eq!(queue.get_priority(str_ref), Some(&5));
}
#[test]
fn test_entry_vacant() {
use super::Entry;
let items = [
("first", 5i32),
("second", 4),
("third", 3),
("fourth", 2),
("fifth", 1),
];
let mut queue = KeyedPriorityQueue::new();
for &(k, v) in items.iter() {
queue.push(k, v);
}
assert_eq!(queue.len(), 5);
match queue.entry("Cotton") {
Entry::Occupied(_) => unreachable!(),
Entry::Vacant(entry) => {
assert_eq!(entry.get_key(), &"Cotton");
entry.set_priority(10);
}
};
assert_eq!(queue.len(), 6);
assert_eq!(queue.get_priority(&"Cotton"), Some(&10));
match queue.entry("Cotton") {
Entry::Occupied(entry) => {
assert_eq!(entry.get_key(), &"Cotton");
assert_eq!(entry.get_priority(), &10);
}
Entry::Vacant(_) => unreachable!(),
};
}
#[test]
fn test_entry_occupied() {
use super::Entry;
let items = [
("first", 5i32),
("second", 4),
("third", 3),
("fourth", 2),
("fifth", 1),
];
let mut queue = KeyedPriorityQueue::new();
for &(k, v) in items.iter() {
queue.push(k, v);
}
assert_eq!(queue.len(), 5);
match queue.entry("third") {
Entry::Occupied(entry) => {
assert_eq!(entry.get_key(), &"third");
assert_eq!(entry.get_priority(), &3);
assert_eq!(entry.set_priority(5), 3);
}
Entry::Vacant(_) => unreachable!(),
};
assert_eq!(queue.len(), 5);
assert_eq!(queue.get_priority(&"third"), Some(&5));
match queue.entry("third") {
Entry::Occupied(entry) => {
assert_eq!(entry.remove(), ("third", 5));
}
Entry::Vacant(_) => unreachable!(),
};
assert_eq!(queue.len(), 4);
assert_eq!(queue.get_priority(&"third"), None);
}
#[test]
fn test_borrow_iter() {
use std::collections::HashMap;
let items = [
("first", 5i32),
("third", 3),
("second", 4),
("fifth", 1),
("fourth", 2),
];
let queue: KeyedPriorityQueue<String, i32> =
items.iter().map(|&(k, p)| (k.to_owned(), p)).collect();
let mut map: HashMap<&str, i32> = HashMap::new();
let mut total_items = 0;
for (key, &value) in queue.iter() {
map.insert(key, value);
total_items += 1;
}
assert_eq!(items.len(), total_items);
assert_eq!(queue.len(), items.len());
let other_map: HashMap<_, _> = items.iter().cloned().collect();
assert_eq!(map, other_map);
}
#[test]
fn test_sync() {
fn assert_sync<T: Sync>() {}
assert_sync::<KeyedPriorityQueue<i32, i32>>();
}
#[test]
fn test_send() {
fn assert_sync<T: Send>() {}
assert_sync::<KeyedPriorityQueue<i32, i32>>();
}
#[test]
fn test_fmt() {
let items = [
("first", 5i32),
("second", 4),
("third", 3),
("fourth", 2),
("fifth", 1),
];
let queue: KeyedPriorityQueue<&str, i32> = items.iter().cloned().collect();
assert_eq!(
format!("{:?}", queue),
"[(\"first\", 5)(\"second\", 4)(\"third\", 3)(\"fourth\", 2)(\"fifth\", 1)]"
);
}
#[test]
fn test_not_clone_works() {
use core::hash::Hash;
#[derive(Hash, PartialEq, Eq)]
struct Key(u32);
let vals = [0u32, 1, 1, 2, 4, 5];
let mut queue: KeyedPriorityQueue<Key, u32> =
vals.iter().copied().map(|v| (Key(v), v)).collect();
queue.set_priority(&Key(1), 10).unwrap();
let mut res = Vec::with_capacity(5);
while let Some((Key(k), p)) = queue.pop() {
res.push((k, p));
}
assert_eq!(&res, &[(1, 10), (5, 5), (4, 4), (2, 2), (0, 0)]);
}
#[test]
fn test_remove_change_tree() {
use std::cmp::Reverse;
let mut queue = KeyedPriorityQueue::new();
queue.push(0, Reverse(300));
queue.push(1, Reverse(500));
queue.push(2, Reverse(400));
queue.push(3, Reverse(400));
queue.push(4, Reverse(600));
queue.push(5, Reverse(100));
queue.push(6, Reverse(200));
queue.remove(&1);
let mut list = Vec::new();
while let Some((_, timestamp)) = queue.pop() {
list.push(timestamp.0);
}
assert_eq!(list, [100, 200, 300, 400, 400, 600])
}
}