semx_bitops 0.1.61

//! 动态大小位图

use alloc::{vec, vec::Vec};

use crate::{
    BITS_PER_LONG, bit_word, bits_to_longs, change_bit, clear_bit, find_first_bit,
    find_first_zero_bit, find_last_bit, find_next_bit, find_next_zero_bit, set_bit,
    test_and_change_bit, test_and_clear_bit, test_and_set_bit, test_bit,
};

pub(super) const fn bitmap_first_word_mask(start: usize) -> usize {
    usize::MAX << (start & (BITS_PER_LONG - 1))
}

pub(super) const fn bitmap_last_word_mask(nbits: usize) -> usize {
    usize::MAX >> (nbits.wrapping_neg() & (BITS_PER_LONG - 1))
}

const fn small_nbits(nbits: usize) -> bool {
    nbits <= BITS_PER_LONG
}

/// 动态大小位图
pub struct BitMap {
    map: Vec<usize>,
    nr: usize,
}

impl Default for BitMap {
    fn default() -> Self {
        Self::new()
    }
}

impl BitMap {
    /// 构造空位图, 需调用 [`init`](Self::init) 后使用
    pub const fn new() -> Self {
        let map = Vec::new();
        Self { map, nr: 0 }
    }

    /// 重置为 `nr` 位的全零位图
    #[inline]
    pub fn init(&mut self, nr: usize) {
        self.nr = nr;
        let sz = bits_to_longs(nr);
        self.map = vec![0; sz];
    }

    /// 创建 `nr` 位的全零位图
    #[inline]
    pub fn create(nr: usize) -> Self {
        let sz = bits_to_longs(nr);
        let map = vec![0; sz];
        Self { map, nr }
    }

    /// 全部清零: `*map = 0`
    #[inline]
    pub fn set_zero(&mut self) {
        self.map.fill(0);
    }

    /// 全部置一: `*map = !0`
    #[inline]
    pub fn set_fill(&mut self) {
        self.map.fill(usize::MAX);
    }

    /// 复制: `*self = *from`
    ///
    /// # Panics
    /// `self.nr != from.nr`
    #[inline]
    pub fn copy_from(&mut self, from: &BitMap) {
        assert_eq!(self.nr, from.nr);
        self.map.copy_from_slice(&from.map);
    }

    /// 与: `*self = *src1 & *src2`
    ///
    /// # Panics
    /// 三者 `nr` 不一致
    #[inline]
    pub fn and(&mut self, src1: &BitMap, src2: &BitMap) {
        assert!(self.nr == src1.nr && self.nr == src2.nr);
        for ((dst, &s1), &s2) in self.map.iter_mut().zip(&src1.map).zip(&src2.map) {
            *dst = s1 & s2;
        }
    }

    /// 或: `*self = *src1 | *src2`
    ///
    /// # Panics
    /// 三者 `nr` 不一致
    #[inline]
    pub fn or(&mut self, src1: &BitMap, src2: &BitMap) {
        assert!(self.nr == src1.nr && self.nr == src2.nr);
        for ((dst, &s1), &s2) in self.map.iter_mut().zip(&src1.map).zip(&src2.map) {
            *dst = s1 | s2;
        }
    }

    /// 异或: `*self = *src1 ^ *src2`
    ///
    /// # Panics
    /// 三者 `nr` 不一致
    #[inline]
    pub fn xor(&mut self, src1: &BitMap, src2: &BitMap) {
        assert!(self.nr == src1.nr && self.nr == src2.nr);
        for ((dst, &s1), &s2) in self.map.iter_mut().zip(&src1.map).zip(&src2.map) {
            *dst = s1 ^ s2;
        }
    }

    /// 与非: `*self = *src1 & !(*src2)`
    ///
    /// # Panics
    /// 三者 `nr` 不一致
    #[inline]
    pub fn andnot(&mut self, src1: &BitMap, src2: &BitMap) {
        assert!(self.nr == src1.nr && self.nr == src2.nr);
        for ((dst, &s1), &s2) in self.map.iter_mut().zip(&src1.map).zip(&src2.map) {
            *dst = s1 & !s2;
        }
    }

    /// 取反: `*self = !(*src)`
    ///
    /// # Panics
    /// `self.nr != src.nr`
    #[inline]
    pub fn complement(&mut self, src: &BitMap) {
        assert_eq!(self.nr, src.nr);
        for (dst, &s) in self.map.iter_mut().zip(&src.map) {
            *dst = !s;
        }
    }

