use fnv::FnvHashSet;
use std::collections::hash_set;
#[derive(Clone, Debug)]
pub enum IdSet {
Zero,
One(usize),
Two([usize; 2]),
Three([usize; 3]),
More(FnvHashSet<usize>),
}
use self::IdSet::*;
impl IdSet {
pub fn new() -> IdSet { IdSet::Zero }
pub fn len(&self) -> usize {
match self {
Zero => 0,
One(_) => 1,
Two(_) => 2,
Three(_) => 3,
More(set) => set.len(),
}
}
pub fn contains(&self, elem: usize) -> bool {
match *self {
Zero => false,
One(a) => a == elem,
Two([a, b]) => a == elem || b == elem,
Three([a, b, c]) => a == elem || b == elem || c == elem,
More(ref set) => set.contains(&elem),
}
}
pub fn remove(&mut self, elem: usize) {
*self = match *self {
Zero => Zero,
One(a) if a == elem => Zero,
One(a) => One(a),
Two([a, b]) if a == elem => One(b),
Two([a, b]) if b == elem => One(a),
Two(x) => Two(x),
Three([a, b, c]) if a == elem => Two([b, c]),
Three([a, b, c]) if b == elem => Two([a, c]),
Three([a, b, c]) if c == elem => Two([a, b]),
Three(x) => Three(x),
More(ref mut set) => {
set.remove(&elem);
return
},
};
}
pub fn insert(&mut self, elem: usize) {
*self = match *self {
Zero => One(elem),
One(a) => if a == elem {
One(a)
} else {
Two([a, elem])
},
Two([a, b]) => if a == elem || b == elem {
Two([a, b])
} else {
Three([a, b, elem])
},
Three([a, b, c]) => if a == elem || b == elem || c ==elem {
Three([a, b, c])
} else {
let mut set = FnvHashSet::default();
set.extend([a, b, c, elem].iter().copied());
More(set)
},
More(ref mut set) => {
set.insert(elem);
return
},
};
}
pub fn iter(&self) -> Iter<std::iter::Copied<hash_set::Iter<usize>>> {
match *self {
Zero => Iter::Small([0; 3], 3),
One(a) => Iter::Small([0, 0, a], 2),
Two([a, b]) => Iter::Small([0, a, b], 1),
Three(x) => Iter::Small(x, 0),
More(ref set) => Iter::Large(set.iter().copied()),
}
}
pub fn drain(&mut self) -> Iter<hash_set::Drain<usize>> {
let out = match *self {
Zero => Iter::Small([0; 3], 3),
One(a) => Iter::Small([0, 0, a], 2),
Two([a, b]) => Iter::Small([0, a, b], 1),
Three(x) => Iter::Small(x, 0),
More(ref mut set) => return Iter::Large(set.drain()),
};
*self = Zero;
out
}
}
pub enum Iter<I> {
Small([usize; 3], usize),
Large(I),
}
impl<'a, I> Iterator for Iter<I> where I: Iterator<Item=usize> {
type Item = usize;
fn next(&mut self) -> Option<usize> {
match self {
Iter::Small(arr, index) => {
let out = arr.get(*index).copied();
*index += 1;
out
},
Iter::Large(iter) => iter.next(),
}
}
}
#[cfg(test)]
const TEST_SETS: &[IdSet] = &[
Zero,
One(1),
One(3459),
Two([1, 2]),
Two([34, 9234]),
Three([3429, 425, 292]),
Three([1, 2, 3]),
Three([234, 340, 912]),
];
#[cfg(test)]
const TEST_ELEMS: &[&[usize]] = &[
&[1],
&[2],
&[3],
&[2, 3],
&[1, 324, 1],
&[1, 2, 3],
&[3405, 319, 430, 209, 425, 43, 319],
&[24, 423, 3, 0],
];
#[cfg(test)]
fn get_pairs() -> impl Iterator<Item=(IdSet, FnvHashSet<usize>)> {
TEST_SETS
.iter()
.map(|set| (set.clone(), set.iter().collect::<_>()))
.chain(TEST_ELEMS.iter().map(|elems| {
let mut idset = IdSet::new();
for x in elems.iter() {
idset.insert(*x);
}
(idset, elems.iter().copied().collect::<_>())
}))
}
#[cfg(test)]
fn assert_eq(a: &IdSet, b: &FnvHashSet<usize>) {
assert_eq!(a.len(), b.len(), "idset is wrong length");
let mut with = b.clone();
with.extend(a.iter());
assert_eq!(with.len(), b.len(), "idset has extra elements");
for x in a.iter() {
with.remove(&x);
}
assert_eq!(with.len(), 0, "idset has missing elements");
}
#[test]
fn test_idset_len() {
for (idset, hashset) in get_pairs() {
assert_eq(&idset, &hashset);
}
}
#[test]
fn test_idset_contains() {
for (idset, hashset) in get_pairs() {
for x in TEST_ELEMS.iter().flat_map(|elems| elems.iter()) {
assert_eq!(idset.contains(*x), hashset.contains(x));
}
}
}
#[test]
fn test_idset_insert() {
for (idset, hashset) in get_pairs() {
for &elems in TEST_ELEMS {
let mut idset = idset.clone();
let mut hashset = hashset.clone();
for &x in elems {
idset.insert(x);
hashset.insert(x);
assert_eq(&idset, &hashset);
}
}
}
}
#[test]
fn test_idset_remove() {
for (idset, hashset) in get_pairs() {
for &elems in TEST_ELEMS {
let mut idset = idset.clone();
let mut hashset = hashset.clone();
for &x in elems {
idset.remove(x);
hashset.remove(&x);
assert_eq(&idset, &hashset);
}
}
}
}
#[test]
fn test_idset_remove_insert() {
for (idset, hashset) in get_pairs() {
for &elems_insert in TEST_ELEMS {
for &elems_remove in TEST_ELEMS {
let mut idset = idset.clone();
let mut hashset = hashset.clone();
for &x in elems_remove {
idset.remove(x);
hashset.remove(&x);
}
for &x in elems_insert {
idset.insert(x);
hashset.insert(x);
}
assert_eq(&idset, &hashset);
}
}
}
}
#[test]
fn test_idset_insert_remove() {
for (idset, hashset) in get_pairs() {
for &elems_insert in TEST_ELEMS {
for &elems_remove in TEST_ELEMS {
let mut idset = idset.clone();
let mut hashset = hashset.clone();
for &x in elems_insert {
idset.insert(x);
hashset.insert(x);
}
for &x in elems_remove {
idset.remove(x);
hashset.remove(&x);
}
assert_eq(&idset, &hashset);
}
}
}
}
#[test]
fn test_idset_drain() {
for (idset, hashset) in get_pairs() {
let mut idset = idset.clone();
let mut drained = IdSet::new();
for x in idset.drain() {
drained.insert(x);
}
assert_eq(&drained, &hashset);
assert!(idset.len() == 0);
}
}