dense_map.rs

use std::{mem::{ManuallyDrop, MaybeUninit, needs_drop, replace}, ptr};

pub struct BitSet {
    unset: usize,
    store: Box<[usize]>,
}

impl BitSet {
    pub fn new(size: usize) -> BitSet {
        let width = usize::BITS as usize;
        let words = size.div_ceil(width);
        let mut store = vec![0; words].into_boxed_slice();
        let rem = size % width;
        if rem > 0 {
            store[words - 1] = (!0usize) << rem;
        }
        BitSet { unset: size, store }
    }

    pub fn unset(&self) -> usize {
        self.unset
    }

    pub fn contains(&self, i: usize) -> bool {
        let (word, mask) = word_mask(i);
        self.store[word] & mask != 0
    }

    pub fn insert(&mut self, x: usize) -> bool {
        let (word, mask) = word_mask(x);
        let inserted = self.store[word] & mask == 0;
        self.unset -= inserted as usize;
        self.store[word] |= mask;
        inserted
    }
}

fn word_mask(x: usize) -> (usize, usize) {
    let width = usize::BITS as usize;
    let pos = x / width;
    let mask = 1 << (x % width);
    (pos, mask)
}

pub struct DenseMap<T>(BitSet, Box<[MaybeUninit<T>]>);

impl<T> DenseMap<T> {
    pub fn new(size: usize) -> DenseMap<T> {
        DenseMap(BitSet::new(size), Box::new_uninit_slice(size))
    }

    pub fn get(&self, i: usize) -> Option<&T> {
        if self.0.contains(i) {
            let slot = &self.1[i];
            Some(unsafe { slot.assume_init_ref() })
        } else {
            None
        }
    }

    pub fn get_mut(&mut self, i: usize) -> &mut T
    where
        T: Default,
    {
        let slot = &mut self.1[i];
        if !self.0.insert(i) {
            slot.write(T::default());
        }
        unsafe { slot.assume_init_mut() }
    }

    pub fn set(&mut self, i: usize, val: T) -> Option<T> {
        let slot = &mut self.1[i];
        let old = self.0.insert(i).then(|| {
            unsafe { slot.assume_init_read() }
        });
        slot.write(val);
        old
    }

    pub fn unwrap(self) -> Box<[T]> {
        assert_eq!(self.0.unset(), 0, "DenseMap has uninitialized entries");
        let mut this = ManuallyDrop::new(self);
        let DenseMap(set, values) = &mut *this;
        unsafe {
            ptr::drop_in_place(set);
            ptr::read(values).assume_init()
        }
    }
}

impl<T> Drop for DenseMap<T> {
    fn drop(&mut self) {
        if needs_drop::<T>() {
            for (i, slot) in self.1.iter_mut().enumerate() {
                if self.0.contains(i) {
                    unsafe { slot.assume_init_drop() };
                }
            }
        }
    }
}