semx_bitops 0.1.6

//! 基于 qemu 的 bitmap 实现

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

use crate::{
    BITS_PER_LONG, bit_word, bits_to_longs, find_next_bit, find_next_zero_bit, set_bit,
    test_and_set_bit,
};

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

#[allow(clippy::cast_sign_loss)]
const fn bitmap_last_word_mask(nbits: isize) -> usize {
    usize::MAX >> (-nbits as usize & (BITS_PER_LONG - 1))
}

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

/// 定义一个 bitmap
pub struct BitMap {
    map: Vec<usize>,
    nr: usize,
}

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

impl BitMap {
    /// 构造 bitmap
    pub const fn new() -> Self {
        let map = Vec::new();
        Self { map, nr: 0 }
    }

    /// 初始化 bitmap
    pub fn init(&mut self, nr: usize) {
        self.nr = nr;
        let sz = bits_to_longs(nr);
        self.map = vec![0; sz];
    }

    /// 创建一个 bitmap
    pub fn create(nr: usize) -> Self {
        let sz = bits_to_longs(nr);
        let map = vec![0; sz];
        Self { map, nr }
    }

    /// 清零 bitmap
    ///
    /// `*map = 0`
    #[inline]
    pub fn set_zero(&mut self) {
        for this in self.map.iter_mut() {
            *this = 0;
        }
    }

    /// 填充 bitmap
    ///
    /// `*map = !0`
    #[inline]
    pub fn set_fill(&mut self) {
        for this in self.map.iter_mut() {
            *this = usize::MAX;
        }
    }

    /// 从另一个复制 bitmap
    ///
    /// `*map = *from`
    ///
    /// # Panics
    /// 当`BitMap`定义`bit`数量不一致时将会 panic.
    #[inline]
    pub fn copy_from(&mut self, from: &BitMap) {
        assert_eq!(self.nr, from.nr);
        for i in 0..bits_to_longs(self.nr) {
            self.map[i] = from.map[i];
        }
    }

    /// 从另外两个进行与操作
    ///
    /// `*map = *src1 & *src2`
    ///
    /// # Panics
    /// 当`BitMap`定义`bit`数量不一致时将会 panic.
    #[inline]
    pub fn and(&mut self, src1: &BitMap, src2: &BitMap) {
        assert!(self.nr == src1.nr && self.nr == src2.nr);
        for i in 0..bits_to_longs(self.nr) {
            self.map[i] = src1.map[i] & src2.map[i];
        }
    }

    /// 从另外两个进行或操作
    ///
    /// `*map = *src1 | *src2`
    ///
    /// # Panics
    /// 当`BitMap`定义`bit`数量不一致时将会 panic.
    #[inline]
    pub fn or(&mut self, src1: &BitMap, src2: &BitMap) {
        assert!(self.nr == src1.nr && self.nr == src2.nr);
        for i in 0..bits_to_longs(self.nr) {
            self.map[i] = src1.map[i] | src2.map[i];
        }
    }

    /// 从另外两个进行异或操作
    ///
    /// `*map = *src1 ^ *src2`
    ///
    /// # Panics
    /// 当`BitMap`定义`bit`数量不一致时将会 panic.
    #[inline]
    pub fn xor(&mut self, src1: &BitMap, src2: &BitMap) {
        assert!(self.nr == src1.nr && self.nr == src2.nr);
        for i in 0..bits_to_longs(self.nr) {
            self.map[i] = src1.map[i] ^ src2.map[i];
        }
    }

    /// 从另外两个进行与非操作
    ///
    /// `*map = *src1 & !(*src2)`
    ///
    /// # Panics
    /// 当`BitMap`定义`bit`数量不一致时将会 panic.
    #[inline]
    pub fn andnot(&mut self, src1: &BitMap, src2: &BitMap) {
        assert!(self.nr == src1.nr && self.nr == src2.nr);
        for i in 0..bits_to_longs(self.nr) {
            self.map[i] = src1.map[i] & !src2.map[i];
        }
    }

    /// `*map = !(*src)`
    ///
    /// # Panics
    /// 当`BitMap`定义`bit`数量不一致时将会 panic.
    #[inline]
    pub fn complement(&mut self, src: &BitMap) {
        assert_eq!(self.nr, src.nr);
        for i in 0..bits_to_longs(self.nr) {
            self.map[i] = !src.map[i];
        }
    }

