aph_disjoint_set 0.1.1

use core::cmp::Ordering;

use crate::storage::Storage;
use crate::tag_type::TagType;
use crate::UnionResult;

/// Initalize roots by values from lower bound to upper bound.
/// If `FILL_RANKS` are true, would also fill ranks with zeros.
/// You MUST provide `FILL_RANKS=true` unless storage filled with zeros already.
pub(crate) fn initialize_storage<Tag: TagType, const FILL_RANKS: bool>(
    storage: &mut impl Storage<Tag>,
) {
    let lower_bound = storage.lower_bound();
    let upper_bound = storage.upper_bound();
    let (roots, ranks) = storage.roots_and_ranks();
    debug_assert_eq!(upper_bound - lower_bound, roots.len());
    debug_assert_eq!(upper_bound - lower_bound, ranks.len());
    debug_assert!(Tag::MAX_VAL.as_u() >= upper_bound - 1);

    let mut current = Tag::from_u(lower_bound);
    for val in roots.iter_mut() {
        *val = current;
        current = current.plus_one();
    }
    if FILL_RANKS {
        ranks.fill(TagType::ZERO);
    }
}

/// # Safety
/// Caller must ensure that Storage are properly initialized storage (see `initialize_storage`).
/// Caller must ensure that `storage.lower_bound() <= idx and idx < storage.upper_bound()`.
/// Caller must ensure that all roots in storage inside is limited to the storage
/// which would not be true only if not freshly initialized storage is splitted.
/// # Algorithm
/// Implemented by using Tarjan's path halving technique.
/// <https://dl.acm.org/doi/10.1145/62.2160/>
pub(crate) unsafe fn get_root<Tag: TagType>(storage: &mut impl Storage<Tag>, idx: usize) -> Tag {
    debug_assert!(storage.lower_bound() <= idx && idx < storage.upper_bound());
    let lower_bound = storage.lower_bound();
    let mut curr = idx - lower_bound;

    let (roots, _) = storage.roots_and_ranks();
    loop {
        unsafe {
            // SAFETY: Validity of original idx is ensured by the caller.
            // For new `curr`s see lower
            let parent = *roots.get_unchecked(curr);
            // SAFETY: Any root value stored in `roots` are valid index + lower_bound
            // if Storage properly implemented and initialized.
            let pp = *roots.get_unchecked(parent.as_u() - lower_bound);
            if parent == pp {
                return parent;
            }
            *roots.get_unchecked_mut(curr) = pp;
            curr = parent.as_u() - lower_bound;
        }
    }
}

/// Joins subsets which contains idx0 and idx1.
/// Implements union-by-rank. If sets have different ranks, smaller joins to bigger one.
/// If sets have same rank, then we choose first as root of second and increment its rank.
/// Rank is upper bound of distance from root to deepest leaf of its tree.
/// # Safety
/// Both indices must be in `[storage.lower_bound()..storage.upper_bound())`.
/// Storage must be properly initialized and modified only by algorithms from this module.
/// Caller must ensure that all roots in storage inside is limited to the storage
/// which would not be true only if not freshly initialized storage is splitted.
pub(crate) unsafe fn union<Tag: TagType>(
    storage: &mut impl Storage<Tag>,
    idx0: usize,
    idx1: usize,
) -> UnionResult {
    if idx0 == idx1 {
        return UnionResult::AlreadyJoined;
    }
    let root0 = unsafe {
        // SAFETY: Argument must be valid by function requirements.
        get_root(storage, idx0)
    };
    let root1 = unsafe {
        // SAFETY: Argument must be valid by function requirements.
        get_root(storage, idx1)
    };
    if root0 == root1 {
        return UnionResult::AlreadyJoined;
    }

    let lower_bound = storage.lower_bound();
    let (all_roots, all_ranks) = storage.roots_and_ranks();

    let r_idx0 = root0.as_u() - lower_bound;
    let r_idx1 = root1.as_u() - lower_bound;

    let rank0 = unsafe {
        // SAFETY: Any (root - lower_bound) < len.
        *all_ranks.get_unchecked(r_idx0)
    };
    let rank1 = unsafe {
        // SAFETY: Any (root - lower_bound) < len.
        *all_ranks.get_unchecked(r_idx1)
    };

    unsafe {
        // SAFETY: Any (root - lower_bound) < len.
        match rank0.cmp(&rank1) {
            Ordering::Less => *all_roots.get_unchecked_mut(r_idx0) = root1,
            Ordering::Equal => {
                *all_roots.get_unchecked_mut(r_idx1) = root0;
                // Increment rank by 1 because subtree starting from root1 have depth rank1,
                // so depth of root0 is rank1 + 1.
                *all_ranks.get_unchecked_mut(r_idx0) = rank1.plus_one();
            }
            Ordering::Greater => *all_roots.get_unchecked_mut(r_idx1) = root0,
        }
    }
    UnionResult::Success
}

