aph_disjoint_set 0.1.0

use crate::algorithm;
use crate::common::{Common, PathsStatus, Splittability};
use crate::macros::{bits_enum, choose_by_size, common_strct, for_current_bitness};
use crate::storage::{BoxStorage, Storage};
use crate::tag_type::TagType;

use super::{ChunkMut, DisjointSetRoView, Root, UnionResult};

/// Disjoint Set with fixed size and integer tags.
/// Guarantees to use single allocation.
#[allow(clippy::len_without_is_empty)]
pub struct DisjointSet(common_strct!(BoxStorage));

#[derive(Clone)]
pub struct DisjointSetRo(bits_enum!(BoxStorage));

#[allow(clippy::semicolon_if_nothing_returned)]
impl DisjointSet {
    /// Creates new Disjoint Set.
    /// # Panics
    /// If provided `size` is too big (currently compared to `usize::MAX / 4 / size_of::<usize>()`).
    #[must_use]
    pub fn new(size: usize) -> Self {
        let storage = choose_by_size!(size, {
            let mut storage = BoxStorage::new(size);
            algorithm::initialize_storage::<_, false>(&mut storage);
            storage
        });
        Self(unsafe {
            // SAFETY: We just initialized storage to default state.
            Common::new(storage, Splittability::Allowed, PathsStatus::Compressed)
        })
    }

    /// Makes every item in disjoint set disjointed.
    /// After this call, for every i != j in range,
    /// `is_united(i, j) == false`.
    /// Equivalent to calling [`new`](Self::new) but reuses allocation.
    /// # Example
    ///
    /// ```
    /// use aph_disjoint_set::DisjointSet;
    /// let mut djs = DisjointSet::new(8);
    /// for i in 1..8{
    ///     djs.union(i-1, i);
    /// }
    /// for i in 0..8{
    ///     for j in 0..8 {
    ///         assert!(djs.is_united(i, j));
    ///     }
    /// }
    /// djs.reset();
    /// for i in 0..8{
    ///     for j in 0..8 {
    ///         assert_eq!(djs.is_united(i, j), i == j);
    ///     }
    /// }
    /// ```
    pub fn reset(&mut self) {
        self.0.reset();
    }

    /// Returns number of values.
    /// Same as argument ot [`new`](Self::new).
    #[must_use = "No need to call this method if result is unused."]
    #[inline]
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Returns market of subset which contains value with idx.
    /// Useful for counting of number of distinct sets.
    /// If you only need to check if 2 values are in same subset,
    /// better to use [`is_united`](Self::is_united) instead.
    ///
    /// # Panics
    /// If provided `idx` is out of bounds.
    ///
    /// # Algorithm
    /// It uses [Tarjan's path halving] optimization.
    ///
    /// [Tarjan's path halving]: https://en.wikipedia.org/wiki/Disjoint-set_data_structure#Finding_set_representatives
    #[must_use]
    pub fn get_root(&mut self, idx: usize) -> Root {
        self.0.get_root(idx)
    }

    /// Unites subsets which contain `idx0` and `idx1`.
    /// Uses union-by-rank algorithm.
    /// # Returns
    /// [`UnionResult::AlreadyJoined`] if `idx0` and `idx1` were already in same subset,
    /// [`UnionResult::Success`] otherwise.
    /// # Panics
    /// If either `idx0` or `idx1` is out of range.
    ///
    /// # Examples
    /// Finding minimal spanning tree using [Kruskal's algorithm].
    /// ```
    /// use aph_disjoint_set::{DisjointSet, UnionResult};
    ///
    /// // (node0, node1, cost)
    /// let mut edges = [
    ///      (0, 1, 4),
    ///      (0, 2, 0),
    ///      (1, 3, 1),
    ///      (1, 4, 1),
    ///      (2, 3, 5),
    ///      (2, 4, 1),
    ///      (2, 5, 0),
    ///      (3, 5, 0),
    ///      (4, 6, 0),
    ///      (5, 6, 2),
    /// ];
    /// edges.sort_by_key(|(_, _, cost)|*cost);
    /// let mut min_spanning_tree = Vec::new();
    /// let mut djs = DisjointSet::new(7);
    /// for (node0, node1, _) in edges {
    ///     if matches!(djs.union(node0, node1), UnionResult::Success) {
    ///         min_spanning_tree.push((node0, node1));
    ///     }
    /// }
    /// assert_eq!(
    ///     &min_spanning_tree,
    ///     &[(0, 2), (2, 5), (3, 5), (4, 6), (1, 3), (1, 4)]
    /// );
    /// ```
    ///
    /// [Kruskal's algorithm]: https://en.wikipedia.org/wiki/Kruskal%27s_algorithm
    pub fn union(&mut self, idx0: usize, idx1: usize) -> UnionResult {
        self.0.union(idx0, idx1)
    }

