Skip to main content

multiversx_sc/storage/mappers/
unique_id_mapper.rs

1use crate::{
2    abi::TypeAbiFrom,
3    codec::{
4        EncodeErrorHandler, TopEncodeMulti, TopEncodeMultiOutput, multi_encode_iter_or_handle_err,
5    },
6    types::ManagedAddress,
7};
8
9use super::{
10    StorageMapper, StorageMapperFromAddress, VecMapper,
11    source::{CurrentStorage, StorageAddress},
12};
13use crate::{
14    abi::{TypeAbi, TypeDescriptionContainer, TypeName},
15    api::{ErrorApiImpl, StorageMapperApi},
16    storage::StorageKey,
17    storage_set,
18    types::{ManagedType, MultiValueEncoded},
19};
20
21pub type UniqueId = usize;
22const EMPTY_ENTRY: UniqueId = 0;
23
24/// A space-optimized storage mapper for managing ID permutations with minimal storage overhead.
25///
26/// # Storage Layout
27///
28/// The `UniqueIdMapper` is optimized for scenarios where most IDs map to themselves (identity mapping).
29/// It uses a `VecMapper` internally but only stores non-identity mappings:
30///
31/// - `base_key + ".len"` → number of slots in the mapper
32/// - `base_key + ".item" + index` → stored value only if `value != index` (stores 0 otherwise)
33///
34/// **Key Optimization**: When `Mapper[i] == i` (identity mapping), it stores `0` instead of `i`,
35/// significantly reducing storage costs when most IDs remain in their original positions.
36///
37/// **Important**: Indexes are 1-based (range: `1..=len()`), consistent with `VecMapper` conventions.
38///
39/// # Main Operations
40///
41/// - **Initialize**: `set_initial_len(n)` - Sets the mapper size (can only be called once).
42/// - **Read**: `get(index)` - Returns the ID at the given index. O(1) with one storage read.
43/// - **Write**: `set(index, id)` - Assigns an ID to an index. O(1) with one storage write.
44/// - **Remove**: `swap_remove(index)` - Removes an index by swapping with the last element. O(1).
45/// - **Iteration**: `iter()` - Iterates over all ID values in index order.
46///
47/// # Trade-offs
48///
49/// - **Pros**: Extremely space-efficient for permutation tracking; zero storage cost for identity mappings.
50/// - **Cons**: Limited to ID permutations (1 to N); length must be set upfront and cannot be extended.
51///
52/// # Use Cases
53///
54/// - NFT position tracking where most NFTs remain in their original slots
55/// - Shuffle/permutation algorithms where most elements stay in place
56/// - ID reassignment systems with minimal changes
57///
58/// # Example
59///
60/// ```rust
61/// # use multiversx_sc::storage::mappers::{StorageMapper, UniqueIdMapper};
62/// # fn example<SA: multiversx_sc::api::StorageMapperApi>() {
63/// # let mut mapper = UniqueIdMapper::<SA>::new(
64/// #     multiversx_sc::storage::StorageKey::new(&b"id_map"[..])
65/// # );
66/// // Initialize with 5 slots (IDs 1-5)
67/// mapper.set_initial_len(5);
68///
69/// // Initially, all IDs map to themselves: [1, 2, 3, 4, 5]
70/// assert_eq!(mapper.get(1), 1);
71/// assert_eq!(mapper.get(3), 3);
72///
73/// // Swap positions 2 and 4
74/// mapper.set(2, 4);
75/// mapper.set(4, 2);
76///
77/// // Now the mapping is: [1, 4, 3, 2, 5]
78/// assert_eq!(mapper.get(2), 4);
79/// assert_eq!(mapper.get(4), 2);
80///
81/// // Remove index 3 (swaps with last)
82/// let removed = mapper.swap_remove(3);
83/// assert_eq!(removed, 3);
84/// assert_eq!(mapper.len(), 4);
85/// # }
86/// ```
87pub struct UniqueIdMapper<SA, A = CurrentStorage>
88where
89    SA: StorageMapperApi,
90    A: StorageAddress<SA>,
91{
92    _address: A,
93    base_key: StorageKey<SA>,
94    vec_mapper: VecMapper<SA, UniqueId, A>,
95}
96
97impl<SA> StorageMapper<SA> for UniqueIdMapper<SA, CurrentStorage>
98where
99    SA: StorageMapperApi,
100{
101    fn new(base_key: StorageKey<SA>) -> Self {
102        Self {
103            _address: CurrentStorage,
104            base_key: base_key.clone(),
105            vec_mapper: VecMapper::new(base_key),
106        }
107    }
108}
109
110impl<SA> StorageMapperFromAddress<SA> for UniqueIdMapper<SA, ManagedAddress<SA>>
111where
112    SA: StorageMapperApi,
113{
114    fn new_from_address(address: ManagedAddress<SA>, base_key: StorageKey<SA>) -> Self {
115        Self {
116            _address: address.clone(),
117            base_key: base_key.clone(),
118            vec_mapper: VecMapper::new_from_address(address, base_key),
119        }
120    }
121}
122
123impl<SA, A> UniqueIdMapper<SA, A>
124where
125    SA: StorageMapperApi,
126    A: StorageAddress<SA>,
127{
128    #[inline]
129    pub fn len(&self) -> usize {
130        self.vec_mapper.len()
131    }
132
133    #[inline]
134    pub fn is_empty(&self) -> bool {
135        self.vec_mapper.is_empty()
136    }
137
138    /// Gets the value for the given `index`. If the entry is empty, `index` is returned.
139    pub fn get(&self, index: usize) -> UniqueId {
140        let id: UniqueId = self.vec_mapper.get(index);
141        if id == EMPTY_ENTRY { index } else { id }
142    }
143
144    /// Provides a forward iterator.
145    pub fn iter(&self) -> Iter<'_, SA, A> {
146        Iter::new(self)
147    }
148}
149
150impl<SA> UniqueIdMapper<SA, CurrentStorage>
151where
152    SA: StorageMapperApi,
153{
154    /// Initializes the mapper's length. This may not be set again afterwards.
155    pub fn set_initial_len(&mut self, len: usize) {
156        if !self.vec_mapper.is_empty() {
157            SA::error_api_impl().signal_error(b"len already set");
158        }
159
160        self.set_internal_mapper_len(len);
161    }
162
163    /// Gets the value from the index and removes it.
164    /// The value is replaced by the last item, and length is decremented.
165    pub fn swap_remove(&mut self, index: usize) -> UniqueId {
166        let last_item_index = self.len();
167        let last_item = self.get(last_item_index);
168
169        let current_item = if index != last_item_index {
170            let item_at_index = self.get(index);
171            self.set(index, last_item);
172
173            item_at_index
174        } else {
175            last_item
176        };
177
178        self.vec_mapper.set(last_item_index, &EMPTY_ENTRY);
179        self.set_internal_mapper_len(last_item_index - 1);
180
181        current_item
182    }
183
184    /// Sets the value at the given index. If index == id, then the entry is cleared.
185    pub fn set(&mut self, index: usize, id: UniqueId) {
186        if index == id {
187            self.vec_mapper.set(index, &EMPTY_ENTRY);
188        } else {
189            self.vec_mapper.set(index, &id);
190        }
191    }
192
193    // Manually sets the internal VecMapper's len value
194    fn set_internal_mapper_len(&mut self, new_len: usize) {
195        let mut len_key = self.base_key.clone();
196        len_key.append_bytes(&b".len"[..]);
197        storage_set(len_key.as_ref(), &new_len);
198    }
199}
200
201impl<'a, SA, A> IntoIterator for &'a UniqueIdMapper<SA, A>
202where
203    SA: StorageMapperApi,
204    A: StorageAddress<SA>,
205{
206    type Item = usize;
207
208    type IntoIter = Iter<'a, SA, A>;
209
210    fn into_iter(self) -> Self::IntoIter {
211        self.iter()
212    }
213}
214
215pub struct Iter<'a, SA, A>
216where
217    SA: StorageMapperApi,
218    A: StorageAddress<SA>,
219{
220    index: usize,
221    len: usize,
222    id_mapper: &'a UniqueIdMapper<SA, A>,
223}
224
225impl<'a, SA, A> Iter<'a, SA, A>
226where
227    SA: StorageMapperApi,
228    A: StorageAddress<SA>,
229{
230    fn new(id_mapper: &'a UniqueIdMapper<SA, A>) -> Iter<'a, SA, A> {
231        Iter {
232            index: 1,
233            len: id_mapper.len(),
234            id_mapper,
235        }
236    }
237}
238
239impl<SA, A> Iterator for Iter<'_, SA, A>
240where
241    SA: StorageMapperApi,
242    A: StorageAddress<SA>,
243{
244    type Item = usize;
245
246    #[inline]
247    fn next(&mut self) -> Option<Self::Item> {
248        let current_index = self.index;
249        if current_index > self.len {
250            return None;
251        }
252
253        self.index += 1;
254        Some(self.id_mapper.get(current_index))
255    }
256}
257
258/// Behaves like a MultiResultVec when an endpoint result.
259impl<SA> TopEncodeMulti for UniqueIdMapper<SA, CurrentStorage>
260where
261    SA: StorageMapperApi,
262{
263    fn multi_encode_or_handle_err<O, H>(&self, output: &mut O, h: H) -> Result<(), H::HandledErr>
264    where
265        O: TopEncodeMultiOutput,
266        H: EncodeErrorHandler,
267    {
268        multi_encode_iter_or_handle_err(self.iter(), output, h)
269    }
270}
271
272impl<SA> TypeAbiFrom<UniqueIdMapper<SA, CurrentStorage>> for MultiValueEncoded<SA, usize> where
273    SA: StorageMapperApi
274{
275}
276
277impl<SA> TypeAbiFrom<Self> for UniqueIdMapper<SA, CurrentStorage> where SA: StorageMapperApi {}
278
279/// Behaves like a MultiResultVec when an endpoint result.
280impl<SA> TypeAbi for UniqueIdMapper<SA, CurrentStorage>
281where
282    SA: StorageMapperApi,
283{
284    type Unmanaged = Self;
285
286    fn type_name() -> TypeName {
287        crate::abi::type_name_variadic::<usize>()
288    }
289
290    fn type_name_rust() -> TypeName {
291        crate::abi::type_name_multi_value_encoded::<usize>()
292    }
293
294    fn provide_type_descriptions<TDC: TypeDescriptionContainer>(accumulator: &mut TDC) {
295        usize::provide_type_descriptions(accumulator);
296    }
297
298    fn is_variadic() -> bool {
299        true
300    }
301}