aph_disjoint_set 0.1.1

use core::hint::unreachable_unchecked;
use core::mem::{align_of, forget, size_of, ManuallyDrop};
use core::ptr::{copy_nonoverlapping, write_bytes, NonNull};
use core::slice::{from_raw_parts, from_raw_parts_mut};

use crate::macros::{bits_enum, choose_by_size};
use crate::tag_type::TagType;

// Guarantees:
// `lower_bound` and `upper_bound` are constant time.
// Guarantees specially for unsafe code:
// 1. roots and ranks do not overlap.
// 2. roots and ranks can be indexed by any `idx - lower_bound()`
// 3. roots and ranks have same length.
// 4. roots can contain every value from half-interval fromlower bound to upper.
// if `idx < upper_bound()` and `idx >= lower_bound()`.
pub(crate) trait Storage<Tag: TagType> {
    // We use bounds instead of 0 and len
    // because we want to be able split storage
    // into chunks for parallel filling.
    #[must_use]
    fn lower_bound(&self) -> usize;
    #[must_use]
    fn upper_bound(&self) -> usize;

    // Max root is `size - 1`.
    // Ranks start from 0, and always less than log2(size)
    // which guaranteed to fit into 0..Tag::MAX.
    #[must_use]
    fn roots_and_ranks(&mut self) -> (&mut [Tag], &mut [Tag]);
    #[must_use]
    fn roots_readonly(&self) -> &[Tag];

    /// **NOTE:** Unlike `slice::split_at_mut`, `split_pos` is relative to
    /// the allocation, e.g. if `self.lower_bound() == 5`,
    /// and `split_pos==6`, it would split self after first item.
    /// # Panics
    /// If `split_pos < self.lower_bound()` or `split_pos > self.upper_bound()`.
    /// # Safety
    /// Caller must ensure that result doesn't outlive allocation.
    #[inline]
    #[must_use]
    unsafe fn split_at<'a>(
        &mut self,
        split_pos: usize,
    ) -> (ViewStorage<'a, Tag>, ViewStorage<'a, Tag>) {
        let lower_bound = self.lower_bound();
        let upper_bound = self.upper_bound();
        assert!(lower_bound <= split_pos && split_pos <= upper_bound);

        let (roots, ranks) = self.roots_and_ranks();
        debug_assert_eq!(roots.len(), upper_bound - lower_bound);
        debug_assert_eq!(ranks.len(), upper_bound - lower_bound);
        let roots_ptr = roots.as_mut_ptr();
        let ranks_ptr = ranks.as_mut_ptr();
        unsafe {
            // SAFETY: We checked bounds.
            // We split roots and ranks into non-overlapping slices.
            let first_roots = from_raw_parts_mut(roots_ptr, split_pos - lower_bound);
            let first_ranks = from_raw_parts_mut(ranks_ptr, split_pos - lower_bound);
            let second_roots = from_raw_parts_mut(
                roots_ptr.add(split_pos - lower_bound),
                upper_bound - split_pos,
            );
            let second_ranks = from_raw_parts_mut(
                ranks_ptr.add(split_pos - lower_bound),
                upper_bound - split_pos,
            );

            (
                ViewStorage {
                    dist_from_alloc: lower_bound,
                    roots: first_roots,
                    ranks: first_ranks,
                },
                ViewStorage {
                    dist_from_alloc: split_pos,
                    roots: second_roots,
                    ranks: second_ranks,
                },
            )
        }
    }
}

// We store both roots and ranks here in one allocation.
// We use same memory block for them because allocations can be costly.
// First half of `data` is roots.
// Second half of `data` is ranks.
#[repr(transparent)]
pub(crate) struct BoxStorage<Tag: TagType>(Box<[Tag]>);

impl<Tag: TagType> BoxStorage<Tag> {
    #[inline]
    pub(crate) fn new(size: usize) -> Self {
        assert!(
            size < usize::MAX / 4 / size_of::<usize>(),
            "Too large size requested"
        );
        Self(vec![Tag::ZERO; size * 2].into_boxed_slice())
    }
}

