aph_disjoint_set 0.1.1

// This module contains common interface implementations.

use crate::macros::{bits_enum, check_bounds, common_strct, for_current_bitness};
use crate::storage::Storage;
use crate::tag_type::TagType;
use crate::{algorithm, ChunkMut, DisjointSetRoView, Root, UnionResult};

pub(crate) const UNSPLITTABLE_ERROR_MSG: &str =
    "It is not possible to split disjoint set after any `union` operation \
    because nodes from different parts can be connected.";

/// This struct maintain invariants:
/// 1. Roots and ranks are initialized.
/// 2. Any root value is valid index for both ranks and roots.
/// 3. If any node is a root for any other, its rank is not zero.
pub(crate) struct Common<
    TU8: Storage<u8>,
    TU16: Storage<u16>,
    #[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))] TU32: Storage<u32>,
    #[cfg(target_pointer_width = "64")] TU64: Storage<u64>,
> {
    storage: bits_enum!(TU8, TU16, TU32, TU64),
    // If it is possible to split it for multithreaded update.
    // Means that any value in `roots` points only to itself.
    is_splittable: Splittability,
    // If not, must compress paths before making readonly view.
    is_compressed: PathsStatus,
}

#[derive(Clone, Copy, PartialEq, Eq)]
#[cfg_attr(test, derive(Debug))]
pub(crate) enum Splittability {
    Disallowed,
    Allowed,
}

#[derive(Clone, Copy, PartialEq, Eq)]
#[cfg_attr(test, derive(Debug))]
pub(crate) enum PathsStatus {
    MaybeNotCompressed,
    Compressed,
}

