MinHeap and MaxHeap

heap.rs

#!/usr/bin/env rust-script

use std::{
    cmp::{Ord, Ordering},
    collections::BinaryHeap,
    ops::{Deref, DerefMut},
};

pub trait Priority<T: ?Sized> {
    fn eq(lhs: &T, rhs: &T) -> bool {
        Self::compare(lhs, rhs).is_eq()
    }
    fn compare(lhs: &T, rhs: &T) -> Ordering;
}
pub struct MaxFirst;
pub struct Reversed<P>(std::marker::PhantomData<fn() -> P>);
pub type MinFirst = Reversed<MaxFirst>;

impl<T: Ord> Priority<T> for MaxFirst {
    fn eq(lhs: &T, rhs: &T) -> bool {
        lhs.eq(rhs)
    }
    fn compare(lhs: &T, rhs: &T) -> Ordering {
        lhs.cmp(rhs)
    }
}
impl<T, P: Priority<T>> Priority<T> for Reversed<P> {
    fn eq(lhs: &T, rhs: &T) -> bool {
        P::eq(lhs, rhs)
    }
    fn compare(lhs: &T, rhs: &T) -> Ordering {
        P::compare(rhs, lhs)
    }
}

#[repr(transparent)]
struct Wrapper<T, P> {
    value: T,
    _priority: std::marker::PhantomData<fn() -> P>,
}
impl<T, P> Wrapper<T, P> {
    fn wrap(value: T) -> Self {
        Self {
            value,
            _priority: std::marker::PhantomData,
        }
    }
    fn unwrap(self) -> T {
        self.value
    }
}

impl<T, P: Priority<T>> PartialEq for Wrapper<T, P> {
    fn eq(&self, other: &Self) -> bool {
        P::eq(&self.value, &other.value)
    }
}
impl<T, P: Priority<T>> Eq for Wrapper<T, P> {}
impl<T, P: Priority<T>> Ord for Wrapper<T, P> {
    fn cmp(&self, other: &Self) -> Ordering {
        P::compare(&self.value, &other.value)
    }
}
impl<T, P: Priority<T>> PartialOrd for Wrapper<T, P> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

pub struct Heap<T, P>(BinaryHeap<Wrapper<T, P>>);
impl<T: Ord, P: Priority<T>> Default for Heap<T, P> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T, P: Priority<T>> Heap<T, P> {
    pub fn new() -> Self {
        Self(BinaryHeap::new())
    }
    pub fn with_capacity(capacity: usize) -> Self {
        Self(BinaryHeap::with_capacity(capacity))
    }
    pub fn peek_mut(&mut self) -> Option<impl '_ + DerefMut<Target = T>> {
        struct PeekMutWrapper<'a, T, P: Priority<T>>(
            std::collections::binary_heap::PeekMut<'a, Wrapper<T, P>>,
        );
        impl<T, P: Priority<T>> Deref for PeekMutWrapper<'_, T, P> {
            type Target = T;

            fn deref(&self) -> &Self::Target {
                &self.0.deref().value
            }
        }
        impl<T, P: Priority<T>> DerefMut for PeekMutWrapper<'_, T, P> {
            fn deref_mut(&mut self) -> &mut Self::Target {
                &mut self.0.deref_mut().value
            }
        }
        self.0.peek_mut().map(PeekMutWrapper)
    }
    pub fn pop(&mut self) -> Option<T> {
        self.0.pop().map(Wrapper::unwrap)
    }
    pub fn push(&mut self, item: T) {
        self.0.push(Wrapper::wrap(item))
    }
    pub fn into_sorted_vec(self) -> Vec<T> {
        self.0
            .into_sorted_vec()
            .into_iter()
            .map(Wrapper::unwrap)
            .collect()
    }
    pub fn append(&mut self, other: &mut Self) {
        self.0.append(&mut other.0)
    }
    pub fn retain<F: FnMut(&T) -> bool>(&mut self, mut f: F) {
        self.0.retain(|v| f(&v.value))
    }
    pub fn iter(&self) -> impl Iterator<Item = &T> {
        self.0.iter().map(|v| &v.value)
    }
    pub fn peek(&self) -> Option<&T> {
        self.0.peek().map(|v| &v.value)
    }
    pub fn capacity(&self) -> usize {
        self.0.capacity()
    }
    pub fn reserve_exact(&mut self, additional: usize) {
        self.0.reserve_exact(additional)
    }
    pub fn reserve(&mut self, additional: usize) {
        self.0.reserve(additional)
    }
    pub fn try_reserve_exact(
        &mut self,
        additional: usize,
    ) -> Result<(), std::collections::TryReserveError> {
        self.0.try_reserve_exact(additional)
    }
    pub fn try_reserve(
        &mut self,
        additional: usize,
    ) -> Result<(), std::collections::TryReserveError> {
        self.0.try_reserve(additional)
    }
    pub fn shrink_to_fit(&mut self) {
        self.0.shrink_to_fit()
    }
    pub fn shrink_to(&mut self, min_capacity: usize) {
        self.0.shrink_to(min_capacity)
    }
    pub fn into_vec(self) -> Vec<T> {
        self.0.into_vec().into_iter().map(Wrapper::unwrap).collect()
    }
    pub fn len(&self) -> usize {
        self.0.len()
    }
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
    pub fn drain(&mut self) -> impl '_ + DoubleEndedIterator<Item = T> {
        self.0.drain().map(Wrapper::unwrap)
    }
    pub fn clear(&mut self) {
        self.0.clear()
    }
}

type MinHeap<T> = Heap<T, MinFirst>;
type MaxHeap<T> = Heap<T, MaxFirst>;

// TODO: Create complete type instead of returning `-> impl ...`.
fn main() {
    let mut a = MinHeap::<u32>::new();
    let mut b = MaxHeap::<u32>::new();
    for i in 0..10 {
        a.push(i);
        b.push(i);
    }
    while let Some(v) = a.pop() {
        println!("{v}");
    }
    while let Some(v) = b.pop() {
        println!("{v}");
    }

    struct MaxFirstIgnoreNaN;
    impl Priority<f64> for MaxFirstIgnoreNaN {
        fn compare(lhs: &f64, rhs: &f64) -> Ordering {
            lhs.partial_cmp(rhs).expect("I hate NaN :(")
        }
    }
    let mut c = Heap::<f64, Reversed<MaxFirstIgnoreNaN>>::new();
    c.push(0.1);
    c.push(0.2);
    while let Some(v) = c.pop() {
        println!("{v}");
    }
}