impl<Tag: TagType> Storage<Tag> for BoxStorage<Tag> {
    #[inline]
    fn lower_bound(&self) -> usize {
        0
    }

    #[inline]
    fn upper_bound(&self) -> usize {
        if 1 == self.0.len() & 1 {
            debug_assert!(false);
            // This is added to make optimizer understand that
            // equality of upper bound imply equality of len.
            unsafe {
                // SAFETY: This cannot happen by construction.
                core::hint::unreachable_unchecked();
            }
        }
        self.0.len() / 2
    }

    #[inline]
    fn roots_and_ranks(&mut self) -> (&mut [Tag], &mut [Tag]) {
        let split_pos = self.upper_bound();
        self.0.split_at_mut(split_pos)
    }

    #[inline]
    fn roots_readonly(&self) -> &[Tag] {
        &self.0[..self.upper_bound()]
    }
}

impl<T: TagType> Clone for BoxStorage<T> {
    #[inline]
    fn clone(&self) -> Self {
        Self(self.0.clone())
    }

    #[inline]
    fn clone_from(&mut self, source: &Self) {
        self.0.clone_from(&source.0);
    }
}

#[derive(Clone, Copy)]
pub(crate) struct ArrayStorage<Tag, const SIZE: usize> {
    roots: [Tag; SIZE],
    ranks: [Tag; SIZE],
}

impl<Tag: TagType, const SIZE: usize> ArrayStorage<Tag, SIZE> {
    pub(crate) const ZEROED: Self = Self {
        roots: [Tag::ZERO; SIZE],
        ranks: [Tag::ZERO; SIZE],
    };
}

impl<Tag: TagType, const SIZE: usize> Storage<Tag> for ArrayStorage<Tag, SIZE> {
    #[inline]
    fn lower_bound(&self) -> usize {
        0
    }
    #[inline]
    fn upper_bound(&self) -> usize {
        SIZE
    }
    #[inline]
    fn roots_and_ranks(&mut self) -> (&mut [Tag], &mut [Tag]) {
        (&mut self.roots, &mut self.ranks)
    }
    #[inline]
    fn roots_readonly(&self) -> &[Tag] {
        &self.roots
    }
}

pub(crate) struct DynamicStorage<Tag: TagType> {
    // How many pairs or root and rank here.
    len: usize,
    // How many pairs or root and rank can be stored here.
    // It is twice smaller than underlying Vec capacity.
    cap: usize,
    // I decided to store pointer instead of Vec
    // because we wouldn't be able to use it len and capacity directly
    // so it would be just a waste of `2*size_of::<usize>` bytes.
    // Data by pointer would be always initialized up to `2*cap` elements.
    ptr: NonNull<Tag>,
}

impl<Tag: TagType> DynamicStorage<Tag> {
    // Note: Result has zeroed valid ranks.
    #[inline]
    pub(crate) fn with_capacity(cap: usize) -> Self {
        if cap == 0 {
            return Self {
                len: 0,
                cap: 0,
                ptr: NonNull::dangling(),
            };
        }
        assert!(cap < usize::MAX / 4, "Too large size requested");
        let v = vec![Tag::ZERO; cap * 2];
        assert_eq!(
            v.capacity(),
            cap * 2,
            "In case of some behaviour of Vec changes"
        );
        let mut v = ManuallyDrop::new(v);
        let v: &mut Vec<_> = &mut *v;
        let len = 0;
        let ptr = v.as_mut_ptr();
        Self {
            len,
            cap,
            ptr: NonNull::new(ptr).expect("Vec ptr cannot be null."),
        }
    }

    #[inline]
    pub(crate) fn capacity(&self) -> usize {
        self.cap
    }

    /// Increases size of storage and returns view
    /// to uninitialized part of self.
    /// # Panics
    /// If capacity is unsufficient. Call `reserve` before it.
    #[inline]
    pub(crate) fn enlarge(&mut self, additional: usize) -> ViewStorage<Tag> {
        assert!(
            self.len.checked_add(additional).unwrap() <= self.cap,
            "Must call `reserve` before enlarge."
        );
        let old_len = self.len;
        self.len += additional;
        unsafe {
            // SAFETY: Returned view lifetime is bound to lifetime of self.
            // Allocation is alive while self is alive.
            // Rust borrow checker wouldn't allow to modify allocation
            // while ViewStorage is alive.
            self.split_at(old_len).1
        }
    }

