fastbit 0.11.1

use std::ptr::NonNull;

use crate::traits::BitWord;
use std::alloc::Layout;

/// Returns the number of words needed to store `len` bits for the given word type `T`.
#[inline(always)]
pub fn word_len<T: BitWord>(len: usize) -> usize {
    (len + T::BITS - 1) / T::BITS
}

/// Returns the word index and bit index within the word for the given bit index `idx`.
#[inline(always)]
pub fn word_bit<T: BitWord>(idx: usize) -> (usize, usize) {
    let word = idx / T::BITS;
    let bit = idx % T::BITS;
    (word, bit)
}

/// Returns the word index, bit index, and a mask for all bits below the given bit index `idx`.
#[inline(always)]
pub fn word_bit_mask<T: BitWord>(len: usize) -> (usize, usize, T) {
    let word = len / T::BITS;
    let bit = len % T::BITS;
    let mask = (T::from(1u8) << bit) - T::from(1u8);
    (word, bit, mask)
}

/// Allocates a zero-initialized bitset of the given length (in bits).
///
/// # Safety
///
/// This function is unsafe because:
/// 1. It allocates raw memory that must be properly freed using `bitset_free`.
/// 2. The caller must ensure the returned pointer is not used after being freed.
/// 3. The caller must not use the pointer after its lifetime has ended.
///
/// # Returns
///
/// A tuple containing:
/// - A non-null pointer to the allocated memory
/// - The capacity in bits (may be larger than requested due to word alignment)
///
/// # Panics
///
/// This function will panic if memory allocation fails.
#[inline(always)]
pub unsafe fn bitset_new<T: BitWord>(len: usize) -> (NonNull<T>, usize) {
    if len == 0 {
        return (NonNull::dangling(), 0);
    }
    let words = word_len::<T>(len);
    let layout = Layout::array::<T>(words).unwrap();
    let raw_ptr = unsafe { std::alloc::alloc_zeroed(layout) } as *mut T;
    if let Some(ptr) = NonNull::new(raw_ptr) {
        return (ptr, words * T::BITS);
    } else {
        std::alloc::handle_alloc_error(layout);
    }
}

/// Allocates an uninitialized bitset of the given length (in bits).
///
/// # Safety
/// Caller must ensure the returned pointer is properly freed using `bitset_free`.
#[inline(always)]
pub unsafe fn bitset_new_uninit<T: BitWord>(len: usize) -> (NonNull<T>, usize) {
    if len == 0 {
        return (NonNull::dangling(), 0);
    }
    let words = word_len::<T>(len);
    let layout = Layout::array::<T>(words).unwrap();
    let raw_ptr = unsafe { std::alloc::alloc(layout) } as *mut T;
    if let Some(ptr) = NonNull::new(raw_ptr) {
        return (ptr, words * T::BITS);
    } else {
        std::alloc::handle_alloc_error(layout);
    }
}

/// Resizes the bitset to a new length (in bits).
///
/// # Safety
/// Caller must ensure the returned pointer is properly freed using `bitset_free`.
#[inline(always)]
pub unsafe fn bitset_resize<T: BitWord>(
    ptr: NonNull<T>,
    old_len: usize,
    new_len: usize,
) -> (NonNull<T>, usize) {
    let new_word_len = word_len::<T>(new_len);
    let old_word_len = word_len::<T>(old_len);
    if new_len == 0 {
        // Free memory if shrinking to zero
        if old_len != 0 {
            let old_layout = Layout::array::<T>(old_word_len).unwrap();
            unsafe { std::alloc::dealloc(ptr.cast().as_ptr(), old_layout) };
        }
        return (NonNull::dangling(), 0);
    }
    if new_word_len != old_word_len {
        // Handle reallocation
        let old_layout = Layout::array::<T>(old_word_len).unwrap();
        let new_layout = Layout::array::<T>(new_word_len).unwrap();
        let raw_ptr = if old_len == 0 {
            // If old length is 0, allocate new memory
            unsafe { std::alloc::alloc(new_layout) }
        } else {
            // Otherwise, reallocate
            unsafe { std::alloc::realloc(ptr.cast().as_ptr(), old_layout, new_layout.size()) }
        } as *mut T;
        if let Some(ptr) = NonNull::new(raw_ptr) {
            // Zero-initialize new words if growing
            if new_word_len > old_word_len {
                unsafe {
                    ptr.add(old_word_len)
                        .write_bytes(0, new_word_len - old_word_len)
                };
            }
            return (ptr, new_word_len * T::BITS);
        } else {
            std::alloc::handle_alloc_error(new_layout);
        }
    }
    // If word length is the same, just update the length
    return (ptr, new_word_len * T::BITS);
}

