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 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///     storage::{
55///         traits::ManualKey,
56///         Mapping,
57///     },
58///     U256,
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 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    #[cfg(feature = "unstable-hostfn")]
205    pub fn try_get<Q>(&self, key: Q) -> Option<ink_env::Result<V>>
206    where
207        Q: scale::EncodeLike<K>,
208    {
209        let key_size = <Q as Encode>::encoded_size(&key);
210
211        if key_size > ink_env::BUFFER_SIZE {
212            return Some(Err(ink_env::Error::BufferTooSmall))
213        }
214
215        let value_size: usize =
216            ink_env::contains_contract_storage(&(&KeyType::KEY, &key))?
217                .try_into()
218                .expect("targets of less than 32bit pointer size are not supported; qed");
219
220        if key_size.saturating_add(value_size) > ink_env::BUFFER_SIZE {
221            return Some(Err(ink_env::Error::BufferTooSmall))
222        }
223
224        self.get(key).map(Ok)
225    }
226
227    /// Removes the `value` at `key`, returning the previous `value` at `key` from
228    /// storage.
229    ///
230    /// Returns `None` if no `value` exists at the given `key`.
231    ///
232    /// # Panics
233    ///
234    /// Traps if the encoded `key` or `value` doesn't fit into the static buffer.
235    ///
236    /// # Warning
237    ///
238    /// This method uses the
239    /// [unstable interface](https://github.com/paritytech/substrate/tree/master/frame/contracts#unstable-interfaces),
240    /// which is unsafe and normally is not available on production chains.
241    #[inline]
242    #[cfg(feature = "unstable-hostfn")]
243    pub fn take<Q>(&self, key: Q) -> Option<V>
244    where
245        Q: scale::EncodeLike<K>,
246    {
247        ink_env::take_contract_storage(&(&KeyType::KEY, key))
248            .unwrap_or_else(|error| panic!("Failed to take value in Mapping: {error:?}"))
249    }
250
251    /// Try to take the `value` at the given `key`.
252    /// On success, this operation will remove the value from the mapping
253    ///
254    /// Returns:
255    /// - `Some(Ok(_))` containing the value if it existed and was decoded successfully.
256    /// - `Some(Err(_))` if either (a) the encoded key doesn't fit into the static buffer
257    ///   or (b) the value existed but its length exceeds the static buffer size.
258    /// - `None` if there was no value under this mapping key.
259    ////
260    /// # Warning
261    ///
262    /// This method uses the
263    /// [unstable interface](https://github.com/paritytech/substrate/tree/master/frame/contracts#unstable-interfaces),
264    /// which is unsafe and normally is not available on production chains.
265    #[inline]
266    #[cfg(feature = "unstable-hostfn")]
267    pub fn try_take<Q>(&self, key: Q) -> Option<ink_env::Result<V>>
268    where
269        Q: scale::EncodeLike<K>,
270    {
271        let key_size = <Q as Encode>::encoded_size(&key);
272
273        if key_size > ink_env::BUFFER_SIZE {
274            return Some(Err(ink_env::Error::BufferTooSmall))
275        }
276
277        let value_size: usize =
278            ink_env::contains_contract_storage(&(&KeyType::KEY, &key))?
279                .try_into()
280                .expect("targets of less than 32bit pointer size are not supported; qed");
281
282        if key_size.saturating_add(value_size) > ink_env::BUFFER_SIZE {
283            return Some(Err(ink_env::Error::BufferTooSmall))
284        }
285
286        self.take(key).map(Ok)
287    }
288
289    /// Get the size in bytes of a value stored at `key` in the contract storage.
290    ///
291    /// Returns `None` if no `value` exists at the given `key`.
292    #[inline]
293    #[cfg(feature = "unstable-hostfn")]
294    pub fn size<Q>(&self, key: Q) -> Option<u32>
295    where
296        Q: scale::EncodeLike<K>,
297    {
298        ink_env::contains_contract_storage(&(&KeyType::KEY, key))
299    }
300
301    /// Checks if a value is stored at the given `key` in the contract storage.
302    ///
303    /// Returns `false` if no `value` exists at the given `key`.
304    #[inline]
305    #[cfg(feature = "unstable-hostfn")]
306    pub fn contains<Q>(&self, key: Q) -> bool
307    where
308        Q: scale::EncodeLike<K>,
309    {
310        ink_env::contains_contract_storage(&(&KeyType::KEY, key)).is_some()
311    }
312
313    /// Clears the value at `key` from storage.
314    #[inline]
315    #[cfg(feature = "unstable-hostfn")]
316    pub fn remove<Q>(&self, key: Q)
317    where
318        Q: scale::EncodeLike<K>,
319    {
320        ink_env::clear_contract_storage(&(&KeyType::KEY, key));
321    }
322}
323
324impl<K, V, KeyType> Storable for Mapping<K, V, KeyType>
325where
326    V: Packed,
327    KeyType: StorageKey,
328{
329    #[inline]
330    fn encode<T: Output + ?Sized>(&self, _dest: &mut T) {}
331
332    #[inline]
333    fn decode<I: Input>(_input: &mut I) -> Result<Self, Error> {
334        Ok(Default::default())
335    }
336
337    #[inline]
338    fn encoded_size(&self) -> usize {
339        0
340    }
341}
342
343impl<K, V, Key, InnerKey> StorableHint<Key> for Mapping<K, V, InnerKey>
344where
345    V: Packed,
346    Key: StorageKey,
347    InnerKey: StorageKey,
348{
349    type Type = Mapping<K, V, Key>;
350    type PreferredKey = InnerKey;
351}
352
353impl<K, V, KeyType> StorageKey for Mapping<K, V, KeyType>
354where
355    V: Packed,
356    KeyType: StorageKey,
357{
358    const KEY: Key = KeyType::KEY;
359}
360
361#[cfg(feature = "std")]
362const _: () = {
363    use crate::traits::StorageLayout;
364    use ink_metadata::layout::{
365        Layout,
366        LayoutKey,
367        RootLayout,
368    };
369
370    impl<K, V, KeyType> StorageLayout for Mapping<K, V, KeyType>
371    where
372        K: scale_info::TypeInfo + 'static,
373        V: Packed + StorageLayout + scale_info::TypeInfo + 'static,
374        KeyType: StorageKey + scale_info::TypeInfo + 'static,
375    {
376        fn layout(_: &Key) -> Layout {
377            Layout::Root(RootLayout::new(
378                LayoutKey::from(&KeyType::KEY),
379                <V as StorageLayout>::layout(&KeyType::KEY),
380                scale_info::meta_type::<Self>(),
381            ))
382        }
383    }
384};
385
386#[cfg(test)]
387mod tests {
388    use super::*;
389    use crate::traits::ManualKey;
390
391    #[test]
392    fn insert_and_get_work() {
393        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
394            let mut mapping: Mapping<u8, _> = Mapping::new();
395            mapping.insert(1, &2);
396            assert_eq!(mapping.get(1), Some(2));
397
398            Ok(())
399        })
400        .unwrap()
401    }
402
403    #[test]
404    fn insert_and_get_work_for_two_mapping_with_same_manual_key() {
405        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
406            let mut mapping: Mapping<u8, u8, ManualKey<123>> = Mapping::new();
407            mapping.insert(1, &2);
408
409            let mapping2: Mapping<u8, u8, ManualKey<123>> = Mapping::new();
410            assert_eq!(mapping2.get(1), Some(2));
411
412            Ok(())
413        })
414        .unwrap()
415    }
416
417    #[test]
418    fn gets_default_if_no_key_set() {
419        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
420            let mapping: Mapping<u8, u8> = Mapping::new();
421            assert_eq!(mapping.get(1), None);
422
423            Ok(())
424        })
425        .unwrap()
426    }
427
428    #[test]
429    fn insert_and_take_work() {
430        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
431            let mut mapping: Mapping<u8, _> = Mapping::new();
432            mapping.insert(1, &2);
433            assert_eq!(mapping.take(1), Some(2));
434            assert!(mapping.get(1).is_none());
435
436            Ok(())
437        })
438        .unwrap()
439    }
440
441    #[test]
442    fn take_empty_value_work() {
443        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
444            let mapping: Mapping<u8, u8> = Mapping::new();
445            assert_eq!(mapping.take(1), None);
446
447            Ok(())
448        })
449        .unwrap()
450    }
451
452    #[test]
453    fn can_clear_entries() {
454        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
455            // Given
456            let mut mapping: Mapping<u8, u8> = Mapping::new();
457
458            mapping.insert(1, &2);
459            assert_eq!(mapping.get(1), Some(2));
460
461            // When
462            mapping.remove(1);
463
464            // Then
465            assert_eq!(mapping.get(1), None);
466
467            Ok(())
468        })
469        .unwrap()
470    }
471
472    #[test]
473    fn can_clear_unexistent_entries() {
474        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
475            // Given
476            let mapping: Mapping<u8, u8> = Mapping::new();
477
478            // When
479            mapping.remove(1);
480
481            // Then
482            assert_eq!(mapping.get(1), None);
483
484            Ok(())
485        })
486        .unwrap()
487    }
488
489    #[test]
490    fn fallible_storage_works_for_fitting_data() {
491        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
492            let mut mapping: Mapping<u8, [u8; ink_env::BUFFER_SIZE - 1]> = Mapping::new();
493
494            let key = 0;
495            let value = [0u8; ink_env::BUFFER_SIZE - 1];
496
497            assert_eq!(mapping.try_insert(key, &value), Ok(None));
498            assert_eq!(mapping.try_get(key), Some(Ok(value)));
499            assert_eq!(mapping.try_take(key), Some(Ok(value)));
500            assert_eq!(mapping.try_get(key), None);
501
502            Ok(())
503        })
504        .unwrap()
505    }
506
507    #[test]
508    fn fallible_storage_fails_gracefully_for_overgrown_data() {
509        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
510            let mut mapping: Mapping<u8, [u8; ink_env::BUFFER_SIZE]> = Mapping::new();
511
512            let key = 0;
513            let value = [0u8; ink_env::BUFFER_SIZE];
514
515            assert_eq!(mapping.try_get(0), None);
516            assert_eq!(
517                mapping.try_insert(key, &value),
518                Err(ink_env::Error::BufferTooSmall)
519            );
520
521            // The off-chain impl conveniently uses a Vec for encoding,
522            // allowing writing values exceeding the static buffer size.
523            ink_env::set_contract_storage(&(&mapping.key(), key), &value);
524            assert_eq!(
525                mapping.try_get(key),
526                Some(Err(ink_env::Error::BufferTooSmall))
527            );
528            assert_eq!(
529                mapping.try_take(key),
530                Some(Err(ink_env::Error::BufferTooSmall))
531            );
532
533            Ok(())
534        })
535        .unwrap()
536    }
537
538    #[test]
539    fn fallible_storage_considers_key_size() {
540        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
541            let mut mapping: Mapping<[u8; ink_env::BUFFER_SIZE + 1], u8> = Mapping::new();
542
543            let key = [0u8; ink_env::BUFFER_SIZE + 1];
544            let value = 0;
545
546            // Key is already too large, so this should fail anyways.
547            assert_eq!(
548                mapping.try_insert(key, &value),
549                Err(ink_env::Error::BufferTooSmall)
550            );
551
552            // The off-chain impl conveniently uses a Vec for encoding,
553            // allowing writing values exceeding the static buffer size.
554            ink_env::set_contract_storage(&(&mapping.key(), key), &value);
555            assert_eq!(
556                mapping.try_get(key),
557                Some(Err(ink_env::Error::BufferTooSmall))
558            );
559            assert_eq!(
560                mapping.try_take(key),
561                Some(Err(ink_env::Error::BufferTooSmall))
562            );
563
564            Ok(())
565        })
566        .unwrap()
567    }
568}