    /// This would not initialize all new values to zero.
    /// # Returns
    /// If values up to new capacity cannot be fit into current `Tag`,
    /// would return new allocation with different `Tag`.
    /// # Panics
    /// If `len + additional` usizes requires more than `isize::MAX / 2` bytes.
    #[allow(clippy::items_after_statements, clippy::type_complexity)]
    #[inline]
    #[must_use]
    pub(crate) fn reserve(&mut self, additional: usize) -> Option<bits_enum!(DynamicStorage)> {
        if self.len + additional <= self.cap {
            // We have enough memory already.
            return None;
        }
        let realloc_res = reallocate(self, additional);
        if realloc_res.is_some() {
            return realloc_res;
        }
        if self.len + additional > self.cap {
            if cfg!(debug_assertions) {
                unreachable!()
            } else {
                unsafe { unreachable_unchecked() };
            }
        }
        return None;

        #[must_use]
        fn allocate_new<Tag: TagType>(
            me: &mut DynamicStorage<Tag>,
            new_cap: usize,
        ) -> bits_enum!(DynamicStorage) {
            assert!(me.cap < new_cap);
            choose_by_size!(new_cap, {
                if size_of::<ChosenTagType>() <= size_of::<Tag>() {
                    unreachable!("We only increase size");
                }

                let mut updated = DynamicStorage::with_capacity(new_cap);
                updated.len = me.len;
                let (old_roots, old_ranks) = me.roots_and_ranks();
                let (new_roots, new_ranks) = updated.roots_and_ranks();
                for (old, new) in old_roots.iter().zip(new_roots.iter_mut()) {
                    *new = ChosenTagType::from_u(old.as_u());
                }
                for (old, new) in old_ranks.iter().zip(new_ranks.iter_mut()) {
                    *new = ChosenTagType::from_u(old.as_u());
                }
                updated
            })
        }

        fn reallocate_inplace<Tag: TagType>(me: &mut DynamicStorage<Tag>, new_cap: usize) {
            let len = me.len;
            let old_cap = me.cap;

            let mut updated = {
                let new_vec_cap = new_cap.checked_mul(2).unwrap();
                let updated: Vec<Tag> = Vec::with_capacity(new_vec_cap);
                assert_eq!(new_vec_cap, updated.capacity());
                updated
            };

            // At this point, this method must never panic.
            let old_ptr = me.ptr.as_ptr();
            let new_ptr = updated.as_mut_ptr();
            // After this unsafe block, all memory in new vec is filled.
            unsafe {
                // SAFETY:
                // 1. memory don't overlap because `target` is newly allocated.
                // 2. old_len <= old_cap
                // 3. old_cap * 2 < new_vec_cap

                // Copy roots
                copy_nonoverlapping(old_ptr, new_ptr, len);
                // Copy ranks
                copy_nonoverlapping(old_ptr.add(old_cap), new_ptr.add(new_cap), len);

                // Initialize missed regions.
                let missed_len = new_cap - len;
                write_bytes(new_ptr.add(len), 0, missed_len);
                write_bytes(new_ptr.add(new_cap).add(len), 0, missed_len);
            }

            forget(updated);
            let old_vec = unsafe {
                // SAFETY:
                // Part values correctness guaranteed by `DynamicStorage` invariant.
                // They wouldn't be accessed anymore because we're being dropped right now.
                Vec::from_raw_parts(old_ptr, old_cap * 2, old_cap * 2)
            };
            me.ptr = NonNull::new(new_ptr).unwrap();
            me.cap = new_cap;
            drop(old_vec);
        }

        // Returns new storage if additional values cannot be encoded by old tag.
        #[cold]
        #[must_use]
        fn reallocate<Tag: TagType>(
            me: &mut DynamicStorage<Tag>,
            additional: usize,
        ) -> Option<bits_enum!(DynamicStorage)> {
            let new_cap = calc_new_cap_dynamic(me.len, me.cap, additional);
            if new_cap > Tag::MAX_VAL.as_u() {
                Some(allocate_new(me, new_cap))
            } else {
                reallocate_inplace(me, new_cap);
                None
            }
        }
    }

