rheaps 0.16.0

use core::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};

use crate::array::InvalidDegree;
use crate::error::{DecreaseKeyError, InvalidHandle};
use crate::{AddressableHeap, DecreaseKeyHeap};

const DEFAULT_HEAP_CAPACITY: usize = 16;
static NEXT_HEAP_ID: AtomicU64 = AtomicU64::new(1);

/// An opaque capability that identifies an entry in an addressable heap.
///
/// Handles are `Copy`, but they are only valid in the heap that created them.
/// They become invalid after their entry is removed or after [`AddressableHeap::clear`].
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct AddressableHandle {
    heap_id: u64,
    slot: usize,
    generation: u64,
}

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct Entry<K, V> {
    key: K,
    value: V,
    slot: usize,
}

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct Slot {
    index: Option<usize>,
    generation: u64,
}

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct AddressableCore<K, V> {
    entries: Vec<Entry<K, V>>,
    slots: Vec<Slot>,
    free_slots: Vec<usize>,
    heap_id: u64,
}

impl<K: Ord, V> AddressableCore<K, V> {
    fn new(capacity: usize) -> Self {
        Self {
            entries: Vec::with_capacity(capacity),
            slots: Vec::with_capacity(capacity),
            free_slots: Vec::new(),
            heap_id: next_heap_id(),
        }
    }

    fn from_vec(entries: Vec<(K, V)>, degree: usize) -> Self {
        let heap_id = next_heap_id();
        let mut slots = Vec::with_capacity(entries.len());
        let mut heap_entries = Vec::with_capacity(entries.len());
        for (index, (key, value)) in entries.into_iter().enumerate() {
            slots.push(Slot {
                index: Some(index),
                generation: 0,
            });
            heap_entries.push(Entry {
                key,
                value,
                slot: index,
            });
        }
        let mut heap = Self {
            entries: heap_entries,
            slots,
            free_slots: Vec::new(),
            heap_id,
        };
        heap.heapify(degree);
        heap
    }

    fn len(&self) -> usize {
        self.entries.len()
    }

    fn handle_for_slot(&self, slot: usize) -> AddressableHandle {
        AddressableHandle {
            heap_id: self.heap_id,
            slot,
            generation: self.slots[slot].generation,
        }
    }

    fn validate(&self, handle: AddressableHandle) -> Result<usize, InvalidHandle> {
        if handle.heap_id != self.heap_id {
            return Err(InvalidHandle::ForeignHeap);
        }
        let Some(slot) = self.slots.get(handle.slot) else {
            return Err(InvalidHandle::Stale);
        };
        if slot.generation != handle.generation {
            return Err(InvalidHandle::Stale);
        }
        slot.index.ok_or(InvalidHandle::Stale)
    }

    fn push(&mut self, key: K, value: V, degree: usize) -> AddressableHandle {
        let slot = match self.free_slots.pop() {
            Some(slot) => {
                self.slots[slot].index = Some(self.entries.len());
                slot
            }
            None => {
                let slot = self.slots.len();
                self.slots.push(Slot {
                    index: Some(self.entries.len()),
                    generation: 0,
                });
                slot
            }
        };
        self.entries.push(Entry { key, value, slot });
        self.sift_up(self.entries.len() - 1, degree);
        self.handle_for_slot(slot)
    }

    fn peek(&self) -> Option<(AddressableHandle, &K, &V)> {
        self.entries
            .first()
            .map(|entry| (self.handle_for_slot(entry.slot), &entry.key, &entry.value))
    }

    fn pop(&mut self, degree: usize) -> Option<(K, V)> {
        if self.entries.is_empty() {
            return None;
        }
        let entry = self.remove_at(0);
        if !self.entries.is_empty() {
            self.sift_down(0, degree);
        }
        Some((entry.key, entry.value))
    }

    fn key(&self, handle: AddressableHandle) -> Result<&K, InvalidHandle> {
        let index = self.validate(handle)?;
        Ok(&self.entries[index].key)
    }

    fn value(&self, handle: AddressableHandle) -> Result<&V, InvalidHandle> {
        let index = self.validate(handle)?;
        Ok(&self.entries[index].value)
    }

    fn value_mut(&mut self, handle: AddressableHandle) -> Result<&mut V, InvalidHandle> {
        let index = self.validate(handle)?;
        Ok(&mut self.entries[index].value)
    }