impl<
        TU8: Storage<u8>,
        TU16: Storage<u16>,
        #[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))] TU32: Storage<u32>,
        #[cfg(target_pointer_width = "64")] TU64: Storage<u64>,
    > common_strct!(TU8, TU16, TU32, TU64)
{
    /// # Safety
    /// `storage` must be initialized.
    /// Caller must provide right flags.
    #[must_use]
    pub(crate) unsafe fn new(
        storage: bits_enum!(TU8, TU16, TU32, TU64),
        is_splittable: Splittability,
        is_compressed: PathsStatus,
    ) -> Self {
        Self {
            storage,
            is_splittable,
            is_compressed,
        }
    }

    /// # Safety
    /// `storage` must be got from [`Self::into_compressed_storage`].
    #[must_use]
    #[inline]
    pub(crate) unsafe fn from_compressed_storage(
        storage: bits_enum!(TU8, TU16, TU32, TU64),
    ) -> Self {
        Self {
            storage,
            // Making RO version is useful only after some union operations done
            // so we can assume that they were done.
            // So disjoint set is not splittable anymore.
            is_splittable: Splittability::Disallowed,
            is_compressed: PathsStatus::Compressed,
        }
    }

    /// # Safety
    /// Caller must maintain struct invariants.
    /// Caller must correctly update flags.
    #[inline]
    pub(crate) unsafe fn access_internals(
        &mut self,
    ) -> (
        &mut bits_enum!(TU8, TU16, TU32, TU64),
        &mut Splittability,
        &mut PathsStatus,
    ) {
        (
            &mut self.storage,
            &mut self.is_splittable,
            &mut self.is_compressed,
        )
    }

    #[inline]
    #[must_use]
    pub(crate) fn get_storage(&self) -> &bits_enum!(TU8, TU16, TU32, TU64) {
        &self.storage
    }

    #[inline]
    #[must_use]
    pub(crate) fn get_flags(&self) -> (Splittability, PathsStatus) {
        (self.is_splittable, self.is_compressed)
    }

    pub fn reset(&mut self) {
        if self.is_splittable == Splittability::Allowed {
            debug_assert!(self.is_compressed == PathsStatus::Compressed);
            return;
        }
        for_current_bitness!(&mut self.storage; s; algorithm::initialize_storage::<_, true>(s));
        self.is_compressed = PathsStatus::Compressed;
        self.is_splittable = Splittability::Allowed;
    }

    #[inline]
    #[must_use]
    pub(crate) fn len(&self) -> usize {
        for_current_bitness!(&self.storage; s; {
            assert_eq!(s.lower_bound(), 0, "len() must be called only for whole storage");
            s.upper_bound()
        })
    }

    #[must_use]
    pub(crate) fn get_root(&mut self, idx: usize) -> Root {
        let s = &mut self.storage;
        check_bounds!(multibit s; idx);
        let root = for_current_bitness!(s; s; unsafe {
            // SAFETY: We just checked idx bounds.
            // Properties of storage are maintained by invariant.
            algorithm::get_root(s, idx).as_u()
        });
        Root(root)
    }

    #[inline]
    #[must_use]
    pub(crate) fn union(&mut self, idx0: usize, idx1: usize) -> UnionResult {
        let s = &mut self.storage;
        check_bounds!(multibit s; idx0, idx1);
        let res = for_current_bitness!(
            s; s;
            unsafe {
                // SAFETY: We just checked idx bounds.
                // Properties of storage are maintained by invariant.
                algorithm::union(s, idx0, idx1)
            }
        );
        if res == UnionResult::Success {
            self.is_compressed = PathsStatus::MaybeNotCompressed;
        }
        self.is_splittable = Splittability::Disallowed;
        res
    }

    #[must_use]
    pub(crate) fn is_united(&mut self, idx0: usize, idx1: usize) -> bool {
        let s = &mut self.storage;
        check_bounds!(multibit s; idx0, idx1);
        for_current_bitness!(s; s; unsafe{
            // SAFETY: We just checked idx bounds.
            // Properties of storage are maintained by invariant.
            algorithm::is_united(s, idx0, idx1)
        })
    }

    pub(crate) fn detach(&mut self, idx: usize) {
        let s = &mut self.storage;
        check_bounds!(multibit s; idx);
        for_current_bitness!(s; s; unsafe{
            // SAFETY: We just checked idx bounds.
            // Properties of storage are maintained by invariant.
            algorithm::detach(s, idx);
        });
    }

    /// # Panics
    /// If any union operation was used.
    /// # Safety
    /// Returned chunks must not outlive allocation.
    #[must_use]
    pub(crate) unsafe fn split_at<'owner>(
        &mut self,
        split_idx: usize,
    ) -> (ChunkMut<'owner>, ChunkMut<'owner>)
    where
        Self: 'owner,
    {
        let is_splittable = self.is_splittable;
        let is_compressed = self.is_compressed;
        assert!(
            is_splittable == Splittability::Allowed,
            "{}",
            UNSPLITTABLE_ERROR_MSG
        );
        let s = &mut self.storage;
        check_bounds!(multibit s; split_idx);
        for_current_bitness!(s; s; {
            let (left, right) = unsafe{
                // SAFETY: Resulting lifetime is bound to self reference.
                s.split_at(split_idx)
            };
            // It is OK to construct this because if `is_splittable`
            // is true, all nodes point to themselves so any split is self-contained.
            (
                ChunkMut(Common{is_splittable, is_compressed, storage: CurrentBitness(left)}),
                ChunkMut(Common{is_splittable, is_compressed, storage: CurrentBitness(right)}),
            )
        })
    }

    #[inline]
    pub(crate) fn compress_paths(&mut self) {
        #[inline(never)]
        fn compress_paths_impl(storage: &mut bits_enum!(impl Storage)) {
            unsafe {
                // SAFETY: Validity of storage maintained by struct invariant.
                for_current_bitness!(storage; s; algorithm::compress_all_paths(s));
            }
        }

        // There are huge chance that this function may be called few times.
        // so we can run check without making function call.
        if self.is_compressed == PathsStatus::MaybeNotCompressed {
            compress_paths_impl(&mut self.storage);
        }
        // We make this assignment in inlined part of the method,
        // to allow optimizer perform Dead Code Elimination
        // if method called multiple times.
        self.is_compressed = PathsStatus::Compressed;
    }

    #[inline]
    #[must_use]
    pub(crate) fn make_ro_view(&mut self) -> DisjointSetRoView {
        debug_assert_eq!(0, for_current_bitness!(&self.storage; s; s.lower_bound()));
        self.compress_paths();
        let slice = for_current_bitness!(&self.storage;s;CurrentBitness(s.roots_readonly()));
        DisjointSetRoView(slice)
    }

    #[inline]
    #[must_use]
    pub(crate) fn into_compressed_storage(mut self) -> bits_enum!(TU8, TU16, TU32, TU64) {
        assert_eq!(0, for_current_bitness!(&self.storage; s; s.lower_bound()));
        self.compress_paths();
        self.storage
    }
}

