meowalloc 0.1.1

use core::{fmt, sync::atomic::{AtomicU32, Ordering}};

pub type Word = AtomicU32;
pub const WORD_BITS: usize = 32;

pub struct AtomicBitSet<const WORDS: usize> {
    words: [Word; WORDS]
}

impl<const WORDS: usize> fmt::Debug for AtomicBitSet<WORDS> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("[\n")?;
        for word in self.words.iter() {
            writeln!(f, "    {:032b},", word.load(Ordering::Relaxed))?
        }
        f.write_str("]")?;
        Ok(())
    }
}

#[derive(Debug)]
pub struct BitSetFull;
#[derive(Debug)]
pub struct BitSetEmpty;

impl<const WORDS: usize> AtomicBitSet<WORDS> {
    pub const fn zeros() -> Self {
        Self {
            words: [const { AtomicU32::new(0) }; WORDS]
        }
    }
    pub const fn ones() -> Self {
        Self {
            words: [const { AtomicU32::new(u32::MAX) }; WORDS]
        }
    }

    pub fn into_inner(self) -> [AtomicU32; WORDS] {
        self.words
    }
    pub fn from_words(words: [AtomicU32; WORDS]) -> Self {
        Self { words: words }
    }

    pub fn is_one(&self, i: usize) -> bool {
        if i >= WORDS * WORD_BITS {
            panic!("the index is {i}, but the length is {}", WORDS * WORD_BITS)
        }
        unsafe { self.is_one_unchecked(i) } 
    }
    
    pub fn is_zero(&self, i: usize) -> bool {
        if i >= WORDS * WORD_BITS {
            panic!("the index is {i}, but the length is {}", WORDS * WORD_BITS)
        }
        unsafe { self.is_zero_unchecked(i) } 
    }

    pub unsafe fn is_one_unchecked(&self, i: usize) -> bool {
        let bit_pos = (WORD_BITS - 1) - (i % WORD_BITS);
        let slot = unsafe { self.words.get_unchecked(i / WORD_BITS) };
        let word = slot.load(Ordering::SeqCst);
        (word & (1 << bit_pos)) != 0
    }
    
    pub unsafe fn is_zero_unchecked(&self, i: usize) -> bool {
        let bit_pos = (WORD_BITS - 1) - (i % WORD_BITS);
        let slot = unsafe { self.words.get_unchecked(i / WORD_BITS) };
        let word = slot.load(Ordering::SeqCst);
        (word & (1 << bit_pos)) == 0
    }

    pub fn set_if_zero(&self, i: usize) -> bool {
        if i >= WORDS * WORD_BITS {
            panic!("the index is {i}, but the length is {}", WORDS * WORD_BITS)
        }
        unsafe { self.set_if_zero_unchecked(i) }
    }
    
    pub fn unset_if_one(&self, i: usize) -> bool {
        if i >= WORDS * WORD_BITS {
            panic!("the index is {i}, but the length is {}", WORDS * WORD_BITS)
        }
        unsafe { self.unset_if_one_unchecked(i) }
    }

    pub fn set(&self, i: usize) {
        if i >= WORDS * WORD_BITS {
            panic!("the index is {i}, but the length is {}", WORDS * WORD_BITS)
        }
        unsafe { self.set_unchecked(i); }
    }
    
    pub fn unset(&self, i: usize) {
        if i >= WORDS * WORD_BITS {
            panic!("the index is {i}, but the length is {}", WORDS * WORD_BITS)
        }
        unsafe { self.unset_unchecked(i); }
    }

    pub unsafe fn set_unchecked(&self, i: usize) {
        let shift = (WORD_BITS - 1) - (i % WORD_BITS);
        let slot = unsafe { self.words.get_unchecked(i / WORD_BITS) };

        slot.update(
            Ordering::SeqCst, Ordering::SeqCst,
            |word| word | (1 << shift)
        );
    }

    pub unsafe fn unset_unchecked(&self, i: usize) {
        let shift = (WORD_BITS - 1) - (i % WORD_BITS);
        let slot = unsafe { self.words.get_unchecked(i / WORD_BITS) };

        slot.update(
            Ordering::SeqCst, Ordering::SeqCst,
            |word| word ^ (1 << shift)
        );
    }

    /// Atomically checks whether the bit at the given index is unset and sets it if so.
    ///
    /// # Safety:
    /// The caller must ensure that the index is within the bounds of the bitset
    /// (i < WORDS * WORD_BITS)
    pub unsafe fn set_if_zero_unchecked(&self, i: usize) -> bool {
        debug_assert!(i < WORDS * WORD_BITS);
        let shift = (WORD_BITS - 1) - (i % WORD_BITS);
        let slot = unsafe { self.words.get_unchecked(i / WORD_BITS) };
        
        slot.try_update(
            Ordering::SeqCst, Ordering::SeqCst,
            |word| {
                ((word & (1 << shift)) == 0)
                    .then_some(word | (1 << shift))
            }
        ).is_ok()
    }
    
    /// Atomically checks whether the bit at the given index is set and unsets it if so.
    ///
    /// # Safety:
    /// The caller must ensure that the index is within the bounds of the bitset
    /// (i < WORDS * WORD_BITS)
    pub unsafe fn unset_if_one_unchecked(&self, i: usize) -> bool { 
        debug_assert!(i < WORDS * WORD_BITS);
        let bit_pos = WORD_BITS - 1 - i % WORD_BITS;
        // Only process one slot at a time, safe
        let slot = unsafe { self.words.get_unchecked(i / WORD_BITS) };
        
        slot.try_update(
            Ordering::SeqCst, Ordering::SeqCst,
            |word| {
                ((word & (1 << bit_pos)) != 0)
                    .then_some(word ^ (1 << bit_pos))
            }
        ).is_ok()
    }