    /// Checks if `idx0` and `idx1` in same subset.
    /// # Panics
    /// If any index are bigger than length.
    #[must_use]
    pub fn is_united(&mut self, idx0: usize, idx1: usize) -> bool {
        self.0.is_united(idx0, idx1)
    }

    /// Detaches node at index to separate subset.
    /// # Panics
    /// If index is out of bounds.
    /// # Time complexity
    /// Detachment of single node can take *O(n)* time if there is any other node
    /// that points to it as parent but would take only *O(1)* time if it is a leaf.
    /// Because of properties of Union-Find algorithm,
    /// at least of half of the nodes would be leaves in worst case (and it is possible that it would be more).
    /// So, with probability 0.5 this operation would take *O(1)* time and with probability 0.5 — *O(n)*.
    /// If you need to detach many nodes, it would probably better to [`reset`](Self::reset) disjoint set
    /// and build it from the beginning, excluding unneeded nodes, which would take *O(n)* time for any realistic n,
    /// compared to _O(detachments * n)_ for consequitive detachments.
    pub fn detach(&mut self, idx: usize) {
        self.0.detach(idx)
    }

    /// Allows to split Disjoint Set into non-overlapping chunks
    /// which can be initialized indepentently.
    /// Useful when you want to run Union-Find algorithm using multiple threads.
    /// **NOTE:** Unlike `<[T]>::split_at`, it takes index relative to start of allocation.
    /// # Panics
    /// If `split_idx` is out of bounds.
    /// If any union operations were run.
    ///
    /// ```should_panic
    /// use aph_disjoint_set::DisjointSet;
    /// let mut djs = DisjointSet::new(10);
    /// djs.union(5,6);
    /// let _ = djs.split_at(3);
    /// ```
    #[must_use]
    pub fn split_at(&mut self, split_idx: usize) -> (ChunkMut, ChunkMut) {
        unsafe {
            // SAFETY: Returned lifetime is bound to reference to self.
            self.0.split_at(split_idx)
        }
    }

    /// Creates read-only view to the disjoint set which can be shared between threads.
    /// Since this method eagerly compress paths between nodes and roots,
    /// it can take *O(n)* time
    /// so it is not recommended to use unless you want to share it between threads.
    /// It would take only *O(1)* time if there weren't any calls to [`union`](Self::union)
    /// after either:
    /// 1. construction by [`new`](Self::new) or [`From::<DisjointSetRo>::from`];
    /// 2. previous call to `make_ro_view`.
    ///
    /// # Example
    /// ```
    /// use aph_disjoint_set::{DisjointSet, DisjointSetRoView};
    /// let mut djs = DisjointSet::new(5);
    /// djs.union(1,3);
    /// djs.union(0,2);
    /// djs.union(0,1);
    /// // This operation takes O(n) time.
    /// let _: DisjointSetRoView = djs.make_ro_view();
    /// // Since we already called `make_ro_view` once,
    /// // this takes *O(1)* time.
    /// let view: DisjointSetRoView = djs.make_ro_view();
    /// std::thread::scope(|sc|{
    ///     let root = view.get_root(0);
    ///     sc.spawn(move||{
    ///         assert_eq!(root, view.get_root(3));
    ///     });
    ///     sc.spawn(move||{
    ///         assert_ne!(root, view.get_root(4));
    ///     });
    /// });
    /// ```
    #[inline]
    #[must_use]
    pub fn make_ro_view(&mut self) -> DisjointSetRoView {
        self.0.make_ro_view()
    }

