Skip to main content

ic_storage/
adapter.rs

1//! Supported upstream collections with governed memory and an explicit API.
2//!
3//! Constructors preserve upstream initialization, panic and persisted-format
4//! contracts. These wrappers expose neither raw memory nor unrestricted
5//! upstream collections.
6//! Use distinct admitted allocations for owning collections. Region views retain
7//! a shared lease for disjoint application frames and reload roots per operation.
8//! Drop collections and all shared lease owners before reopening.
9//! IC Memory enforces one live lease per allocation within each runtime;
10//! hosts retain exclusive custody of the backing across runtimes.
11use crate::substrate::{self, DefaultMemoryImpl};
12use ic_memory::RuntimeMemory;
13use std::{borrow::Cow, ops::RangeBounds};
14/// Exact substrate interfaces for backing bounds and record codecs.
15/// Implementing `Memory` grants no allocation authority; constructors require
16/// committed `RuntimeMemory` leases.
17pub 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
25/// Governed upstream `BTreeMap`; initialization and failures follow upstream contracts.
26pub 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    /// Delegates to upstream `BTreeMap::new` with the governed allocation.
31    pub fn new(memory: RuntimeMemory<M>) -> Self {
32        Self {
33            inner: substrate::BTreeMap::new(memory),
34        }
35    }
36    /// Delegates to upstream `BTreeMap::init` with the governed allocation.
37    pub fn init(memory: RuntimeMemory<M>) -> Self {
38        Self {
39            inner: substrate::BTreeMap::init(memory),
40        }
41    }
42    /// Delegates to upstream `BTreeMap::load` with the governed allocation.
43    pub fn load(memory: RuntimeMemory<M>) -> Self {
44        Self {
45            inner: substrate::BTreeMap::load(memory),
46        }
47    }
48    /// Clear entries while retaining this allocation's exclusive lease.
49    pub fn clear(&mut self) {
50        self.inner.clear_new();
51    }
52    /// Reserve a minimum virtual extent without exposing or duplicating the lease.
53    ///
54    /// Returns the map on both success and ordinary growth refusal. Entries are
55    /// preserved; reloading the upstream header refreshes its cached allocator.
56    /// Backing traps retain the upstream panic contract.
57    ///
58    /// `minimum` is an absolute virtual extent in 65,536-byte Wasm pages, not an
59    /// entry count or additional pages. Existing extents never shrink. A request
60    /// already covered performs no growth or writes. This reserves backing only:
61    /// it does not promise that any particular number or encoding of entries fits.
62    /// Callers own their byte/page budget and must qualify it for their codecs;
63    /// upstream node-layout constants are not a supported capacity contract.
64    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    /// Observe the virtual extent in 65,536-byte pages while retaining the lease.
71    ///
72    /// Consumes and returns the map because upstream exposes memory only through
73    /// ownership transfer. Reloading reads the header without writing or growing.
74    /// This is allocated extent, not free space or guaranteed entry capacity.
75    #[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    /// Delegates to upstream `BTreeMap::insert` with the governed allocation.
82    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
83        self.inner.insert(key, value)
84    }
85    /// Delegates to upstream `BTreeMap::get` with the governed allocation.
86    pub fn get(&self, key: &K) -> Option<V> {
87        self.inner.get(key)
88    }
89    /// Delegates to upstream `BTreeMap::contains_key` with the governed allocation.
90    pub fn contains_key(&self, key: &K) -> bool {
91        self.inner.contains_key(key)
92    }
93    /// Delegates to upstream `BTreeMap::is_empty` with the governed allocation.
94    pub fn is_empty(&self) -> bool {
95        self.inner.is_empty()
96    }
97    /// Delegates to upstream `BTreeMap::len` with the governed allocation.
98    pub fn len(&self) -> u64 {
99        self.inner.len()
100    }
101    /// Delegates to upstream `BTreeMap::remove` with the governed allocation.
102    pub fn remove(&mut self, key: &K) -> Option<V> {
103        self.inner.remove(key)
104    }
105    /// Delegates to upstream `BTreeMap::first_key_value` with the governed allocation.
106    pub fn first_key_value(&self) -> Option<(K, V)> {
107        self.inner.first_key_value()
108    }
109    /// Delegates to upstream `BTreeMap::last_key_value` with the governed allocation.
110    pub fn last_key_value(&self) -> Option<(K, V)> {
111        self.inner.last_key_value()
112    }
113    /// Delegates to upstream `BTreeMap::pop_first` with the governed allocation.
114    pub fn pop_first(&mut self) -> Option<(K, V)> {
115        self.inner.pop_first()
116    }
117    /// Delegates to upstream `BTreeMap::pop_last` with the governed allocation.
118    pub fn pop_last(&mut self) -> Option<(K, V)> {
119        self.inner.pop_last()
120    }
121    /// Delegates to upstream `BTreeMap::iter` with the governed allocation.
122    pub fn iter(&self) -> MapIter<'_, K, V, M> {
123        self.inner.iter()
124    }
125    /// Delegates to upstream `BTreeMap::range` with the governed allocation.
126    pub fn range(&self, key_range: impl RangeBounds<K>) -> MapIter<'_, K, V, M> {
127        self.inner.range(key_range)
128    }
129    /// Delegates to upstream `BTreeMap::keys` with the governed allocation.
130    pub fn keys(&self) -> impl DoubleEndedIterator<Item = K> + '_ {
131        self.inner.keys()
132    }
133    /// Delegates to upstream `BTreeMap::values` with the governed allocation.
134    pub fn values(&self) -> impl DoubleEndedIterator<Item = V> + '_ {
135        self.inner.values()
136    }
137}
138
139/// Governed upstream `BTreeSet`; initialization and failures follow upstream contracts.
140pub 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    /// Delegates to upstream `BTreeSet::new` with the governed allocation.
145    pub fn new(memory: RuntimeMemory<M>) -> Self {
146        Self {
147            inner: substrate::BTreeSet::new(memory),
148        }
149    }
150    /// Delegates to upstream `BTreeSet::init` with the governed allocation.
151    pub fn init(memory: RuntimeMemory<M>) -> Self {
152        Self {
153            inner: substrate::BTreeSet::init(memory),
154        }
155    }
156    /// Delegates to upstream `BTreeSet::load` with the governed allocation.
157    pub fn load(memory: RuntimeMemory<M>) -> Self {
158        Self {
159            inner: substrate::BTreeSet::load(memory),
160        }
161    }
162    /// Delegates to upstream `BTreeSet::insert` with the governed allocation.
163    pub fn insert(&mut self, key: K) -> bool {
164        self.inner.insert(key)
165    }
166    /// Delegates to upstream `BTreeSet::contains` with the governed allocation.
167    pub fn contains(&self, key: &K) -> bool {
168        self.inner.contains(key)
169    }
170    /// Delegates to upstream `BTreeSet::is_empty` with the governed allocation.
171    pub fn is_empty(&self) -> bool {
172        self.inner.is_empty()
173    }
174    /// Delegates to upstream `BTreeSet::len` with the governed allocation.
175    pub fn len(&self) -> u64 {
176        self.inner.len()
177    }
178    /// Delegates to upstream `BTreeSet::remove` with the governed allocation.
179    pub fn remove(&mut self, key: &K) -> bool {
180        self.inner.remove(key)
181    }
182    /// Delegates to upstream `BTreeSet::first` with the governed allocation.
183    pub fn first(&self) -> Option<K> {
184        self.inner.first()
185    }
186    /// Delegates to upstream `BTreeSet::last` with the governed allocation.
187    pub fn last(&self) -> Option<K> {
188        self.inner.last()
189    }
190    /// Delegates to upstream `BTreeSet::pop_first` with the governed allocation.
191    pub fn pop_first(&mut self) -> Option<K> {
192        self.inner.pop_first()
193    }
194    /// Delegates to upstream `BTreeSet::pop_last` with the governed allocation.
195    pub fn pop_last(&mut self) -> Option<K> {
196        self.inner.pop_last()
197    }
198    /// Delegates to upstream `BTreeSet::iter` with the governed allocation.
199    pub fn iter(&self) -> SetIter<'_, K, M> {
200        self.inner.iter()
201    }
202    /// Delegates to upstream `BTreeSet::range` with the governed allocation.
203    pub fn range(&self, key_range: impl RangeBounds<K>) -> SetIter<'_, K, M> {
204        self.inner.range(key_range)
205    }
206    /// Delegates to upstream `BTreeSet::clear` with the governed allocation.
207    pub fn clear(&mut self) {
208        self.inner.clear();
209    }
210}
211
212/// Governed upstream `Cell`; initialization and failures follow upstream contracts.
213pub struct StableCell<T: Storable, M: Memory = DefaultMemoryImpl> {
214    inner: substrate::Cell<T, RuntimeMemory<M>>,
215}
216impl<T: Storable, M: Memory> StableCell<T, M> {
217    /// Delegates to upstream `Cell::new` with the governed allocation.
218    pub fn new(memory: RuntimeMemory<M>, value: T) -> Self {
219        Self {
220            inner: substrate::Cell::new(memory, value),
221        }
222    }
223    /// Delegates to upstream `Cell::init` with the governed allocation.
224    pub fn init(memory: RuntimeMemory<M>, default_value: T) -> Self {
225        Self {
226            inner: substrate::Cell::init(memory, default_value),
227        }
228    }
229    /// Delegates to upstream `Cell::get` with the governed allocation.
230    pub fn get(&self) -> &T {
231        self.inner.get()
232    }
233    /// Delegates to upstream `Cell::set` with the governed allocation.
234    pub fn set(&mut self, value: T) -> T {
235        self.inner.set(value)
236    }
237}
238
239/// Governed upstream `Vec`; initialization and failures follow upstream contracts.
240pub struct StableVec<T: Storable, M: Memory = DefaultMemoryImpl> {
241    inner: substrate::Vec<T, RuntimeMemory<M>>,
242}
243impl<T: Storable, M: Memory> StableVec<T, M> {
244    /// Delegates to upstream `Vec::new` with the governed allocation.
245    pub fn new(memory: RuntimeMemory<M>) -> Self {
246        Self {
247            inner: substrate::Vec::new(memory),
248        }
249    }
250    /// Delegates to upstream `Vec::init` with the governed allocation.
251    pub fn init(memory: RuntimeMemory<M>) -> Self {
252        Self {
253            inner: substrate::Vec::init(memory),
254        }
255    }
256    /// Delegates to upstream `Vec::get` with the governed allocation.
257    pub fn get(&self, index: u64) -> Option<T> {
258        self.inner.get(index)
259    }
260    /// Delegates to upstream `Vec::set` with the governed allocation.
261    pub fn set(&self, index: u64, item: &T) {
262        self.inner.set(index, item);
263    }
264    /// Delegates to upstream `Vec::push` with the governed allocation.
265    pub fn push(&self, item: &T) {
266        self.inner.push(item);
267    }
268    /// Delegates to upstream `Vec::pop` with the governed allocation.
269    pub fn pop(&self) -> Option<T> {
270        self.inner.pop()
271    }
272    /// Reset to an empty vector while retaining this allocation's exclusive lease.
273    /// Existing payload bytes follow upstream `Vec::new` semantics.
274    #[must_use]
275    pub fn reset(self) -> Self {
276        Self::new(self.inner.into_memory())
277    }
278    /// Delegates to upstream `Vec::len` with the governed allocation.
279    pub fn len(&self) -> u64 {
280        self.inner.len()
281    }
282    /// Delegates to upstream `Vec::is_empty` with the governed allocation.
283    pub fn is_empty(&self) -> bool {
284        self.inner.is_empty()
285    }
286    /// Delegates to upstream `Vec::iter` with the governed allocation.
287    pub fn iter(&self) -> impl DoubleEndedIterator<Item = T> + '_ {
288        self.inner.iter()
289    }
290}
291
292/// Governed upstream `MinHeap`; initialization and failures follow upstream contracts.
293pub 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    /// Delegates to upstream `MinHeap::new` with the governed allocation.
298    pub fn new(memory: RuntimeMemory<M>) -> Self {
299        Self {
300            inner: substrate::MinHeap::new(memory),
301        }
302    }
303    /// Delegates to upstream `MinHeap::init` with the governed allocation.
304    pub fn init(memory: RuntimeMemory<M>) -> Self {
305        Self {
306            inner: substrate::MinHeap::init(memory),
307        }
308    }
309    /// Delegates to upstream `MinHeap::push` with the governed allocation.
310    pub fn push(&mut self, item: &T) {
311        self.inner.push(item);
312    }
313    /// Delegates to upstream `MinHeap::pop` with the governed allocation.
314    pub fn pop(&mut self) -> Option<T> {
315        self.inner.pop()
316    }
317    /// Delegates to upstream `MinHeap::peek` with the governed allocation.
318    pub fn peek(&self) -> Option<T> {
319        self.inner.peek()
320    }
321    /// Delegates to upstream `MinHeap::len` with the governed allocation.
322    pub fn len(&self) -> u64 {
323        self.inner.len()
324    }
325    /// Delegates to upstream `MinHeap::is_empty` with the governed allocation.
326    pub fn is_empty(&self) -> bool {
327        self.inner.is_empty()
328    }
329    /// Delegates to upstream `MinHeap::iter` with the governed allocation.
330    pub fn iter(&self) -> impl Iterator<Item = T> + '_ {
331        self.inner.iter()
332    }
333}
334
335/// Governed upstream `Log`; initialization and failures follow upstream contracts.
336pub 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    /// Delegates to upstream `Log::new` with the governed allocation.
341    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    /// Delegates to upstream `Log::init` with the governed allocation.
347    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    /// Delegates to upstream `Log::first` with the governed allocation.
353    pub fn first(&self) -> Option<T> {
354        self.inner.first()
355    }
356    /// Delegates to upstream `Log::last` with the governed allocation.
357    pub fn last(&self) -> Option<T> {
358        self.inner.last()
359    }
360    /// Delegates to upstream `Log::get` with the governed allocation.
361    pub fn get(&self, idx: u64) -> Option<T> {
362        self.inner.get(idx)
363    }
364    /// Delegates to upstream `Log::len` with the governed allocation.
365    pub fn len(&self) -> u64 {
366        self.inner.len()
367    }
368    /// Delegates to upstream `Log::is_empty` with the governed allocation.
369    pub fn is_empty(&self) -> bool {
370        self.inner.is_empty()
371    }
372    /// Delegates to upstream `Log::append` with the governed allocation.
373    /// # Errors
374    /// Returns the upstream write error if memory cannot grow.
375    pub fn append(&self, item: &T) -> Result<u64, LogWriteError> {
376        self.inner.append(item)
377    }
378    /// Delegates to upstream `Log::iter` with the governed allocation.
379    pub fn iter(&self) -> LogIter<'_, T, M> {
380        self.inner.iter()
381    }
382}
383
384/// Upstream map iterator with governed memory.
385pub type MapIter<'a, K, V, M = DefaultMemoryImpl> =
386    substrate::btreemap::Iter<'a, K, V, RuntimeMemory<M>>;
387/// Upstream set iterator with governed memory.
388pub type SetIter<'a, K, M = DefaultMemoryImpl> = substrate::btreeset::Iter<'a, K, RuntimeMemory<M>>;
389/// Upstream log iterator with governed index and data memories.
390pub 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}
420/// A stable cell whose writes enforce the record's declared `Storable` bound.
421pub 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    /// Open a current cell or initialize a fresh bounded record.
427    ///
428    /// # Panics
429    /// Panics for unbounded records or encodings that violate the declared bound.
430    #[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    /// Read the cached current record without serializing or cloning it.
439    #[must_use]
440    pub fn get(&self) -> &T {
441        &self.cell.get().0
442    }
443
444    /// Validate and persist a complete replacement, returning the prior record.
445    ///
446    /// # Panics
447    /// Panics before writing if the encoding violates its declared bound.
448    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}