    fn decrease_key(
        &mut self,
        handle: AddressableHandle,
        key: K,
        degree: usize,
    ) -> Result<(), DecreaseKeyError> {
        let index = self
            .validate(handle)
            .map_err(DecreaseKeyError::InvalidHandle)?;
        if key > self.entries[index].key {
            return Err(DecreaseKeyError::NotDecreased);
        }
        self.entries[index].key = key;
        self.sift_up(index, degree);
        Ok(())
    }

    fn delete(
        &mut self,
        handle: AddressableHandle,
        degree: usize,
    ) -> Result<(K, V), InvalidHandle> {
        let index = self.validate(handle)?;
        let entry = self.remove_at(index);
        if index < self.entries.len() {
            self.restore_at(index, degree);
        }
        Ok((entry.key, entry.value))
    }

    fn clear(&mut self) {
        while let Some(entry) = self.entries.pop() {
            self.invalidate_slot(entry.slot);
        }
    }

    fn handles(&self) -> impl Iterator<Item = AddressableHandle> + '_ {
        self.entries
            .iter()
            .map(|entry| self.handle_for_slot(entry.slot))
    }

    fn remove_at(&mut self, index: usize) -> Entry<K, V> {
        let entry = self.entries.swap_remove(index);
        self.invalidate_slot(entry.slot);
        if let Some(moved) = self.entries.get(index) {
            self.slots[moved.slot].index = Some(index);
        }
        entry
    }

    fn invalidate_slot(&mut self, slot_index: usize) {
        let slot = &mut self.slots[slot_index];
        slot.index = None;
        slot.generation = slot.generation.wrapping_add(1);
        self.free_slots.push(slot_index);
    }

    fn heapify(&mut self, degree: usize) {
        if self.entries.len() < 2 {
            return;
        }
        let last_parent = (self.entries.len() - 2) / degree;
        for index in (0..=last_parent).rev() {
            self.sift_down(index, degree);
        }
    }

    fn restore_at(&mut self, index: usize, degree: usize) {
        if index > 0 {
            let parent = (index - 1) / degree;
            if self.entries[index].key < self.entries[parent].key {
                self.sift_up(index, degree);
                return;
            }
        }
        self.sift_down(index, degree);
    }

    fn sift_up(&mut self, mut index: usize, degree: usize) {
        while index > 0 {
            let parent = (index - 1) / degree;
            if self.entries[parent].key <= self.entries[index].key {
                break;
            }
            self.swap_entries(parent, index);
            index = parent;
        }
    }

    fn sift_down(&mut self, mut index: usize, degree: usize) {
        loop {
            let first_child = index
                .checked_mul(degree)
                .and_then(|value| value.checked_add(1))
                .unwrap_or(self.entries.len());
            if first_child >= self.entries.len() {
                return;
            }
            let end = first_child.saturating_add(degree).min(self.entries.len());
            let mut smallest = first_child;
            for child in first_child + 1..end {
                if self.entries[child].key < self.entries[smallest].key {
                    smallest = child;
                }
            }
            if self.entries[index].key <= self.entries[smallest].key {
                return;
            }
            self.swap_entries(index, smallest);
            index = smallest;
        }
    }

    fn swap_entries(&mut self, left: usize, right: usize) {
        self.entries.swap(left, right);
        self.slots[self.entries[left].slot].index = Some(left);
        self.slots[self.entries[right].slot].index = Some(right);
    }
}

fn next_heap_id() -> u64 {
    let id = NEXT_HEAP_ID.fetch_add(1, AtomicOrdering::Relaxed);
    if id == 0 {
        NEXT_HEAP_ID.fetch_add(1, AtomicOrdering::Relaxed)
    } else {
        id
    }
}

/// An array-backed binary min-heap with stable, checked entry handles.
///
/// Insertion, removal, deletion by handle, and key decreases are `O(log n)`;
/// inspecting the minimum is `O(1)`.
///
/// ```
/// use rheaps::AddressableHeap;
/// use rheaps::array::BinaryArrayAddressableHeap;
///
/// let mut heap = BinaryArrayAddressableHeap::new();
/// let task = heap.insert(10, "compile report");
/// heap.insert(5, "answer mail");
///
/// assert_eq!(heap.delete(task), Ok((10, "compile report")));
/// assert_eq!(heap.pop(), Some((5, "answer mail")));
/// ```
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct BinaryArrayAddressableHeap<K, V> {
    inner: AddressableCore<K, V>,
}