    pub fn unset_first_set(&self) -> Result<usize, BitSetEmpty> {
        // Only process one atomic at a time, safe
        let mut bit_pos = 0;
        for (idx, slot) in self.words.iter().enumerate() {
            if slot.try_update(
                Ordering::SeqCst, Ordering::SeqCst,
                |word| {
                    if word == 0 {
                        return None
                    }
                    for shift in (0..WORD_BITS).rev() {
                        if (word & (1 << shift)) != 0 {
                            bit_pos = WORD_BITS - 1 - shift;
                            return Some(word ^ (1 << shift))
                        }
                    }
                    None
                }
            ).is_ok() {
                return Ok(idx * WORD_BITS + bit_pos)
            }
        }
        Err(BitSetEmpty)
    }
    
    pub fn set_first_unset(&self) -> Result<usize, BitSetFull> {
        let mut bit_pos = 0;
        for (idx, slot) in self.words.iter().enumerate() {
            if slot.try_update(
                Ordering::SeqCst, Ordering::SeqCst,
                |word| {
                    if word == 0 {
                        return None
                    }
                    for shift in (0..WORD_BITS).rev() {
                        if (word & (1 << shift)) == 0 {
                            bit_pos = WORD_BITS - 1 - shift;
                            return Some(word | (1 << shift))
                        }
                    }
                    None
                }
            ).is_ok() {
                return Ok(idx * WORD_BITS + bit_pos)
            }
        }
        Err(BitSetFull)
    }
}

impl<const WORDS: usize> From<AtomicBitSet<WORDS>> for [Word; WORDS] {
    fn from(value: AtomicBitSet<WORDS>) -> Self { value.into_inner() }
}
impl<const WORDS: usize> From<[Word; WORDS]> for AtomicBitSet<WORDS> {
    fn from(value: [Word; WORDS]) -> Self { Self::from_words(value) }
}
impl<const WORDS: usize> From<[u32; WORDS]> for AtomicBitSet<WORDS> {
    fn from(value: [u32; WORDS]) -> Self {
        value.map(|n| AtomicU32::new(n)).into()
    }
}
impl<const WORDS: usize> From<AtomicBitSet<WORDS>> for [u32; WORDS] {
    fn from(value: AtomicBitSet<WORDS>) -> Self {
        value.into_inner().map(|n| n.into_inner())
    }
}
impl<const WORDS: usize> From<&AtomicBitSet<WORDS>> for [u32; WORDS] {
    fn from(value: &AtomicBitSet<WORDS>) -> Self {
        let mut arr = [0; WORDS];
        let iter = value.words.iter().map(|n| n.load(Ordering::Relaxed)).enumerate();
        for (i, elem) in iter {
            arr[i] = elem;
        }
        arr

    }
}

#[cfg(test)]
fn atomic_bitset_fixture() -> AtomicBitSet<4> {
    let set = AtomicBitSet::ones();
    set.unset(10);
    set.unset(126);
    set.unset(22);
    set.unset(25);
    set
}

#[test]
fn set_if_zero() {
    let set = atomic_bitset_fixture();
    assert!(!set.set_if_zero(0));
    assert!(!set.set_if_zero(9));
    assert!(!set.set_if_zero(127));

    assert!(set.set_if_zero(10));
    assert!(set.set_if_zero(22));
    assert!(set.set_if_zero(25));
    assert!(set.set_if_zero(126));
    
    assert!(!set.set_if_zero(10));
    assert!(!set.set_if_zero(22));
    assert!(!set.set_if_zero(25));
    assert!(!set.set_if_zero(126));
}

#[test]
fn unset_if_one() {
    let set = atomic_bitset_fixture();
    assert!(!set.unset_if_one(10));
    assert!(!set.unset_if_one(126));
    assert!(!set.unset_if_one(22));
    assert!(!set.unset_if_one(25));

    assert!(set.unset_if_one(0));
    assert!(set.unset_if_one(9));
    assert!(set.unset_if_one(127));
    assert!(set.unset_if_one(54));
    
    assert!(!set.unset_if_one(0));
    assert!(!set.unset_if_one(9));
    assert!(!set.unset_if_one(127));
    assert!(!set.unset_if_one(54));
}

#[test]
fn set_first_unset() {
    AtomicBitSet::<4>::ones().set_first_unset().unwrap_err();
    let set = atomic_bitset_fixture();
    
    assert_eq!(set.set_first_unset().unwrap(), 10);
    assert_eq!(set.set_first_unset().unwrap(), 22);
    assert_eq!(set.set_first_unset().unwrap(), 25);
    assert_eq!(set.set_first_unset().unwrap(), 126);
}

#[test]
fn unset_first_set() {
    AtomicBitSet::<4>::zeros().unset_first_set().unwrap_err();
    let set = atomic_bitset_fixture();
    set.unset(2);
    set.unset(3);
    set.unset(5);
    
    assert_eq!(set.unset_first_set().unwrap(), 0);
    assert_eq!(set.unset_first_set().unwrap(), 1);
    assert_eq!(set.unset_first_set().unwrap(), 4);
    assert_eq!(set.unset_first_set().unwrap(), 6);
}