ink_storage/lazy/
mapping.rs

1// Copyright (C) Use Ink (UK) Ltd.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! A simple mapping to contract storage.
16//!
17//! # Note
18//!
19//! This mapping doesn't actually "own" any data.
20//! Instead, it is just a simple wrapper around the contract storage facilities.
21
22use crate::traits::{
23    AutoKey,
24    Packed,
25    StorableHint,
26    StorageKey,
27};
28use core::marker::PhantomData;
29use ink_primitives::Key;
30use ink_storage_traits::Storable;
31use scale::{
32    Encode,
33    Error,
34    Input,
35    Output,
36};
37
38/// A mapping of key-value pairs directly into contract storage.
39///
40/// # Important
41///
42/// The mapping requires its own pre-defined storage key where to store values. By
43/// default, the key is automatically calculated using [`AutoKey`](crate::traits::AutoKey)
44/// during compilation. However, anyone can specify a storage key using
45/// [`ManualKey`](crate::traits::ManualKey). Specifying the storage key can be helpful for
46/// upgradeable contracts or you want to be resistant to future changes of storage key
47/// calculation strategy.
48///
49/// This is an example of how you can do this:
50/// ```rust
51/// # #[ink::contract]
52/// # mod my_module {
53/// use ink::{
54///     U256,
55///     storage::{
56///         Mapping,
57///         traits::ManualKey,
58///     },
59/// };
60///
61/// #[ink(storage)]
62/// #[derive(Default)]
63/// pub struct MyContract {
64///     balances: Mapping<Address, U256, ManualKey<123>>,
65/// }
66///
67/// impl MyContract {
68///     #[ink(constructor)]
69///     pub fn new() -> Self {
70///         let mut instance = Self::default();
71///         let caller = Self::env().caller();
72///         let value: U256 = Default::default();
73///         instance.balances.insert(&caller, &value);
74///         instance
75///     }
76///
77/// #   #[ink(message)]
78/// #   pub fn my_message(&self) { }
79/// }
80/// # }
81/// ```
82///
83/// More usage examples can be found [in the ink! examples](https://github.com/use-ink/ink-examples).
84#[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
85pub struct Mapping<K, V: Packed, KeyType: StorageKey = AutoKey> {
86    #[allow(clippy::type_complexity)]
87    _marker: PhantomData<fn() -> (K, V, KeyType)>,
88}
89
90/// We implement this manually because the derived implementation adds trait bounds.
91impl<K, V, KeyType> Default for Mapping<K, V, KeyType>
92where
93    V: Packed,
94    KeyType: StorageKey,
95{
96    fn default() -> Self {
97        Self::new()
98    }
99}
100
101impl<K, V, KeyType> Mapping<K, V, KeyType>
102where
103    V: Packed,
104    KeyType: StorageKey,
105{
106    /// Creates a new empty `Mapping`.
107    pub const fn new() -> Self {
108        Self {
109            _marker: PhantomData,
110        }
111    }
112}
113
114impl<K, V, KeyType> ::core::fmt::Debug for Mapping<K, V, KeyType>
115where
116    V: Packed,
117    KeyType: StorageKey,
118{
119    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
120        f.debug_struct("Mapping")
121            .field("key", &KeyType::KEY)
122            .finish()
123    }
124}
125
126impl<K, V, KeyType> Mapping<K, V, KeyType>
127where
128    K: Encode,
129    V: Packed,
130    KeyType: StorageKey,
131{
132    /// Insert the given `value` to the contract storage.
133    ///
134    /// Returns the size in bytes of the pre-existing value at the specified key if any.
135    ///
136    /// # Panics
137    ///
138    /// Traps if encoding the `key` together with the `value` doesn't fit into the static
139    /// buffer.
140    #[inline]
141    pub fn insert<Q, R>(&mut self, key: Q, value: &R) -> Option<u32>
142    where
143        Q: scale::EncodeLike<K>,
144        R: Storable + scale::EncodeLike<V>,
145    {
146        ink_env::set_contract_storage(&(&KeyType::KEY, key), value)
147    }
148
149    /// Try to insert the given `value` into the mapping under given `key`.
150    ///
151    /// Fails if `key` or `value` exceeds the static buffer size.
152    ///
153    /// Returns:
154    /// - `Ok(Some(_))` if the value was inserted successfully, containing the size in
155    ///   bytes of the pre-existing value at the specified key if any.
156    /// - `Ok(None)` if the insert was successful but there was no pre-existing value.
157    /// - `Err(_)` if encoding the `key` together with the `value` exceeds the static
158    ///   buffer size.
159    #[inline]
160    pub fn try_insert<Q, R>(&mut self, key: Q, value: &R) -> ink_env::Result<Option<u32>>
161    where
162        Q: scale::EncodeLike<K>,
163        R: Storable + scale::EncodeLike<V>,
164    {
165        let key_size = <Q as Encode>::encoded_size(&key);
166
167        if key_size > ink_env::BUFFER_SIZE {
168            return Err(ink_env::Error::BufferTooSmall)
169        }
170
171        let value_size = <R as Storable>::encoded_size(value);
172
173        if key_size.saturating_add(value_size) > ink_env::BUFFER_SIZE {
174            return Err(ink_env::Error::BufferTooSmall)
175        }
176
177        Ok(self.insert(key, value))
178    }
179
180    /// Get the `value` at `key` from the contract storage.
181    ///
182    /// Returns `None` if no `value` exists at the given `key`.
183    ///
184    /// # Panics
185    ///
186    /// Traps if the encoded `key` or `value` doesn't fit into the static buffer.
187    #[inline]
188    pub fn get<Q>(&self, key: Q) -> Option<V>
189    where
190        Q: scale::EncodeLike<K>,
191    {
192        ink_env::get_contract_storage(&(&KeyType::KEY, key))
193            .unwrap_or_else(|error| panic!("Failed to get value in Mapping: {error:?}"))
194    }
195
196    /// Try to get the `value` at the given `key`.
197    ///
198    /// Returns:
199    /// - `Some(Ok(_))` containing the value if it existed and was decoded successfully.
200    /// - `Some(Err(_))` if either (a) the encoded key doesn't fit into the static buffer
201    ///   or (b) the value existed but its length exceeds the static buffer size.
202    /// - `None` if there was no value under this mapping key.
203    #[inline]
204    pub fn try_get<Q>(&self, key: Q) -> Option<ink_env::Result<V>>
205    where
206        Q: scale::EncodeLike<K>,
207    {
208        let key_size = <Q as Encode>::encoded_size(&key);
209
210        if key_size > ink_env::remaining_buffer() {
211            return Some(Err(ink_env::Error::BufferTooSmall))
212        }
213
214        let value_size: usize =
215            ink_env::contains_contract_storage(&(&KeyType::KEY, &key))?
216                .try_into()
217                .expect("targets of less than 32bit pointer size are not supported; qed");
218
219        if key_size.saturating_add(value_size) > ink_env::remaining_buffer() {
220            return Some(Err(ink_env::Error::BufferTooSmall))
221        }
222
223        self.get(key).map(Ok)
224    }
225
226    /// Removes the `value` at `key`, returning the previous `value` at `key` from
227    /// storage.
228    ///
229    /// Returns `None` if no `value` exists at the given `key`.
230    ///
231    /// # Panics
232    ///
233    /// Traps if the encoded `key` or `value` doesn't fit into the static buffer.
234    #[inline]
235    pub fn take<Q>(&self, key: Q) -> Option<V>
236    where
237        Q: scale::EncodeLike<K>,
238    {
239        ink_env::take_contract_storage(&(&KeyType::KEY, key))
240            .unwrap_or_else(|error| panic!("Failed to take value in Mapping: {error:?}"))
241    }
242
243    /// Try to take the `value` at the given `key`.
244    /// On success, this operation will remove the value from the mapping
245    ///
246    /// Returns:
247    /// - `Some(Ok(_))` containing the value if it existed and was decoded successfully.
248    /// - `Some(Err(_))` if either (a) the encoded key doesn't fit into the static buffer
249    ///   or (b) the value existed but its length exceeds the static buffer size.
250    /// - `None` if there was no value under this mapping key.
251    #[inline]
252    pub fn try_take<Q>(&self, key: Q) -> Option<ink_env::Result<V>>
253    where
254        Q: scale::EncodeLike<K>,
255    {
256        let key_size = <Q as Encode>::encoded_size(&key);
257
258        // todo @cmichi explain calculation, move it into `contains_contract_storage
259        let required_buffer =
260            key_size.saturating_add(4 + 32 + 32 + 64 + key_size + 32 + 32);
261        if required_buffer > ink_env::remaining_buffer() {
262            return Some(Err(ink_env::Error::BufferTooSmall))
263        }
264
265        let value_size: usize =
266            ink_env::contains_contract_storage(&(&KeyType::KEY, &key))?
267                .try_into()
268                .expect("targets of less than 32bit pointer size are not supported; qed");
269
270        let required_buffer = key_size
271            .saturating_add(4 + 32 + 32 + 64 + key_size + 32 + 32)
272            .saturating_add(value_size)
273            .saturating_add(4 + 32 + 32 + 64 + key_size + 64 + value_size);
274        if required_buffer > ink_env::remaining_buffer() {
275            return Some(Err(ink_env::Error::BufferTooSmall))
276        }
277
278        self.take(key).map(Ok)
279    }
280
281    /// Get the size in bytes of a value stored at `key` in the contract storage.
282    ///
283    /// Returns `None` if no `value` exists at the given `key`.
284    #[inline]
285    pub fn size<Q>(&self, key: Q) -> Option<u32>
286    where
287        Q: scale::EncodeLike<K>,
288    {
289        ink_env::contains_contract_storage(&(&KeyType::KEY, key))
290    }
291
292    /// Checks if a value is stored at the given `key` in the contract storage.
293    ///
294    /// Returns `false` if no `value` exists at the given `key`.
295    #[inline]
296    pub fn contains<Q>(&self, key: Q) -> bool
297    where
298        Q: scale::EncodeLike<K>,
299    {
300        ink_env::contains_contract_storage(&(&KeyType::KEY, key)).is_some()
301    }
302
303    /// Clears the value at `key` from storage.
304    #[inline]
305    pub fn remove<Q>(&self, key: Q)
306    where
307        Q: scale::EncodeLike<K>,
308    {
309        ink_env::clear_contract_storage(&(&KeyType::KEY, key));
310    }
311}
312
313impl<K, V, KeyType> Storable for Mapping<K, V, KeyType>
314where
315    V: Packed,
316    KeyType: StorageKey,
317{
318    #[inline]
319    fn encode<T: Output + ?Sized>(&self, _dest: &mut T) {}
320
321    #[inline]
322    fn decode<I: Input>(_input: &mut I) -> Result<Self, Error> {
323        Ok(Default::default())
324    }
325
326    #[inline]
327    fn encoded_size(&self) -> usize {
328        0
329    }
330}
331
332impl<K, V, Key, InnerKey> StorableHint<Key> for Mapping<K, V, InnerKey>
333where
334    V: Packed,
335    Key: StorageKey,
336    InnerKey: StorageKey,
337{
338    type Type = Mapping<K, V, Key>;
339    type PreferredKey = InnerKey;
340}
341
342impl<K, V, KeyType> StorageKey for Mapping<K, V, KeyType>
343where
344    V: Packed,
345    KeyType: StorageKey,
346{
347    const KEY: Key = KeyType::KEY;
348}
349
350#[cfg(feature = "std")]
351const _: () = {
352    use crate::traits::StorageLayout;
353    use ink_metadata::layout::{
354        Layout,
355        LayoutKey,
356        RootLayout,
357    };
358
359    impl<K, V, KeyType> StorageLayout for Mapping<K, V, KeyType>
360    where
361        K: scale_info::TypeInfo + 'static,
362        V: Packed + StorageLayout + scale_info::TypeInfo + 'static,
363        KeyType: StorageKey + scale_info::TypeInfo + 'static,
364    {
365        fn layout(_: &Key) -> Layout {
366            Layout::Root(RootLayout::new(
367                LayoutKey::from(&KeyType::KEY),
368                <V as StorageLayout>::layout(&KeyType::KEY),
369                scale_info::meta_type::<Self>(),
370            ))
371        }
372    }
373};
374
375#[cfg(test)]
376mod tests {
377    use super::*;
378    use crate::traits::ManualKey;
379
380    #[test]
381    fn insert_and_get_work() {
382        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
383            let mut mapping: Mapping<u8, _> = Mapping::new();
384            mapping.insert(1, &2);
385            assert_eq!(mapping.get(1), Some(2));
386
387            Ok(())
388        })
389        .unwrap()
390    }
391
392    #[test]
393    fn insert_and_get_work_for_two_mapping_with_same_manual_key() {
394        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
395            let mut mapping: Mapping<u8, u8, ManualKey<123>> = Mapping::new();
396            mapping.insert(1, &2);
397
398            let mapping2: Mapping<u8, u8, ManualKey<123>> = Mapping::new();
399            assert_eq!(mapping2.get(1), Some(2));
400
401            Ok(())
402        })
403        .unwrap()
404    }
405
406    #[test]
407    fn gets_default_if_no_key_set() {
408        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
409            let mapping: Mapping<u8, u8> = Mapping::new();
410            assert_eq!(mapping.get(1), None);
411
412            Ok(())
413        })
414        .unwrap()
415    }
416
417    #[test]
418    fn insert_and_take_work() {
419        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
420            let mut mapping: Mapping<u8, _> = Mapping::new();
421            mapping.insert(1, &2);
422            assert_eq!(mapping.take(1), Some(2));
423            assert!(mapping.get(1).is_none());
424
425            Ok(())
426        })
427        .unwrap()
428    }
429
430    #[test]
431    fn take_empty_value_work() {
432        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
433            let mapping: Mapping<u8, u8> = Mapping::new();
434            assert_eq!(mapping.take(1), None);
435
436            Ok(())
437        })
438        .unwrap()
439    }
440
441    #[test]
442    fn can_clear_entries() {
443        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
444            // Given
445            let mut mapping: Mapping<u8, u8> = Mapping::new();
446
447            mapping.insert(1, &2);
448            assert_eq!(mapping.get(1), Some(2));
449
450            // When
451            mapping.remove(1);
452
453            // Then
454            assert_eq!(mapping.get(1), None);
455
456            Ok(())
457        })
458        .unwrap()
459    }
460
461    #[test]
462    fn can_clear_unexistent_entries() {
463        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
464            // Given
465            let mapping: Mapping<u8, u8> = Mapping::new();
466
467            // When
468            mapping.remove(1);
469
470            // Then
471            assert_eq!(mapping.get(1), None);
472
473            Ok(())
474        })
475        .unwrap()
476    }
477
478    #[test]
479    #[ignore]
480    fn fallible_storage_works_for_fitting_data() {
481        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
482            let mut mapping: Mapping<u8, [u8; ink_env::BUFFER_SIZE - 1]> = Mapping::new();
483
484            let key = 0;
485            let value = [0u8; ink_env::BUFFER_SIZE - 1];
486
487            assert_eq!(mapping.try_insert(key, &value), Ok(None));
488            assert_eq!(mapping.try_get(key), Some(Ok(value)));
489            assert_eq!(mapping.try_take(key), Some(Ok(value)));
490            assert_eq!(mapping.try_get(key), None);
491
492            Ok(())
493        })
494        .unwrap()
495    }
496
497    #[test]
498    fn fallible_storage_fails_gracefully_for_overgrown_data() {
499        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
500            let mut mapping: Mapping<u8, [u8; ink_env::BUFFER_SIZE]> = Mapping::new();
501
502            let key = 0;
503            let value = [0u8; ink_env::BUFFER_SIZE];
504
505            assert_eq!(mapping.try_get(0), None);
506            assert_eq!(
507                mapping.try_insert(key, &value),
508                Err(ink_env::Error::BufferTooSmall)
509            );
510
511            // The off-chain impl conveniently uses a Vec for encoding,
512            // allowing writing values exceeding the static buffer size.
513            ink_env::set_contract_storage(&(&mapping.key(), key), &value);
514            assert_eq!(
515                mapping.try_get(key),
516                Some(Err(ink_env::Error::BufferTooSmall))
517            );
518            assert_eq!(
519                mapping.try_take(key),
520                Some(Err(ink_env::Error::BufferTooSmall))
521            );
522
523            Ok(())
524        })
525        .unwrap()
526    }
527
528    #[test]
529    fn fallible_storage_considers_key_size() {
530        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
531            let mut mapping: Mapping<[u8; ink_env::BUFFER_SIZE + 1], u8> = Mapping::new();
532
533            let key = [0u8; ink_env::BUFFER_SIZE + 1];
534            let value = 0;
535
536            // Key is already too large, so this should fail anyways.
537            assert_eq!(
538                mapping.try_insert(key, &value),
539                Err(ink_env::Error::BufferTooSmall)
540            );
541
542            // The off-chain impl conveniently uses a Vec for encoding,
543            // allowing writing values exceeding the static buffer size.
544            ink_env::set_contract_storage(&(&mapping.key(), key), &value);
545            assert_eq!(
546                mapping.try_get(key),
547                Some(Err(ink_env::Error::BufferTooSmall))
548            );
549            assert_eq!(
550                mapping.try_take(key),
551                Some(Err(ink_env::Error::BufferTooSmall))
552            );
553
554            Ok(())
555        })
556        .unwrap()
557    }
558}