#[allow(clippy::expl_impl_clone_on_copy)]
impl<
        TU8: Storage<u8> + Clone,
        TU16: Storage<u16> + Clone,
        #[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))] TU32: Storage<u32> + Clone,
        #[cfg(target_pointer_width = "64")] TU64: Storage<u64> + Clone,
    > Clone for common_strct!(TU8, TU16, TU32, TU64)
{
    #[inline]
    fn clone(&self) -> Self {
        Self {
            storage: self.storage.clone(),
            is_splittable: self.is_splittable,
            is_compressed: self.is_compressed,
        }
    }

    fn clone_from(&mut self, _: &Self) {
        unreachable!("Must not be called");
    }
}

impl<
        TU8: Storage<u8> + Copy,
        TU16: Storage<u16> + Copy,
        #[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))] TU32: Storage<u32> + Copy,
        #[cfg(target_pointer_width = "64")] TU64: Storage<u64> + Copy,
    > Copy for common_strct!(TU8, TU16, TU32, TU64)
{
}

#[allow(clippy::let_underscore_untyped, clippy::items_after_statements)]
#[cfg(test)]
mod tests {
    use rstest::rstest;

    use crate::algorithm;
    use crate::bits_enum::BitsEnum;
    use crate::macros::{common_strct, for_current_bitness};
    use crate::ro_view::DisjointSetRoView;
    use crate::storage::{ArrayStorage, BoxStorage, DynamicStorage, Storage};

    use super::Common;
    use super::PathsStatus::Compressed;
    use super::Splittability::Allowed;

    type DynStorage<T> = Box<dyn Storage<T>>;

    fn make_dyn_storage(size: usize) -> DynamicStorage<u8> {
        let mut r = DynamicStorage::with_capacity(size);
        r.enlarge(size);
        r
    }

    fn make_common_strt(mut d: DynStorage<u8>) -> common_strct!(DynStorage) {
        algorithm::initialize_storage::<_, false>(&mut d);
        unsafe { Common::new(BitsEnum::U8(d), Allowed, Compressed) }
    }

    #[rstest]
    fn test_readonly_view(
        #[values(
            Box::new(BoxStorage::new(7)),
            Box::new(make_dyn_storage(7)),
            Box::new(ArrayStorage::<_, 7>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);

        #[allow(unused_must_use)]
        {
            djs.union(0, 2);
            djs.union(3, 4);
            djs.union(0, 3);
            djs.union(1, 5);
            djs.union(1, 6);
        }

        let view = djs.make_ro_view();
        for (i, j) in [
            (0, 2),
            (3, 4),
            (0, 3),
            (0, 4),
            (2, 3),
            (2, 4),
            (1, 5),
            (1, 6),
            (5, 6),
        ] {
            assert!(view.is_united(i, j), "view.is_united({}, {})", i, j);
            assert!(view.is_united(j, i), "view.is_united({}, {})", j, i);
            assert_eq!(view.get_root(i), view.get_root(j));
        }

        for (i, j) in itertools::Itertools::cartesian_product([0, 2, 3, 4].into_iter(), [1, 5, 6]) {
            assert!(!view.is_united(i, j), "not view.is_united({}, {})", i, j);
            assert!(!view.is_united(j, i,), "not view.is_united({}, {})", j, i,);
            assert_ne!(view.get_root(i), view.get_root(j));
        }

        fn test_sync<T: Sync + Send>(_: T) {}
        test_sync(view);
    }

    #[rstest]
    fn test_into_readonly(
        #[values(
            Box::new(BoxStorage::new(7)),
            Box::new(make_dyn_storage(7)),
            Box::new(ArrayStorage::<_, 7>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);

        #[allow(unused_must_use)]
        {
            djs.union(0, 2);
            djs.union(3, 4);
            djs.union(0, 3);
            djs.union(1, 5);
            djs.union(1, 6);
        }

        const UNITED_INDICES: &[(usize, usize)] = &[
            (0, 2),
            (3, 4),
            (0, 3),
            (0, 4),
            (2, 3),
            (2, 4),
            (1, 5),
            (1, 6),
            (5, 6),
        ];

        let djs = djs.into_compressed_storage();
        let view = DisjointSetRoView(
            for_current_bitness!(&djs; djs; CurrentBitness(djs.roots_readonly())),
        );
        for &(i, j) in UNITED_INDICES {
            assert!(view.is_united(i, j), "view.is_united({}, {})", i, j);
            assert!(view.is_united(j, i), "view.is_united({}, {})", j, i);
            assert_eq!(view.get_root(i), view.get_root(j));
        }

        for (i, j) in itertools::Itertools::cartesian_product([0, 2, 3, 4].into_iter(), [1, 5, 6]) {
            assert!(!view.is_united(i, j), "not view.is_united({}, {})", i, j);
            assert!(!view.is_united(j, i,), "not view.is_united({}, {})", j, i,);
            assert_ne!(view.get_root(i), view.get_root(j));
        }

        fn test_sync<T: Sync + Send>(_: T) {}
        test_sync(view);
    }

