1use 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#[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
90impl<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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 let mut mapping: Mapping<u8, u8> = Mapping::new();
457
458 mapping.insert(1, &2);
459 assert_eq!(mapping.get(1), Some(2));
460
461 mapping.remove(1);
463
464 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 let mapping: Mapping<u8, u8> = Mapping::new();
477
478 mapping.remove(1);
480
481 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 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 assert_eq!(
548 mapping.try_insert(key, &value),
549 Err(ink_env::Error::BufferTooSmall)
550 );
551
552 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}