popset.rs

#[inline(always)]
fn most_significant_bit(n: u64) -> usize {
    use std::intrinsics::ctlz;
    (63u64 - unsafe {ctlz(n)}) as usize
}

pub struct PopSet {
    upper: u64,
    lower: Box<[u64; 64]>,
}


impl PopSet {

    pub fn new() -> PopSet {
        PopSet {
            upper: 0u64,
            lower: Box::new([0u64; 64]),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.upper == 0
    }

    pub fn insert(&mut self, el: usize) {
        debug_assert!(el < 4096);
        let up = el / 64;
        self.upper |= 1 << up;
        let low = el % 64;
        self.lower[up] |= 1 << low;
    }

    pub fn pop(&mut self) -> Option<usize> {
        if self.is_empty() {
            None
        }
        else {
            let up = most_significant_bit(self.upper);
            let low = most_significant_bit(self.lower[up]);
            self.lower[up] ^= 1 << low;
            if self.lower[up] == 0 {
                self.upper ^= 1 << up;
            }
            Some(up * 64 + low)
        }
    }
}


#[cfg(test)]
mod tests {

    use super::PopSet;

    #[test]
    fn test_popset() {
        let mut popset = PopSet::new();
        assert_eq!(popset.pop(), None);
        popset.insert(10);
        popset.insert(11);
        popset.insert(4002);
        assert_eq!(popset.pop(), Some(4002));
        assert_eq!(popset.pop(), Some(11));
        assert_eq!(popset.pop(), Some(10));
        assert_eq!(popset.pop(), None);
        assert_eq!(popset.pop(), None);
        popset.insert(1);
        popset.insert(4095);
        assert_eq!(popset.pop(), Some(4095));
        assert_eq!(popset.pop(), Some(1));
        assert_eq!(popset.pop(), None);
    }
}