use super::error::{CollectionError as Error, Result, ensure};
use super::{MapIter, Memory, StableBTreeMap, Storable};
use crate::DefaultMemoryImpl;
use crate::RuntimeMemory;
pub struct StableBTreeSet<K: Storable + Ord + Clone, M: Memory = DefaultMemoryImpl> {
inner: StableBTreeMap<K, (), M>,
}
impl<K: Storable + Ord + Clone, M: Memory> StableBTreeSet<K, M> {
pub fn create(memory: RuntimeMemory<M>) -> Result<Self> {
ensure(memory.size() == 0, Error::NotEmpty)?;
Self::reset(memory)
}
pub fn reset(memory: RuntimeMemory<M>) -> Result<Self> {
Ok(Self {
inner: StableBTreeMap::reset(memory)?,
})
}
pub fn load(memory: RuntimeMemory<M>) -> Result<Self> {
Ok(Self {
inner: StableBTreeMap::load(memory)?,
})
}
pub fn insert(&mut self, key: K) -> Result<bool> {
Ok(self.inner.insert(key, ())?.is_none())
}
pub fn contains(&self, key: &K) -> Result<bool> {
self.inner.contains_key(key)
}
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
pub fn len(&self) -> u64 {
self.inner.len()
}
pub fn remove(&mut self, key: &K) -> Result<bool> {
Ok(self.inner.remove(key)?.is_some())
}
pub fn iter(&self) -> SetIter<'_, K, M> {
SetIter(self.inner.iter())
}
}
impl<'a, K: Storable + Ord + Clone, M: Memory> IntoIterator for &'a StableBTreeSet<K, M> {
type Item = Result<K>;
type IntoIter = SetIter<'a, K, M>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
pub struct SetIter<'a, K: Storable + Ord + Clone, M: Memory>(MapIter<'a, K, (), M>);
impl<K: Storable + Ord + Clone, M: Memory> Iterator for SetIter<'_, K, M> {
type Item = Result<K>;
fn next(&mut self) -> Option<Result<K>> {
self.0
.next()
.map(|entry| entry.map(super::MapEntry::into_key))
}
}
impl<K: Storable + Ord + Clone, M: Memory> DoubleEndedIterator for SetIter<'_, K, M> {
fn next_back(&mut self) -> Option<Result<K>> {
self.0
.next_back()
.map(|entry| entry.map(super::MapEntry::into_key))
}
}