/// Frees the memory allocated for the bitset.
///
/// # Safety
/// The pointer must have been allocated by `bitset_new`, `bitset_new_uninit`, or `bitset_resize`.
#[inline(always)]
pub unsafe fn bitset_free<T: BitWord>(ptr: NonNull<T>, len: usize) {
    if len > 0 {
        let word_len = word_len::<T>(len);
        let layout = Layout::array::<T>(word_len).unwrap();
        unsafe { std::alloc::dealloc(ptr.cast().as_ptr(), layout) };
    }
}

/// Gets the value of the bit at `idx` in the bitset pointed to by `ptr` (with `len` bits).
///
/// # Safety
/// Caller must ensure `idx < len`.
#[inline(always)]
pub unsafe fn bitset_test<T: BitWord>(ptr: NonNull<T>, len: usize, idx: usize) -> bool {
    let (word, bit) = word_bit::<T>(idx);
    if idx < len {
        *unsafe { ptr.add(word).as_ref() } & T::from(1u8) << bit != T::from(0u8)
    } else {
        false
    }
}

/// Sets the bit at `idx` in the bitset pointed to by `ptr` (with `len` bits).
///
/// # Safety
/// Caller must ensure `idx < len`.
#[inline(always)]
pub unsafe fn bitset_set<T: BitWord>(ptr: NonNull<T>, len: usize, idx: usize) {
    debug_assert!(idx < len, "Index out of bounds: {} >= {}", idx, len);
    let (word, bit) = word_bit::<T>(idx);
    *unsafe { ptr.add(word).as_mut() } |= T::from(1u8) << bit;
}

/// Clears the bit at `idx` in the bitset pointed to by `ptr` (with `len` bits).
///
/// # Safety
/// Caller must ensure `idx < len`.
#[inline(always)]
pub unsafe fn bitset_reset<T: BitWord>(ptr: NonNull<T>, len: usize, idx: usize) {
    debug_assert!(idx < len, "Index out of bounds: {} >= {}", idx, len);
    let (word, bit) = word_bit::<T>(idx);
    if idx < len {
        *unsafe { ptr.add(word).as_mut() } &= !(T::from(1u8) << bit);
    }
}

/// Flips the bit at `idx` in the bitset pointed to by `ptr` (with `len` bits).
///
/// # Safety
/// Caller must ensure `idx < len`.
#[inline(always)]
pub unsafe fn bitset_flip<T: BitWord>(ptr: NonNull<T>, len: usize, idx: usize) {
    debug_assert!(idx < len, "Index out of bounds: {} >= {}", idx, len);
    let (word, bit) = word_bit::<T>(idx);
    if idx < len {
        *unsafe { ptr.add(word).as_mut() } ^= T::from(1u8) << bit;
    }
}

/// Tests and sets the bit at `idx`. Returns the previous value.
///
/// # Safety
/// Caller must ensure `idx < len`.
#[inline(always)]
pub unsafe fn bitset_test_and_set<T: BitWord>(ptr: NonNull<T>, len: usize, idx: usize) -> bool {
    debug_assert!(idx < len, "Index out of bounds: {} >= {}", idx, len);
    let prev = unsafe { bitset_test(ptr, len, idx) };
    unsafe { bitset_set(ptr, len, idx) };
    prev
}

/// Tests and resets the bit at `idx`. Returns the previous value.
///
/// # Safety
/// Caller must ensure `idx < len`.
#[inline(always)]
pub unsafe fn bitset_test_and_reset<T: BitWord>(ptr: NonNull<T>, len: usize, idx: usize) -> bool {
    debug_assert!(idx < len, "Index out of bounds: {} >= {}", idx, len);
    let prev = unsafe { bitset_test(ptr, len, idx) };
    unsafe { bitset_reset(ptr, len, idx) };
    prev
}

/// Tests and resets the bit at `idx`, returning the previous value.
///
/// # Safety
/// Caller must ensure `idx < len`.
pub unsafe fn bitset_take<T: BitWord>(ptr: NonNull<T>, len: usize, idx: usize) -> bool {
    debug_assert!(idx < len, "Index out of bounds: {} >= {}", idx, len);
    unsafe { bitset_test_and_reset(ptr, len, idx) }
}

