Skip to main content

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
11/// Key type for [`BigMap`]: must be convertible to/from `u64`.
12pub trait BigMapKey: Clone + Copy + Eq + PartialEq + Hash {
13    /// Converts the key to its `u64` representation.
14    fn to_u64(&self) -> u64;
15    /// Reconstructs a key from its `u64` representation.
16    fn from_u64(value: u64) -> Self;
17}
18
19/// Auto-incrementing dense map that generates monotone `u64`-backed typed keys.
20pub struct BigMap<K: BigMapKey, V> {
21    inner: HashMap<u64, V>,
22    current_index: u64,
23    phantom_k: PhantomData<K>,
24}
25
26impl<K: BigMapKey, V> Default for BigMap<K, V> {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl<K: BigMapKey, V> BigMap<K, V> {
33    /// Creates an empty `BigMap`.
34    pub fn new() -> Self {
35        Self {
36            inner: HashMap::new(),
37            current_index: 0,
38            phantom_k: PhantomData,
39        }
40    }
41
42    /// Returns a reference to the value for `key`, or `None` if not present.
43    pub fn get(&self, key: &K) -> Option<&V> {
44        self.inner.get(&key.to_u64())
45    }
46
47    /// Returns a mutable reference to the value for `key`, or `None` if not present.
48    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
49        self.inner.get_mut(&key.to_u64())
50    }
51
52    /// Inserts a value and returns its newly generated key. Panics on `u64` overflow.
53    pub fn insert(&mut self, value: V) -> K {
54        // [entity-replication-11] GlobalEntity rollover is a terminal error
55        if self.current_index == u64::MAX {
56            panic!(
57                "BigMap counter overflow: cannot allocate new key (current_index = u64::MAX). \
58                 This is a terminal error per entity-replication-11 spec."
59            );
60        }
61
62        let old_index = self.current_index;
63        self.current_index = self.current_index.wrapping_add(1);
64
65        self.inner.insert(old_index, value);
66
67        K::from_u64(old_index)
68    }
69
70    /// Removes and returns the value for `key`, or `None` if not present.
71    pub fn remove(&mut self, key: &K) -> Option<V> {
72        self.inner.remove(&key.to_u64())
73    }
74
75    /// Returns `true` if the map contains an entry for `key`.
76    pub fn contains_key(&self, key: &K) -> bool {
77        self.inner.contains_key(&key.to_u64())
78    }
79
80    /// Returns an iterator over `(K, &V)` pairs in arbitrary order.
81    #[allow(clippy::type_complexity)]
82    pub fn iter<'a>(&'a self) -> Map<Iter<'a, u64, V>, fn((&'a u64, &'a V)) -> (K, &'a V)> {
83        self
84            .inner
85            .iter()
86            .map(|(key, value)| (K::from_u64(*key), value))
87    }
88
89    /// Returns a mutable iterator over `(K, &mut V)` pairs in arbitrary order.
90    #[allow(clippy::type_complexity)]
91    pub fn iter_mut<'a>(
92        &'a mut self,
93    ) -> Map<IterMut<'a, u64, V>, fn((&'a u64, &'a mut V)) -> (K, &'a mut V)> {
94        self
95            .inner
96            .iter_mut()
97            .map(|(key, value)| (K::from_u64(*key), value))
98    }
99
100    /// Returns the number of entries currently in the map.
101    pub fn len(&self) -> usize {
102        self.inner.len()
103    }
104
105    /// Returns `true` if the map contains no entries.
106    pub fn is_empty(&self) -> bool {
107        self.inner.is_empty()
108    }
109
110    #[cfg(feature = "test_utils")]
111    #[doc(hidden)]
112    pub fn set_current_index_for_test(&mut self, index: u64) {
113        self.current_index = index;
114    }
115}