    pub(crate) fn make_copied(&self) -> bits_enum!(DynamicStorage) {
        choose_by_size!(self.len, {
            let mut res: DynamicStorage<ChosenTagType> = DynamicStorage::with_capacity(self.len);
            res.copy_data_from(self);
            res
        })
    }

    /// # Panics
    /// If other have more values than our capacity.
    pub(crate) fn copy_data_from<OtherT: TagType>(&mut self, source: &DynamicStorage<OtherT>) {
        // This struct used only to zero len in case of panic.
        struct ZeroLen<'a, T: TagType>(&'a mut DynamicStorage<T>);

        impl<T: TagType> Drop for ZeroLen<'_, T> {
            fn drop(&mut self) {
                self.0.len = 0;
            }
        }

        assert!(self.cap >= source.len);

        let zero_len = ZeroLen(self);
        let dest = &mut *zero_len.0;

        let len = source.len;

        dest.len = len;
        let (dest_roots, dest_ranks) = dest.roots_and_ranks();
        let (source_roots, source_ranks) = source.get_ro_roots_ranks();

        assert_eq!(dest_roots.len(), len);
        assert_eq!(dest_ranks.len(), len);
        assert_eq!(source_roots.len(), len);
        assert_eq!(source_ranks.len(), len);

        if size_of::<OtherT>() == size_of::<Tag>() {
            assert_eq!(align_of::<OtherT>(), align_of::<Tag>());
            unsafe {
                // SAFETY:
                // Slices point to valid data.
                // Mutable and immutable references cannot alias.
                // We checked lengths validity.
                // Tag types has unique sizes so in that branch those types are same.
                let source = source_roots.as_ptr();
                let target = dest_roots.as_mut_ptr();
                copy_nonoverlapping(source, target.cast(), len);
            }
            unsafe {
                // SAFETY: Same as above.
                let source = source_ranks.as_ptr();
                let target = dest_ranks.as_mut_ptr();
                copy_nonoverlapping(source, target.cast(), len);
            }
        } else {
            for (d, s) in [(dest_roots, source_roots), (dest_ranks, source_ranks)] {
                for (dest, src) in d.iter_mut().zip(s.iter()) {
                    *dest = TagType::from_u(src.as_u());
                }
            }
        }

        // Copy succeeded so we can destroy our guard.
        forget(zero_len);
    }

    #[inline]
    pub(crate) fn clear(&mut self) {
        self.len = 0;
    }

    #[inline]
    #[must_use]
    pub(crate) fn get_ro_roots_ranks(&self) -> (&[Tag], &[Tag]) {
        unsafe {
            // SAFETY:
            // 1. Allocated memory size has self.cap * 2 elements
            // 2. self.cap >= self.len.
            // 3. All memory is initialized.
            let roots = from_raw_parts(self.ptr.as_ptr(), self.len);
            let ranks = from_raw_parts(self.ptr.as_ptr().add(self.cap), self.len);
            (roots, ranks)
        }
    }
}

#[allow(clippy::items_after_statements)]
fn calc_new_cap_dynamic(len: usize, old_cap: usize, additional: usize) -> usize {
    let min_new_vec_cap = len
        .checked_add(additional)
        .and_then(|x| x.checked_mul(2))
        .unwrap();
    assert!(
        min_new_vec_cap > old_cap * 2,
        "This method is used for calculation of new cap for reallocation."
    );

    let new_vec_cap = min_new_vec_cap
        .max(
            min_new_vec_cap
                .checked_next_power_of_two()
                .unwrap_or_default(),
        )
        .max(
            old_cap
                .checked_next_power_of_two()
                .and_then(|x| x.checked_mul(4))
                .unwrap_or_default(),
        )
        .max(old_cap.checked_mul(4).unwrap_or_default())
        .max(8);

    // Standard library don't allocate more than that.
    const STD_LIMIT: usize = (isize::MAX as usize) / size_of::<usize>();
    assert!(min_new_vec_cap < STD_LIMIT, "Too large allocation request");
    new_vec_cap.min(STD_LIMIT) / 2
}

