refactor: refine small set

This commit is contained in:
Zixuan Chen 2022-11-09 19:37:47 +08:00
parent 1653f8d109
commit 6eff5ddd3f

View file

@ -1,12 +1,17 @@
use std::{hash::Hash, mem::MaybeUninit}; use std::{collections::hash_set::IntoIter, hash::Hash, mem::MaybeUninit};
use fxhash::FxHashSet; use fxhash::FxHashSet;
pub enum SmallSet<T, const SIZE: usize> { pub enum SmallSet<T, const SIZE: usize> {
Arr([Option<T>; SIZE]), Arr([Option<T>; SIZE], usize),
Set(FxHashSet<T>), Set(FxHashSet<T>),
} }
pub enum SmallSetIter<T, const SIZE: usize> {
Arr([Option<T>; SIZE], usize),
Set(IntoIter<T>),
}
impl<T: Eq + Hash, const SIZE: usize> SmallSet<T, SIZE> { impl<T: Eq + Hash, const SIZE: usize> SmallSet<T, SIZE> {
pub fn new() -> Self { pub fn new() -> Self {
let a: MaybeUninit<[Option<T>; SIZE]> = MaybeUninit::zeroed(); let a: MaybeUninit<[Option<T>; SIZE]> = MaybeUninit::zeroed();
@ -15,16 +20,23 @@ impl<T: Eq + Hash, const SIZE: usize> SmallSet<T, SIZE> {
for i in a.iter_mut() { for i in a.iter_mut() {
*i = None; *i = None;
} }
SmallSet::Arr(a) SmallSet::Arr(a, 0)
} }
pub fn is_empty(&self) -> bool { pub fn is_empty(&self) -> bool {
match self { match self {
SmallSet::Arr(a) => a.is_empty(), SmallSet::Arr(_, size) => *size == 0,
SmallSet::Set(s) => s.is_empty(), SmallSet::Set(s) => s.is_empty(),
} }
} }
pub fn len(&self) -> usize {
match self {
SmallSet::Arr(_, size) => *size,
SmallSet::Set(s) => s.len(),
}
}
/// Adds a value to the set. /// Adds a value to the set.
/// ///
/// Returns whether the value was newly inserted. That is: /// Returns whether the value was newly inserted. That is:
@ -34,7 +46,8 @@ impl<T: Eq + Hash, const SIZE: usize> SmallSet<T, SIZE> {
/// ///
pub fn insert(&mut self, v: T) -> bool { pub fn insert(&mut self, v: T) -> bool {
match self { match self {
SmallSet::Arr(a) => { SmallSet::Arr(a, i) => {
*i += 1;
for i in a.iter_mut() { for i in a.iter_mut() {
if let Some(i) = i { if let Some(i) = i {
if *i == v { if *i == v {
@ -61,7 +74,7 @@ impl<T: Eq + Hash, const SIZE: usize> SmallSet<T, SIZE> {
pub fn contains(&mut self, v: &T) -> bool { pub fn contains(&mut self, v: &T) -> bool {
match self { match self {
SmallSet::Arr(a) => { SmallSet::Arr(a, _) => {
for i in a.iter_mut() { for i in a.iter_mut() {
if let Some(i) = i { if let Some(i) = i {
if i == v { if i == v {
@ -80,7 +93,8 @@ impl<T: Eq + Hash, const SIZE: usize> SmallSet<T, SIZE> {
pub fn remove(&mut self, v: &T) -> bool { pub fn remove(&mut self, v: &T) -> bool {
match self { match self {
SmallSet::Arr(a) => { SmallSet::Arr(a, i) => {
*i -= 1;
for i in a.iter_mut() { for i in a.iter_mut() {
if i.as_ref() == Some(v) { if i.as_ref() == Some(v) {
*i = None; *i = None;
@ -100,3 +114,36 @@ impl<T: Eq + Hash, const SIZE: usize> Default for SmallSet<T, SIZE> {
Self::new() Self::new()
} }
} }
impl<T, const SIZE: usize> Iterator for SmallSetIter<T, SIZE> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
match self {
SmallSetIter::Arr(arr, i) => {
#[allow(clippy::needless_range_loop)]
for index in *i..arr.len() {
if let Some(v) = std::mem::take(&mut arr[index]) {
*i += 1;
return Some(v);
}
}
None
}
SmallSetIter::Set(set) => set.next(),
}
}
}
impl<T: Eq + Hash, const SIZE: usize> IntoIterator for SmallSet<T, SIZE> {
type Item = T;
type IntoIter = SmallSetIter<T, SIZE>;
fn into_iter(self) -> Self::IntoIter {
match self {
SmallSet::Arr(arr, _) => SmallSetIter::Arr(arr, 0),
SmallSet::Set(arr) => SmallSetIter::Set(arr.into_iter()),
}
}
}