path: root/crates/core/src/egraph/union_find.rs

//! Union-find / disjoint-set data structure implementations.

use alloc::{
    alloc::{Allocator, Global},
    collections::TryReserveError,
    vec::Vec,
};
use core::cmp;

use crate::id::Identifier;

//NOTE: maybe we should look into finding a way to bind identifiers to the
//      isize cap of vector sizes to avoid the error condition that would
//      result from saturating this bound?

/// Simple union-find implementation using union by min-id.
///
/// All [`Identifier`]s are assumed to be their own parents by default.
#[derive(Clone)]
#[cfg_attr(feature = "core-fmt", derive(Debug))]
pub struct UnionFind<Id, A = Global>
where
    A: Allocator + Clone,
{
    /// Vector mapping an e-class to its parent e-class.
    ///
    /// # Invariants
    ///
    /// For every stored parent `p`, `p.into_usize() < self.inner.len()`; and
    /// for every `i < self.inner.len()`, `Id::from_usize(i)` succeeds.
    pub(crate) inner: Vec<Id, A>,

    /// Allocator used by the [`UnionFind`].
    pub(crate) alloc: A,
}

impl<Id> UnionFind<Id, Global> {
    /// Constructs an empty union-find data structure.
    #[must_use]
    #[inline]
    pub fn new() -> Self {
        Self::new_in(Global)
    }
}

impl<Id, A> UnionFind<Id, A>
where
    A: Allocator + Clone,
{
    /// Constructs an empty union-find data structure using the provided
    /// [`Allocator`].
    #[must_use]
    #[inline]
    pub fn new_in(alloc: A) -> Self {
        Self {
            inner: Vec::new_in(alloc.clone()),
            alloc,
        }
    }
}

impl<Id, A> UnionFind<Id, A>
where
    A: Allocator + Clone,
    Id: Identifier,
{
    /// Helper method for converting a `usize` into an [`Identifier`],
    /// encoding our core assumption.
    #[must_use]
    #[inline(always)]
    fn index_to_id(id: usize) -> Id {
        debug_assert!(Id::from_usize(id).is_ok());

        //SAFETY: the e-class wouldn't be representable if a usize didn't fit
        //        into its id
        unsafe { Id::from_usize(id).unwrap_unchecked() }
    }

    /// Makes all represented e-classes their own parents.
    ///
    /// This does not free up any space; it only edits e-classes already
    /// represented.
    #[inline]
    pub fn reset(&mut self) {
        for (child, parent) in self.inner.iter_mut().enumerate() {
            *parent = Self::index_to_id(child);
        }
    }

    /// Removes all represented e-classes from the data structure.
    ///
    /// This does not return any used memory to the operating system; it only
    /// calls [`Vec::clear`].
    #[inline(always)]
    pub fn clear(&mut self) {
        self.inner.clear();
    }

    /// Shrinks the internal data structure so that the minimum amount of
    /// memory is used.
    #[inline(always)]
    pub fn shrink_to_fit(&mut self) {
        self.inner.shrink_to_fit();
    }

    /// Returns the id of the represented e-class with the greatest id.
    #[inline]
    pub fn greatest_represented(&self) -> Option<Id> {
        if let Some(id) = self.inner.len().checked_sub(1) {
            return Some(Self::index_to_id(id));
        }
        None
    }

    /// Indicates if no e-classes are represented.
    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }

    /// Reserves sufficient space in the union-find data structure to contain
    /// all of the given e-class ids.
    ///
    /// # Errors
    ///
    /// An error is returned if reserving space for new ids fails.
    #[inline(always)]
    pub fn ensure_contains_all(
        &mut self,
        ids: impl IntoIterator<Item = Id>,
    ) -> Result<(), TryReserveError> {
        if let Some(id) = ids.into_iter().max() {
            self.ensure_contains(id)?;
        }
        Ok(())
    }

    /// Reserves sufficient space in the union-find data structure to contain
    /// the given e-class id.
    ///
    /// # Errors
    ///
    /// An error is returned if reserving space for new ids fails.
    #[inline]
    pub fn ensure_contains(&mut self, id: Id) -> Result<(), TryReserveError> {
        let index = id.into_usize();

        if index < self.inner.len() {
            return Ok(());
        }

        let current_size = self.inner.len();

        //NOTE: we just established that index is at least equal to the
        //      length of the parent vector and `saturating_add` will prevent
        //      roll-over to 0 from `usize::MAX`
        self.inner.try_reserve(
            index.wrapping_sub(current_size).saturating_add(1),
        )?;

        for (parent, child) in self
            .inner
            .spare_capacity_mut()
            .iter_mut()
            .zip(current_size..=index)
        {
            parent.write(Self::index_to_id(child));
        }

        //NOTE: we know by the requirement that vectors (containing non-ZSTs)
        //      may have at most a size of `isize::MAX` that our index is
        //      nowhere near `usize::MAX`
        //SAFETY: we just wrote to the unused capacity
        unsafe {
            self.inner.set_len(index.wrapping_add(1));
        }

        Ok(())
    }

    /// Locates the canonical representation of the given e-class without
    /// performing path compression.
    #[inline]
    pub fn find(&self, id: Id) -> Id {
        if id.into_usize() >= self.inner.len() {
            return id;
        }

        //NOTE: we know all further indices are valid from this point on

        let mut current = id;
        loop {
            //SAFETY: refer to the above note
            let parent =
                *unsafe { self.inner.get_unchecked(current.into_usize()) };
            if current == parent {
                break current;
            }
            current = parent;
        }
    }

    /// Indicates if two e-class ids share the same parent.
    #[inline(always)]
    pub fn equivalent(&self, a: Id, b: Id) -> bool {
        self.find(a) == self.find(b)
    }

    /// Indicates if two e-class ids share the same parent, performing path
    /// compression.
    #[inline(always)]
    pub fn equivalent_mut(&mut self, a: Id, b: Id) -> bool {
        self.find_mut(a) == self.find_mut(b)
    }

    /// Locates the canonical representation of the given e-class, performing
    /// path compression.
    #[inline]
    pub fn find_mut(&mut self, id: Id) -> Id {
        if id.into_usize() >= self.inner.len() {
            return id;
        }

        let mut current = id;
        let base_ptr = self.inner.as_mut_ptr();

        //SAFETY: we know these indices are in bounds after the check above
        //        and due to invariants on the contents of the vector
        unsafe {
            loop {
                let parent_ptr = base_ptr.add(current.into_usize());
                let parent = *parent_ptr;

                if current == parent {
                    break current;
                }

                let grandparent = *base_ptr.add(parent.into_usize());
                parent_ptr.write(grandparent);
                current = grandparent;
            }
        }
    }

    /* TODO: implement this and the data structure it will use
    /// Locates the canonical representation of several e-classes, performing path compression.
    #[inline]
    pub fn find_mut_all(&mut self, ids: impl IntoIterator<Item = Id>) -> Result<FindMutAll<Id>, TryReserveError> {
        //NOTE: needs a try_extend or something to fallibly collect the iterator into a new vec and sort it. we can do this but it'd be more work in the utilities file
    }
    */

    /// Merges the given e-classes and returns their representatives,
    /// indicating which is the parent of the other.
    ///
    /// If either is not present and they differ, space is reserved for them
    /// both.
    ///
    /// # Errors
    ///
    /// An error is returned if reserving space for new ids fails.
    #[inline]
    pub fn merge(
        &mut self,
        a: Id,
        b: Id,
    ) -> Result<Union<Id>, TryReserveError> {
        let a = self.find_mut(a);
        let b = self.find_mut(b);

        let [parent, child] = cmp::minmax(a, b);

        if parent != child {
            self.ensure_contains(child)?;

            //SAFETY: because of invariants on the inner vector, and due to
            //        us reserving space if the id is not already contained,
            //        this is safe
            *unsafe { self.inner.get_unchecked_mut(child.into_usize()) } =
                parent;
        }

        Ok(Union { parent, child })
    }
}