impl<Tag: TagType> Drop for DynamicStorage<Tag> {
    fn drop(&mut self) {
        let v = unsafe {
            // SAFETY:
            // Part values correctness guaranteed by `DynamicStorage` invariant.
            // They wouldn't be accessed anymore because we're being dropped right now.
            Vec::from_raw_parts(self.ptr.as_ptr(), self.cap * 2, self.cap * 2)
        };
        drop(v);
    }
}

impl<Tag: TagType> Storage<Tag> for DynamicStorage<Tag> {
    #[inline]
    fn lower_bound(&self) -> usize {
        0
    }

    #[inline]
    fn upper_bound(&self) -> usize {
        self.len
    }

    #[inline]
    fn roots_and_ranks(&mut self) -> (&mut [Tag], &mut [Tag]) {
        unsafe {
            // SAFETY:
            // 1. Allocated memory size has self.cap * 2 elements
            // 2. self.cap >= self.len.
            // 3. All memory is initialized.
            let roots = from_raw_parts_mut(self.ptr.as_ptr(), self.len);
            let ranks = from_raw_parts_mut(self.ptr.as_ptr().add(self.cap), self.len);
            (roots, ranks)
        }
    }

    #[inline]
    fn roots_readonly(&self) -> &[Tag] {
        self.get_ro_roots_ranks().0
    }
}

/// This type is view to some part of another storage.
/// It is intended to be used for parallel initialization.
pub(crate) struct ViewStorage<'owner, Tag: TagType> {
    dist_from_alloc: usize,
    roots: &'owner mut [Tag],
    ranks: &'owner mut [Tag],
}

impl<'owner, Tag: TagType> Storage<Tag> for ViewStorage<'owner, Tag> {
    #[inline]
    fn lower_bound(&self) -> usize {
        self.dist_from_alloc
    }

    #[inline]
    fn upper_bound(&self) -> usize {
        self.dist_from_alloc + self.roots.len()
    }

    #[inline]
    fn roots_and_ranks(&mut self) -> (&mut [Tag], &mut [Tag]) {
        (self.roots, self.ranks)
    }

    #[inline]
    fn roots_readonly(&self) -> &[Tag] {
        self.roots
    }
}

#[cfg(test)]
impl<Tag: TagType> Storage<Tag> for Box<dyn Storage<Tag>> {
    fn lower_bound(&self) -> usize {
        (**self).lower_bound()
    }

    fn upper_bound(&self) -> usize {
        (**self).upper_bound()
    }

    fn roots_and_ranks(&mut self) -> (&mut [Tag], &mut [Tag]) {
        (**self).roots_and_ranks()
    }

    fn roots_readonly(&self) -> &[Tag] {
        (**self).roots_readonly()
    }
}

#[allow(clippy::items_after_statements)]
#[cfg(test)]
mod tests {
    use core::mem::size_of;

    use rstest::rstest;

    use crate::bits_enum::BitsEnum;

    use super::{calc_new_cap_dynamic, ArrayStorage, BoxStorage, DynamicStorage, Storage};

