tempo_precompiles/storage/types/
mapping.rs

1//! Type-safe wrapper for EVM storage mappings (hash-based key-value storage).
2
3use alloy::primitives::{Address, U256};
4use std::{
5    hash::Hash,
6    ops::{Index, IndexMut},
7};
8
9use crate::storage::{Layout, LayoutCtx, StorableType, StorageKey, types::HandlerCache};
10
11/// Type-safe access wrapper for EVM storage mappings (hash-based key-value storage).
12///
13/// This struct does not store data itself. Instead, it provides a zero-cost abstraction
14/// for accessing mapping storage slots using Solidity's hash-based layout. It wraps a
15/// base slot number and provides methods to compute the actual storage slots for keys.
16///
17/// # Type Parameters
18///
19/// - `K`: Key type (must implement `StorageKey`)
20/// - `V`: Value type (must implement `StorableType`)
21///
22/// `Mapping<K, V>` is essentially a slot computation helper. The `[key]` method
23/// performs the keccak256 hash to compute the actual storage slot and returns a
24/// `Handler` that can be used for read/write operations.
25///
26/// # Storage Layout
27///
28/// Mappings use Solidity's storage layout:
29/// - Base slot: stored in `base_slot` field (never accessed directly)
30/// - Actual slot for key `k`: `keccak256(k || base_slot)`
31///
32/// # Usage Pattern
33///
34/// The typical usage follows a composable pattern:
35/// 1. Create a `Mapping<K, V>` with a base slot (usually from generated constants)
36/// 2. Call `[key]` to compute and obtain a `Handler` for that key
37/// 3. Use `.read()`, `.write()`, or `.delete()` on the resulting slot
38///
39/// # Accessing Mapping Fields Within Structs
40///
41/// When a mapping is a field within a struct stored in another mapping, use the static
42/// `at_offset` method to compute the slot without creating a `Mapping` instance:
43#[derive(Debug, Clone)]
44pub struct Mapping<K, V: StorableType> {
45    base_slot: U256,
46    address: Address,
47    cache: HandlerCache<K, V::Handler>,
48}
49
50impl<K, V: StorableType> Mapping<K, V> {
51    /// Creates a new `Mapping` with the given base slot number and address.
52    ///
53    /// This is typically called with slot constants generated by the `#[contract]` macro.
54    #[inline]
55    pub fn new(base_slot: U256, address: Address) -> Self {
56        Self {
57            base_slot,
58            address,
59            cache: HandlerCache::new(),
60        }
61    }
62
63    /// Returns the U256 base storage slot number for this mapping.
64    #[inline]
65    pub const fn slot(&self) -> U256 {
66        self.base_slot
67    }
68
69    /// Returns a `Handler` for the given key.
70    ///
71    /// This enables the composable pattern: `mapping.at(&key).read()`
72    /// where the mapping slot computation happens once, and the resulting slot
73    /// can be used for multiple operations.
74    ///
75    /// The handler is computed on first access and cached for subsequent accesses.
76    /// Takes a reference to avoid cloning on cache hits.
77    pub fn at(&self, key: &K) -> &V::Handler
78    where
79        K: StorageKey + Hash + Eq + Clone,
80    {
81        let (base_slot, address) = (self.base_slot, self.address);
82        self.cache.get_or_insert(key, || {
83            V::handle(key.mapping_slot(base_slot), LayoutCtx::FULL, address)
84        })
85    }
86
87    /// Returns a mutable `Handler` for the given key.
88    ///
89    /// Use this when you need to call mutable methods like `write()` or `delete()`.
90    ///
91    /// The handler is computed on first access and cached for subsequent accesses.
92    /// Takes a reference to avoid cloning on cache hits.
93    pub fn at_mut(&mut self, key: &K) -> &mut V::Handler
94    where
95        K: StorageKey + Hash + Eq + Clone,
96    {
97        let (base_slot, address) = (self.base_slot, self.address);
98        self.cache.get_or_insert_mut(key, || {
99            V::handle(key.mapping_slot(base_slot), LayoutCtx::FULL, address)
100        })
101    }
102}
103
104impl<K, V: StorableType> Default for Mapping<K, V> {
105    fn default() -> Self {
106        Self::new(U256::ZERO, Address::ZERO)
107    }
108}
109
110impl<K, V: StorableType> Index<K> for Mapping<K, V>
111where
112    K: StorageKey + Hash + Eq + Clone,
113{
114    type Output = V::Handler;
115
116    /// Returns a reference to the cached handler for the given key.
117    ///
118    /// The handler is computed on first access and cached for subsequent accesses.
119    fn index(&self, key: K) -> &Self::Output {
120        let (base_slot, address) = (self.base_slot, self.address);
121        self.cache.get_or_insert(&key, || {
122            V::handle(key.mapping_slot(base_slot), LayoutCtx::FULL, address)
123        })
124    }
125}
126
127impl<K, V: StorableType> IndexMut<K> for Mapping<K, V>
128where
129    K: StorageKey + Hash + Eq + Clone,
130{
131    /// Returns a mutable reference to the cached handler for the given key.
132    fn index_mut(&mut self, key: K) -> &mut Self::Output {
133        let (base_slot, address) = (self.base_slot, self.address);
134        self.cache.get_or_insert_mut(&key, || {
135            V::handle(key.mapping_slot(base_slot), LayoutCtx::FULL, address)
136        })
137    }
138}
139
140// Mappings occupy a full 32-byte slot in the layout (used as a base for hashing),
141// even though they don't store data in that slot directly.
142//
143// **NOTE:** Necessary to allow it to participate in struct layout calculations.
144impl<K, V> StorableType for Mapping<K, V>
145where
146    V: StorableType,
147{
148    const LAYOUT: Layout = Layout::Slots(1);
149    type Handler = Self;
150
151    fn handle(slot: U256, _ctx: LayoutCtx, address: Address) -> Self::Handler {
152        Self::new(slot, address)
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159    use crate::storage::StorageKey;
160    use alloy::primitives::{Address, B256, keccak256};
161
162    // Backward compatibility helper to verify the trait impl.
163    fn old_mapping_slot<K: AsRef<[u8]>>(key: K, slot: U256) -> U256 {
164        let key = key.as_ref();
165        let mut buf = [0u8; 64];
166        buf[32 - key.len()..32].copy_from_slice(key);
167        buf[32..].copy_from_slice(&slot.to_be_bytes::<32>());
168        U256::from_be_bytes(keccak256(buf).0)
169    }
170
171    #[test]
172    fn test_mapping_slot_encoding() {
173        let key = Address::random();
174        let base_slot = U256::random();
175
176        // Manual computation to validate
177        let mut buf = [0u8; 64];
178        // Left-pad the address to 32 bytes
179        buf[12..32].copy_from_slice(key.as_ref());
180        // Slot in big-endian
181        buf[32..].copy_from_slice(&base_slot.to_be_bytes::<32>());
182
183        let expected = U256::from_be_bytes(keccak256(buf).0);
184        let computed = key.mapping_slot(base_slot);
185
186        assert_eq!(computed, expected, "mapping_slot encoding mismatch");
187    }
188
189    #[test]
190    fn test_mapping_slot_matches_old_impl() {
191        let slot = U256::random();
192
193        let addr = Address::random();
194        assert_eq!(
195            addr.mapping_slot(slot),
196            old_mapping_slot(addr.as_slice(), slot),
197        );
198
199        let b256 = B256::random();
200        assert_eq!(
201            b256.mapping_slot(slot),
202            old_mapping_slot(b256.as_slice(), slot),
203        );
204
205        let u256 = U256::random();
206        assert_eq!(
207            u256.mapping_slot(slot),
208            old_mapping_slot(u256.to_be_bytes::<32>(), slot),
209        );
210    }
211
212    #[test]
213    fn test_mapping_basic_properties() {
214        let address = Address::random();
215        let base_slot = U256::random();
216        let mapping = Mapping::<Address, U256>::new(base_slot, address);
217
218        // Property 1: Determinism - same key always produces same slot
219        let key = Address::random();
220        let slot1 = &mapping[key];
221        let slot2 = &mapping[key];
222        assert_eq!(
223            slot1.slot(),
224            slot2.slot(),
225            "same key should produce same slot"
226        );
227
228        // Property 2: Different keys produce different slots
229        let key1 = Address::random();
230        let key2 = Address::random();
231        let slot_a = &mapping[key1];
232        let slot_b = &mapping[key2];
233        assert_ne!(
234            slot_a.slot(),
235            slot_b.slot(),
236            "different keys should produce different slots"
237        );
238
239        // Property 3: Derived slot matches manual computation
240        let test_key = Address::random();
241        let derived_slot = &mapping[test_key];
242        let expected_slot = test_key.mapping_slot(base_slot);
243        assert_eq!(derived_slot.slot(), expected_slot);
244    }
245
246    #[test]
247    fn test_nested_mapping_basic_properties() {
248        let address = Address::random();
249        let base_slot = U256::random();
250        // Nested mappings use recursive Mapping<K, Mapping<K2, V>> type
251        let nested = Mapping::<Address, Mapping<B256, U256>>::new(base_slot, address);
252
253        let key1 = Address::random();
254        let key2 = B256::random();
255
256        // Property 1: Chaining - first .at() returns intermediate Mapping with correct slot
257        let intermediate = &nested[key1];
258        let expected_intermediate_slot = key1.mapping_slot(base_slot);
259        assert_eq!(
260            intermediate.slot(),
261            expected_intermediate_slot,
262            "intermediate mapping should have correct slot"
263        );
264
265        // Property 2: Double-hash - second .at() returns final Slot with correct double-derived slot
266        let final_slot = &intermediate[key2];
267        let expected_final_slot = key2.mapping_slot(expected_intermediate_slot);
268        assert_eq!(
269            final_slot.slot(),
270            expected_final_slot,
271            "final slot should use double-hash"
272        );
273
274        // Property 3: Determinism - same keys always produce same slot
275        let slot_a = &nested[key1][key2];
276        let slot_b = &nested[key1][key2];
277        assert_eq!(
278            slot_a.slot(),
279            slot_b.slot(),
280            "same keys should produce same slot"
281        );
282
283        // Property 4: Different first-level keys produce different final slots
284        let different_key1 = Address::random();
285        let different_slot = &nested[different_key1][key2];
286        assert_ne!(
287            final_slot.slot(),
288            different_slot.slot(),
289            "different first-level keys should produce different slots"
290        );
291
292        // Property 5: Different second-level keys produce different final slots
293        let different_key2 = B256::random();
294        let another_slot = &nested[key1][different_key2];
295        assert_ne!(
296            final_slot.slot(),
297            another_slot.slot(),
298            "different second-level keys should produce different slots"
299        );
300    }
301
302    #[test]
303    fn test_mapping_slot_boundaries() {
304        let address = Address::random();
305
306        // Test .slot() getter with ZERO boundary
307        let zero_mapping = Mapping::<Address, U256>::new(U256::ZERO, address);
308        assert_eq!(zero_mapping.slot(), U256::ZERO);
309        let user = Address::random();
310        let slot = &zero_mapping[user];
311        assert_eq!(slot.slot(), user.mapping_slot(U256::ZERO));
312
313        // Test .slot() getter with MAX boundary
314        let max_mapping = Mapping::<Address, U256>::new(U256::MAX, address);
315        assert_eq!(max_mapping.slot(), U256::MAX);
316        let user2 = Address::random();
317        let slot2 = &max_mapping[user2];
318        assert_eq!(slot2.slot(), user2.mapping_slot(U256::MAX));
319
320        // Test .slot() getter with arbitrary values
321        let random_slot = U256::random();
322        let arbitrary_mapping = Mapping::<Address, U256>::new(random_slot, address);
323        assert_eq!(arbitrary_mapping.slot(), random_slot);
324    }
325}
tempo_precompiles/storage/types/mapping.rs

tempo_precompiles/storage/types/
mapping.rs