    /// See also [`make_ro_view`](Self::make_ro_view).
    /// This moves existing disjoint set into read-only form.
    /// Useful when you want to make disjoint set part of readonly structure after initialization.
    /// For other usecases, [`make_ro_view`](Self::make_ro_view) would probably better.
    /// This operation can take up to *O(n)* time,
    /// unless it is called on disjoint set which was constructed from [`DisjointSetRo`]
    /// which never was updated using [`union`](Self::union).
    ///
    /// # Example
    /// ```
    /// use aph_disjoint_set::{DisjointSet, DisjointSetRo};
    /// struct ContainsDisjointSet{
    ///    stored_ro: DisjointSetRo,
    /// }
    ///
    /// let mut djs0 = DisjointSet::new(5);
    /// djs0.union(1,3);
    /// djs0.union(0,2);
    /// djs0.union(0,1);
    /// // This operation takes O(n) time.
    /// let djs1: DisjointSetRo = djs0.into_readonly();
    /// // This operation takes O(1) time.
    /// let djs2: DisjointSet = djs1.into();
    /// // This operation takes O(1) time because we never updated djs2.
    /// let djs3 = djs2.into_readonly();
    /// let bigger_struct = ContainsDisjointSet{stored_ro:djs3};
    ///
    /// // We can `get_root` using immutable reference!
    /// assert_eq!(
    ///     bigger_struct.stored_ro.view().get_root(0),
    ///     bigger_struct.stored_ro.view().get_root(3),
    /// );
    /// // And we can share it between multiple threads:
    /// std::thread::scope(|sc|{
    ///     let root = bigger_struct.stored_ro.view().get_root(0);
    ///     let sro = &bigger_struct.stored_ro;
    ///     sc.spawn(move||{
    ///         assert_eq!(root, sro.view().get_root(3));
    ///     });
    ///     sc.spawn(move||{
    ///         assert_ne!(root, sro.view().get_root(4));
    ///     });
    /// });
    /// ```
    #[inline]
    #[must_use]
    pub fn into_readonly(self) -> DisjointSetRo {
        let s = self.0.into_compressed_storage();
        DisjointSetRo(s)
    }

    /// Eagerly apply path compression optimization.
    /// Has complexity _O(n)_.
    #[inline]
    pub fn compress_paths(&mut self) {
        self.0.compress_paths();
    }
}

impl Clone for DisjointSet {
    fn clone(&self) -> Self {
        Self(self.0.clone())
    }

    fn clone_from(&mut self, source: &Self) {
        #[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
        use crate::bits_enum::BitsEnum::U32;
        #[cfg(target_pointer_width = "64")]
        use crate::bits_enum::BitsEnum::U64;
        use crate::bits_enum::BitsEnum::{U16, U8};

        // This function needed only to make `match` look prettier
        #[inline]
        fn ss<T: TagType>(a: &BoxStorage<T>, b: &BoxStorage<T>) -> bool {
            (a.lower_bound() == b.lower_bound()) & (a.upper_bound() == b.upper_bound())
        }

        let (s, is_splittable, is_compressed) = unsafe {
            // SAFETY: We fill storage with copied bytes from original.
            // Since original storage is valid, copied would be valid too.
            // We update flags correctly.
            self.0.access_internals()
        };

        // Conservatively fill flags.
        // If we panic in the middle of storage copy,
        // this flags would prevent UB in other operations.
        *is_splittable = Splittability::Disallowed;
        *is_compressed = PathsStatus::MaybeNotCompressed;

        // Try to copy storage.
        match (s, source.0.get_storage()) {
            (U8(dest), U8(src)) if ss(dest, src) => dest.clone_from(src),
            (U16(dest), U16(src)) if ss(dest, src) => dest.clone_from(src),
            #[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
            (U32(dest), U32(src)) if ss(dest, src) => dest.clone_from(src),
            #[cfg(target_pointer_width = "64")]
            (U64(dest), U64(src)) if ss(dest, src) => dest.clone_from(src),
            _ => {
                // If we need reallocation, it is simpler to just clone everything.
                *self = source.clone();
                return;
            }
        }
        let (source_splittable, source_compressed) = source.0.get_flags();
        *is_splittable = source_splittable;
        *is_compressed = source_compressed;
    }
}

