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::{DynamicStorage, Storage};
use crate::{ChunkMut, DisjointSetRoView, Root, UnionResult};

/// Disjoint Set which allows dynamic addition of new items.
/// If you can estimate upper limit of number of elements in disjoint set,
/// it is probably better to use [`DisjointSet`](crate::DisjointSet)
/// because it doesn't have complexity for dynamic size changing.
pub struct DisjointSetDynamic(common_strct!(DynamicStorage));

#[allow(clippy::module_name_repetitions)]
pub struct DisjointSetDynamicRo(bits_enum!(DynamicStorage));

impl DisjointSetDynamic {
    /// Creates empty structure.
    #[must_use]
    pub fn new() -> Self {
        Self::with_capacity(0)
    }

    /// Creates empty structure with memory enough to store `cap` nodes without reallocation.
    #[must_use]
    pub fn with_capacity(cap: usize) -> Self {
        let storage = choose_by_size!(cap, {
            let s = DynamicStorage::with_capacity(cap);
            debug_assert_eq!(s.upper_bound(), 0);
            s
        });
        Self(unsafe {
            // SAFETY: Our storage is empty.
            Common::new(storage, Splittability::Allowed, PathsStatus::Compressed)
        })
    }

    /// Current amount of items.
    #[inline]
    #[must_use = "No need to call this method if result is unused."]
    pub fn len(&self) -> usize {
        self.0.len()
    }

    #[inline]
    #[must_use]
    pub fn is_empty(&self) -> bool {
        0 == self.len()
    }

    /// Total number of elements which structure can hold without reallocation.
    #[inline]
    #[must_use = "No need to call this method if result is unused."]
    pub fn capacity(&self) -> usize {
        for_current_bitness!(self.0.get_storage(); s; s.capacity())
    }

    /// Same as [`DisjointSet::get_root`](crate::DisjointSet::get_root).
    /// # Panics
    /// If index is out of bounds.
    #[must_use]
    pub fn get_root(&mut self, idx: usize) -> Root {
        self.0.get_root(idx)
    }

    /// Same as [`DisjointSet::union`](crate::DisjointSet::union).
    /// # Panics
    /// If any index is out of bounds.
    pub fn union(&mut self, idx0: usize, idx1: usize) -> UnionResult {
        self.0.union(idx0, idx1)
    }

    /// Same as [`DisjointSet::is_united`](crate::DisjointSet::is_united).
    /// # Panics
    /// If any index is out of bounds.
    #[must_use]
    pub fn is_united(&mut self, idx0: usize, idx1: usize) -> bool {
        self.0.is_united(idx0, idx1)
    }

    /// Same as [`DisjointSet::detach`](crate::DisjointSet::detach).
    /// # Panics
    /// If index is out of bounds.
    pub fn detach(&mut self, idx: usize) {
        self.0.detach(idx);
    }

    /// Same as [`DisjointSet::split_at`](crate::DisjointSet::split_at).
    #[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)
        }
    }

    /// Same as [`DisjointSet::make_ro_view`](crate::DisjointSet::make_ro_view).
    #[inline]
    #[must_use]
    pub fn make_ro_view(&mut self) -> DisjointSetRoView {
        self.0.make_ro_view()
    }

    /// Same as [`DisjointSet::into_readonly`](crate::DisjointSet::into_readonly).
    #[inline]
    #[must_use]
    pub fn into_readonly(self) -> DisjointSetDynamicRo {
        let s = self.0.into_compressed_storage();
        DisjointSetDynamicRo(s)
    }

    /// Same as [`DisjointSet::compress_paths`](crate::DisjointSet::compress_paths).
    #[inline]
    pub fn compress_paths(&mut self) {
        self.0.compress_paths();
    }

    /// Same as [`DisjointSet::reset`](crate::DisjointSet::reset).
    pub fn reset(&mut self) {
        self.0.reset();
    }

    /// Removes all values.
    pub fn clear(&mut self) {
        unsafe {
            // SAFETY: We correctly update flags.
            let (s, is_splittable, is_compressed) = self.0.access_internals();
            for_current_bitness!(s; s; s.clear());
            // This properties are always true for empty.
            *is_compressed = PathsStatus::Compressed;
            *is_splittable = Splittability::Allowed;
        }
    }

    /// Reserves memory so at least `additional`
    /// values can be added without reallocation.
    /// # Panics
    /// If `len + additional` usizes requires more than `isize::MAX / 2` bytes.
    #[inline]
    pub fn reserve(&mut self, additional: usize) {
        let s = unsafe {
            // SAFETY: We don't modify internal structure of roots and ranks,
            // but only update capacity.
            self.0.access_internals().0
        };
        let with_changed_bitness = for_current_bitness!(s; s; s.reserve(additional));
        if let Some(newly_allocated) = with_changed_bitness {
            *s = newly_allocated;
        }
    }

    /// Adds `added_size` new single-item subsets.
    /// If new size is larger than capacity, reallocates.
    /// New items would be unconnected to others.
    /// # Panics
    /// If `len + additional` usizes requires more than `isize::MAX / 2` bytes.
    #[inline]
    pub fn make_sets(&mut self, added_size: usize) {
        self.reserve(added_size);

        let s = unsafe {
            // SAFETY: We add new values to storage.
            // Those values would be splittable and compressed
            // those they don't affect flags for existing elements.
            self.0.access_internals().0
        };

        for_current_bitness!(s;s;{
            let mut newly_added = s.enlarge(added_size);
            // Our new memory can contain garbage because of clone_from.
            algorithm::initialize_storage::<_,true>(&mut newly_added);
        });
    }
}

