multiversx_sc/storage/mappers/
address_to_id_mapper.rs

1use core::marker::PhantomData;
2
3use super::{
4    StorageMapper, StorageMapperFromAddress,
5    source::{CurrentStorage, StorageAddress},
6};
7use crate::{
8    api::{ErrorApiImpl, StorageMapperApi},
9    storage::{StorageKey, storage_clear, storage_set},
10    types::{ManagedAddress, ManagedType},
11};
12
13static ID_SUFFIX: &[u8] = b"addrId";
14static ADDRESS_SUFFIX: &[u8] = b"addr";
15static LAST_ID_SUFFIX: &[u8] = b"lastId";
16
17static UNKNOWN_ADDR_ERR_MSG: &[u8] = b"Unknown address";
18
19pub type AddressId = u64;
20pub const NULL_ID: AddressId = 0;
21
22pub struct AddressToIdMapper<SA, A = CurrentStorage>
23where
24    SA: StorageMapperApi,
25    A: StorageAddress<SA>,
26{
27    _phantom_api: PhantomData<SA>,
28    address: A,
29    base_key: StorageKey<SA>,
30}
31
32impl<SA> StorageMapper<SA> for AddressToIdMapper<SA>
33where
34    SA: StorageMapperApi,
35{
36    fn new(base_key: StorageKey<SA>) -> Self {
37        AddressToIdMapper {
38            _phantom_api: PhantomData,
39            address: CurrentStorage,
40            base_key,
41        }
42    }
43}
44
45impl<SA> StorageMapperFromAddress<SA> for AddressToIdMapper<SA, ManagedAddress<SA>>
46where
47    SA: StorageMapperApi,
48{
49    fn new_from_address(address: ManagedAddress<SA>, base_key: StorageKey<SA>) -> Self {
50        AddressToIdMapper {
51            _phantom_api: PhantomData,
52            address,
53            base_key,
54        }
55    }
56}
57
58impl<SA, A> AddressToIdMapper<SA, A>
59where
60    SA: StorageMapperApi,
61    A: StorageAddress<SA>,
62{
63    pub fn contains_id(&self, id: AddressId) -> bool {
64        let key = self.id_to_address_key(id);
65        self.address.address_storage_get_len(key.as_ref()) != 0
66    }
67
68    pub fn get_id(&self, address: &ManagedAddress<SA>) -> AddressId {
69        let key = self.address_to_id_key(address);
70        self.address.address_storage_get(key.as_ref())
71    }
72
73    pub fn get_id_non_zero(&self, address: &ManagedAddress<SA>) -> AddressId {
74        let id = self.get_id(address);
75        if id == NULL_ID {
76            SA::error_api_impl().signal_error(UNKNOWN_ADDR_ERR_MSG);
77        }
78
79        id
80    }
81
82    pub fn get_address(&self, id: AddressId) -> Option<ManagedAddress<SA>> {
83        let key = self.id_to_address_key(id);
84        if self.address.address_storage_get_len(key.as_ref()) == 0 {
85            return None;
86        }
87
88        let addr = self.address.address_storage_get(key.as_ref());
89        Some(addr)
90    }
91
92    fn id_to_address_key(&self, id: AddressId) -> StorageKey<SA> {
93        let mut item_key = self.base_key.clone();
94        item_key.append_bytes(ID_SUFFIX);
95        item_key.append_item(&id);
96
97        item_key
98    }
99
100    fn address_to_id_key(&self, address: &ManagedAddress<SA>) -> StorageKey<SA> {
101        let mut item_key = self.base_key.clone();
102        item_key.append_bytes(ADDRESS_SUFFIX);
103        item_key.append_item(address);
104
105        item_key
106    }
107
108    fn last_id_key(&self) -> StorageKey<SA> {
109        let mut item_key = self.base_key.clone();
110        item_key.append_bytes(LAST_ID_SUFFIX);
111
112        item_key
113    }
114
115    pub fn get_last_id(&self) -> AddressId {
116        self.address
117            .address_storage_get(self.last_id_key().as_ref())
118    }
119}
120
121impl<SA> AddressToIdMapper<SA, CurrentStorage>
122where
123    SA: StorageMapperApi,
124{
125    pub fn get_id_or_insert(&self, address: &ManagedAddress<SA>) -> AddressId {
126        let current_id = self
127            .address
128            .address_storage_get(self.address_to_id_key(address).as_ref());
129        if current_id != 0 {
130            return current_id;
131        }
132
133        self.insert_address(address)
134    }
135
136    pub fn insert_new(&self, address: &ManagedAddress<SA>) -> AddressId {
137        let existing_id = self.get_id(address);
138        if existing_id != NULL_ID {
139            SA::error_api_impl().signal_error(b"Address already registered");
140        }
141
142        self.insert_address(address)
143    }
144
145    pub fn remove_by_id(&self, id: AddressId) -> Option<ManagedAddress<SA>> {
146        let address = self.get_address(id)?;
147        self.remove_entry(id, &address);
148
149        Some(address)
150    }
151
152    pub fn remove_by_address(&self, address: &ManagedAddress<SA>) -> AddressId {
153        let current_id = self.get_id(address);
154        if current_id != NULL_ID {
155            self.remove_entry(current_id, address);
156        }
157
158        current_id
159    }
160
161    fn insert_address(&self, address: &ManagedAddress<SA>) -> AddressId {
162        let new_id = self.get_last_id() + 1;
163        storage_set(self.address_to_id_key(address).as_ref(), &new_id);
164        storage_set(self.id_to_address_key(new_id).as_ref(), address);
165
166        self.set_last_id(new_id);
167
168        new_id
169    }
170
171    fn set_last_id(&self, last_id: AddressId) {
172        if last_id == 0 {
173            SA::error_api_impl().signal_error(b"ID Overflow");
174        }
175
176        storage_set(self.last_id_key().as_ref(), &last_id);
177    }
178
179    fn remove_entry(&self, id: AddressId, address: &ManagedAddress<SA>) {
180        storage_clear(self.address_to_id_key(address).as_ref());
181        storage_clear(self.id_to_address_key(id).as_ref());
182    }
183}
multiversx_sc/storage/mappers/address_to_id_mapper.rs

multiversx_sc/storage/mappers/
address_to_id_mapper.rs