impl<K: Ord, V> BinaryArrayAddressableHeap<K, V> {
    /// Creates an empty heap.
    #[must_use]
    pub fn new() -> Self {
        Self::with_capacity(DEFAULT_HEAP_CAPACITY)
    }

    /// Creates an empty heap with storage for at least `capacity` entries.
    #[must_use]
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            inner: AddressableCore::new(capacity),
        }
    }

    /// Builds a heap from key-value pairs in linear time.
    #[must_use]
    pub fn from_vec(entries: Vec<(K, V)>) -> Self {
        Self {
            inner: AddressableCore::from_vec(entries, 2),
        }
    }
}

impl<K: Ord, V> Default for BinaryArrayAddressableHeap<K, V> {
    fn default() -> Self {
        Self::new()
    }
}

impl<K: Ord, V> FromIterator<(K, V)> for BinaryArrayAddressableHeap<K, V> {
    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
        Self::from_vec(iter.into_iter().collect())
    }
}

impl<K: Ord, V> Extend<(K, V)> for BinaryArrayAddressableHeap<K, V> {
    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
        for (key, value) in iter {
            self.insert(key, value);
        }
    }
}

impl<K: Ord, V> BinaryArrayAddressableHeap<K, V> {
    /// Inserts an entry and returns a handle that addresses it while live.
    pub fn insert(&mut self, key: K, value: V) -> AddressableHandle {
        self.inner.push(key, value, 2)
    }

    /// Returns the handle, key, and value of a minimum entry.
    #[must_use]
    pub fn peek_entry(&self) -> Option<(AddressableHandle, &K, &V)> {
        self.inner.peek()
    }

    /// Removes and returns a minimum key-value pair.
    pub fn pop_entry(&mut self) -> Option<(K, V)> {
        self.inner.pop(2)
    }

    /// Returns the key associated with `handle`.
    pub fn key(&self, handle: AddressableHandle) -> Result<&K, InvalidHandle> {
        self.inner.key(handle)
    }

    /// Returns the value associated with `handle`.
    pub fn value(&self, handle: AddressableHandle) -> Result<&V, InvalidHandle> {
        self.inner.value(handle)
    }

    /// Returns mutable access to the value associated with `handle`.
    pub fn value_mut(&mut self, handle: AddressableHandle) -> Result<&mut V, InvalidHandle> {
        self.inner.value_mut(handle)
    }

    /// Decreases an entry's key and restores heap order.
    pub fn decrease_key(
        &mut self,
        handle: AddressableHandle,
        key: K,
    ) -> Result<(), DecreaseKeyError> {
        self.inner.decrease_key(handle, key, 2)
    }

    /// Removes the entry associated with `handle`.
    pub fn delete(&mut self, handle: AddressableHandle) -> Result<(K, V), InvalidHandle> {
        self.inner.delete(handle, 2)
    }

    /// Returns handles for all live entries in unspecified heap order.
    pub fn handles(&self) -> impl Iterator<Item = AddressableHandle> + '_ {
        self.inner.handles()
    }

    /// Returns the number of live entries.
    #[must_use]
    pub fn len(&self) -> usize {
        self.inner.len()
    }

    /// Returns whether the heap contains no entries.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.inner.len() == 0
    }

    /// Removes all entries and invalidates every outstanding handle.
    pub fn clear(&mut self) {
        self.inner.clear();
    }
}

impl<K: Ord, V> AddressableHeap<K, V> for BinaryArrayAddressableHeap<K, V> {
    type Handle = AddressableHandle;

    fn insert(&mut self, key: K, value: V) -> Self::Handle {
        Self::insert(self, key, value)
    }

    fn peek(&self) -> Option<(Self::Handle, &K, &V)> {
        Self::peek_entry(self)
    }

    fn pop(&mut self) -> Option<(K, V)> {
        Self::pop_entry(self)
    }

    fn key(&self, handle: Self::Handle) -> Result<&K, InvalidHandle> {
        Self::key(self, handle)
    }