/// Replaces the bit at `idx` with `value`, returning the previous value.
///
/// # Safety
/// Caller must ensure `idx < len`.
#[inline(always)]
pub unsafe fn bitset_replace<T: BitWord>(
    ptr: NonNull<T>,
    len: usize,
    idx: usize,
    value: bool,
) -> bool {
    debug_assert!(idx < len, "Index out of bounds: {} >= {}", idx, len);
    let prev = unsafe { bitset_test(ptr, len, idx) };
    if value {
        unsafe { bitset_set(ptr, len, idx) };
    } else {
        unsafe { bitset_reset(ptr, len, idx) };
    }
    prev
}

/// Fills all bits in the bitset pointed to by `ptr` (with `len` bits) with 1.
///
/// # Safety
/// Caller must ensure the pointer is valid for `len` bits.
#[inline(always)]
pub unsafe fn bitset_fill<T: BitWord>(ptr: NonNull<T>, len: usize) {
    let (word, bit) = word_bit::<T>(len);
    unsafe { ptr.write_bytes(u8::MAX, word) };
    // Clear unused bits in the last word
    if bit != 0 {
        unsafe { ptr.add(word).write((T::from(1u8) << bit) - T::from(1u8)) };
    }
}

/// Clears all bits in the bitset pointed to by `ptr` (with `len` bits).
///
/// # Safety
/// Caller must ensure the pointer is valid for `len` bits.
#[inline(always)]
pub unsafe fn bitset_clear<T: BitWord>(ptr: NonNull<T>, len: usize) {
    let words = word_len::<T>(len);
    unsafe { ptr.write_bytes(0, words) };
}

/// Returns true if all bits are zero in the bitset pointed to by `ptr` (with `len` bits).
///
/// # Safety
/// Caller must ensure the pointer is valid for `len` bits.
#[inline(always)]
pub unsafe fn bitset_is_empty<T: BitWord>(ptr: NonNull<T>, len: usize) -> bool {
    let (word, bit) = word_bit::<T>(len);
    for i in 0..word {
        if *unsafe { ptr.add(i).as_ref() } != T::from(0u8) {
            return false;
        }
    }
    if bit != 0 {
        let mask = (T::from(1u8) << bit) - T::from(1u8);
        if *unsafe { ptr.add(word).as_ref() } & mask != T::from(0u8) {
            return false;
        }
    }
    true
}

/// Counts the number of set bits in the bitset pointed to by `ptr` (with `len` bits).
///
/// # Safety
/// Caller must ensure the pointer is valid for `len` bits.
#[inline(always)]
pub unsafe fn bitset_count_ones<T: BitWord>(ptr: NonNull<T>, len: usize) -> usize {
    let (word, bit) = word_bit::<T>(len);
    let words = unsafe { NonNull::slice_from_raw_parts(ptr, word).as_ref() };
    let mut count = words.iter().fold(0, |count, w| count + w.count_ones());
    if bit != 0 {
        count += (*unsafe { ptr.add(word).as_ref() } & (T::from(1u8) << bit) - T::from(1u8))
            .count_ones() as usize;
    }
    count
}

/// Returns true if all bits are set in the bitset pointed to by `ptr` (with `len` bits).
///
/// # Safety
/// Caller must ensure the pointer is valid for `len` bits.
#[inline(always)]
pub unsafe fn bitset_all<T: BitWord>(ptr: NonNull<T>, len: usize) -> bool {
    let (word, bit) = word_bit::<T>(len);
    for i in 0..word {
        if *unsafe { ptr.add(i).as_ref() } != T::MAX {
            return false;
        }
    }
    if bit != 0 {
        let mask = (T::from(1u8) << bit) - T::from(1u8);
        if *unsafe { ptr.add(word).as_ref() } & mask != mask {
            return false;
        }
    }
    true
}

/// Returns true if any bit is set in the bitset pointed to by `ptr` (with `len` bits).
///
/// # Safety
/// Caller must ensure the pointer is valid for `len` bits.
#[inline(always)]
pub unsafe fn bitset_any<T: BitWord>(ptr: NonNull<T>, len: usize) -> bool {
    let (word, bit) = word_bit::<T>(len);
    for i in 0..word {
        if *unsafe { ptr.add(i).as_ref() } != T::from(0u8) {
            return true;
        }
    }
    if bit != 0 {
        let mask = (T::from(1u8) << bit) - T::from(1u8);
        if *unsafe { ptr.add(word).as_ref() } & mask != T::from(0u8) {
            return true;
        }
    }
    false
}