    /// 判断两个位图有效位是否完全相同
    ///
    /// # Panics
    /// `self.nr != cmp.nr`
    #[inline]
    pub fn is_equal(&self, cmp: &BitMap) -> bool {
        assert_eq!(self.nr, cmp.nr);
        if self.nr == 0 {
            return true;
        }

        if small_nbits(self.nr) {
            ((self.map[0] ^ cmp.map[0]) & bitmap_last_word_mask(self.nr)) == 0
        } else {
            let lim = self.nr / BITS_PER_LONG;
            for i in 0..lim {
                if self.map[i] != cmp.map[i] {
                    return false;
                }
            }

            if !self.nr.is_multiple_of(BITS_PER_LONG)
                && (self.map[lim] ^ cmp.map[lim]) & bitmap_last_word_mask(self.nr) != 0
            {
                return false;
            }
            true
        }
    }

    /// 是否所有有效位均为零
    #[inline]
    pub fn is_empty(&self) -> bool {
        if self.nr == 0 {
            return true;
        }

        if small_nbits(self.nr) {
            (self.map[0] & bitmap_last_word_mask(self.nr)) == 0
        } else {
            let lim = self.nr / BITS_PER_LONG;
            for i in 0..lim {
                if self.map[i] != 0 {
                    return false;
                }
            }
            if !self.nr.is_multiple_of(BITS_PER_LONG)
                && self.map[lim] & bitmap_last_word_mask(self.nr) != 0
            {
                return false;
            }
            true
        }
    }

    /// 是否所有有效位均为一
    #[inline]
    pub fn is_full(&self) -> bool {
        if self.nr == 0 {
            return true;
        }

        if small_nbits(self.nr) {
            return (!self.map[0] & bitmap_last_word_mask(self.nr)) == 0;
        }
        let lim = self.nr / BITS_PER_LONG;
        for i in 0..lim {
            if self.map[i] != usize::MAX {
                return false;
            }
        }
        if !self.nr.is_multiple_of(BITS_PER_LONG)
            && (!self.map[lim] & bitmap_last_word_mask(self.nr)) != 0
        {
            return false;
        }
        true
    }

    /// 判断两个位图是否存在同时为一的位
    ///
    /// # Panics
    /// `self.nr != src.nr`
    #[inline]
    pub fn intersects(&self, src: &BitMap) -> bool {
        assert_eq!(self.nr, src.nr);
        if self.nr == 0 {
            return false;
        }

        if small_nbits(self.nr) {
            (self.map[0] & src.map[0]) & bitmap_last_word_mask(self.nr) != 0
        } else {
            let lim = self.nr / BITS_PER_LONG;
            for i in 0..lim {
                if self.map[i] & src.map[i] != 0 {
                    return true;
                }
            }
            if !self.nr.is_multiple_of(BITS_PER_LONG)
                && (self.map[lim] & src.map[lim]) & bitmap_last_word_mask(self.nr) != 0
            {
                return true;
            }
            false
        }
    }

    /// 统计置一的位数（popcount）
    #[inline]
    pub fn count_ones(&self) -> usize {
        if self.nr == 0 {
            return 0;
        }

        if small_nbits(self.nr) {
            (self.map[0] & bitmap_last_word_mask(self.nr)).count_ones() as usize
        } else {
            let lim = self.nr / BITS_PER_LONG;
            let mut result = 0;
            for i in 0..lim {
                result += self.map[i].count_ones();
            }
            if !self.nr.is_multiple_of(BITS_PER_LONG) {
                result += (self.map[lim] & bitmap_last_word_mask(self.nr)).count_ones();
            }
            result as usize
        }
    }

    /// 将 `[start, start + nr)` 区间内的位全部置一
    ///
    /// # Panics
    /// `start + nr > self.nr`
    #[inline]
    pub fn set_range(&mut self, start: usize, mut nr: usize) {
        assert!(self.nr >= start + nr);
        let mut pos = bit_word(start);
        let size = start + nr;
        let mut bits_to_set = BITS_PER_LONG - (start % BITS_PER_LONG);
        let mut mask_to_set = bitmap_first_word_mask(start);

        while nr >= bits_to_set {
            self.map[pos] |= mask_to_set;
            nr -= bits_to_set;
            bits_to_set = BITS_PER_LONG;
            mask_to_set = usize::MAX;
            pos += 1;
        }
        if nr != 0 {
            mask_to_set &= bitmap_last_word_mask(size);
            self.map[pos] |= mask_to_set;
        }
    }

