rustsol/types/
mapping.rs

1use std::fmt::Debug;
2use std::marker::PhantomData;
3use std::sync::Arc;
4use alloy_primitives::U256;
5use derivative::Derivative;
6use crate::utils::{bytes32_to_u256, ceil_div, index_to_position, keccak256_concat, u256_to_bytes32, u256_to_u64};
7use crate::types::{Value, DynamicArray};
8use crate::types::{Position, SlotsGetter, SlotsGetterSetter};
9use crate::types::keys::Key;
10
11// The value corresponding to a mapping key k is located at keccak256(h(k) . p) where . is
12// concatenation and h is a function that is applied to the key depending on its type:
13// - for value types, h pads the value to 32 bytes in the same way as when storing the value in memory.
14// - for strings and byte arrays, h computes the keccak256 hash of the unpadded data.
15#[derive(Derivative)]
16#[derivative(Debug)]
17pub struct Mapping<KeyNativeType, ElementType> {
18    __slot: U256,
19    __marker: PhantomData<(KeyNativeType, ElementType)>,
20    #[derivative(Debug = "ignore")]
21    __slots_getter: Option<Arc<dyn SlotsGetter>>,
22}
23
24impl<KeyNativeType: Key, ElementType> Mapping<KeyNativeType, ElementType> {
25    pub fn slot(&self) -> U256 {
26        self.__slot
27    }
28
29    pub fn position(&self) -> (U256, usize, usize) {
30        (self.__slot, 0, 32)
31    }
32
33    fn new_element(&self, slot: U256, offset: usize) -> ElementType
34        where ElementType: Position + SlotsGetterSetter,
35    {
36        let mut element = ElementType::from_position(slot, offset);
37        match &self.__slots_getter {
38            None => {
39                // No slots getter to pass to children.
40            }
41            Some(getter) => {
42                // Set child's slots getter.
43                element.set_slots_getter(getter.clone());
44            }
45        }
46        element
47    }
48    fn storage_at_bytes(&self, key: [u8; 32]) -> U256 {
49        let value_slot_bytes = keccak256_concat(key, u256_to_bytes32(self.__slot));
50        bytes32_to_u256(value_slot_bytes)
51    }
52
53    pub fn get_value_at(&self, key: KeyNativeType) -> Result<<ElementType as Value>::ValueType, String>
54        where
55            KeyNativeType: Key,
56            ElementType: Position + Value + SlotsGetterSetter,
57    {
58        let getter = self.__slots_getter.as_ref().expect("No slots getter");
59        let element_slot = self.storage_at_bytes(key.to_bytes());
60        let element_size_slots = ceil_div(ElementType::size(), 32);
61        let element_slot_values = getter.get_slots(element_slot, element_size_slots)
62            .map_err(|err| format!("Failed to get slot values: {}", err))?;
63        self.new_element(element_slot, 0).get_value_from_slots_content(element_slot_values)
64    }
65}
66
67impl<KeyNativeType, ElementType> Position for Mapping<KeyNativeType, ElementType> {
68    fn from_position(slot: U256, _: usize) -> Self {
69        Mapping::<KeyNativeType, ElementType> { __slot: slot, __marker: PhantomData, __slots_getter: None }
70    }
71
72    fn size() -> usize {
73        32
74    }
75}
76
77
78impl<KeyNativeType, ElementType> Mapping<KeyNativeType, ElementType> {
79    pub fn at(&self, key: KeyNativeType) -> ElementType
80        where
81            KeyNativeType: Key,
82            ElementType: Position + SlotsGetterSetter,
83    {
84        let element_slot = self.storage_at_bytes(key.to_bytes());
85        self.new_element(element_slot, 0)
86    }
87}
88
89impl<KeyNativeType: Debug, ElementType: Debug> SlotsGetterSetter for Mapping<KeyNativeType, ElementType> {
90    fn set_slots_getter(&mut self, getter: Arc<dyn SlotsGetter>) {
91        self.__slots_getter = Some(getter);
92    }
93}
94
95impl<KeyNativeType, ElementType> Value for Mapping<KeyNativeType, ElementType> {
96    type ValueType = Self;
97
98    fn get_value_from_slots_content(&self, _: Vec<U256>) -> Result<Self::ValueType, String> {
99        // Mapping base slot is always empty and is used only for element storage calculation.
100        // Clone self.
101        Ok(Mapping {
102            __slot: self.__slot,
103            __marker: self.__marker,
104            __slots_getter: self.__slots_getter.clone(),
105        })
106    }
107}