impl DisjointSetRo {
    #[inline]
    #[must_use]
    pub fn view(&self) -> DisjointSetRoView {
        let slice = for_current_bitness!(&self.0;s;CurrentBitness(s.roots_readonly()));
        DisjointSetRoView(slice)
    }
}

impl From<DisjointSetRo> for DisjointSet {
    #[inline]
    fn from(value: DisjointSetRo) -> Self {
        Self(unsafe {
            // SAFETY: `DisjointSetRo` can contain only compressed storage.
            Common::from_compressed_storage(value.0)
        })
    }
}

#[allow(clippy::let_underscore_untyped, clippy::items_after_statements)]
#[cfg(test)]
mod tests {
    use crate::bits_enum::BitsEnum;
    use crate::common::{PathsStatus, Splittability};
    use crate::DisjointSet;

    #[test]
    fn test_new() {
        let mut djs;

        djs = DisjointSet::new(255);
        assert_eq!(djs.len(), 255);
        assert!(matches!(djs.0.get_storage(), &BitsEnum::U8(_)));

        djs = DisjointSet::new(256);
        assert_eq!(djs.len(), 256);
        assert!(matches!(djs.0.get_storage(), &BitsEnum::U16(_)));

        assert_eq!(
            djs.0.get_flags(),
            (Splittability::Allowed, PathsStatus::Compressed)
        );
        for i in 0..djs.len() {
            assert_eq!(djs.get_root(i).into_inner(), i);
        }
    }

    #[test]
    fn test_clone() {
        fn make_from_params(size: usize, split: Splittability, comp: PathsStatus) -> DisjointSet {
            let mut sets = DisjointSet::new(size);
            if split == Splittability::Allowed {
                return sets;
            }
            sets.union(3, 4);
            sets.union(1, 2);
            sets.union(4, 2);
            sets.union(5, 6);
            if comp == PathsStatus::Compressed {
                assert_ne!(sets.0.get_flags().1, PathsStatus::Compressed);
                sets.compress_paths();
                assert_eq!(sets.0.get_flags().1, PathsStatus::Compressed);
            }
            sets
        }

        let params = [
            (7, Splittability::Allowed, PathsStatus::Compressed),
            (10, Splittability::Allowed, PathsStatus::Compressed),
            (7, Splittability::Disallowed, PathsStatus::Compressed),
            (10, Splittability::Disallowed, PathsStatus::Compressed),
            (7, Splittability::Allowed, PathsStatus::MaybeNotCompressed),
            (10, Splittability::Allowed, PathsStatus::MaybeNotCompressed),
            (300, Splittability::Allowed, PathsStatus::MaybeNotCompressed),
        ];

        let ops: [fn(&mut DisjointSet, &DisjointSet); 2] =
            [|a, b| *a = b.clone(), Clone::clone_from];
        for &left_param in params.iter() {
            for &right_param in params.iter() {
                for op in ops.iter() {
                    let mut left = make_from_params(left_param.0, left_param.1, left_param.2);
                    let mut right = make_from_params(right_param.0, right_param.1, right_param.2);

                    op(&mut left, &right);

                    assert_eq!(left.len(), right.len());
                    assert_eq!(left.0.get_flags(), right.0.get_flags());
                    for i in 0..left.len() {
                        assert_eq!(left.get_root(i), right.get_root(i));
                    }
                }
            }
        }
    }
}