/// Returns true if no bits are set in the bitset pointed to by `ptr` (with `len` bits).
///
/// # Safety
/// Caller must ensure the pointer is valid for `len` bits.
#[inline(always)]
pub unsafe fn bitset_none<T: BitWord>(ptr: NonNull<T>, len: usize) -> bool {
    !unsafe { bitset_any(ptr, len) }
}

#[cfg(test)]
mod tests {

    use std::fmt::Debug;

    use super::*;

    #[allow(dead_code)]
    fn bitset_print<T: Debug + BitWord>(ptr: NonNull<T>, len: usize) {
        let words = word_len::<T>(len);
        let store = unsafe { std::slice::from_raw_parts(ptr.as_ptr(), words) };
        eprintln!("len: {len}, store: {store:?}");
    }

    #[test]
    fn test_word_len() {
        assert_eq!(word_len::<usize>(0), 0);
        assert_eq!(word_len::<u8>(1), 1);
        assert_eq!(word_len::<u8>(7), 1);
        assert_eq!(word_len::<u8>(8), 1);
        assert_eq!(word_len::<u8>(9), 2);
        assert_eq!(word_len::<u8>(15), 2);
        assert_eq!(word_len::<u8>(16), 2);
        assert_eq!(word_len::<u8>(17), 3);
    }

    #[test]
    fn test_word_bit() {
        assert_eq!(word_bit::<u8>(0), (0, 0));
        assert_eq!(word_bit::<u8>(1), (0, 1));
        assert_eq!(word_bit::<u8>(7), (0, 7));
        assert_eq!(word_bit::<u8>(8), (1, 0));
        assert_eq!(word_bit::<u8>(9), (1, 1));
        assert_eq!(word_bit::<u8>(15), (1, 7));
        assert_eq!(word_bit::<u8>(16), (2, 0));
        assert_eq!(word_bit::<u8>(17), (2, 1));
    }

