use std::{fmt, iter::FusedIterator, marker::PhantomData};
use serde::{Deserialize, Serialize};
use smol_bitmap::SmolBitmap;
use crate::sparse_index::TrySparseIndex;
#[repr(transparent)]
pub struct SparseSet<K> {
bits: SmolBitmap,
_phantom: PhantomData<K>,
}
impl<K> Clone for SparseSet<K> {
fn clone(&self) -> Self {
Self { bits: self.bits.clone(), _phantom: PhantomData }
}
}
impl<K: TrySparseIndex + fmt::Debug> fmt::Debug for SparseSet<K> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut s = f.debug_set();
for key in self {
s.entry(&key);
}
s.finish()
}
}
impl<K: TrySparseIndex + Serialize> Serialize for SparseSet<K> {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
use serde::ser::SerializeSeq;
if serializer.is_human_readable() {
let mut seq = serializer.serialize_seq(Some(self.len()))?;
for key in self {
seq.serialize_element(&key)?;
}
seq.end()
} else {
self.bits.serialize(serializer)
}
}
}
impl<'de, K: TrySparseIndex + Deserialize<'de>> Deserialize<'de> for SparseSet<K> {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
if deserializer.is_human_readable() {
use std::fmt;
use serde::de::{SeqAccess, Visitor};
struct SparseSetVisitor<K> {
_phantom: PhantomData<K>,
}
impl<'de, K: TrySparseIndex + Deserialize<'de>> Visitor<'de> for SparseSetVisitor<K> {
type Value = SparseSet<K>;
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
formatter.write_str("a sequence of indices")
}
fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
where
A: SeqAccess<'de>,
{
let mut set = SparseSet::new();
if let Some(len) = seq.size_hint() {
set.reserve(len);
}
let mut prev_index = None;
while let Some(key) = seq.next_element::<K>()? {
let index = key.index();
if prev_index.is_some_and(|prev| prev >= index) {
return Err(serde::de::Error::invalid_value(
serde::de::Unexpected::Unsigned(index as u64),
&"indices must be in ascending order",
));
}
prev_index = Some(index);
set.bits.insert(index);
}
Ok(set)
}
}
deserializer.deserialize_seq(SparseSetVisitor { _phantom: PhantomData })
} else {
let bits = Deserialize::deserialize(deserializer)?;
let set = Self { bits, _phantom: PhantomData };
K::validate_sorted(set.bits.iter()).map_err(serde::de::Error::custom)?;
Ok(set)
}
}
}
impl<K> Default for SparseSet<K> {
fn default() -> Self {
Self::new()
}
}
impl<K> PartialEq for SparseSet<K> {
fn eq(&self, other: &Self) -> bool {
self.bits == other.bits
}
}
impl<K> Eq for SparseSet<K> {}
impl<K> SparseSet<K> {
pub const fn new() -> Self {
Self { bits: SmolBitmap::new(), _phantom: PhantomData }
}
pub fn with_capacity(capacity: usize) -> Self {
Self { bits: SmolBitmap::with_capacity(capacity), _phantom: PhantomData }
}
#[inline]
pub fn len(&self) -> usize {
self.bits.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.bits.is_empty()
}
#[inline]
pub const fn capacity(&self) -> usize {
self.bits.capacity()
}
pub fn clear(&mut self) {
self.bits.clear();
}
pub fn shrink_to_fit(&mut self) {
self.bits.shrink_to_fit();
}
pub fn reserve(&mut self, additional: usize) {
self.bits.reserve(additional);
}
#[inline]
pub fn into_parts(self) -> SmolBitmap {
self.bits
}
#[inline]
pub const fn from_parts(bits: SmolBitmap) -> Self {
Self { bits, _phantom: PhantomData }
}
pub fn is_subset(&self, other: &Self) -> bool {
self.bits.is_subset(&other.bits)
}
pub fn is_superset(&self, other: &Self) -> bool {
self.bits.is_superset(&other.bits)
}
pub fn is_disjoint(&self, other: &Self) -> bool {
self.bits.is_disjoint(&other.bits)
}
pub fn union(&self, other: &Self) -> Self {
Self { bits: self.bits.union(&other.bits), _phantom: PhantomData }
}
pub fn intersection(&self, other: &Self) -> Self {
Self { bits: self.bits.intersection(&other.bits), _phantom: PhantomData }
}
pub fn difference(&self, other: &Self) -> Self {
Self { bits: self.bits.difference(&other.bits), _phantom: PhantomData }
}
pub fn symmetric_difference(&self, other: &Self) -> Self {
Self { bits: self.bits.symmetric_difference(&other.bits), _phantom: PhantomData }
}
}
impl<K: TrySparseIndex> SparseSet<K> {
#[inline]
pub fn contains(&self, key: K) -> bool {
self.bits.get(key.index())
}
pub fn insert(&mut self, key: K) -> bool {
self.bits.insert(key.index())
}
pub fn remove(&mut self, key: K) -> bool {
self.bits.remove(key.index())
}
pub fn retain<F>(&mut self, mut f: F)
where
F: FnMut(K) -> bool,
{
self.bits.retain(|idx| f(K::from_index(idx)));
}
#[define_opaque(Iter)]
pub fn iter(&self) -> Iter<'_, K> {
self.bits.iter().map(K::from_index)
}
pub fn first(&self) -> Option<K> {
self.bits.first().map(K::from_index)
}
pub fn last(&self) -> Option<K> {
self.bits.last().map(K::from_index)
}
pub fn is_sparse(&self) -> bool {
match (self.bits.first(), self.bits.last()) {
(Some(first), Some(last)) => {
self.len() < (last - first + 1)
},
_ => false, }
}
}
pub type Iter<'a, K: TrySparseIndex> =
impl DoubleEndedIterator<Item = K> + ExactSizeIterator + FusedIterator + Clone;
pub type IntoIter<K: TrySparseIndex> =
impl DoubleEndedIterator<Item = K> + ExactSizeIterator + FusedIterator + Clone;
impl<'a, K: TrySparseIndex> IntoIterator for &'a SparseSet<K> {
type IntoIter = Iter<'a, K>;
type Item = K;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<K: TrySparseIndex> IntoIterator for SparseSet<K> {
type IntoIter = IntoIter<K>;
type Item = K;
#[define_opaque(IntoIter)]
fn into_iter(self) -> Self::IntoIter {
self.bits.into_iter().map(K::from_index)
}
}
impl<K: TrySparseIndex> FromIterator<K> for SparseSet<K> {
fn from_iter<T: IntoIterator<Item = K>>(iter: T) -> Self {
let iter = iter.into_iter();
let mut set = Self::with_capacity(iter.size_hint().0);
for key in iter {
set.insert(key);
}
set
}
}
impl<K: TrySparseIndex> Extend<K> for SparseSet<K> {
fn extend<T: IntoIterator<Item = K>>(&mut self, iter: T) {
let iter = iter.into_iter();
self.reserve(iter.size_hint().0);
for key in iter {
self.insert(key);
}
}
}