    #[test]
    fn test_enlarge() {
        let mut storage: DynamicStorage<u8> = DynamicStorage::with_capacity(0);

        assert!(storage.reserve(1).is_none());
        assert_eq!(storage.capacity(), 4);
        assert_eq!(storage.upper_bound(), 0);
        assert!(storage.roots_and_ranks().0.is_empty());
        assert!(storage.roots_and_ranks().1.is_empty());

        let uninit_part = storage.enlarge(1);
        assert_eq!(
            (uninit_part.lower_bound(), uninit_part.upper_bound()),
            (0, 1)
        );
        assert_eq!(storage.upper_bound(), 1);
        assert_eq!(storage.roots_and_ranks().0, &[0]);
        assert_eq!(storage.roots_and_ranks().1, &[0]);

        assert!(storage.reserve(10).is_none());
        assert_eq!(storage.capacity(), 16);
        let uninit_part = storage.enlarge(10);
        assert_eq!(
            (uninit_part.lower_bound(), uninit_part.upper_bound()),
            (1, 11)
        );
        assert_eq!(storage.upper_bound(), 11);
        assert_eq!(storage.roots_and_ranks().0, &[0; 11]);
        assert_eq!(storage.roots_and_ranks().1, &[0; 11]);

        let updated = storage.reserve(200);
        assert!(updated.is_some());
        let mut updated = if let Some(BitsEnum::U16(u)) = updated {
            u
        } else {
            panic!("Must to be u16")
        };
        assert_eq!(updated.capacity(), 256);
        let uninit_part = updated.enlarge(200);
        assert_eq!(
            (uninit_part.lower_bound(), uninit_part.upper_bound()),
            (11, 211)
        );
        assert_eq!(updated.upper_bound(), 211);
        assert_eq!(updated.roots_and_ranks().0, &[0; 211]);
        assert_eq!(updated.roots_and_ranks().1, &[0; 211]);
    }

    #[rstest]
    #[case(0, 1, 4)]
    #[case(5, 1, 16)]
    #[case(250, 500, 1024)]
    #[case(usize::MAX / 4 / size_of::<usize>() - 1000 * size_of::<usize>(), 500, usize::MAX / 4 / size_of::<usize>())]
    fn test_new_capacity(
        #[case] old_len: usize,
        #[case] added: usize,
        #[case] expected_cap: usize,
    ) {
        assert_eq!(calc_new_cap_dynamic(old_len, old_len, added), expected_cap);
    }

    #[test]
    fn test_split_at() {
        let mut arr = ArrayStorage {
            roots: [0, 1, 2, 3, 4, 5, 6, 7],
            ranks: [7, 6, 5, 4, 3, 2, 1, 0],
        };

        let mut dynamic = DynamicStorage::with_capacity(8);
        dynamic.enlarge(8);
        dynamic.roots_and_ranks().0.copy_from_slice(&arr.roots);
        dynamic.roots_and_ranks().1.copy_from_slice(&arr.ranks);

        let mut fixed = BoxStorage::new(8);
        fixed.roots_and_ranks().0.copy_from_slice(&arr.roots);
        fixed.roots_and_ranks().1.copy_from_slice(&arr.ranks);

        let all: [&mut dyn Storage<u8>; 3] = [&mut arr, &mut dynamic, &mut fixed];

        for s in all {
            let (mut left, mut right) = unsafe { s.split_at(3) };
            assert_eq!(left.lower_bound(), 0);
            assert_eq!(left.upper_bound(), 3);
            assert_eq!(left.roots_and_ranks().0, &[0, 1, 2]);
            assert_eq!(left.roots_and_ranks().1, &[7, 6, 5]);

            assert_eq!(right.lower_bound(), 3);
            assert_eq!(right.upper_bound(), 8);
            assert_eq!(right.roots_and_ranks().0, &[3, 4, 5, 6, 7]);
            assert_eq!(right.roots_and_ranks().1, &[4, 3, 2, 1, 0]);

            let (mut middle, mut most_right) = unsafe { right.split_at(6) };

            assert_eq!(middle.lower_bound(), 3);
            assert_eq!(middle.upper_bound(), 6);
            assert_eq!(middle.roots_and_ranks().0, &[3, 4, 5,]);
            assert_eq!(middle.roots_and_ranks().1, &[4, 3, 2,]);

            assert_eq!(most_right.lower_bound(), 6);
            assert_eq!(most_right.upper_bound(), 8);
            assert_eq!(most_right.roots_and_ranks().0, &[6, 7]);
            assert_eq!(most_right.roots_and_ranks().1, &[1, 0]);
        }
    }
}