hierarchical bitset

hibitset.rs

#[derive(Debug, PartialEq, Eq)]
enum State {
    /// There are `0`s only.
    Empty,

    /// There may be `1`s and `0`s.
    Partial,

    /// There are `1`s only.
    Full,
}

trait BitSet {

    const INDEX_BITS: u32;
    const SIZE: u32;

    fn get(&self, index: u32) -> bool;
    fn set(&mut self, index: u32) -> bool;
    fn unset(&mut self, index: u32) -> bool;
    fn next_0(&self) -> Option<u32>;
    fn next_1(&self) -> Option<u32>;
    fn state(&self, range: Range<u32>) -> State;
}

impl BitSet for u32 {

    const INDEX_BITS: u32 = 5;
    const SIZE: u32 = 32;

    fn get(&self, index: u32) -> bool {
        debug_assert!(index < Self::SIZE);
        ((1 << index) & *self) != 0
    }

    fn set(&mut self, index: u32) -> bool {
        debug_assert!(index < Self::SIZE);
        *self |= 1 << index;
        *self == !0
    }

    fn unset(&mut self, index: u32) -> bool {
        debug_assert!(index < Self::SIZE);
        *self &= !(1 << index);
        *self == 0
    }

    fn next_0(&self) -> Option<u32> {
        let index = (!*self).trailing_zeros();
        debug_assert!(index <= Self::SIZE);
        if index == Self::SIZE {
            None
        } else {
            Some(index)
        }
    }

    fn next_1(&self) -> Option<u32> {
        let index = (self).trailing_zeros();
        debug_assert!(index <= Self::SIZE);
        if index == Self::SIZE {
            None
        } else {
            Some(index)
        }
    }

    fn state(&self, range: Range<u32>) -> State {
        debug_assert!(range.start <= range.end && range.end <= Self::SIZE);
        let pattern = (((1u64 << (range.end - range.start)) - 1) << range.start) as u32;
        let value = *self;
        if (value & pattern) == pattern {
            State::Full
        } else if (value & !pattern) == value {
            State::Empty
        } else {
            State::Partial
        }
    }
}

struct Level<T> {
    // Each pair of bits is associated with index.
    // Each pattern represent higher level knowledge of the `T`
    // `00` - `T` is empty. There are `0`s only.
    // `01` - `T` may be empty. There are `0`s.
    // `10` - Invalid pattern. Indicates and error.
    // `11` - `T` is full. There are `1`s only.
    mask: u64,
    bits: Vec<T>,
}

impl<T> Level<T> {
    fn split_index(index: u32) -> (u32, u32) {
        let mask_index = index >> T::INDEX_BITS;
        let next_index = index & ((1u32 << T::INDEX_BITS) - 1);
        (mask_index, next_index)
    }

    fn merge_index(mask_index: u32, next_index: u32) -> u32 {
        mask_index << T::INDEX_BITS | next_index
    }

    fn state(&self, mask_index: u32) -> State {
        match ((self.mask >> mask_index * 2) & 3) {
            0 => State::Empty,
            1 => State::Partial,
            3 => State::Full,
            _ => unreachable!(),
        }
    }

    fn set_state(&self, mask_index: u32, state: State) {
        match state {
            State::Empty => {
                self.mask &= !(3 << mask_index);
            }
            State::Partial => {
                self.mask &= !(2 << mask_index);
                self.mask |= 1 << mask_index;
            }
            State::Full => {
                self.mask |= 3 << mask_index;
            }
        };
    }
}

impl<T> BitSet for Level<T>
where
    T: BitSet,
{
    const INDEX_BITS: u32 = T::INDEX_BITS + 5;
    const SIZE: u32 = 32 * T::SIZE;

    fn get(&self, index: u32) -> bool {
        debug_assert!(index < Self::SIZE);
        let (mask_index, next_index) = Self::split_index(index);
        if self.state(mask_index) == State::Empty {
            false
        } else {
            self.bits[mask_index].get(next_index)
        }
    }

    fn set(&mut self, index: u32) -> bool {
        debug_assert!(index < Self::SIZE);
        let (mask_index, next_index) = Self::split_index(index);
        if self.state(mask_index) != State::Full {
            if self.bits[mask_index].set(next_index) {
                self.set_state(mask_index, State::Full);
            }
        }
        self.mask == !0
    }

    fn unset(&mut self, index: u32) -> bool {
        debug_assert!(index < Self::SIZE);
        let (mask_index, next_index) = Self::split_index(index);
        if self.state(mask_index) != State::Empty {
            if self.bits[mask_index].unset(next_index) {
                self.set_state(mask_index, State::Empty);
            }
        }
        self.mask == 0
    }

    fn next_0(&self) -> Option<u32> {
        let mask_index = (!self.mask).trailing_zeros() / 2;
        debug_assert!(mask_index <= 32);
        if mask_index == 32 {
            None
        } else {
            match self.state(mask_index) {
                State::Full => unreachable!(),
                State::Partial => self.bits[mask_index].next_0().map(|next_index| Self::merge_index(mask_index, next_index)),
                State::Empty => Some(Self::merge_index(mask_index, 0)),
            }
        }
    }

    fn next_1(&self) -> Option<u32> {
        let mask_index = self.mask.trailing_zeros() / 2;
        debug_assert!(mask_index <= 32);
        if mask_index == 32 {
            None
        } else {
            match self.state(mask_index) {
                State::Empty => unreachable!(),
                State::Partial => self.bits[mask_index].next_1().map(|next_index| Self::merge_index(mask_index, next_index)),
                State::Full => Some(Self::merge_index(mask_index, 0)),
            }
        }
    }

    fn state(&self, range: Range<u32>) -> State {
        debug_assert!(range.start <= range.end && range.end <= Self::SIZE);

        let (mask_start, next_start) = Self::split_index(range.start);
        let (mask_end, next_end) = Self::split_index(range.end);

        if mask_start != mask_end {
            let mask_range = mask_start * 2 .. (mask_end + (next_end > 0) as u32) * 2;
            let pattern = (((1u64 << (mask_range.end - mask_range.start)) - 1) << mask_range.start) as u32;
            let value = *self;
            if (value & pattern) == pattern {
                State::Full
            } else if (value & !pattern) == value {
                State::Empty
            } else {
                State::Partial
            }
        } else {
            self.bits[mask_start].state(next.start .. next.end)
        }
    }
}

// Bit set that holds up to 32^4 bits.
// It can be used to find index with `1` or `0` in it.
// Also it supports faster query for all bits in range to be `1`s or `0`s.
struct BitSetImpl {
    levels: Level<Level<Level<u32>>>,
}

It requires few more functions to be suitable for for fast resource allocation.