    /// 将 `[start, start + nr)` 区间内的位全部清零
    ///
    /// # Panics
    /// `start + nr > self.nr`
    #[inline]
    pub fn clear_range(&mut self, start: usize, mut nr: usize) {
        assert!(self.nr >= start + nr);
        let mut pos = bit_word(start);
        let size = start + nr;
        let mut bits_to_clear = BITS_PER_LONG - (start % BITS_PER_LONG);
        let mut mask_to_clear = bitmap_first_word_mask(start);

        while nr >= bits_to_clear {
            self.map[pos] &= !mask_to_clear;
            nr -= bits_to_clear;
            bits_to_clear = BITS_PER_LONG;
            mask_to_clear = usize::MAX;
            pos += 1;
        }
        if nr != 0 {
            mask_to_clear &= bitmap_last_word_mask(size);
            self.map[pos] &= !mask_to_clear;
        }
    }

    /// 从 `start` 起查找连续 `nr` 个零位的对齐区域, 返回起始索引
    ///
    /// 找不到时返回值 > `self.nr`
    #[inline]
    pub fn find_next_zero_area(&self, mut start: usize, nr: usize, align_mask: usize) -> usize {
        loop {
            let mut index = find_next_zero_bit(&self.map, self.nr, start);
            index = (index + align_mask) & !align_mask;
            let end = index + nr;
            if end > self.nr {
                return end;
            }
            let i = find_next_bit(&self.map, end, index);
            if i < end {
                start = i + 1;
                continue;
            }
            return index;
        }
    }

    /// 在尾部扩展 `add` 个零位
    #[inline]
    pub fn extend(&mut self, add: usize) {
        let tail = self.nr & (BITS_PER_LONG - 1);
        if tail != 0 {
            let idx = self.nr / BITS_PER_LONG;
            self.map[idx] &= bitmap_last_word_mask(self.nr);
        }
        self.map.resize(bits_to_longs(self.nr + add), 0);
        self.nr += add;
    }

    /// 返回位图总位数
    #[inline]
    pub fn nr(&self) -> usize {
        self.nr
    }

    /// 置位第 `nr` 位
    #[inline]
    pub fn set_bit(&mut self, nr: usize) {
        set_bit(nr, &mut self.map);
    }

    /// 清除第 `nr` 位
    #[inline]
    pub fn clear_bit(&mut self, nr: usize) {
        clear_bit(nr, &mut self.map);
    }

    /// 翻转第 `nr` 位
    #[inline]
    pub fn change_bit(&mut self, nr: usize) {
        change_bit(nr, &mut self.map);
    }

    /// 测试第 `nr` 位是否为一
    #[inline]
    pub fn test_bit(&self, nr: usize) -> bool {
        test_bit(nr, &self.map)
    }

    /// 置位第 `nr` 位, 返回该位的旧值
    #[inline]
    pub fn test_and_set_bit(&mut self, nr: usize) -> bool {
        test_and_set_bit(nr, &mut self.map)
    }

    /// 清除第 `nr` 位, 返回该位的旧值
    #[inline]
    pub fn test_and_clear_bit(&mut self, nr: usize) -> bool {
        test_and_clear_bit(nr, &mut self.map)
    }

    /// 翻转第 `nr` 位, 返回该位的旧值
    #[inline]
    pub fn test_and_change_bit(&mut self, nr: usize) -> bool {
        test_and_change_bit(nr, &mut self.map)
    }

    /// 查找第一个为一的位, 找不到时返回 `self.nr`
    #[inline]
    pub fn find_first_bit(&self) -> usize {
        find_first_bit(&self.map, self.nr)
    }

    /// 查找第一个为零的位, 找不到时返回 `self.nr`
    #[inline]
    pub fn find_first_zero_bit(&self) -> usize {
        find_first_zero_bit(&self.map, self.nr)
    }

    /// 从 `offset` 起查找下一个为一的位, 找不到时返回 `self.nr`
    #[inline]
    pub fn find_next_bit(&self, offset: usize) -> usize {
        find_next_bit(&self.map, self.nr, offset)
    }

    /// 从 `offset` 起查找下一个为零的位, 找不到时返回 `self.nr`
    #[inline]
    pub fn find_next_zero_bit(&self, offset: usize) -> usize {
        find_next_zero_bit(&self.map, self.nr, offset)
    }

    /// 查找最后一个为一的位, 找不到时返回 `self.nr`
    #[inline]
    pub fn find_last_bit(&self) -> usize {
        find_last_bit(&self.map, self.nr)
    }
}