/// # Safety
/// Both arguments must be in bounds of storage.
/// Storage must be properly initialized and modified only by algorithms from this module.
/// Caller must ensure that all roots in storage inside is limited to the storage
/// which would not be true only if not freshly initialized storage is splitted.
pub(crate) unsafe fn is_united<Tag: TagType>(
    storage: &mut impl Storage<Tag>,
    idx0: usize,
    idx1: usize,
) -> bool {
    idx0 == idx1
        || unsafe {
            // SAFETY: Caller must ensure that indices are valid.
            get_root(storage, idx0) == get_root(storage, idx1)
        }
}

/// # Safety
/// `idx` must be in-bounds of storage.
/// Caller must ensure that all roots in storage inside is limited to the storage
/// which would not be true only if not freshly initialized storage is splitted.
pub(crate) unsafe fn detach<Tag: TagType>(storage: &mut impl Storage<Tag>, idx: usize) {
    debug_assert!(idx >= storage.lower_bound() && idx < storage.upper_bound());
    let lower_bound = storage.lower_bound();
    let detach_tag = Tag::from_u(idx);
    let idx = idx - lower_bound;
    let (all_roots, all_ranks) = storage.roots_and_ranks();
    let (prev_root, prev_rank) = unsafe {
        // SAFETY: Caller must ensure valid idx.
        let prev_root = *all_roots.get_unchecked(idx);
        let prev_rank = *all_ranks.get_unchecked(idx);
        if prev_rank == Tag::ZERO {
            // Lucky O(1) case.
            // We were leaf so there is no need to update any other node.
            // We can be root of subset with exactly 1 element but in that
            // case we are already detached.
            *all_roots.get_unchecked_mut(idx) = detach_tag;
            return;
        }
        (prev_root, prev_rank)
    };

    // We can reach here only if our node was not a leaf.
    // There is a possibility that we are root of some subtree in graph.

    // So let's change all roots of any nodes which target ourselves.
    let new_root = if prev_root == detach_tag {
        // `idx` happened to be the root of our tree
        // so we need to find some our child to replace
        // idx as root of this tree.
        let mut new_root_idx = 0;
        unsafe {
            // SAFETY:
            // 1. All roots are valid indices.
            // 2. Since our rank is more than 0, there must be some other node
            // which points to `idx` as its root.
            // 3. So we would find such root before hitting end of the slice here.
            while *all_roots.get_unchecked(new_root_idx) != detach_tag || new_root_idx == idx {
                new_root_idx += 1;
                debug_assert!(new_root_idx < all_roots.len());
            }
            debug_assert!(
                all_ranks[new_root_idx] < prev_rank,
                "Rank of the child must be smaller",
            );
            *all_ranks.get_unchecked_mut(new_root_idx) = prev_rank;
        }
        Tag::from_u(new_root_idx + lower_bound)
    } else {
        // Our own root is best solution.
        // We need to connect our children to our parent.
        prev_root
    };

    // Unfortunate O(n) loop here.
    for v in all_roots.iter_mut() {
        if *v == detach_tag {
            *v = new_root;
        }
    }
    // Actual detachment.
    // We do it only after whole loops to prevent data modification
    // in case if some debug assert fires before.
    // At this point, it is guaranteed that the node is a leaf.
    unsafe {
        // SAFETY: idx validity guaranteed by caller.
        *all_roots.get_unchecked_mut(idx) = detach_tag;
        // We are detached now so our rank must be zero.
        *all_ranks.get_unchecked_mut(idx) = Tag::ZERO;
    }
}