    #[rstest]
    #[should_panic]
    fn test_bounds_check_get_root(
        #[values(8, 9, 10)] idx: usize,
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        let _ = djs.get_root(idx);
    }

    #[rstest]
    #[case(0, 8)]
    #[case(0, 9)]
    #[case(8, 0)]
    #[case(9, 0)]
    #[case(8, 9)]
    #[case(9, 8)]
    #[should_panic]
    fn test_bounds_check_union_djs(
        #[case] idx0: usize,
        #[case] idx1: usize,
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        let _ = djs.union(idx0, idx1);
    }

    #[rstest]
    #[case(0, 8)]
    #[case(0, 9)]
    #[case(8, 0)]
    #[case(9, 0)]
    #[case(8, 9)]
    #[case(9, 8)]
    #[should_panic]
    fn test_bounds_check_is_united_djs(
        #[case] idx0: usize,
        #[case] idx1: usize,
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED)
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        let _ = djs.is_united(idx0, idx1);
    }

    #[rstest]
    #[case(0, 8)]
    #[case(0, 9)]
    #[case(8, 0)]
    #[case(9, 0)]
    #[case(8, 9)]
    #[case(9, 8)]
    #[should_panic]
    fn test_bounds_check_union_view(
        #[case] idx0: usize,
        #[case] idx1: usize,
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED)
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        let view = djs.make_ro_view();
        let _ = view.is_united(idx0, idx1);
    }

    #[rstest]
    #[should_panic]
    fn test_bounds_check_get_root_view(
        #[values(8, 9, 10)] idx: usize,
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        let view = djs.make_ro_view();
        let _ = view.get_root(idx);
    }

    #[rstest]
    fn test_split(
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        for i in 0..8 {
            let (left, right) = unsafe { djs.split_at(i) };
            assert_eq!(left.lower_bound(), 0);
            assert_eq!(left.upper_bound(), i);
            assert_eq!(right.lower_bound(), i);
            assert_eq!(right.upper_bound(), 8);
        }
    }

    #[rstest]
    fn test_split_multithread(
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        let (mut left, mut right) = unsafe { djs.split_at(4) };
        std::thread::scope(|s| {
            s.spawn(|| left.union(1, 3));
            s.spawn(|| right.union(4, 5));
        });
        assert!(djs.is_united(1, 3));
        assert!(djs.is_united(4, 5));
    }

    #[rstest]
    #[should_panic]
    fn test_cannot_split_if_modified(
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        let _ = djs.union(5, 7);
        let _ = unsafe { djs.split_at(4) };
    }

    #[rstest]
    fn test_can_split_if_reset(
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        let _ = djs.union(5, 7);
        djs.reset();
        let _ = unsafe { djs.split_at(4) };
    }

    #[rstest]
    #[should_panic]
    fn test_split_bounds_check(
        #[values(9, 10, usize::MAX)] idx: usize,
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        let _ = unsafe { djs.split_at(idx) };
    }

    #[rstest]
    fn test_compress(
        #[values(
            Box::new(BoxStorage::new(8)),
            Box::new(make_dyn_storage(8)),
            Box::new(ArrayStorage::<_, 8>::ZEROED),
        )]
        storage: DynStorage<u8>,
    ) {
        let mut djs = make_common_strt(storage);
        #[allow(unused_must_use)]
        {
            djs.union(1, 2);
            djs.union(3, 4);
            djs.union(6, 7);
            djs.union(2, 4);
        }
        assert_ne!(djs.is_compressed, Compressed);
        assert_ne!(djs.is_splittable, Allowed);

        djs.compress_paths();
        assert_eq!(djs.is_compressed, Compressed);
        assert_ne!(djs.is_splittable, Allowed);

        djs.reset();
        assert_eq!(djs.is_compressed, Compressed);
        assert_eq!(djs.is_splittable, Allowed);
    }
}