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 {
13 fn to_u64(&self) -> u64;
15 fn from_u64(value: u64) -> Self;
17}
18
19pub 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 pub fn new() -> Self {
35 Self {
36 inner: HashMap::new(),
37 current_index: 0,
38 phantom_k: PhantomData,
39 }
40 }
41
42 pub fn get(&self, key: &K) -> Option<&V> {
44 self.inner.get(&key.to_u64())
45 }
46
47 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
49 self.inner.get_mut(&key.to_u64())
50 }
51
52 pub fn insert(&mut self, value: V) -> K {
54 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 pub fn remove(&mut self, key: &K) -> Option<V> {
72 self.inner.remove(&key.to_u64())
73 }
74
75 pub fn contains_key(&self, key: &K) -> bool {
77 self.inner.contains_key(&key.to_u64())
78 }
79
80 #[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 #[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 pub fn len(&self) -> usize {
102 self.inner.len()
103 }
104
105 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}