use fxhash::FxHashMap;
use std::cell::Cell;
const SMALL_ELEMS: usize = 12;
#[derive(Clone, Debug)]
enum AdaptiveMap {
Small {
len: u32,
keys: [u32; SMALL_ELEMS],
values: [u64; SMALL_ELEMS],
},
Large(FxHashMap<u32, u64>),
}
const INVALID: u32 = 0xffff_ffff;
impl AdaptiveMap {
fn new() -> Self {
Self::Small {
len: 0,
keys: [INVALID; SMALL_ELEMS],
values: [0; SMALL_ELEMS],
}
}
#[inline(always)]
fn get_or_insert<'a>(&'a mut self, key: u32) -> &'a mut u64 {
let small_mode_idx = match self {
&mut Self::Small {
len,
ref mut keys,
ref values,
} => {
if let Some(i) = keys[..len as usize].iter().position(|&k| k == key) {
Some(i)
} else if len != SMALL_ELEMS as u32 {
debug_assert!(len < SMALL_ELEMS as u32);
None
} else if let Some(i) = values.iter().position(|&v| v == 0) {
keys[i] = key;
Some(i)
} else {
*self = Self::Large(keys.iter().copied().zip(values.iter().copied()).collect());
None
}
}
_ => None,
};
match self {
Self::Small { len, keys, values } => {
if let Some(i) = small_mode_idx {
return &mut values[i];
}
debug_assert!(*len < SMALL_ELEMS as u32);
let idx = *len as usize;
*len += 1;
keys[idx] = key;
values[idx] = 0;
&mut values[idx]
}
Self::Large(map) => map.entry(key).or_insert(0),
}
}
#[inline(always)]
fn get_mut(&mut self, key: u32) -> Option<&mut u64> {
match self {
&mut Self::Small {
len,
ref keys,
ref mut values,
} => {
for i in 0..len {
if keys[i as usize] == key {
return Some(&mut values[i as usize]);
}
}
None
}
&mut Self::Large(ref mut map) => map.get_mut(&key),
}
}
#[inline(always)]
fn get(&self, key: u32) -> Option<u64> {
match self {
&Self::Small {
len,
ref keys,
ref values,
} => {
for i in 0..len {
if keys[i as usize] == key {
let value = values[i as usize];
return Some(value);
}
}
None
}
&Self::Large(ref map) => {
let value = map.get(&key).cloned();
value
}
}
}
fn iter<'a>(&'a self) -> AdaptiveMapIter<'a> {
match self {
&Self::Small {
len,
ref keys,
ref values,
} => AdaptiveMapIter::Small(&keys[0..len as usize], &values[0..len as usize]),
&Self::Large(ref map) => AdaptiveMapIter::Large(map.iter()),
}
}
fn is_empty(&self) -> bool {
match self {
AdaptiveMap::Small { values, .. } => values.iter().all(|&value| value == 0),
AdaptiveMap::Large(m) => m.values().all(|&value| value == 0),
}
}
}
enum AdaptiveMapIter<'a> {
Small(&'a [u32], &'a [u64]),
Large(std::collections::hash_map::Iter<'a, u32, u64>),
}
impl<'a> std::iter::Iterator for AdaptiveMapIter<'a> {
type Item = (u32, u64);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
match self {
&mut Self::Small(ref mut keys, ref mut values) => {
if keys.is_empty() {
None
} else {
let (k, v) = ((*keys)[0], (*values)[0]);
*keys = &(*keys)[1..];
*values = &(*values)[1..];
Some((k, v))
}
}
&mut Self::Large(ref mut it) => it.next().map(|(&k, &v)| (k, v)),
}
}
}
#[derive(Clone)]
pub struct IndexSet {
elems: AdaptiveMap,
cache: Cell<(u32, u64)>,
}
const BITS_PER_WORD: usize = 64;
impl IndexSet {
pub fn new() -> Self {
Self {
elems: AdaptiveMap::new(),
cache: Cell::new((INVALID, 0)),
}
}
#[inline(always)]
fn elem(&mut self, bit_index: usize) -> &mut u64 {
let word_index = (bit_index / BITS_PER_WORD) as u32;
if self.cache.get().0 == word_index {
self.cache.set((INVALID, 0));
}
self.elems.get_or_insert(word_index)
}
#[inline(always)]
fn maybe_elem_mut(&mut self, bit_index: usize) -> Option<&mut u64> {
let word_index = (bit_index / BITS_PER_WORD) as u32;
if self.cache.get().0 == word_index {
self.cache.set((INVALID, 0));
}
self.elems.get_mut(word_index)
}
#[inline(always)]
fn maybe_elem(&self, bit_index: usize) -> Option<u64> {
let word_index = (bit_index / BITS_PER_WORD) as u32;
if self.cache.get().0 == word_index {
Some(self.cache.get().1)
} else {
self.elems.get(word_index)
}
}
#[inline(always)]
pub fn set(&mut self, idx: usize, val: bool) {
let bit = idx % BITS_PER_WORD;
if val {
*self.elem(idx) |= 1 << bit;
} else if let Some(word) = self.maybe_elem_mut(idx) {
*word &= !(1 << bit);
}
}
pub fn assign(&mut self, other: &Self) {
self.elems = other.elems.clone();
self.cache = other.cache.clone();
}
#[inline(always)]
pub fn get(&self, idx: usize) -> bool {
let bit = idx % BITS_PER_WORD;
if let Some(word) = self.maybe_elem(idx) {
(word & (1 << bit)) != 0
} else {
false
}
}
pub fn union_with(&mut self, other: &Self) -> bool {
let mut changed = 0;
for (word_idx, bits) in other.elems.iter() {
if bits == 0 {
continue;
}
let word_idx = word_idx as usize;
let self_word = self.elem(word_idx * BITS_PER_WORD);
changed |= bits & !*self_word;
*self_word |= bits;
}
changed != 0
}
pub fn iter<'a>(&'a self) -> impl Iterator<Item = usize> + 'a {
self.elems.iter().flat_map(|(word_idx, bits)| {
let word_idx = word_idx as usize;
SetBitsIter(bits).map(move |i| BITS_PER_WORD * word_idx + i)
})
}
pub(crate) fn is_small(&self) -> bool {
match &self.elems {
&AdaptiveMap::Small { .. } => true,
_ => false,
}
}
pub(crate) fn is_empty(&self) -> bool {
self.elems.is_empty()
}
}
pub struct SetBitsIter(u64);
impl Iterator for SetBitsIter {
type Item = usize;
#[inline]
fn next(&mut self) -> Option<usize> {
std::num::NonZeroU64::new(self.0).map(|nz| {
let bitidx = nz.trailing_zeros();
self.0 &= self.0 - 1; bitidx as usize
})
}
}
impl std::fmt::Debug for IndexSet {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
let vals = self.iter().collect::<Vec<_>>();
write!(f, "{:?}", vals)
}
}
#[cfg(test)]
mod test {
use super::IndexSet;
#[test]
fn test_set_bits_iter() {
let mut vec = IndexSet::new();
let mut sum = 0;
for i in 0..1024 {
if i % 17 == 0 {
vec.set(i, true);
sum += i;
}
}
let mut checksum = 0;
for bit in vec.iter() {
debug_assert!(bit % 17 == 0);
checksum += bit;
}
debug_assert_eq!(sum, checksum);
}
#[test]
fn test_expand_remove_zero_elems() {
let mut vec = IndexSet::new();
for i in 0..12 {
vec.set(64 * i, true);
}
vec.set(64 * 5, false);
vec.set(64 * 100, true);
debug_assert!(vec.is_small());
}
}