    /// 是否相等
    ///
    /// # Panics
    /// 当`BitMap`定义`bit`数量不一致时将会 panic.
    pub fn is_equal(&self, cmp: &BitMap) -> bool {
        assert_eq!(self.nr, cmp.nr);

        if small_nbits(self.nr) {
            ((self.map[0] ^ cmp.map[0]) & bitmap_last_word_mask(self.nr as isize)) == 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 & BITS_PER_LONG != 0
                && (self.map[lim] ^ cmp.map[lim]) & bitmap_last_word_mask(self.nr as isize) != 0
            {
                return false;
            }
            true
        }
    }

    /// 是否为空
    pub fn is_empty(&self) -> bool {
        if small_nbits(self.nr) {
            (self.map[0] & bitmap_last_word_mask(self.nr as isize)) == 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 as isize) != 0
            {
                return false;
            }
            true
        }
    }

    /// 是否满了
    pub fn is_full(&self) -> bool {
        if small_nbits(self.nr) {
            return (!self.map[0] & bitmap_last_word_mask(self.nr as isize)) == 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 as isize) != 0
            {
                return false;
            }
        }
        true
    }

    /// 是否相交
    ///
    /// # Panics
    /// 当`BitMap`定义`bit`数量不一致时将会 panic.
    pub fn is_intersects(&self, src: &BitMap) -> bool {
        assert_eq!(self.nr, src.nr);
        if small_nbits(self.nr) {
            (self.map[0] & src.map[0]) & bitmap_last_word_mask(self.nr as isize) != 0
        } else {
            let lim = self.nr / BITS_PER_LONG;
            for i in 0..lim {
                if (self.map[i] & src.map[i]) & bitmap_last_word_mask(self.nr as isize) != 0 {
                    return true;
                }
            }
            false
        }
    }

    /// 计算 bit1 的位数
    pub fn count_one(&self) -> usize {
        if small_nbits(self.nr) {
            (self.map[0] & bitmap_last_word_mask(self.nr as isize)).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 as isize)).count_ones();
            }
            result as usize
        }
    }

    /// 设置特定的 bit 域
    ///
    /// # Panics
    /// 当`start + nr`大于`BitMap`最大容量数将会 panic.
    pub fn bitmap_set(&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 as isize - bits_to_set as isize) >= 0 {
            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 as isize);
            self.map[pos] |= mask_to_set;
        }
    }

    /// 清除特定的 bit 域
    ///
    /// # Panics
    /// 当`start + nr`大于`BitMap`最大容量数将会 panic.
    pub fn bitmap_clear(&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 as isize - bits_to_clear as isize) >= 0 {
            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 as isize);
            self.map[pos] &= !mask_to_clear;
        }
    }

    /// 寻找下一个空闲域
    pub fn find_next_zero_area(&self, mut start: usize, nr: usize, align_mask: usize) -> usize {
        loop {
            let index = find_next_zero_bit(&self.map, self.nr, start);
            let 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;
        }
    }

    /// 扩大 bitmap
    pub fn extend_nr(&mut self, add: usize) {
        for _ in 0..(bits_to_longs(self.nr + add) - bits_to_longs(self.nr)) {
            self.map.push(0);
        }
        let start = self.nr;
        self.nr += add;
        self.bitmap_clear(start, add);
    }

    /// Bitmap nr 最大数量
    pub fn bitmap_nr(&self) -> usize {
        self.nr
    }

    /// 设置一个 bit
    pub fn bitmap_set_bit(&mut self, nr: usize) {
        let vec = &mut self.map;
        set_bit(nr, vec);
    }

    /// 测试并且设置一个 bit
    pub fn bitmap_test_and_set_bit(&mut self, nr: usize) -> bool {
        let vec = &mut self.map;
        test_and_set_bit(nr, vec)
    }

    /// 寻找下一个为零 bit
    pub fn bitmap_find_next_zero_bit(&self, size: usize, offset: usize) -> usize {
        let vec = &self.map;
        find_next_zero_bit(vec, size, offset)
    }
}