Skip to main content

ic_storage/adapter/
map.rs

1//! Governed map operations over the copied ICSS B-tree.
2use super::bridge::{self, Codec, TreeMemory};
3use super::error::{CollectionError as Error, Result, ensure};
4use super::{MapIter, Memory, Storable};
5use crate::{DefaultMemoryImpl, RuntimeMemory};
6use std::ops::RangeBounds;
7
8///
9/// `StableBTreeMap`
10///
11/// ICSS-baseline ordered map owning one governed allocation. Stored codec and
12/// engine failures trap. Native callers must discard the runtime after a mutation panic.
13///
14pub struct StableBTreeMap<K: Storable + Ord + Clone, V: Storable, M: Memory = DefaultMemoryImpl> {
15    inner: bridge::Map<K, V, M>,
16}
17impl<K: Storable + Ord + Clone, V: Storable, M: Memory> StableBTreeMap<K, V, M> {
18    /// Create only in empty admitted memory.
19    /// # Errors
20    /// Refuses nonempty memory or initial growth refusal.
21    pub fn create(memory: RuntimeMemory<M>) -> Result<Self> {
22        ensure(memory.size() == 0, Error::NotEmpty)?;
23        Self::reset(memory)
24    }
25    /// Explicitly replace this allocation's collection image.
26    /// # Errors
27    /// Refuses initial growth before creating the map.
28    pub fn reset(memory: RuntimeMemory<M>) -> Result<Self> {
29        bridge::reserve(&memory)?;
30        #[cfg(not(target_arch = "wasm32"))]
31        let _guard = memory.poison_on_unwind();
32        Ok(Self {
33            inner: bridge::Map::new(TreeMemory::Allocation(memory)),
34        })
35    }
36    /// Load the current BTR version without initialization or migration.
37    /// # Errors
38    /// Refuses empty or foreign framing.
39    pub fn load(memory: RuntimeMemory<M>) -> Result<Self> {
40        Ok(Self {
41            inner: bridge::Map::load(TreeMemory::Allocation(memory)).ok_or(Error::Format)?,
42        })
43    }
44    /// Clear entries and reuse the map's existing backing capacity.
45    pub fn clear(&mut self) {
46        #[cfg(not(target_arch = "wasm32"))]
47        let _guard = self.inner.memory().allocation().poison_on_unwind();
48        self.inner.clear_new();
49    }
50    /// Reserve backing pages without promising entry capacity.
51    pub fn reserve_pages(
52        self,
53        minimum: u64,
54    ) -> (Self, std::result::Result<(), crate::RuntimeGrowError>) {
55        let additional = minimum.saturating_sub(self.inner.memory().allocation().size());
56        let result = self
57            .inner
58            .memory()
59            .allocation()
60            .grow(additional)
61            .map(|_| ());
62        (self, result)
63    }
64    /// Observe allocated pages while retaining exclusive custody.
65    #[must_use]
66    pub fn allocated_pages(self) -> (Self, u64) {
67        let pages = self.inner.memory().allocation().size();
68        (self, pages)
69    }
70    /// Insert or replace. Invalid codecs and allocation failures trap.
71    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
72        #[cfg(not(target_arch = "wasm32"))]
73        let _guard = self.inner.memory().allocation().poison_on_unwind();
74        self.inner.insert(Codec(key), Codec(value)).map(|v| v.0)
75    }
76    /// Read a matching value; malformed stored values trap.
77    pub fn get(&self, key: &K) -> Option<V> {
78        self.inner.get(Codec::borrowed(key)).map(|v| v.0)
79    }
80    /// Test key membership without decoding values.
81    pub fn contains_key(&self, key: &K) -> bool {
82        self.inner.contains_key(Codec::borrowed(key))
83    }
84    /// Whether no entries are committed.
85    #[must_use]
86    pub const fn is_empty(&self) -> bool {
87        self.inner.is_empty()
88    }
89    /// Number of entries.
90    #[must_use]
91    pub const fn len(&self) -> u64 {
92        self.inner.len()
93    }
94    /// Remove a matching value; codec and engine failures trap.
95    pub fn remove(&mut self, key: &K) -> Option<V> {
96        #[cfg(not(target_arch = "wasm32"))]
97        let _guard = self.inner.memory().allocation().poison_on_unwind();
98        self.inner.remove(Codec::borrowed(key)).map(|v| v.0)
99    }
100    /// First ordered key/value pair.
101    #[must_use]
102    pub fn first_key_value(&self) -> Option<(K, V)> {
103        self.inner.first_key_value().map(|(k, v)| (k.0, v.0))
104    }
105    /// Last ordered key/value pair.
106    #[must_use]
107    pub fn last_key_value(&self) -> Option<(K, V)> {
108        self.inner.last_key_value().map(|(k, v)| (k.0, v.0))
109    }
110    /// Iterate in key order.
111    #[must_use]
112    pub fn iter(&self) -> MapIter<'_, K, V, M> {
113        MapIter(self.inner.iter())
114    }
115    /// Iterate a bounded key range.
116    pub fn range(&self, bounds: impl RangeBounds<K>) -> MapIter<'_, K, V, M> {
117        MapIter(self.inner.range(codec_bounds(&bounds)))
118    }
119    /// Iterate keys without reading value payloads.
120    pub fn keys(&self) -> impl DoubleEndedIterator<Item = K> + '_ {
121        self.inner.keys().map(|k| k.0)
122    }
123    /// Iterate bounded keys without reading value payloads.
124    pub fn keys_range(
125        &self,
126        bounds: impl RangeBounds<K>,
127    ) -> impl DoubleEndedIterator<Item = K> + '_ {
128        self.inner.keys_range(codec_bounds(&bounds)).map(|k| k.0)
129    }
130    /// Iterate decoded values.
131    pub fn values(&self) -> impl DoubleEndedIterator<Item = V> + '_ {
132        self.inner.values().map(|v| v.0)
133    }
134}
135
136pub(super) fn codec_bounds<K: Clone>(
137    range: &impl RangeBounds<K>,
138) -> (std::ops::Bound<Codec<K>>, std::ops::Bound<Codec<K>>) {
139    (
140        range.start_bound().map(|k| Codec(k.clone())),
141        range.end_bound().map(|k| Codec(k.clone())),
142    )
143}
144impl<'a, K: Storable + Ord + Clone, V: Storable, M: Memory> IntoIterator
145    for &'a StableBTreeMap<K, V, M>
146{
147    type Item = <MapIter<'a, K, V, M> as Iterator>::Item;
148    type IntoIter = MapIter<'a, K, V, M>;
149    fn into_iter(self) -> Self::IntoIter {
150        self.iter()
151    }
152}