luallaby 0.1.0-alpha.3

/// A bit set struct specifically made for tracking reserved registers in the compilation phase.
pub struct BitSet {
    bits: Vec<usize>,
    count: usize,
}

impl BitSet {
    pub fn new() -> Self {
        Self {
            bits: vec![0],
            count: 0,
        }
    }

    pub fn reserve(&mut self) -> usize {
        // Check if there is space in currently stored bits
        for (i, bits) in self.bits.iter_mut().enumerate() {
            // All reserved, continue
            let trailing = bits.trailing_ones();
            if trailing >= usize::BITS {
                continue;
            }

            // Set the new bit and return index
            let ind = i * usize::BITS as usize + trailing as usize;
            self.count = self.count.max(ind + 1);
            let mask = 0x1 << trailing;
            *bits |= mask;
            return ind;
        }

        // No space, append new bits
        self.count += 1;
        let res = self.bits.len() * usize::BITS as usize;
        self.bits.push(0x1);
        res
    }

    pub fn free(&mut self, ind: usize) {
        let bits = self
            .bits
            .get_mut(ind / usize::BITS as usize)
            .expect("invalid register index");
        let ind = ind % usize::BITS as usize;
        if (*bits >> ind) & 0x1 == 0 {
            panic!("register was not reserved")
        }
        *bits &= !(0x1 << ind);
    }

    pub fn is_empty(&self) -> bool {
        for bits in &self.bits {
            if bits != &0 {
                return false;
            }
        }
        true
    }

    /// Returns the number of registers that have to be allocated to fit all of the reserved registers that have been
    /// reserved during the lifetime of this bit set.
    pub fn count(&self) -> usize {
        self.count
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn bitset() {
        let mut set = BitSet::new();
        assert!(set.is_empty());

        assert_eq!(0, set.reserve());
        assert_eq!(vec![1], set.bits);
        assert!(!set.is_empty());

        assert_eq!(1, set.reserve());
        assert_eq!(vec![3], set.bits);

        assert_eq!(2, set.reserve());
        assert_eq!(vec![7], set.bits);

        set.free(1);
        assert_eq!(vec![5], set.bits);

        assert_eq!(1, set.reserve());
        assert_eq!(vec![7], set.bits);

        set.free(0);
        assert_eq!(vec![6], set.bits);

        set.free(1);
        assert_eq!(vec![4], set.bits);

        set.free(2);
        assert_eq!(vec![0], set.bits);
        assert!(set.is_empty());

        assert_eq!(0, set.reserve());
        assert_eq!(vec![1], set.bits);
        assert!(!set.is_empty());
        assert_eq!(3, set.count());
    }

    #[test]
    fn bitset_expand() {
        let mut set = BitSet::new();

        for i in 0..usize::BITS as usize {
            assert_eq!(i, set.reserve());
        }
        assert_eq!(vec![usize::MAX], set.bits);

        assert_eq!(usize::BITS as usize, set.reserve());
        assert_eq!(vec![usize::MAX, 1], set.bits);

        assert_eq!(usize::BITS as usize + 1, set.reserve());
        assert_eq!(vec![usize::MAX, 3], set.bits);

        set.free(4);
        assert_eq!(vec![!(0x1 << 4), 3], set.bits);

        assert_eq!(usize::BITS as usize + 2, set.count());
    }
}