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