multiversx_sc/storage/mappers/
user_mapper.rs

1use core::marker::PhantomData;
2
3use crate::{
4    abi::TypeAbiFrom,
5    codec::{
6        multi_encode_iter_or_handle_err, EncodeErrorHandler, TopEncodeMulti, TopEncodeMultiOutput,
7    },
8};
9
10use super::{
11    source::{CurrentStorage, StorageAddress},
12    StorageMapper, StorageMapperFromAddress,
13};
14use crate::{
15    abi::{TypeAbi, TypeName},
16    api::StorageMapperApi,
17    storage::{storage_set, StorageKey},
18    types::{ManagedAddress, ManagedType, ManagedVec, MultiValueEncoded},
19};
20
21const ADDRESS_TO_ID_SUFFIX: &[u8] = b"_address_to_id";
22const ID_TO_ADDRESS_SUFFIX: &[u8] = b"_id_to_address";
23const COUNT_SUFFIX: &[u8] = b"_count";
24
25/// Very widely used mapper, that manages the users of a smart contract.
26/// It holds a bi-directional map, from addresses to ids and viceversa.
27/// This is so we can easily iterate over all users, using their ids.
28/// Also holds the user count in sync. This is also necessary for iteration.
29///
30/// This particular implementation of a user mapper doesn't contain any additional
31/// user data other than address/id.
32///
33/// It also doesn't allow removing users. Once in, their ids are reserved forever.
34pub struct UserMapper<SA, A = CurrentStorage>
35where
36    SA: StorageMapperApi,
37    A: StorageAddress<SA>,
38{
39    _phantom_api: PhantomData<SA>,
40    address: A,
41    base_key: StorageKey<SA>,
42}
43
44impl<SA> StorageMapper<SA> for UserMapper<SA>
45where
46    SA: StorageMapperApi,
47{
48    fn new(base_key: StorageKey<SA>) -> Self {
49        UserMapper {
50            _phantom_api: PhantomData,
51            address: CurrentStorage,
52            base_key,
53        }
54    }
55}
56
57impl<SA> StorageMapperFromAddress<SA> for UserMapper<SA, ManagedAddress<SA>>
58where
59    SA: StorageMapperApi,
60{
61    fn new_from_address(address: ManagedAddress<SA>, base_key: StorageKey<SA>) -> Self {
62        UserMapper {
63            _phantom_api: PhantomData,
64            address,
65            base_key,
66        }
67    }
68}
69
70impl<SA, A> UserMapper<SA, A>
71where
72    SA: StorageMapperApi,
73    A: StorageAddress<SA>,
74{
75    fn get_user_id_key(&self, address: &ManagedAddress<SA>) -> StorageKey<SA> {
76        let mut user_id_key = self.base_key.clone();
77        user_id_key.append_bytes(ADDRESS_TO_ID_SUFFIX);
78        user_id_key.append_item(address);
79        user_id_key
80    }
81
82    fn get_user_address_key(&self, id: usize) -> StorageKey<SA> {
83        let mut user_address_key = self.base_key.clone();
84        user_address_key.append_bytes(ID_TO_ADDRESS_SUFFIX);
85        user_address_key.append_item(&id);
86        user_address_key
87    }
88
89    fn get_user_count_key(&self) -> StorageKey<SA> {
90        let mut user_count_key = self.base_key.clone();
91        user_count_key.append_bytes(COUNT_SUFFIX);
92        user_count_key
93    }
94
95    /// Yields the user id for a given address.
96    /// Will return 0 if the address is not known to the contract.
97    pub fn get_user_id(&self, address: &ManagedAddress<SA>) -> usize {
98        self.address
99            .address_storage_get(self.get_user_id_key(address).as_ref())
100    }
101
102    /// Yields the user address for a given id, if the id is valid.
103    pub fn get_user_address(&self, id: usize) -> Option<ManagedAddress<SA>> {
104        let key = self.get_user_address_key(id);
105        // TODO: optimize, storage_load_managed_buffer_len is currently called twice
106
107        if self.address.address_storage_get_len(key.as_ref()) > 0 {
108            Some(self.address.address_storage_get(key.as_ref()))
109        } else {
110            None
111        }
112    }
113
114    /// Yields the user address for a given id.
115    /// Will cause a deserialization error if the id is invalid.
116    pub fn get_user_address_unchecked(&self, id: usize) -> ManagedAddress<SA> {
117        self.address
118            .address_storage_get(self.get_user_address_key(id).as_ref())
119    }
120
121    /// Yields the user address for a given id, if the id is valid.
122    /// Otherwise returns the zero address (0x000...)
123    pub fn get_user_address_or_zero(&self, id: usize) -> ManagedAddress<SA> {
124        let key = self.get_user_address_key(id);
125        // TODO: optimize, storage_load_managed_buffer_len is currently called twice
126        if self.address.address_storage_get_len(key.as_ref()) > 0 {
127            self.address.address_storage_get(key.as_ref())
128        } else {
129            ManagedAddress::zero()
130        }
131    }
132
133    /// Number of users.
134    pub fn get_user_count(&self) -> usize {
135        self.address
136            .address_storage_get(self.get_user_count_key().as_ref())
137    }
138
139    /// Loads all addresses from storage and places them in a ManagedVec.
140    /// Can easily consume a lot of gas.
141    pub fn get_all_addresses(&self) -> ManagedVec<SA, ManagedAddress<SA>> {
142        let user_count = self.get_user_count();
143        let mut result = ManagedVec::new();
144        for i in 1..=user_count {
145            result.push(self.get_user_address_or_zero(i));
146        }
147        result
148    }
149}
150
151impl<SA> UserMapper<SA, CurrentStorage>
152where
153    SA: StorageMapperApi,
154{
155    fn set_user_id(&self, address: &ManagedAddress<SA>, id: usize) {
156        storage_set(self.get_user_id_key(address).as_ref(), &id);
157    }
158
159    fn set_user_address(&self, id: usize, address: &ManagedAddress<SA>) {
160        storage_set(self.get_user_address_key(id).as_ref(), address);
161    }
162
163    fn set_user_count(&self, user_count: usize) {
164        storage_set(self.get_user_count_key().as_ref(), &user_count);
165    }
166
167    /// Tries to insert a number of addresses.
168    /// Calls a lambda function for each, with the new user id and whether of nor the user was already present.
169    pub fn get_or_create_users<AddressIter, F>(
170        &self,
171        address_iter: AddressIter,
172        mut user_id_lambda: F,
173    ) where
174        AddressIter: Iterator<Item = ManagedAddress<SA>>,
175        F: FnMut(usize, bool),
176    {
177        let mut user_count = self.get_user_count();
178        for address in address_iter {
179            let user_id = self.get_user_id(&address);
180            if user_id > 0 {
181                user_id_lambda(user_id, false);
182            } else {
183                user_count += 1;
184                let new_user_id = user_count;
185                self.set_user_id(&address, new_user_id);
186                self.set_user_address(new_user_id, &address);
187                user_id_lambda(new_user_id, true);
188            }
189        }
190        self.set_user_count(user_count);
191    }
192
193    /// Yields the user id for a given address, or creates a new user id if there isn't one.
194    /// Will safely keep the user count in sync.
195    pub fn get_or_create_user(&self, address: &ManagedAddress<SA>) -> usize {
196        let mut user_id = self.get_user_id(address);
197        if user_id == 0 {
198            let next_user_count = self.get_user_count() + 1;
199            self.set_user_count(next_user_count);
200            user_id = next_user_count;
201            self.set_user_id(address, user_id);
202            self.set_user_address(user_id, address);
203        }
204        user_id
205    }
206}
207
208/// Behaves like a MultiResultVec<Address> when an endpoint result,
209/// and lists all users addresses.
210impl<SA> TopEncodeMulti for UserMapper<SA, CurrentStorage>
211where
212    SA: StorageMapperApi,
213{
214    fn multi_encode_or_handle_err<O, H>(&self, output: &mut O, h: H) -> Result<(), H::HandledErr>
215    where
216        O: TopEncodeMultiOutput,
217        H: EncodeErrorHandler,
218    {
219        let all_addresses = self.get_all_addresses();
220        multi_encode_iter_or_handle_err(all_addresses.into_iter(), output, h)
221    }
222}
223
224impl<SA> TypeAbiFrom<UserMapper<SA, CurrentStorage>> for MultiValueEncoded<SA, ManagedAddress<SA>> where
225    SA: StorageMapperApi
226{
227}
228
229impl<SA> TypeAbiFrom<Self> for UserMapper<SA, CurrentStorage> where SA: StorageMapperApi {}
230
231/// Behaves like a MultiResultVec when an endpoint result.
232impl<SA> TypeAbi for UserMapper<SA, CurrentStorage>
233where
234    SA: StorageMapperApi,
235{
236    type Unmanaged = Self;
237
238    fn type_name() -> TypeName {
239        crate::abi::type_name_variadic::<ManagedAddress<SA>>()
240    }
241
242    fn type_name_rust() -> TypeName {
243        crate::abi::type_name_multi_value_encoded::<ManagedAddress<SA>>()
244    }
245
246    fn is_variadic() -> bool {
247        true
248    }
249}
multiversx_sc/storage/mappers/user_mapper.rs

multiversx_sc/storage/mappers/
user_mapper.rs