impl<Id> Default for UnionFind<Id, Global> {
    fn default() -> Self {
        Self::new()
    }
}

/// Result of merging two e-classes.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct Union<Id> {
    /// The parent e-class.
    pub parent: Id,

    /// The child e-class.
    pub child: Id,
}

impl<Id> Union<Id>
where
    Id: Identifier,
{
    /// Indicates if an e-class merge was made in this union operation.
    #[must_use]
    #[inline(always)]
    pub fn changed(&self) -> bool {
        self.parent != self.child
    }
}

/// Iterator over canonical representations of e-classes.
pub struct FindMutAll<'a, Id, A>
where
    A: Allocator + Clone,
{
    /// The union find we are locating canonical representations within.
    pub(crate) union_find: &'a mut UnionFind<Id, A>,

    /// The `Id`s we are finding canonical representations of.
    pub(crate) ids: Vec<Id, A>,
}

#[cfg(all(test, feature = "core-error"))]
mod test {
    use super::*;

    #[test]
    fn union_find_reserve() {
        let mut union_find = UnionFind::default();

        union_find
            .ensure_contains(5usize)
            .expect("unable to reserve space to contain an identifier");

        assert_eq!(&[0, 1, 2, 3, 4, 5], union_find.inner.as_slice());

        union_find
            .ensure_contains(6usize)
            .expect("unable to reserve space to contain an identifier");

        assert_eq!(&[0, 1, 2, 3, 4, 5, 6], union_find.inner.as_slice());

        union_find
            .ensure_contains(5usize)
            .expect("unable to reserve space to contain an identifier");

        assert_eq!(&[0, 1, 2, 3, 4, 5, 6], union_find.inner.as_slice());

        union_find
            .ensure_contains(0usize)
            .expect("unable to reserve space to contain an identifier");

        assert_eq!(&[0, 1, 2, 3, 4, 5, 6], union_find.inner.as_slice());

        union_find
            .ensure_contains(usize::MAX)
            .expect_err("reserving greater than `isize::MAX` should fail");
    }
}