gmsol_utils/
fixed_map.rs

1use anchor_lang::solana_program::hash::hashv;
2
3/// Convert to fixed-size key.
4pub fn to_key(key: &str) -> [u8; 32] {
5    hashv(&[key.as_bytes()]).to_bytes()
6}
7
8#[macro_export]
9macro_rules! fixed_map {
10    ($map:ident, $value:ty, $len:expr, $padding:expr) => {
11        $crate::fixed_map!($map, str, $crate::fixed_map::to_key, $value, $len, $padding);
12    };
13
14    ($map:ident, $key:ty, $to_key:path, $value:ty, $len:expr, $padding:expr) => {
15        $crate::fixed_map!($map, 32, $key, $to_key, $value, $len, $padding);
16    };
17
18    ($map:ident, $key_len:expr, $key:ty, $to_key:path, $value:ty, $len:expr, $padding:expr) => {
19        $crate::paste::paste! {
20            /// Entry.
21            #[anchor_lang::zero_copy]
22            #[cfg_attr(feature = "debug", derive(Debug))]
23            struct [<$map Entry>] {
24                key: [u8; $key_len],
25                value: $value,
26            }
27
28            impl Default for [<$map Entry>] {
29                fn default() -> Self {
30                    Self {
31                        key: Default::default(),
32                        value: Default::default(),
33                    }
34                }
35            }
36
37            /// Fixed size map generated by the macro.
38            #[anchor_lang::zero_copy]
39            #[cfg_attr(feature = "debug", derive(Debug))]
40            pub struct $map {
41                data: [[<$map Entry>]; $len],
42                padding: [u8; $padding],
43                count: u32,
44            }
45
46            $crate::impl_fixed_map!($map, $key_len, $key, $to_key, $value, $len);
47        }
48    };
49}
50
51#[macro_export]
52macro_rules! impl_fixed_map {
53    ($map:ident, $value:ty, $len:expr) => {
54        $crate::impl_fixed_map!($map, str, $crate::fixed_map::to_key, $value, $len);
55    };
56
57    ($map:ident, $key:ty, $to_key:path, $value:ty, $len:expr) => {
58        $crate::impl_fixed_map!($map, 32, $key, $to_key, $value, $len);
59    };
60
61    ($map:ident, $key_len:expr, $key:ty, $to_key:path, $value:ty, $len:expr) => {
62        $crate::paste::paste! {
63            impl $crate::InitSpace for $map {
64                const INIT_SPACE: usize = std::mem::size_of::<$map>();
65            }
66
67            #[cfg(test)]
68            $crate::static_assertions::const_assert_eq!(
69                std::mem::size_of::<$map>(),
70                <$map as $crate::InitSpace>::INIT_SPACE
71            );
72
73            impl Default for $map {
74                fn default() -> Self {
75                    bytemuck::Zeroable::zeroed()
76                }
77            }
78
79            #[allow(dead_code)]
80            impl $map {
81                fn binary_search(&self, key: &[u8; $key_len]) -> std::result::Result<usize, usize> {
82                    self.data[..self.len()].binary_search_by(|entry| entry.key.cmp(key))
83                }
84
85                /// Get.
86                pub fn get(&self, key: &$key) -> Option<&$value> {
87                    let key = $to_key(key);
88                    self
89                        .binary_search(&key)
90                        .ok()
91                        .map(|index| &self.data[index].value)
92                }
93
94                /// Get entry by index.
95                pub fn get_entry_by_index(&self, idx: usize) -> Option<(&[u8; $key_len], &$value)> {
96                    if idx < self.len() {
97                        let data = &self.data[idx];
98                        Some((&data.key, &data.value))
99                    } else {
100                        None
101                    }
102                }
103
104                /// Get mutable reference to the corresponding value.
105                pub fn get_mut(&mut self, key: &$key) -> Option<&mut $value> {
106                    let key = $to_key(key);
107                    self
108                        .binary_search(&key)
109                        .ok()
110                        .map(|index| &mut self.data[index].value)
111                }
112
113                /// Insert.
114                pub fn insert(&mut self, key: &$key, value: $value) -> Option<$value> {
115                    self.insert_with_options(key, value, false).expect("must be success")
116                }
117
118                /// Insert with options.
119                pub fn insert_with_options(
120                    &mut self,
121                    key: &$key,
122                    value: $value,
123                    new: bool,
124                ) -> std::result::Result<Option<$value>, anchor_lang::error::Error> {
125                    let key = $to_key(key);
126                    match self.binary_search(&key) {
127                        Ok(index) => {
128                            if new {
129                                anchor_lang::err!($crate::GeneralError::AlreadyExist)
130                            } else {
131                                let previous = std::mem::replace(&mut self.data[index].value, value);
132                                Ok(Some(previous))
133                            }
134                        }
135                        Err(index) => {
136                            if self.len() >= $len {
137                                anchor_lang::err!($crate::GeneralError::ExceedMaxLengthLimit)
138                            } else {
139                                for i in (index..self.len()).rev() {
140                                    self.data[i + 1] = self.data[i];
141                                }
142                                self.data[index] = [<$map Entry>] { key, value };
143                                self.count += 1;
144                                Ok(None)
145                            }
146                        }
147                    }
148                }
149
150                /// Remove.
151                pub fn remove(&mut self, key: &$key) -> Option<$value> {
152                    let key = $to_key(key);
153                    self.binary_search(&key).ok().map(|index| {
154                        let value = std::mem::take(&mut self.data[index].value);
155                        let len = self.len();
156                        for i in index..len {
157                            self.data[i] = self.data[i + 1];
158                        }
159                        self.data[len - 1] = [<$map Entry>]::default();
160                        self.count -= 1;
161                        value
162                    })
163                }
164
165                /// Get length.
166                pub fn len(&self) -> usize {
167                    self.count as usize
168                }
169
170                /// Is empty.
171                pub fn is_empty(&self) -> bool {
172                    self.count == 0
173                }
174
175                /// Entries.
176                pub fn entries(&self) -> impl Iterator<Item = (&[u8; $key_len], &$value)> {
177                    self.data.iter().take(self.len()).map(|entry| {
178                        (&entry.key, &entry.value)
179                    })
180                }
181
182                /// Entries with mutable access.
183                pub fn entries_mut(&mut self) -> impl Iterator<Item = (&[u8; $key_len], &mut $value)> {
184                    let len = self.len();
185                    self.data.iter_mut().take(len).map(|entry| {
186                        (&entry.key, &mut entry.value)
187                    })
188                }
189
190                /// Clear.
191                pub fn clear(&mut self) {
192                    let len = self.len();
193                    for i in 0..len {
194                        self.data[i] = [<$map Entry>]::default();
195                    }
196                    self.count = 0;
197                }
198            }
199        }
200    };
201}
202
203#[cfg(test)]
204mod tests {
205    use anchor_lang::solana_program::pubkey::Pubkey;
206
207    fixed_map!(FixedFactorMap, u128, 32, 12);
208
209    #[test]
210    fn test_insert_and_get() {
211        let mut map = FixedFactorMap::default();
212
213        assert!(map.is_empty());
214
215        assert_eq!(map.insert("key1", 123), None);
216        assert_eq!(map.insert("key1", 234), Some(123));
217
218        assert_eq!(map.insert("key2", 345), None);
219        assert_eq!(map.insert("key2", 456), Some(345));
220
221        assert_eq!(map.insert("key1", 789), Some(234));
222        assert_eq!(map.get("key1"), Some(&789));
223
224        *map.get_mut("key2").unwrap() = 42;
225        assert_eq!(map.get("key2"), Some(&42));
226
227        assert_eq!(map.len(), 2);
228    }
229
230    #[test]
231    fn test_insert_and_remove() {
232        let mut map = FixedFactorMap::default();
233
234        assert_eq!(map.insert("key1", 123), None);
235        assert_eq!(map.insert("key1", 234), Some(123));
236
237        assert_eq!(map.insert("key2", 345), None);
238        assert_eq!(map.insert("key2", 456), Some(345));
239
240        assert_eq!(map.remove("key1"), Some(234));
241        assert_eq!(map.insert("key1", 789), None);
242
243        assert_eq!(map.len(), 2);
244    }
245
246    fn to_bytes(key: &Pubkey) -> [u8; 32] {
247        key.to_bytes()
248    }
249
250    fixed_map!(RolesMap, Pubkey, to_bytes, u64, 32, 4);
251
252    #[test]
253    fn test_insert_and_get_for_roles_map() {
254        let mut map = RolesMap::default();
255
256        let address_1 = Pubkey::new_unique();
257        let address_2 = Pubkey::new_unique();
258
259        assert!(map.is_empty());
260
261        assert_eq!(map.insert(&address_1, 123), None);
262        assert_eq!(map.insert(&address_1, 234), Some(123));
263
264        assert_eq!(map.insert(&address_2, 345), None);
265        assert_eq!(map.insert(&address_2, 456), Some(345));
266
267        assert_eq!(map.insert(&address_1, 789), Some(234));
268        assert_eq!(map.get(&address_1), Some(&789));
269
270        *map.get_mut(&address_2).unwrap() = 42;
271        assert_eq!(map.get(&address_2), Some(&42));
272
273        assert_eq!(map.len(), 2);
274    }
275
276    #[test]
277    fn test_insert_and_remove_for_roles_map() {
278        let mut map = RolesMap::default();
279
280        let address_1 = Pubkey::new_unique();
281        let address_2 = Pubkey::new_unique();
282
283        assert_eq!(map.insert(&address_1, 123), None);
284        assert_eq!(map.insert(&address_1, 234), Some(123));
285
286        assert_eq!(map.insert(&address_2, 345), None);
287        assert_eq!(map.insert(&address_2, 456), Some(345));
288
289        assert_eq!(map.remove(&address_1), Some(234));
290        assert_eq!(map.insert(&address_1, 789), None);
291
292        assert_eq!(map.len(), 2);
293    }
294}
gmsol_utils/fixed_map.rs

gmsol_utils/
fixed_map.rs
