1use crate::substrate::{self, DefaultMemoryImpl};
10use ic_memory::RuntimeMemory;
11use std::{borrow::Cow, ops::RangeBounds};
12pub use substrate::{Memory, Storable, storable, storable::Bound};
16
17#[cfg(all(test, not(target_arch = "wasm32")))]
18mod tests;
19
20pub struct StableBTreeMap<K: Storable + Ord + Clone, V: Storable, M: Memory = DefaultMemoryImpl> {
22 inner: substrate::BTreeMap<K, V, RuntimeMemory<M>>,
23}
24impl<K: Storable + Ord + Clone, V: Storable, M: Memory> StableBTreeMap<K, V, M> {
25 pub fn new(memory: RuntimeMemory<M>) -> Self {
27 Self {
28 inner: substrate::BTreeMap::new(memory),
29 }
30 }
31 pub fn init(memory: RuntimeMemory<M>) -> Self {
33 Self {
34 inner: substrate::BTreeMap::init(memory),
35 }
36 }
37 pub fn load(memory: RuntimeMemory<M>) -> Self {
39 Self {
40 inner: substrate::BTreeMap::load(memory),
41 }
42 }
43 pub fn clear(&mut self) {
45 self.inner.clear_new();
46 }
47 pub fn reserve_pages(self, minimum: u64) -> (Self, Result<(), ic_memory::RuntimeGrowError>) {
60 let memory = self.inner.into_memory();
61 let additional = minimum.saturating_sub(memory.size());
62 let result = memory.grow(additional).map(|_| ());
63 (Self::load(memory), result)
64 }
65 #[must_use]
71 pub fn allocated_pages(self) -> (Self, u64) {
72 let memory = self.inner.into_memory();
73 let pages = memory.size();
74 (Self::load(memory), pages)
75 }
76 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
78 self.inner.insert(key, value)
79 }
80 pub fn get(&self, key: &K) -> Option<V> {
82 self.inner.get(key)
83 }
84 pub fn contains_key(&self, key: &K) -> bool {
86 self.inner.contains_key(key)
87 }
88 pub fn is_empty(&self) -> bool {
90 self.inner.is_empty()
91 }
92 pub fn len(&self) -> u64 {
94 self.inner.len()
95 }
96 pub fn remove(&mut self, key: &K) -> Option<V> {
98 self.inner.remove(key)
99 }
100 pub fn first_key_value(&self) -> Option<(K, V)> {
102 self.inner.first_key_value()
103 }
104 pub fn last_key_value(&self) -> Option<(K, V)> {
106 self.inner.last_key_value()
107 }
108 pub fn pop_first(&mut self) -> Option<(K, V)> {
110 self.inner.pop_first()
111 }
112 pub fn pop_last(&mut self) -> Option<(K, V)> {
114 self.inner.pop_last()
115 }
116 pub fn iter(&self) -> MapIter<'_, K, V, M> {
118 self.inner.iter()
119 }
120 pub fn range(&self, key_range: impl RangeBounds<K>) -> MapIter<'_, K, V, M> {
122 self.inner.range(key_range)
123 }
124 pub fn keys(&self) -> impl DoubleEndedIterator<Item = K> + '_ {
126 self.inner.keys()
127 }
128 pub fn values(&self) -> impl DoubleEndedIterator<Item = V> + '_ {
130 self.inner.values()
131 }
132}
133
134pub struct StableBTreeSet<K: Storable + Ord + Clone, M: Memory = DefaultMemoryImpl> {
136 inner: substrate::BTreeSet<K, RuntimeMemory<M>>,
137}
138impl<K: Storable + Ord + Clone, M: Memory> StableBTreeSet<K, M> {
139 pub fn new(memory: RuntimeMemory<M>) -> Self {
141 Self {
142 inner: substrate::BTreeSet::new(memory),
143 }
144 }
145 pub fn init(memory: RuntimeMemory<M>) -> Self {
147 Self {
148 inner: substrate::BTreeSet::init(memory),
149 }
150 }
151 pub fn load(memory: RuntimeMemory<M>) -> Self {
153 Self {
154 inner: substrate::BTreeSet::load(memory),
155 }
156 }
157 pub fn insert(&mut self, key: K) -> bool {
159 self.inner.insert(key)
160 }
161 pub fn contains(&self, key: &K) -> bool {
163 self.inner.contains(key)
164 }
165 pub fn is_empty(&self) -> bool {
167 self.inner.is_empty()
168 }
169 pub fn len(&self) -> u64 {
171 self.inner.len()
172 }
173 pub fn remove(&mut self, key: &K) -> bool {
175 self.inner.remove(key)
176 }
177 pub fn first(&self) -> Option<K> {
179 self.inner.first()
180 }
181 pub fn last(&self) -> Option<K> {
183 self.inner.last()
184 }
185 pub fn pop_first(&mut self) -> Option<K> {
187 self.inner.pop_first()
188 }
189 pub fn pop_last(&mut self) -> Option<K> {
191 self.inner.pop_last()
192 }
193 pub fn iter(&self) -> SetIter<'_, K, M> {
195 self.inner.iter()
196 }
197 pub fn range(&self, key_range: impl RangeBounds<K>) -> SetIter<'_, K, M> {
199 self.inner.range(key_range)
200 }
201 pub fn clear(&mut self) {
203 self.inner.clear();
204 }
205}
206
207pub struct StableCell<T: Storable, M: Memory = DefaultMemoryImpl> {
209 inner: substrate::Cell<T, RuntimeMemory<M>>,
210}
211impl<T: Storable, M: Memory> StableCell<T, M> {
212 pub fn new(memory: RuntimeMemory<M>, value: T) -> Self {
214 Self {
215 inner: substrate::Cell::new(memory, value),
216 }
217 }
218 pub fn init(memory: RuntimeMemory<M>, default_value: T) -> Self {
220 Self {
221 inner: substrate::Cell::init(memory, default_value),
222 }
223 }
224 pub fn get(&self) -> &T {
226 self.inner.get()
227 }
228 pub fn set(&mut self, value: T) -> T {
230 self.inner.set(value)
231 }
232}
233
234pub struct StableVec<T: Storable, M: Memory = DefaultMemoryImpl> {
236 inner: substrate::Vec<T, RuntimeMemory<M>>,
237}
238impl<T: Storable, M: Memory> StableVec<T, M> {
239 pub fn new(memory: RuntimeMemory<M>) -> Self {
241 Self {
242 inner: substrate::Vec::new(memory),
243 }
244 }
245 pub fn init(memory: RuntimeMemory<M>) -> Self {
247 Self {
248 inner: substrate::Vec::init(memory),
249 }
250 }
251 pub fn get(&self, index: u64) -> Option<T> {
253 self.inner.get(index)
254 }
255 pub fn set(&self, index: u64, item: &T) {
257 self.inner.set(index, item);
258 }
259 pub fn push(&self, item: &T) {
261 self.inner.push(item);
262 }
263 pub fn pop(&self) -> Option<T> {
265 self.inner.pop()
266 }
267 #[must_use]
270 pub fn reset(self) -> Self {
271 Self::new(self.inner.into_memory())
272 }
273 pub fn len(&self) -> u64 {
275 self.inner.len()
276 }
277 pub fn is_empty(&self) -> bool {
279 self.inner.is_empty()
280 }
281 pub fn iter(&self) -> impl DoubleEndedIterator<Item = T> + '_ {
283 self.inner.iter()
284 }
285}
286
287pub struct StableMinHeap<T: Storable + Ord, M: Memory = DefaultMemoryImpl> {
289 inner: substrate::MinHeap<T, RuntimeMemory<M>>,
290}
291impl<T: Storable + Ord, M: Memory> StableMinHeap<T, M> {
292 pub fn new(memory: RuntimeMemory<M>) -> Self {
294 Self {
295 inner: substrate::MinHeap::new(memory),
296 }
297 }
298 pub fn init(memory: RuntimeMemory<M>) -> Self {
300 Self {
301 inner: substrate::MinHeap::init(memory),
302 }
303 }
304 pub fn push(&mut self, item: &T) {
306 self.inner.push(item);
307 }
308 pub fn pop(&mut self) -> Option<T> {
310 self.inner.pop()
311 }
312 pub fn peek(&self) -> Option<T> {
314 self.inner.peek()
315 }
316 pub fn len(&self) -> u64 {
318 self.inner.len()
319 }
320 pub fn is_empty(&self) -> bool {
322 self.inner.is_empty()
323 }
324 pub fn iter(&self) -> impl Iterator<Item = T> + '_ {
326 self.inner.iter()
327 }
328}
329
330pub struct StableLog<T: Storable, M: Memory = DefaultMemoryImpl> {
332 inner: substrate::Log<T, RuntimeMemory<M>, RuntimeMemory<M>>,
333}
334impl<T: Storable, M: Memory> StableLog<T, M> {
335 pub fn new(index_memory: RuntimeMemory<M>, data_memory: RuntimeMemory<M>) -> Self {
337 Self {
338 inner: substrate::Log::new(index_memory, data_memory),
339 }
340 }
341 pub fn init(index_memory: RuntimeMemory<M>, data_memory: RuntimeMemory<M>) -> Self {
343 Self {
344 inner: substrate::Log::init(index_memory, data_memory),
345 }
346 }
347 pub fn first(&self) -> Option<T> {
349 self.inner.first()
350 }
351 pub fn last(&self) -> Option<T> {
353 self.inner.last()
354 }
355 pub fn get(&self, idx: u64) -> Option<T> {
357 self.inner.get(idx)
358 }
359 pub fn len(&self) -> u64 {
361 self.inner.len()
362 }
363 pub fn is_empty(&self) -> bool {
365 self.inner.is_empty()
366 }
367 pub fn append(&self, item: &T) -> Result<u64, LogWriteError> {
371 self.inner.append(item)
372 }
373 pub fn iter(&self) -> LogIter<'_, T, M> {
375 self.inner.iter()
376 }
377}
378
379pub type MapIter<'a, K, V, M = DefaultMemoryImpl> =
381 substrate::btreemap::Iter<'a, K, V, RuntimeMemory<M>>;
382pub type SetIter<'a, K, M = DefaultMemoryImpl> = substrate::btreeset::Iter<'a, K, RuntimeMemory<M>>;
384pub type LogIter<'a, T, M = DefaultMemoryImpl> =
386 substrate::log::Iter<'a, T, RuntimeMemory<M>, RuntimeMemory<M>>;
387
388pub use substrate::log::WriteError as LogWriteError;
389
390impl<'a, K: Storable + Ord + Clone, V: Storable, M: Memory> IntoIterator
391 for &'a StableBTreeMap<K, V, M>
392{
393 type Item = <MapIter<'a, K, V, M> as Iterator>::Item;
394 type IntoIter = MapIter<'a, K, V, M>;
395 fn into_iter(self) -> Self::IntoIter {
396 self.iter()
397 }
398}
399
400impl<'a, K: Storable + Ord + Clone, M: Memory> IntoIterator for &'a StableBTreeSet<K, M> {
401 type Item = K;
402 type IntoIter = SetIter<'a, K, M>;
403 fn into_iter(self) -> Self::IntoIter {
404 self.iter()
405 }
406}
407
408impl<'a, T: Storable, M: Memory> IntoIterator for &'a StableLog<T, M> {
409 type Item = T;
410 type IntoIter = LogIter<'a, T, M>;
411 fn into_iter(self) -> Self::IntoIter {
412 self.iter()
413 }
414}
415pub struct BoundedCell<T: Storable, M: Memory = DefaultMemoryImpl> {
417 cell: substrate::Cell<BoundedValue<T>, RuntimeMemory<M>>,
418}
419
420impl<T: Storable, M: Memory> BoundedCell<T, M> {
421 #[must_use]
426 pub fn init(memory: RuntimeMemory<M>, value: T) -> Self {
427 assert!(matches!(T::BOUND, Bound::Bounded { .. }));
428 Self {
429 cell: substrate::Cell::init(memory, BoundedValue(value)),
430 }
431 }
432
433 #[must_use]
435 pub fn get(&self) -> &T {
436 &self.cell.get().0
437 }
438
439 pub fn set(&mut self, value: T) -> T {
444 self.cell.set(BoundedValue(value)).0
445 }
446}
447
448struct BoundedValue<T>(T);
449
450impl<T: Storable> Storable for BoundedValue<T> {
451 const BOUND: Bound = T::BOUND;
452
453 fn to_bytes(&self) -> Cow<'_, [u8]> {
454 self.0.to_bytes_checked()
455 }
456
457 fn into_bytes(self) -> Vec<u8> {
458 self.0.into_bytes_checked()
459 }
460
461 fn from_bytes(bytes: Cow<'_, [u8]>) -> Self {
462 let Bound::Bounded {
463 max_size,
464 is_fixed_size,
465 } = T::BOUND
466 else {
467 unreachable!("bounded cell record must declare a bound");
468 };
469 if is_fixed_size {
470 assert_eq!(bytes.len(), max_size as usize);
471 } else {
472 assert!(bytes.len() <= max_size as usize);
473 }
474 Self(T::from_bytes(bytes))
475 }
476}