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 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 #[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 #[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 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 #[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 #[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 #[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 let mut mapping: Mapping<u8, u8> = Mapping::new();
446
447 mapping.insert(1, &2);
448 assert_eq!(mapping.get(1), Some(2));
449
450 mapping.remove(1);
452
453 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 let mapping: Mapping<u8, u8> = Mapping::new();
466
467 mapping.remove(1);
469
470 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 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 assert_eq!(
538 mapping.try_insert(key, &value),
539 Err(ink_env::Error::BufferTooSmall)
540 );
541
542 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}