impl Default for DisjointSetDynamic {
    fn default() -> Self {
        Self::new()
    }
}

impl Clone for DisjointSetDynamic {
    fn clone(&self) -> Self {
        let (is_splittable, is_compressed) = self.0.get_flags();
        let src = self.0.get_storage();

        let storage = for_current_bitness!(src; src; src.make_copied());

        Self(unsafe {
            // SAFETY: Because our storage contains same data as original,
            // maybe with different capacity and bitness,
            // its correctness and flags must not be changed.
            Common::new(storage, is_splittable, is_compressed)
        })
    }

    fn clone_from(&mut self, source: &Self) {
        if self.capacity() < source.len() {
            *self = source.clone();
            return;
        }

        let (src_splittable, src_compressed) = source.0.get_flags();
        let src = source.0.get_storage();

        let (dest, dst_splittable, dst_compressed) = unsafe { self.0.access_internals() };

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

        for_current_bitness!(
            dest; dest;
            for_current_bitness!(src; src; {
                dest.copy_data_from(src);
            })
        );

        *dst_splittable = src_splittable;
        *dst_compressed = src_compressed;
    }
}

impl From<DisjointSetDynamicRo> for DisjointSetDynamic {
    fn from(value: DisjointSetDynamicRo) -> Self {
        Self(unsafe {
            // SAFETY: `DisjointSetDynamicRo` cannot have uncompressed storage.
            Common::from_compressed_storage(value.0)
        })
    }
}

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

impl Clone for DisjointSetDynamicRo {
    fn clone(&self) -> Self {
        for_current_bitness!(&self.0;s;Self(s.make_copied()))
    }
}

#[cfg(test)]
mod tests {
    use crate::common::{PathsStatus, Splittability};
    use crate::{DisjointSetDynamic, Root};

    fn roots_vec(djs: &mut DisjointSetDynamic) -> Vec<Root> {
        (0..djs.len()).map(|i| djs.get_root(i)).collect()
    }

    #[test]
    fn test_enlarge() {
        let mut djs = DisjointSetDynamic::new();
        djs.make_sets(2);
        assert_eq!(djs.len(), 2);
        assert_eq!(djs.capacity(), 4);
        assert_eq!(djs.get_root(0), Root(0));
        assert_eq!(djs.get_root(1), Root(1));
        djs.union(0, 1);
        assert_eq!(djs.get_root(0), djs.get_root(1));

        let mut old_roots = roots_vec(&mut djs);
        djs.make_sets(3);
        assert_eq!(djs.len(), 5);
        assert_eq!(&old_roots, &roots_vec(&mut djs)[..2]);
        assert_eq!(djs.get_root(2), Root(2));
        assert_eq!(djs.get_root(3), Root(3));
        assert_eq!(djs.get_root(4), Root(4));

        old_roots = roots_vec(&mut djs);
        djs.make_sets(255);
        assert_eq!(&old_roots, &roots_vec(&mut djs)[..old_roots.len()]);
        for i in old_roots.len()..djs.len() {
            assert_eq!(Root(i), djs.get_root(i));
        }
    }

    #[test]
    fn test_old_values_are_removed() {
        let mut current = DisjointSetDynamic::with_capacity(10);
        current.make_sets(10);
        assert_eq!(current.len(), 10);

        for i in 0..current.len() {
            current.union(0, i);
            assert_eq!(current.get_root(i), Root(0));
        }

        let mut other = DisjointSetDynamic::with_capacity(7);
        other.make_sets(7);
        for i in 0..other.len() {
            other.union(6, i);
            assert_eq!(other.get_root(i), Root(6));
        }

        current.clone_from(&other);
        assert_eq!(current.len(), 7);
        assert_eq!(current.capacity(), 10);
        for i in 0..current.len() {
            assert_eq!(current.get_root(i), Root(6));
        }

        current.make_sets(3);
        for i in 7..current.len() {
            assert_eq!(current.get_root(i), Root(i));
        }

        current.clear();
        current.make_sets(10);
        for i in 0..current.len() {
            assert_eq!(current.get_root(i), Root(i));
        }
    }

    #[test]
    fn test_clone() {
        use PathsStatus::{Compressed, MaybeNotCompressed};
        use Splittability::{Allowed, Disallowed};

        fn make_from_params(
            size: usize,
            cap: usize,
            split: Splittability,
            comp: PathsStatus,
        ) -> DisjointSetDynamic {
            let mut sets = DisjointSetDynamic::with_capacity(cap);
            sets.make_sets(size);
            if split == Allowed {
                return sets;
            }
            sets.union(3, 4);
            sets.union(1, 2);
            sets.union(4, 2);
            sets.union(5, 6);
            if comp == Compressed {
                assert_ne!(sets.0.get_flags().1, Compressed);
                sets.compress_paths();
                assert_eq!(sets.0.get_flags().1, Compressed);
            }
            sets
        }

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

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

                    op(&mut left, &right);

                    if op_idx == 0 {
                        assert_eq!(left.capacity(), right_param.0);
                    } else {
                        assert_eq!(left.capacity(), std::cmp::max(left_param.1, right_param.0));
                    }
                    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));
                    }
                }
            }
        }
    }
}