aph_disjoint_set 0.1.1

use core::mem::{align_of, size_of, MaybeUninit};
use core::ptr::write_bytes;

use crate::algorithm;
use crate::bits_enum::BitsEnum;
use crate::common::{Common, PathsStatus, Splittability};
use crate::macros::{common_strct, for_current_bitness};
use crate::storage::{ArrayStorage, Storage};
use crate::tag_type::TagType;
use crate::{ChunkMut, DisjointSetRoView, Root, UnionResult};

// We need this for combination of 2 reasons:
// 1. We want to reuse implementation of `Common`.
// 2. We know exactly the size and bitness of ArrayStorage.
// By using inconstructible variant in enum, we allow Rust to eliminate all overhead of the enum
// while satisfying interface requirements of `Common`.
#[derive(Clone, Copy)]
enum Impossible {}

macro_rules! implement_inconstructible_storage {
    ($tag:ty) => {
        impl Storage<$tag> for Impossible {
            #[inline]
            fn lower_bound(&self) -> usize {
                unreachable!("Impossible cannot be constructed!")
            }

            #[inline]
            fn upper_bound(&self) -> usize {
                unreachable!("Impossible cannot be constructed!")
            }

            #[inline]
            fn roots_and_ranks(&mut self) -> (&mut [$tag], &mut [$tag]) {
                unreachable!("Impossible cannot be constructed!")
            }

            #[inline]
            fn roots_readonly(&self) -> &[$tag] {
                unreachable!("Impossible cannot be constructed!")
            }
        }
    };
}

implement_inconstructible_storage!(u8);
implement_inconstructible_storage!(u16);
implement_inconstructible_storage!(u32);
implement_inconstructible_storage!(u64);

// TODO: Switch to compile test in const block when they would be stabilized.
trait CompileTest {
    const COMPILE_TEST: ();
}

const fn test_struct_size<T, Stor>() {
    let extra_place_for_flags = if align_of::<Stor>() > 1 {
        align_of::<Stor>()
    } else {
        // For `is_splittable` and `is_compressed`.
        2
    };
    assert!(
        size_of::<T>() ==
        // We must not have extra place for enum impl.
        size_of::<Stor>() + extra_place_for_flags
    );
}

impl<T: TagType, const SIZE: usize> CompileTest for [T; SIZE] {
    #[allow(clippy::cast_possible_truncation)]
    const COMPILE_TEST: () = {
        let max_val = match T::BIT_SIZE {
            8 => u8::MAX as usize,
            16 => u16::MAX as usize,
            32 => u32::MAX as usize,
            64 => u64::MAX as usize,
            _ => unreachable!(),
        };
        assert!(
            SIZE <= max_val,
            "SIZE must not be bigger than tag type max."
        );

        test_struct_size::<DisjointSetArrayU8<128>, ArrayStorage<u8, 128>>();
        test_struct_size::<DisjointSetArrayU16<128>, ArrayStorage<u16, 128>>();
        test_struct_size::<DisjointSetArrayU32<128>, ArrayStorage<u32, 128>>();
        test_struct_size::<DisjointSetArrayU64<128>, ArrayStorage<u64, 128>>();
    };
}

macro_rules! impl_array_djs {
    ($name:ident,$ro_type:ident,
        $common_decl:ty,$storage:ty,
        $tag:ty,$enm:path
    ) => {

#[allow(clippy::module_name_repetitions)]
#[derive(Clone, Copy)]
pub struct $name<const SIZE: usize>(
    $common_decl
);

#[allow(clippy::module_name_repetitions)]
#[derive(Clone, Copy)]
pub struct $ro_type<const SIZE: usize>($storage);

impl<const SIZE: usize> $name<SIZE> {
    /// Creates new disjoint set.
    /// Does not compile if SIZE too big:
    /// ```compile_fail
    /// use aph_disjoint_set::DisjointSetArrayU8;
    /// let _: DisjointSetArrayU8<500> = DisjointSetArrayU8::new();
    /// #[cfg(miri)] // Miri doesn't catch compile error because it doesn't evaluate code in such tests.
    /// let _: u32 = 0u64;
    /// ```
    #[must_use]
    pub fn new() -> Self {
        #[allow(clippy::let_unit_value)]
        let _: () = <[$tag; SIZE] as CompileTest>::COMPILE_TEST;

        let mut storage = ArrayStorage::ZEROED;
        algorithm::initialize_storage::<_, false>(&mut storage);
        Self(unsafe {
            // SAFETY: We just initialized storage.
            Common::new(
                $enm(storage),
                Splittability::Allowed,
                PathsStatus::Compressed,
            )
        })
    }

    /// Initializes some memory with default state of disjoint set
    /// and returns reference to it.
    /// Useful if you want hard guarantee that compiler doesn't allocate some place on stack
    /// for temporary value during moves.
    ///
    /// Does not compile if SIZE too big:
    /// ```compile_fail
    /// use core::mem::MaybeUninit;
    /// use aph_disjoint_set::DisjointSetArrayU8;
    ///
    /// let mut value: MaybeUninit<DisjointSetArrayU8<500>> = MaybeUninit::uninit();
    /// DisjointSetArrayU8::initialize_inplace(&mut value);
    /// #[cfg(miri)] // Miri doesn't catch compile error because it doesn't evaluate code in such tests.
    /// let _: u32 = 0u64;
    /// ```
    /// # Example
    /// ```
    /// use core::mem::MaybeUninit;
    /// use aph_disjoint_set::DisjointSetArrayU8;
    ///
    /// let mut value: MaybeUninit<DisjointSetArrayU8<10>> = MaybeUninit::uninit();
    /// let djs = DisjointSetArrayU8::initialize_inplace(&mut value);
    /// djs.union(9, 1);
    /// assert!(djs.is_united(1, 9));
    /// ```
    pub fn initialize_inplace(memory: &mut MaybeUninit<Self>) -> &mut Self {
        #[allow(clippy::let_unit_value)]
        let _: () = <[$tag; SIZE] as CompileTest>::COMPILE_TEST;
        unsafe {
            let ptr = memory.as_mut_ptr();
            // SAFETY: Safe because we have mutable reference to memory.
            write_bytes(ptr, 0, 1);
            // SAFETY: It is safe to dereference here because we filled it with zeros.
            let me = &mut *ptr;

            // SAFETY: We would initialize it just now.
            let (s, is_splittable, is_compressed) = me.0.access_internals();
            // Since bitness enum contains only single variant with constructible type,
            // it is guaranteed to match it.
            for_current_bitness!(s; s; algorithm::initialize_storage::<$tag, true>(s));
            // We just initialized our storage.
            *is_splittable = Splittability::Allowed;
            *is_compressed = PathsStatus::Compressed;

            me
        }
    }
}

impl<const SIZE: usize> $name<SIZE> {
    /// 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) -> $ro_type<SIZE> {
        let s = self.0.into_compressed_storage();
        match s {
            $enm(s) => $ro_type(s),
            _ => unreachable!(),
        }
    }

    /// 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();
    }
}

