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