1use crate::substrate::{self, DefaultMemoryImpl};
12use ic_memory::RuntimeMemory;
13use std::{borrow::Cow, ops::RangeBounds};
14pub use substrate::{Memory, Storable, storable, storable::Bound};
18
19#[cfg(all(test, not(target_arch = "wasm32")))]
20mod tests;
21
22mod region;
23pub use region::{BTreeMapRegion, RegionGrowError};
24
25pub struct StableBTreeMap<K: Storable + Ord + Clone, V: Storable, M: Memory = DefaultMemoryImpl> {
27 inner: substrate::BTreeMap<K, V, RuntimeMemory<M>>,
28}
29impl<K: Storable + Ord + Clone, V: Storable, M: Memory> StableBTreeMap<K, V, M> {
30 pub fn new(memory: RuntimeMemory<M>) -> Self {
32 Self {
33 inner: substrate::BTreeMap::new(memory),
34 }
35 }
36 pub fn init(memory: RuntimeMemory<M>) -> Self {
38 Self {
39 inner: substrate::BTreeMap::init(memory),
40 }
41 }
42 pub fn load(memory: RuntimeMemory<M>) -> Self {
44 Self {
45 inner: substrate::BTreeMap::load(memory),
46 }
47 }
48 pub fn clear(&mut self) {
50 self.inner.clear_new();
51 }
52 pub fn reserve_pages(self, minimum: u64) -> (Self, Result<(), ic_memory::RuntimeGrowError>) {
65 let memory = self.inner.into_memory();
66 let additional = minimum.saturating_sub(memory.size());
67 let result = memory.grow(additional).map(|_| ());
68 (Self::load(memory), result)
69 }
70 #[must_use]
76 pub fn allocated_pages(self) -> (Self, u64) {
77 let memory = self.inner.into_memory();
78 let pages = memory.size();
79 (Self::load(memory), pages)
80 }
81 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
83 self.inner.insert(key, value)
84 }
85 pub fn get(&self, key: &K) -> Option<V> {
87 self.inner.get(key)
88 }
89 pub fn contains_key(&self, key: &K) -> bool {
91 self.inner.contains_key(key)
92 }
93 pub fn is_empty(&self) -> bool {
95 self.inner.is_empty()
96 }
97 pub fn len(&self) -> u64 {
99 self.inner.len()
100 }
101 pub fn remove(&mut self, key: &K) -> Option<V> {
103 self.inner.remove(key)
104 }
105 pub fn first_key_value(&self) -> Option<(K, V)> {
107 self.inner.first_key_value()
108 }
109 pub fn last_key_value(&self) -> Option<(K, V)> {
111 self.inner.last_key_value()
112 }
113 pub fn pop_first(&mut self) -> Option<(K, V)> {
115 self.inner.pop_first()
116 }
117 pub fn pop_last(&mut self) -> Option<(K, V)> {
119 self.inner.pop_last()
120 }
121 pub fn iter(&self) -> MapIter<'_, K, V, M> {
123 self.inner.iter()
124 }
125 pub fn range(&self, key_range: impl RangeBounds<K>) -> MapIter<'_, K, V, M> {
127 self.inner.range(key_range)
128 }
129 pub fn keys(&self) -> impl DoubleEndedIterator<Item = K> + '_ {
131 self.inner.keys()
132 }
133 pub fn values(&self) -> impl DoubleEndedIterator<Item = V> + '_ {
135 self.inner.values()
136 }
137}
138
139pub struct StableBTreeSet<K: Storable + Ord + Clone, M: Memory = DefaultMemoryImpl> {
141 inner: substrate::BTreeSet<K, RuntimeMemory<M>>,
142}
143impl<K: Storable + Ord + Clone, M: Memory> StableBTreeSet<K, M> {
144 pub fn new(memory: RuntimeMemory<M>) -> Self {
146 Self {
147 inner: substrate::BTreeSet::new(memory),
148 }
149 }
150 pub fn init(memory: RuntimeMemory<M>) -> Self {
152 Self {
153 inner: substrate::BTreeSet::init(memory),
154 }
155 }
156 pub fn load(memory: RuntimeMemory<M>) -> Self {
158 Self {
159 inner: substrate::BTreeSet::load(memory),
160 }
161 }
162 pub fn insert(&mut self, key: K) -> bool {
164 self.inner.insert(key)
165 }
166 pub fn contains(&self, key: &K) -> bool {
168 self.inner.contains(key)
169 }
170 pub fn is_empty(&self) -> bool {
172 self.inner.is_empty()
173 }
174 pub fn len(&self) -> u64 {
176 self.inner.len()
177 }
178 pub fn remove(&mut self, key: &K) -> bool {
180 self.inner.remove(key)
181 }
182 pub fn first(&self) -> Option<K> {
184 self.inner.first()
185 }
186 pub fn last(&self) -> Option<K> {
188 self.inner.last()
189 }
190 pub fn pop_first(&mut self) -> Option<K> {
192 self.inner.pop_first()
193 }
194 pub fn pop_last(&mut self) -> Option<K> {
196 self.inner.pop_last()
197 }
198 pub fn iter(&self) -> SetIter<'_, K, M> {
200 self.inner.iter()
201 }
202 pub fn range(&self, key_range: impl RangeBounds<K>) -> SetIter<'_, K, M> {
204 self.inner.range(key_range)
205 }
206 pub fn clear(&mut self) {
208 self.inner.clear();
209 }
210}
211
212pub struct StableCell<T: Storable, M: Memory = DefaultMemoryImpl> {
214 inner: substrate::Cell<T, RuntimeMemory<M>>,
215}
216impl<T: Storable, M: Memory> StableCell<T, M> {
217 pub fn new(memory: RuntimeMemory<M>, value: T) -> Self {
219 Self {
220 inner: substrate::Cell::new(memory, value),
221 }
222 }
223 pub fn init(memory: RuntimeMemory<M>, default_value: T) -> Self {
225 Self {
226 inner: substrate::Cell::init(memory, default_value),
227 }
228 }
229 pub fn get(&self) -> &T {
231 self.inner.get()
232 }
233 pub fn set(&mut self, value: T) -> T {
235 self.inner.set(value)
236 }
237}
238
239pub struct StableVec<T: Storable, M: Memory = DefaultMemoryImpl> {
241 inner: substrate::Vec<T, RuntimeMemory<M>>,
242}
243impl<T: Storable, M: Memory> StableVec<T, M> {
244 pub fn new(memory: RuntimeMemory<M>) -> Self {
246 Self {
247 inner: substrate::Vec::new(memory),
248 }
249 }
250 pub fn init(memory: RuntimeMemory<M>) -> Self {
252 Self {
253 inner: substrate::Vec::init(memory),
254 }
255 }
256 pub fn get(&self, index: u64) -> Option<T> {
258 self.inner.get(index)
259 }
260 pub fn set(&self, index: u64, item: &T) {
262 self.inner.set(index, item);
263 }
264 pub fn push(&self, item: &T) {
266 self.inner.push(item);
267 }
268 pub fn pop(&self) -> Option<T> {
270 self.inner.pop()
271 }
272 #[must_use]
275 pub fn reset(self) -> Self {
276 Self::new(self.inner.into_memory())
277 }
278 pub fn len(&self) -> u64 {
280 self.inner.len()
281 }
282 pub fn is_empty(&self) -> bool {
284 self.inner.is_empty()
285 }
286 pub fn iter(&self) -> impl DoubleEndedIterator<Item = T> + '_ {
288 self.inner.iter()
289 }
290}
291
292pub struct StableMinHeap<T: Storable + Ord, M: Memory = DefaultMemoryImpl> {
294 inner: substrate::MinHeap<T, RuntimeMemory<M>>,
295}
296impl<T: Storable + Ord, M: Memory> StableMinHeap<T, M> {
297 pub fn new(memory: RuntimeMemory<M>) -> Self {
299 Self {
300 inner: substrate::MinHeap::new(memory),
301 }
302 }
303 pub fn init(memory: RuntimeMemory<M>) -> Self {
305 Self {
306 inner: substrate::MinHeap::init(memory),
307 }
308 }
309 pub fn push(&mut self, item: &T) {
311 self.inner.push(item);
312 }
313 pub fn pop(&mut self) -> Option<T> {
315 self.inner.pop()
316 }
317 pub fn peek(&self) -> Option<T> {
319 self.inner.peek()
320 }
321 pub fn len(&self) -> u64 {
323 self.inner.len()
324 }
325 pub fn is_empty(&self) -> bool {
327 self.inner.is_empty()
328 }
329 pub fn iter(&self) -> impl Iterator<Item = T> + '_ {
331 self.inner.iter()
332 }
333}
334
335pub struct StableLog<T: Storable, M: Memory = DefaultMemoryImpl> {
337 inner: substrate::Log<T, RuntimeMemory<M>, RuntimeMemory<M>>,
338}
339impl<T: Storable, M: Memory> StableLog<T, M> {
340 pub fn new(index_memory: RuntimeMemory<M>, data_memory: RuntimeMemory<M>) -> Self {
342 Self {
343 inner: substrate::Log::new(index_memory, data_memory),
344 }
345 }
346 pub fn init(index_memory: RuntimeMemory<M>, data_memory: RuntimeMemory<M>) -> Self {
348 Self {
349 inner: substrate::Log::init(index_memory, data_memory),
350 }
351 }
352 pub fn first(&self) -> Option<T> {
354 self.inner.first()
355 }
356 pub fn last(&self) -> Option<T> {
358 self.inner.last()
359 }
360 pub fn get(&self, idx: u64) -> Option<T> {
362 self.inner.get(idx)
363 }
364 pub fn len(&self) -> u64 {
366 self.inner.len()
367 }
368 pub fn is_empty(&self) -> bool {
370 self.inner.is_empty()
371 }
372 pub fn append(&self, item: &T) -> Result<u64, LogWriteError> {
376 self.inner.append(item)
377 }
378 pub fn iter(&self) -> LogIter<'_, T, M> {
380 self.inner.iter()
381 }
382}
383
384pub type MapIter<'a, K, V, M = DefaultMemoryImpl> =
386 substrate::btreemap::Iter<'a, K, V, RuntimeMemory<M>>;
387pub type SetIter<'a, K, M = DefaultMemoryImpl> = substrate::btreeset::Iter<'a, K, RuntimeMemory<M>>;
389pub type LogIter<'a, T, M = DefaultMemoryImpl> =
391 substrate::log::Iter<'a, T, RuntimeMemory<M>, RuntimeMemory<M>>;
392
393pub use substrate::log::WriteError as LogWriteError;
394
395impl<'a, K: Storable + Ord + Clone, V: Storable, M: Memory> IntoIterator
396 for &'a StableBTreeMap<K, V, M>
397{
398 type Item = <MapIter<'a, K, V, M> as Iterator>::Item;
399 type IntoIter = MapIter<'a, K, V, M>;
400 fn into_iter(self) -> Self::IntoIter {
401 self.iter()
402 }
403}
404
405impl<'a, K: Storable + Ord + Clone, M: Memory> IntoIterator for &'a StableBTreeSet<K, M> {
406 type Item = K;
407 type IntoIter = SetIter<'a, K, M>;
408 fn into_iter(self) -> Self::IntoIter {
409 self.iter()
410 }
411}
412
413impl<'a, T: Storable, M: Memory> IntoIterator for &'a StableLog<T, M> {
414 type Item = T;
415 type IntoIter = LogIter<'a, T, M>;
416 fn into_iter(self) -> Self::IntoIter {
417 self.iter()
418 }
419}
420pub struct BoundedCell<T: Storable, M: Memory = DefaultMemoryImpl> {
422 cell: substrate::Cell<BoundedValue<T>, RuntimeMemory<M>>,
423}
424
425impl<T: Storable, M: Memory> BoundedCell<T, M> {
426 #[must_use]
431 pub fn init(memory: RuntimeMemory<M>, value: T) -> Self {
432 assert!(matches!(T::BOUND, Bound::Bounded { .. }));
433 Self {
434 cell: substrate::Cell::init(memory, BoundedValue(value)),
435 }
436 }
437
438 #[must_use]
440 pub fn get(&self) -> &T {
441 &self.cell.get().0
442 }
443
444 pub fn set(&mut self, value: T) -> T {
449 self.cell.set(BoundedValue(value)).0
450 }
451}
452
453struct BoundedValue<T>(T);
454
455impl<T: Storable> Storable for BoundedValue<T> {
456 const BOUND: Bound = T::BOUND;
457
458 fn to_bytes(&self) -> Cow<'_, [u8]> {
459 self.0.to_bytes_checked()
460 }
461
462 fn into_bytes(self) -> Vec<u8> {
463 self.0.into_bytes_checked()
464 }
465
466 fn from_bytes(bytes: Cow<'_, [u8]>) -> Self {
467 let Bound::Bounded {
468 max_size,
469 is_fixed_size,
470 } = T::BOUND
471 else {
472 unreachable!("bounded cell record must declare a bound");
473 };
474 if is_fixed_size {
475 assert_eq!(bytes.len(), max_size as usize);
476 } else {
477 assert!(bytes.len() <= max_size as usize);
478 }
479 Self(T::from_bytes(bytes))
480 }
481}