    fn value(&self, handle: Self::Handle) -> Result<&V, InvalidHandle> {
        Self::value(self, handle)
    }

    fn value_mut(&mut self, handle: Self::Handle) -> Result<&mut V, InvalidHandle> {
        Self::value_mut(self, handle)
    }

    fn delete(&mut self, handle: Self::Handle) -> Result<(K, V), InvalidHandle> {
        Self::delete(self, handle)
    }

    fn len(&self) -> usize {
        Self::len(self)
    }

    fn clear(&mut self) {
        Self::clear(self);
    }
}

impl<K: Ord, V> DecreaseKeyHeap<K, V> for BinaryArrayAddressableHeap<K, V> {
    fn decrease_key(&mut self, handle: Self::Handle, key: K) -> Result<(), DecreaseKeyError> {
        Self::decrease_key(self, handle, key)
    }
}

impl<K: Ord> BinaryArrayAddressableHeap<K, ()> {
    /// Inserts a key into this value-less heap and returns a checked handle.
    pub fn push(&mut self, key: K) -> AddressableHandle {
        self.insert(key, ())
    }

    /// Returns the minimum key, if present.
    #[must_use]
    pub fn peek(&self) -> Option<&K> {
        self.peek_entry().map(|(_, key, _)| key)
    }

    /// Removes and returns the minimum key, if present.
    pub fn pop(&mut self) -> Option<K> {
        self.pop_entry().map(|(key, ())| key)
    }
}

crate::impl_heap_via_addressable!(BinaryArrayAddressableHeap);

/// An array-backed d-ary min-heap with stable, checked entry handles.
///
/// The degree must be at least two. Larger degrees reduce insertion height but
/// increase comparisons during removal.
///
/// ```
/// use rheaps::AddressableHeap;
/// use rheaps::array::DaryArrayAddressableHeap;
///
/// let mut heap = DaryArrayAddressableHeap::new(4).unwrap();
/// let task = heap.insert(10, "compile report");
/// heap.insert(5, "answer mail");
///
/// assert_eq!(heap.degree(), 4);
/// assert_eq!(heap.delete(task), Ok((10, "compile report")));
/// assert_eq!(heap.pop(), Some((5, "answer mail")));
/// ```
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct DaryArrayAddressableHeap<K, V> {
    inner: AddressableCore<K, V>,
    degree: usize,
}

impl<K: Ord, V> DaryArrayAddressableHeap<K, V> {
    /// Creates an empty heap with `degree` children per node.
    pub fn new(degree: usize) -> Result<Self, InvalidDegree> {
        Self::with_capacity(degree, DEFAULT_HEAP_CAPACITY)
    }

    /// Creates an empty heap with storage for at least `capacity` entries.
    pub fn with_capacity(degree: usize, capacity: usize) -> Result<Self, InvalidDegree> {
        validate_degree(degree)?;
        Ok(Self {
            inner: AddressableCore::new(capacity),
            degree,
        })
    }

    /// Builds a heap from key-value pairs in linear time.
    pub fn from_vec(degree: usize, entries: Vec<(K, V)>) -> Result<Self, InvalidDegree> {
        validate_degree(degree)?;
        Ok(Self {
            inner: AddressableCore::from_vec(entries, degree),
            degree,
        })
    }
}

impl<K: Ord, V> Default for DaryArrayAddressableHeap<K, V> {
    fn default() -> Self {
        Self::new(2).expect("binary degree is valid")
    }
}

impl<K: Ord, V> FromIterator<(K, V)> for DaryArrayAddressableHeap<K, V> {
    /// Builds a binary (`degree = 2`) addressable d-ary heap from an iterator.
    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
        Self::from_vec(2, iter.into_iter().collect()).expect("binary degree is valid")
    }
}

impl<K: Ord, V> Extend<(K, V)> for DaryArrayAddressableHeap<K, V> {
    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
        for (key, value) in iter {
            self.insert(key, value);
        }
    }
}

impl<K: Ord, V> DaryArrayAddressableHeap<K, V> {
    /// Returns the number of children per node.
    #[must_use]
    pub const fn degree(&self) -> usize {
        self.degree
    }

    /// Inserts an entry and returns a handle that addresses it while live.
    pub fn insert(&mut self, key: K, value: V) -> AddressableHandle {
        self.inner.push(key, value, self.degree)
    }

