use super::bridge::{self, TreeMemory};
use super::error::{CollectionError as Error, Result, ensure};
use super::{MapIter, Memory, Storable};
use crate::{DefaultMemoryImpl, RuntimeMemory};
use std::ops::RangeBounds;
pub struct StableBTreeMap<K: Storable + Ord + Clone, V: Storable, M: Memory = DefaultMemoryImpl> {
inner: bridge::Map<K, V, M>,
}
impl<K: Storable + Ord + Clone, V: Storable, M: Memory> StableBTreeMap<K, V, 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> {
bridge::reserve(&memory)?;
#[cfg(not(target_arch = "wasm32"))]
let _guard = memory.poison_on_unwind();
Ok(Self {
inner: bridge::Map::new(TreeMemory::Allocation(memory)),
})
}
pub fn load(memory: RuntimeMemory<M>) -> Result<Self> {
Ok(Self {
inner: bridge::Map::load(TreeMemory::Allocation(memory)).ok_or(Error::Format)?,
})
}
pub fn clear(&mut self) {
#[cfg(not(target_arch = "wasm32"))]
let _guard = self.inner.memory().allocation().poison_on_unwind();
self.inner.clear();
}
pub fn reserve_pages(
self,
minimum: u64,
) -> (Self, std::result::Result<(), crate::RuntimeGrowError>) {
let additional = minimum.saturating_sub(self.inner.memory().allocation().size());
let result = self
.inner
.memory()
.allocation()
.grow(additional)
.map(|_| ());
(self, result)
}
#[must_use]
pub fn allocated_pages(self) -> (Self, u64) {
let pages = self.inner.memory().allocation().size();
(self, pages)
}
pub fn insert(&mut self, key: K, value: V) -> Option<V> {
#[cfg(not(target_arch = "wasm32"))]
let _guard = self.inner.memory().allocation().poison_on_unwind();
self.inner.insert(key, value)
}
pub fn get(&self, key: &K) -> Option<V> {
self.inner.get(key)
}
pub fn contains_key(&self, key: &K) -> bool {
self.inner.contains_key(key)
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.inner.is_empty()
}
#[must_use]
pub const fn len(&self) -> u64 {
self.inner.len()
}
pub fn remove(&mut self, key: &K) -> Option<V> {
#[cfg(not(target_arch = "wasm32"))]
let _guard = self.inner.memory().allocation().poison_on_unwind();
self.inner.remove(key)
}
#[must_use]
pub fn first_key_value(&self) -> Option<(K, V)> {
self.inner.first_key_value()
}
#[must_use]
pub fn last_key_value(&self) -> Option<(K, V)> {
self.inner.last_key_value()
}
#[must_use]
pub const fn iter(&self) -> MapIter<'_, K, V, M> {
MapIter(self.inner.iter())
}
pub fn range(&self, bounds: impl RangeBounds<K>) -> MapIter<'_, K, V, M> {
MapIter(self.inner.range(bounds))
}
pub fn keys(&self) -> impl DoubleEndedIterator<Item = K> + '_ {
self.inner.keys()
}
pub fn keys_range(
&self,
bounds: impl RangeBounds<K>,
) -> impl DoubleEndedIterator<Item = K> + '_ {
self.inner.keys_range(bounds)
}
pub fn values(&self) -> impl DoubleEndedIterator<Item = V> + '_ {
self.inner.values()
}
}
impl<'a, K: Storable + Ord + Clone, V: Storable, M: Memory> IntoIterator
for &'a StableBTreeMap<K, V, M>
{
type Item = <MapIter<'a, K, V, M> as Iterator>::Item;
type IntoIter = MapIter<'a, K, V, M>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}