naia_shared/
bigmap.rs

1use std::{
2    collections::{
3        hash_map::{Iter, IterMut},
4        HashMap,
5    },
6    hash::Hash,
7    iter::Map,
8    marker::PhantomData,
9};
10
11pub trait BigMapKey: Clone + Copy + Eq + PartialEq + Hash {
12    fn to_u64(&self) -> u64;
13    fn from_u64(value: u64) -> Self;
14}
15
16pub struct BigMap<K: BigMapKey, V> {
17    inner: HashMap<u64, V>,
18    current_index: u64,
19    phantom_k: PhantomData<K>,
20}
21
22impl<K: BigMapKey, V> BigMap<K, V> {
23    pub fn new() -> Self {
24        Self {
25            inner: HashMap::new(),
26            current_index: 0,
27            phantom_k: PhantomData,
28        }
29    }
30
31    pub fn get(&self, key: &K) -> Option<&V> {
32        self.inner.get(&key.to_u64())
33    }
34
35    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
36        self.inner.get_mut(&key.to_u64())
37    }
38
39    pub fn insert(&mut self, value: V) -> K {
40        let old_index = self.current_index;
41        self.current_index = self.current_index.wrapping_add(1);
42
43        self.inner.insert(old_index, value);
44
45        K::from_u64(old_index)
46    }
47
48    pub fn remove(&mut self, key: &K) -> Option<V> {
49        self.inner.remove(&key.to_u64())
50    }
51
52    pub fn contains_key(&self, key: &K) -> bool {
53        self.inner.contains_key(&key.to_u64())
54    }
55
56    #[allow(clippy::type_complexity)]
57    pub fn iter<'a>(&'a self) -> Map<Iter<'a, u64, V>, fn((&'a u64, &'a V)) -> (K, &'a V)> {
58        return self
59            .inner
60            .iter()
61            .map(|(key, value)| (K::from_u64(*key), value));
62    }
63
64    #[allow(clippy::type_complexity)]
65    pub fn iter_mut<'a>(
66        &'a mut self,
67    ) -> Map<IterMut<'a, u64, V>, fn((&'a u64, &'a mut V)) -> (K, &'a mut V)> {
68        return self
69            .inner
70            .iter_mut()
71            .map(|(key, value)| (K::from_u64(*key), value));
72    }
73
74    pub fn len(&self) -> usize {
75        self.inner.len()
76    }
77
78    pub fn is_empty(&self) -> bool {
79        self.inner.is_empty()
80    }
81}