    /// Returns the handle, key, and value of a minimum entry.
    #[must_use]
    pub fn peek_entry(&self) -> Option<(AddressableHandle, &K, &V)> {
        self.inner.peek()
    }

    /// Removes and returns a minimum key-value pair.
    pub fn pop_entry(&mut self) -> Option<(K, V)> {
        self.inner.pop(self.degree)
    }

    /// Returns the key associated with `handle`.
    pub fn key(&self, handle: AddressableHandle) -> Result<&K, InvalidHandle> {
        self.inner.key(handle)
    }

    /// Returns the value associated with `handle`.
    pub fn value(&self, handle: AddressableHandle) -> Result<&V, InvalidHandle> {
        self.inner.value(handle)
    }

    /// Returns mutable access to the value associated with `handle`.
    pub fn value_mut(&mut self, handle: AddressableHandle) -> Result<&mut V, InvalidHandle> {
        self.inner.value_mut(handle)
    }

    /// Decreases an entry's key and restores heap order.
    pub fn decrease_key(
        &mut self,
        handle: AddressableHandle,
        key: K,
    ) -> Result<(), DecreaseKeyError> {
        self.inner.decrease_key(handle, key, self.degree)
    }

    /// Removes the entry associated with `handle`.
    pub fn delete(&mut self, handle: AddressableHandle) -> Result<(K, V), InvalidHandle> {
        self.inner.delete(handle, self.degree)
    }

    /// Returns handles for all live entries in unspecified heap order.
    pub fn handles(&self) -> impl Iterator<Item = AddressableHandle> + '_ {
        self.inner.handles()
    }

    /// Returns the number of live entries.
    #[must_use]
    pub fn len(&self) -> usize {
        self.inner.len()
    }

    /// Returns whether the heap contains no entries.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.inner.len() == 0
    }

    /// Removes all entries and invalidates every outstanding handle.
    pub fn clear(&mut self) {
        self.inner.clear();
    }
}

impl<K: Ord, V> AddressableHeap<K, V> for DaryArrayAddressableHeap<K, V> {
    type Handle = AddressableHandle;

    fn insert(&mut self, key: K, value: V) -> Self::Handle {
        Self::insert(self, key, value)
    }

    fn peek(&self) -> Option<(Self::Handle, &K, &V)> {
        Self::peek_entry(self)
    }

    fn pop(&mut self) -> Option<(K, V)> {
        Self::pop_entry(self)
    }

    fn key(&self, handle: Self::Handle) -> Result<&K, InvalidHandle> {
        Self::key(self, handle)
    }

    fn value(&self, handle: Self::Handle) -> Result<&V, InvalidHandle> {
        Self::value(self, handle)
    }

    fn value_mut(&mut self, handle: Self::Handle) -> Result<&mut V, InvalidHandle> {
        Self::value_mut(self, handle)
    }

    fn delete(&mut self, handle: Self::Handle) -> Result<(K, V), InvalidHandle> {
        Self::delete(self, handle)
    }

    fn len(&self) -> usize {
        Self::len(self)
    }

    fn clear(&mut self) {
        Self::clear(self);
    }
}

impl<K: Ord, V> DecreaseKeyHeap<K, V> for DaryArrayAddressableHeap<K, V> {
    fn decrease_key(&mut self, handle: Self::Handle, key: K) -> Result<(), DecreaseKeyError> {
        Self::decrease_key(self, handle, key)
    }
}

impl<K: Ord> DaryArrayAddressableHeap<K, ()> {
    /// Inserts a key into this value-less heap and returns a checked handle.
    pub fn push(&mut self, key: K) -> AddressableHandle {
        self.insert(key, ())
    }

    /// Returns the minimum key, if present.
    #[must_use]
    pub fn peek(&self) -> Option<&K> {
        self.peek_entry().map(|(_, key, _)| key)
    }

    /// Removes and returns the minimum key, if present.
    pub fn pop(&mut self) -> Option<K> {
        self.pop_entry().map(|(key, ())| key)
    }
}

crate::impl_heap_via_addressable!(DaryArrayAddressableHeap);

fn validate_degree(degree: usize) -> Result<(), InvalidDegree> {
    if degree < 2 {
        Err(InvalidDegree(degree))
    } else {
        Ok(())
    }
}