impl<const SIZE: usize> Default for $name<SIZE> {
    fn default() -> Self {
        Self::new()
    }
}

impl<const SIZE: usize> From<$ro_type<SIZE>> for $name<SIZE> {
    fn from(value: $ro_type<SIZE>) -> Self {
        Self(unsafe {
            // SAFETY: Only way to produce ro version is to use
            // valid mutable version.
            Common::new(
                $enm(value.0),
                // We probably had union operation before conversion to RO.
                Splittability::Disallowed,
                // Any RO struct is compressed.
                PathsStatus::Compressed,
            )
        })
    }
}

impl<const SIZE: usize> $ro_type<SIZE> {
    #[inline]
    #[must_use]
    pub fn view(&self)->DisjointSetRoView{
        DisjointSetRoView($enm(self.0.roots_readonly()))
    }
}
};
}

impl_array_djs!(
    DisjointSetArrayU8,
    RoDisjointSetArrayU8,
    common_strct!(ArrayStorage<u8, SIZE>, Impossible, Impossible,Impossible),
    ArrayStorage<u8, SIZE>,
    u8,
    BitsEnum::U8
);
impl_array_djs!(
    DisjointSetArrayU16,
    RoDisjointSetArrayU16,
    common_strct!(Impossible, ArrayStorage<u16, SIZE>, Impossible,Impossible),
    ArrayStorage<u16, SIZE>,
    u16,
    BitsEnum::U16
);
#[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
impl_array_djs!(
    DisjointSetArrayU32,
    RoDisjointSetArrayU32,
    common_strct!(Impossible,Impossible, ArrayStorage<u32, SIZE>, Impossible),
    ArrayStorage<u32, SIZE>,
    u32,
    BitsEnum::U32
);
#[cfg(target_pointer_width = "64")]
impl_array_djs!(
    DisjointSetArrayU64,
    RoDisjointSetArrayU64,
    common_strct!(Impossible,Impossible,Impossible, ArrayStorage<u64, SIZE>),
    ArrayStorage<u64, SIZE>,
    u64,
    BitsEnum::U64
);

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

    use crate::{DisjointSetArrayU16, DisjointSetArrayU8};

    #[test]
    fn test_initialize() {
        let mut djs: DisjointSetArrayU8<20> = DisjointSetArrayU8::new();
        for i in 0..20 {
            assert_eq!(djs.get_root(i).into_inner(), i);
        }
        let _ = djs.union(5, 7);
        assert!(djs.is_united(7, 5));

        let mut mem: MaybeUninit<DisjointSetArrayU8<20>> = MaybeUninit::uninit();
        let djs = DisjointSetArrayU8::initialize_inplace(&mut mem);
        for i in 0..20 {
            assert_eq!(djs.get_root(i).into_inner(), i);
        }
        let _ = djs.union(5, 7);
        assert!(djs.is_united(7, 5));
    }

    #[test]
    fn test_initialize2() {
        const NUM: usize = 500;
        let mut djs: DisjointSetArrayU16<NUM> = DisjointSetArrayU16::new();
        for i in 0..NUM {
            assert_eq!(djs.get_root(i).into_inner(), i);
        }
        let _ = djs.union(5, 7);
        assert!(djs.is_united(7, 5));

        let mut mem: MaybeUninit<DisjointSetArrayU16<NUM>> = MaybeUninit::uninit();
        let djs = DisjointSetArrayU16::initialize_inplace(&mut mem);
        for i in 0..NUM {
            assert_eq!(djs.get_root(i).into_inner(), i);
        }
        let _ = djs.union(5, 7);
        assert!(djs.is_united(7, 5));
    }
}