    #[test]
    fn test_mem() {
        let mut len = 7;
        unsafe {
            let (ptr, cap) = bitset_new::<u32>(0);
            assert_eq!(ptr, NonNull::dangling());
            assert_eq!(cap, 0);
            let (ptr, cap) = bitset_resize(ptr, cap, 0);
            assert_eq!(ptr, NonNull::dangling());
            assert_eq!(cap, 0);
            let (ptr, cap) = bitset_resize(ptr, cap, len);
            assert_ne!(ptr, NonNull::dangling());
            assert_eq!(
                cap,
                (len + u32::BITS as usize - 1) / u32::BITS as usize * u32::BITS as usize
            );
            assert!(!bitset_test(ptr, len, len - 1));
            bitset_set(ptr, len, len - 1);
            assert!(bitset_test(ptr, len, len - 1));
            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                bitset_set::<u32>(ptr, len, len)
            }));
            assert!(result.is_err());
            let (ptr, _cap) = bitset_resize(ptr, len, len + 1);
            len += 1;
            assert!(!bitset_test(ptr, len, len - 1));
            bitset_set(ptr, len, len - 1);
            assert!(bitset_test(ptr, len, len - 1));
            let result = std::panic::catch_unwind(|| bitset_set::<u32>(ptr, len, len));
            assert!(result.is_err());
            let cap = (len + 17) * 31;
            let (ptr, _cap) = bitset_resize(ptr, len, cap);
            assert_ne!(ptr, NonNull::dangling());
            bitset_free(ptr, 0);
        }
    }

    #[test]
    fn test_set_and_bitset_test() {
        let len = 100;
        unsafe {
            let (ptr, words) = bitset_new::<u8>(len);
            for i in 0..len {
                assert!(!bitset_test(ptr, len, i));
                bitset_set(ptr, len, i);
                assert!(bitset_test(ptr, len, i));
                bitset_reset(ptr, len, i);
                assert!(!bitset_test(ptr, len, i));
            }
            bitset_free(ptr, words);
        }
    }

    #[test]
    fn test_fill_and_bitset_clear() {
        let len = 77;
        unsafe {
            let (ptr, words) = bitset_new::<u8>(len);
            bitset_fill(ptr, len);
            for i in 0..len {
                assert!(bitset_test(ptr, len, i));
            }
            bitset_clear(ptr, len);
            for i in 0..len {
                assert!(!bitset_test(ptr, len, i));
            }
            bitset_free(ptr, words);
        }
    }

    #[test]
    fn test_bitset_count_ones() {
        let len = 65;
        unsafe {
            let (ptr, words) = bitset_new::<u8>(len);
            assert_eq!(bitset_count_ones(ptr, len), 0);
            bitset_set(ptr, len, 0);
            bitset_set(ptr, len, 1);
            bitset_set(ptr, len, 64);
            assert_eq!(bitset_count_ones(ptr, len), 3);
            bitset_fill(ptr, len);
            assert_eq!(bitset_count_ones(ptr, len), len);
            bitset_clear(ptr, len);
            assert_eq!(bitset_count_ones(ptr, len), 0);
            bitset_free(ptr, words);
        }
    }

    #[test]
    fn test_bitset_all_any_none() {
        let len = 33;
        unsafe {
            let (ptr, words) = bitset_new::<u8>(len);
            assert!(bitset_none(ptr, len));
            assert!(!bitset_any(ptr, len));
            assert!(!bitset_all(ptr, len));

            bitset_set(ptr, len, 0);
            assert!(!bitset_none(ptr, len));
            assert!(bitset_any(ptr, len));
            assert!(!bitset_all(ptr, len));

            bitset_fill(ptr, len);
            assert!(!bitset_none(ptr, len));
            assert!(bitset_any(ptr, len));
            assert!(bitset_all(ptr, len));

            bitset_reset(ptr, len, 0);
            assert!(!bitset_none(ptr, len));
            assert!(bitset_any(ptr, len));
            assert!(!bitset_all(ptr, len));
            bitset_free(ptr, words);
        }
    }

    #[test]
    fn test_bitset_flip() {
        let len = 10;
        unsafe {
            let (ptr, words) = bitset_new::<u8>(len);
            for i in 0..len {
                assert!(!bitset_test(ptr, len, i));
                bitset_flip(ptr, len, i);
                assert!(bitset_test(ptr, len, i));
                bitset_flip(ptr, len, i);
                assert!(!bitset_test(ptr, len, i));
            }
            bitset_free(ptr, words);
        }
    }

    #[test]
    fn test_bitset_take() {
        let len = 10;
        unsafe {
            let (ptr, words) = bitset_new::<u8>(len);
            bitset_fill(ptr, len);
            for i in 0..len {
                assert!(bitset_take(ptr, len, i));
                assert!(!bitset_take(ptr, len, i));
            }
            bitset_free(ptr, words);
        }
    }

    #[test]
    fn test_bitset_replace() {
        let len = 10;
        unsafe {
            let (ptr, words) = bitset_new::<u8>(len);
            for i in 0..len {
                assert!(!bitset_replace(ptr, len, i, true));
                assert!(bitset_test(ptr, len, i));
                assert!(bitset_replace(ptr, len, i, false));
                assert!(!bitset_test(ptr, len, i));
            }
            bitset_free(ptr, words);
        }
    }

    #[test]
    fn test_bitset_test_and_set() {
        let len = 10;
        unsafe {
            let (ptr, words) = bitset_new::<u8>(len);
            for i in 0..len {
                assert!(!bitset_test_and_set(ptr, len, i));
                assert!(bitset_test_and_set(ptr, len, i));
            }
            bitset_free(ptr, words);
        }
    }

    #[test]
    fn test_bitset_fill_and_clear() {
        let len = 100;
        unsafe {
            let (ptr, words) = bitset_new::<u8>(len);
            bitset_fill(ptr, len);
            for i in 0..len {
                assert!(bitset_test(ptr, len, i));
            }
            bitset_clear(ptr, len);
            for i in 0..len {
                assert!(!bitset_test(ptr, len, i));
            }
            bitset_free(ptr, words);
        }
    }

    #[test]
    #[should_panic]
    fn test_out_of_bounds() {
        let len = 10;
        unsafe {
            let (ptr, words) = bitset_new::<u8>(len);
            // Should not panic or set anything
            bitset_set(ptr, len, len + 5);
            assert_eq!(bitset_count_ones(ptr, len), 0);
            bitset_reset(ptr, len, len + 5);
            assert_eq!(bitset_count_ones(ptr, len), 0);
            assert!(!bitset_test(ptr, len, len + 5));
            bitset_free(ptr, words);
        }
    }
}