/// Must be called before using any readonly functions.
/// Have complexity O(n).
/// It would still be compressed until first call to `Self::union`.
/// Disjoint set would be compressed after calling `initialize_storage` as well.
/// # Safety
/// Caller must ensure that all roots in storage inside is limited to the storage
/// which would not be true only if not freshly initialized storage is splitted.
pub(crate) unsafe fn compress_all_paths<Tag: TagType>(storage: &mut impl Storage<Tag>) {
    let lower_bound = storage.lower_bound();
    let upper_bound = storage.upper_bound();
    for idx in lower_bound..upper_bound {
        // SAFETY: Caller must pass valid storage.
        // idx is in bounds by loop construction.
        let root = unsafe { get_root(storage, idx) };
        let (all_roots, _) = storage.roots_and_ranks();
        let root_ref = unsafe {
            // SAFETY: idx is in bounds by loop construction.
            all_roots.get_unchecked_mut(idx - lower_bound)
        };
        if *root_ref != root {
            *root_ref = root;
        }
    }

    // Reset ranks of leaves to make `detach` operation fast more often.
    let (all_roots, all_ranks) = storage.roots_and_ranks();
    for ((root_tag, &real_tag), rank) in Iterator::zip(
        Iterator::zip(
            (lower_bound..upper_bound).map(Tag::from_u),
            all_roots.iter(),
        ),
        all_ranks.iter_mut(),
    ) {
        // After compression of paths, all nodes are either leaf or root,
        // so if node's root is not itself, it is a leaf.
        if root_tag != real_tag {
            *rank = Tag::ZERO;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{detach, get_root, initialize_storage, is_united, union};
    use crate::algorithm::compress_all_paths;
    use crate::storage::{BoxStorage, Storage};
    use crate::UnionResult;

    fn make_storage() -> impl Storage<u8> {
        let mut storage = BoxStorage::new(8);
        initialize_storage::<_, false>(&mut storage);
        storage
    }

    #[test]
    fn test_initialize() {
        {
            let mut storage = make_storage();
            let (roots, ranks) = storage.roots_and_ranks();
            assert_eq!(roots, &[0, 1, 2, 3, 4, 5, 6, 7]);
            assert_eq!(ranks, &[0; 8]);
        }

        {
            let mut storage = BoxStorage::new(8);
            let (roots, ranks) = storage.roots_and_ranks();
            roots.fill(128);
            ranks.fill(50);
            initialize_storage::<u8, true>(&mut storage);
            let (roots, ranks) = storage.roots_and_ranks();
            assert_eq!(roots, &[0, 1, 2, 3, 4, 5, 6, 7]);
            assert_eq!(ranks, &[0; 8]);
        }
    }

    #[test]
    fn test_get_root() {
        for key in [0, 1, 2, 3, 4, 5, 6, 7] {
            let mut storage = make_storage();

            let (roots, ranks) = storage.roots_and_ranks();
            assert_eq!(roots.len(), 8);
            assert_eq!(ranks.len(), 8);
            roots.copy_from_slice(&[0, 0, 1, 2, 3, 4, 5, 6]);

            unsafe {
                assert_eq!(get_root(&mut storage, key), 0);
                assert_eq!(get_root(&mut storage, key), 0);
            }
        }

        for (key, value) in [
            (0, 2),
            (1, 2),
            (2, 2),
            (3, 3),
            (4, 3),
            (5, 3),
            (6, 6),
            (7, 7),
        ] {
            let mut storage = make_storage();

            let (roots, ranks) = storage.roots_and_ranks();
            assert_eq!(roots.len(), 8);
            assert_eq!(ranks.len(), 8);
            roots.copy_from_slice(&[1, 2, 2, 3, 3, 4, 6, 7]);

            unsafe {
                assert_eq!(get_root(&mut storage, key), value);
                assert_eq!(get_root(&mut storage, key), value);
            }
        }
    }

    #[test]
    fn test_ranks() {
        let mut storage = make_storage();

        for i in 0..8 {
            unsafe {
                union(&mut storage, 0, i);
            }
        }
        assert_eq!(storage.roots_and_ranks().1, &[1, 0, 0, 0, 0, 0, 0, 0]);

        storage = make_storage();
        for i in 0..4 {
            unsafe {
                union(&mut storage, i * 2, i * 2 + 1);
            }
        }
        assert_eq!(storage.roots_and_ranks().1, &[1, 0, 1, 0, 1, 0, 1, 0]);
        for i in (0..7).step_by(4) {
            unsafe {
                union(&mut storage, i, i + 2);
            }
        }
        assert_eq!(storage.roots_and_ranks().1, &[2, 0, 1, 0, 2, 0, 1, 0]);
        unsafe {
            union(&mut storage, 0, 4);
        }
        assert_eq!(storage.roots_and_ranks().1, &[3, 0, 1, 0, 2, 0, 1, 0]);
    }

    #[test]
    fn test_unions() {
        unsafe {
            let mut storage = make_storage();
            for f in (0..8).step_by(2) {
                let s = f + 1;
                assert!(!is_united(&mut storage, f, s));
                assert!(matches!(union(&mut storage, f, s), UnionResult::Success));
                assert!(is_united(&mut storage, f, s));
            }

            for f in (0..8).step_by(2) {
                let s = f + 1;
                assert!(is_united(&mut storage, f, s));
                assert!(matches!(
                    union(&mut storage, f, s),
                    UnionResult::AlreadyJoined
                ));
            }

            assert!(!is_united(&mut storage, 1, 2));
            assert!(matches!(union(&mut storage, 0, 3), UnionResult::Success));
            assert!(matches!(
                union(&mut storage, 1, 2),
                UnionResult::AlreadyJoined
            ));
        }
    }

    #[test]
    fn test_detach() {
        const ALL_CONNS: &[(usize, usize)] = &[
            (0, 1),
            (0, 2),
            (0, 3),
            (1, 2),
            (1, 3),
            (2, 3),
            (4, 5),
            (4, 6),
            (5, 6),
        ];
        fn check_no_unexpected_conn(djs: &mut impl Storage<u8>, case: &str) {
            for i in djs.lower_bound()..djs.upper_bound() {
                for j in i + 1..djs.upper_bound() {
                    if ALL_CONNS.contains(&(i, j)) {
                        continue;
                    }
                    unsafe {
                        assert!(
                            !is_united(djs, i, j),
                            "Unexpected connection after {}",
                            case
                        );
                    }
                }
            }
        }
        let make_djs = || unsafe {
            let mut djs = make_storage();
            for (l, r) in [(0, 1), (2, 3), (0, 2), (4, 5), (4, 6)] {
                union(&mut djs, l, r);
            }
            djs
        };

        // Detach already detached.
        unsafe {
            let mut djs = make_djs();
            detach(&mut djs, 7);

            for i in (djs.lower_bound()..djs.upper_bound()).filter(|&i| i != 7) {
                assert!(!is_united(&mut djs, i, 7));
            }
            for &(i, j) in ALL_CONNS {
                assert!(is_united(&mut djs, i, j));
            }
            check_no_unexpected_conn(&mut djs, "already detached");
        }
        // Detach leaf.
        unsafe {
            let mut djs = make_djs();
            detach(&mut djs, 3);
            for i in (djs.lower_bound()..djs.upper_bound()).filter(|&i| i != 3) {
                assert!(!is_united(&mut djs, i, 3));
            }
            for &(i, j) in ALL_CONNS {
                if i == 3 || j == 3 {
                    continue;
                }
                assert!(is_united(&mut djs, i, j));
            }
            check_no_unexpected_conn(&mut djs, "leaf");
        }

        // Detach in the middle
        unsafe {
            let mut djs = make_djs();
            detach(&mut djs, 2);
            for i in (djs.lower_bound()..djs.upper_bound()).filter(|&i| i != 2) {
                assert!(!is_united(&mut djs, i, 2));
            }
            for &(i, j) in ALL_CONNS {
                if i == 2 || j == 2 {
                    continue;
                }
                assert!(is_united(&mut djs, i, j));
            }
            check_no_unexpected_conn(&mut djs, "middle");
        }
        // Detach root.
        unsafe {
            let mut djs = make_djs();
            detach(&mut djs, 0);
            for i in (djs.lower_bound()..djs.upper_bound()).filter(|&i| i != 0) {
                assert!(!is_united(&mut djs, i, 0));
            }
            for &(i, j) in ALL_CONNS {
                if i == 0 || j == 0 {
                    continue;
                }
                assert!(is_united(&mut djs, i, j));
            }
            check_no_unexpected_conn(&mut djs, "root");
        }
    }

    #[test]
    fn test_compress_paths() {
        let mut djs = make_storage();
        unsafe {
            union(&mut djs, 0, 2);
            union(&mut djs, 3, 4);
            union(&mut djs, 0, 3);
            union(&mut djs, 1, 5);
            union(&mut djs, 1, 6);
        }
        // It is unexpected to get compressed storage right after unions.
        assert_ne!(djs.roots_and_ranks().0, &[0, 1, 0, 0, 0, 1, 1, 7]);
        assert_ne!(djs.roots_and_ranks().1, &[2, 1, 0, 0, 0, 0, 0, 0]);

        unsafe {
            compress_all_paths(&mut djs);
        }
        assert_eq!(djs.roots_and_ranks().0, &[0, 1, 0, 0, 0, 1, 1, 7]);
        assert_eq!(djs.roots_and_ranks().1, &[2, 1, 0, 0, 0, 0, 0, 0]);
    }
}

#[cfg(test)]
mod splitted_storage_tests {
    use super::{detach, get_root, initialize_storage, is_united, union};
    use crate::algorithm::compress_all_paths;
    use crate::storage::{BoxStorage, Storage};
    use crate::UnionResult;

    #[test]
    fn test_initialize() {
        unsafe {
            let mut storage: BoxStorage<u8> = BoxStorage::new(8);
            let (mut left, mut right) = storage.split_at(5);
            initialize_storage::<u8, true>(&mut left);
            initialize_storage::<u8, true>(&mut right);
            let (roots, ranks) = storage.roots_and_ranks();
            assert_eq!(roots, &[0, 1, 2, 3, 4, 5, 6, 7]);
            assert_eq!(ranks, &[0; 8]);
        }

        unsafe {
            let mut storage = BoxStorage::new(8);
            let (roots, ranks) = storage.roots_and_ranks();
            roots.fill(128);
            ranks.fill(50);

            initialize_storage::<u8, true>(&mut storage.split_at(4).0);
            let (roots, ranks) = storage.roots_and_ranks();
            assert_eq!(roots, &[0, 1, 2, 3, 128, 128, 128, 128]);
            assert_eq!(ranks, &[0, 0, 0, 0, 50, 50, 50, 50]);

            initialize_storage::<u8, false>(&mut storage.split_at(4).1);
            let (roots, ranks) = storage.roots_and_ranks();
            assert_eq!(roots, &[0, 1, 2, 3, 4, 5, 6, 7]);
            assert_eq!(ranks, &[0, 0, 0, 0, 50, 50, 50, 50]);
        }
    }

    fn make_storage() -> impl Storage<u8> {
        let mut storage = BoxStorage::new(8);
        initialize_storage::<_, false>(&mut storage);
        storage
    }

    #[test]
    fn test_get_root() {
        for key in [4, 5, 6, 7] {
            let mut storage = make_storage();

            let (roots, ranks) = storage.roots_and_ranks();
            assert_eq!(roots.len(), 8);
            assert_eq!(ranks.len(), 8);
            roots.copy_from_slice(&[0, 1, 2, 3, 4, 4, 5, 6]);
            let (_, mut part) = unsafe { storage.split_at(4) };

            unsafe {
                assert_eq!(get_root(&mut part, key), 4);
                assert_eq!(get_root(&mut part, key), 4);
            }
        }

        for (key, value) in [(4, 5), (5, 5), (6, 5), (7, 7)] {
            let mut storage = make_storage();

            let (roots, ranks) = storage.roots_and_ranks();
            assert_eq!(roots.len(), 8);
            assert_eq!(ranks.len(), 8);
            roots.copy_from_slice(&[0, 1, 2, 3, /* */ 5, 5, 4, 7]);
            let (_, mut part) = unsafe { storage.split_at(4) };

            unsafe {
                assert_eq!(get_root(&mut part, key), value);
                assert_eq!(get_root(&mut part, key), value);
            }
        }
    }

    #[test]
    fn test_ranks() {
        let mut storage = make_storage();
        let mut part = unsafe { storage.split_at(4).1 };

        for i in 4..8 {
            unsafe {
                union(&mut part, 4, i);
            }
        }
        assert_eq!(storage.roots_and_ranks().1, &[0, 0, 0, 0, 1, 0, 0, 0]);

        storage = make_storage();
        let mut part = unsafe { storage.split_at(4).1 };
        for i in (part.lower_bound()..part.upper_bound()).step_by(2) {
            unsafe {
                union(&mut part, i, i + 1);
            }
        }
        assert_eq!(part.roots_and_ranks().1, &[1, 0, 1, 0]);
        unsafe {
            union(&mut part, 5, 7);
        }
        assert_eq!(part.roots_and_ranks().1, &[2, 0, 1, 0]);
        assert_eq!(storage.roots_and_ranks().1, &[0, 0, 0, 0, 2, 0, 1, 0]);
    }

    #[test]
    fn test_unions() {
        unsafe {
            let mut st = make_storage();
            let (_, mut right) = st.split_at(2);
            for f in (right.lower_bound()..right.upper_bound()).step_by(2) {
                let s = f + 1;
                assert!(!is_united(&mut right, f, s));
                assert!(matches!(union(&mut right, f, s), UnionResult::Success));
                assert!(is_united(&mut right, f, s));
            }

            for f in (2..8).step_by(2) {
                let s = f + 1;
                assert!(is_united(&mut right, f, s));
                assert!(matches!(
                    union(&mut right, f, s),
                    UnionResult::AlreadyJoined
                ));
            }

            assert!(!is_united(&mut right, 3, 4));
            assert!(matches!(union(&mut right, 2, 5), UnionResult::Success));
            assert!(matches!(
                union(&mut right, 3, 4),
                UnionResult::AlreadyJoined
            ));
        }
    }

    #[test]
    fn test_detach() {
        const ALL_CONNS: &[(usize, usize)] = &[(1, 2), (1, 3), (2, 3), (4, 5), (4, 6), (5, 6)];
        fn check_no_unexpected_conn(djs: &mut impl Storage<u8>, case: &str) {
            for i in djs.lower_bound()..djs.upper_bound() {
                for j in i + 1..djs.upper_bound() {
                    if ALL_CONNS.contains(&(i, j)) {
                        continue;
                    }
                    unsafe {
                        assert!(
                            !is_united(djs, i, j),
                            "Unexpected connection after {}",
                            case
                        );
                    }
                }
            }
        }
        let make_djs = || unsafe {
            let mut djs = make_storage();
            let mut right = djs.split_at(1).1;
            for (l, r) in [(1, 2), (2, 3), (4, 5), (4, 6)] {
                union(&mut right, l, r);
            }
            djs
        };

        // Detach already detached.
        unsafe {
            let mut djs = make_djs();
            let mut right = djs.split_at(1).1;
            detach(&mut right, 7);

            for i in (djs.lower_bound()..djs.upper_bound()).filter(|&i| i != 7) {
                assert!(!is_united(&mut djs, i, 7));
            }
            for &(i, j) in ALL_CONNS {
                assert!(is_united(&mut djs, i, j));
            }
            check_no_unexpected_conn(&mut djs, "already detached");
        }
        // Detach leaf.
        unsafe {
            let mut djs = make_djs();
            let mut right = djs.split_at(1).1;
            detach(&mut right, 3);
            for i in (djs.lower_bound()..djs.upper_bound()).filter(|&i| i != 3) {
                assert!(!is_united(&mut djs, i, 3));
            }
            for &(i, j) in ALL_CONNS {
                if i == 3 || j == 3 {
                    continue;
                }
                assert!(is_united(&mut djs, i, j));
            }
            check_no_unexpected_conn(&mut djs, "leaf");
        }

        // Detach in the middle
        unsafe {
            let mut djs = make_djs();
            let mut right = djs.split_at(1).1;
            detach(&mut right, 2);
            for i in (djs.lower_bound()..djs.upper_bound()).filter(|&i| i != 2) {
                assert!(!is_united(&mut djs, i, 2));
            }
            for &(i, j) in ALL_CONNS {
                if i == 2 || j == 2 {
                    continue;
                }
                assert!(is_united(&mut djs, i, j));
            }
            check_no_unexpected_conn(&mut djs, "middle");
        }
        // Detach root.
        unsafe {
            let mut djs = make_djs();
            let mut right = djs.split_at(1).1;
            detach(&mut right, 1);
            for i in (djs.lower_bound()..djs.upper_bound()).filter(|&i| i != 0) {
                assert!(!is_united(&mut djs, i, 0));
            }
            for &(i, j) in ALL_CONNS {
                if i == 1 || j == 1 {
                    continue;
                }
                assert!(is_united(&mut djs, i, j));
            }
            check_no_unexpected_conn(&mut djs, "root");
        }
    }

    #[test]
    fn test_compress_paths() {
        let mut storage = BoxStorage::new(9);
        initialize_storage::<u8, false>(&mut storage);
        let mut right = unsafe { storage.split_at(1).1 };
        unsafe {
            union(&mut right, 1, 3);
            union(&mut right, 4, 5);
            union(&mut right, 1, 4);
            union(&mut right, 2, 6);
            union(&mut right, 2, 7);
        }
        // It is unexpected to get compressed storage right after unions.
        assert_ne!(right.roots_and_ranks().0, &[1, 2, 1, 1, 1, 2, 2, 8]);
        assert_ne!(right.roots_and_ranks().1, &[2, 1, 0, 0, 0, 0, 0, 0]);

        unsafe {
            compress_all_paths(&mut right);
        }
        assert_eq!(right.roots_and_ranks().0, &[1, 2, 1, 1, 1, 2, 2, 8]);
        assert_eq!(right.roots_and_ranks().1, &[2, 1, 0, 0, 0, 0, 0, 0]);
        assert_eq!(storage.roots_and_ranks().0, &[0, 1, 2, 1, 1, 1, 2, 2, 8]);
        assert_eq!(storage.roots_and_ranks().1, &[0, 2, 1, 0, 0, 0, 0, 0, 0]);
    }
}