1use std::{marker::PhantomData, collections::BTreeMap};
9use borsh::{BorshSerialize, BorshDeserialize};
10use crate::{storage::{self}, Storable, StoragePath};
11
12pub struct FastMap<K, V>
50 where K: BorshSerialize,
51 V: Insertable {
52 parent_key: Vec<u8>,
53 write_set: BTreeMap<Vec<u8>, UpdateOperation<V>>,
54 _marker: PhantomData<Box<(K, V)>>
55}
56
57impl<K, V> FastMap<K, V>
58 where K: BorshSerialize,
59 V: Insertable {
60
61 pub fn new() -> Self {
69 Self { parent_key: vec![], write_set: BTreeMap::default(), _marker: PhantomData }
70 }
71
72 pub fn get(&self, key: &K) -> Option<V> {
85 let key_bs = key.try_to_vec().unwrap();
86
87 match self.write_set.get(&key_bs) {
88 Some(UpdateOperation::Delete) => return None,
89 Some(UpdateOperation::Insert(v,_)) => {
90 let v_serialized = v.try_to_vec().unwrap();
91 return Some(V::deserialize(&mut v_serialized.as_slice()).unwrap());
92 }
93 None => {},
94 }
95
96 let ws_key = self.child_key(key_bs);
98 V::load(ws_key)
99 }
100
101 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
114 match self.get(key) {
115 Some(v) => {
116 self.insert_inner(key, v, false);
117 let key_bs = key.try_to_vec().unwrap();
118 match self.write_set.get_mut(&key_bs) {
119 Some(UpdateOperation::Insert(mut_v, _)) => Some(mut_v),
120 _=> None
121 }
122 },
123 None => None,
124 }
125 }
126
127 pub fn insert(&mut self, key: &K, value: V) {
133 self.insert_inner(key, value, true);
134 }
135
136 fn insert_inner(&mut self, key: &K, value: V, new_record: bool) -> Option<&mut V> {
137 let key_bs = key.try_to_vec().unwrap();
138 self.write_set.insert(key_bs.clone(), UpdateOperation::Insert(value, new_record));
139 match self.write_set.get_mut(&key_bs) {
140 Some(UpdateOperation::Insert(mut_insertable, _)) => Some(mut_insertable),
141 _=> None
142 }
143 }
144
145 pub fn remove(&mut self, key: &K) {
151 let key_bs = key.try_to_vec().unwrap();
152 self.write_set.insert(key_bs, UpdateOperation::Delete);
153 }
154
155 fn child_key(&self, key: Vec<u8>) -> Vec<u8> {
156 let edition = Self::edition(&self.parent_key);
157 Self::make_child_key(self.parent_key.to_vec(), edition, key)
158 }
159 fn make_child_key(parent_key: Vec<u8>, edition: u32, key: Vec<u8>) -> Vec<u8> {
160 [
161 parent_key.to_vec(),
162 edition.to_le_bytes().to_vec(),
163 key
164 ].concat()
165 }
166
167}
168
169impl<K, V> Insertable for FastMap<K, V>
170 where K: BorshSerialize,
171 V: Insertable {
172 fn save(&mut self, key: Vec<u8>, is_new: bool){
174 if self.parent_key.is_empty() {
175 self.parent_key = key;
176 }
177
178 let edition = match storage::get(&self.parent_key) {
179 Some(bytes) => {
180 match Cell::deserialize(&mut bytes.as_slice()) {
181 Ok(c) => c.edition + u32::from(is_new),
182 Err(_) => 0,
183 }
184 },
185 None => 0
186 };
187
188 let c = Cell { edition, data: Some(self.parent_key.try_to_vec().unwrap()) };
189 storage::set(&self.parent_key, c.try_to_vec().unwrap().as_slice());
190
191 self.write_set.iter_mut().for_each(|(k, v)| {
192 let vkey = Self::make_child_key(self.parent_key.to_vec(), edition, k.clone());
193 match v {
194 UpdateOperation::Insert(v, is_new) => {
195 v.save(vkey, *is_new);
196 },
197 UpdateOperation::Delete => {
198 V::delete(vkey);
199 },
200 }
201 });
202 }
203}
204
205impl<K, V> BorshSerialize for FastMap<K, V>
206 where K: BorshSerialize,
207 V: Insertable {
208 fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
209 self.parent_key.serialize(writer)
211 }
212}
213
214impl<K, V> BorshDeserialize for FastMap<K, V>
215 where K: BorshSerialize,
216 V: Insertable {
217 fn deserialize_reader<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
218 let parent_key = Vec::<u8>::deserialize_reader(reader)?;
219 Ok(Self{
220 parent_key,
221 write_set: BTreeMap::default(),
222 _marker: PhantomData,
223 })
224 }
225}
226
227impl<K, V> Storable for FastMap<K, V>
228 where K: BorshSerialize,
229 V: Insertable {
230
231 fn __load_storage(field: &StoragePath) -> Self {
233 Self {
234 parent_key: field.get_path().to_vec(),
235 write_set: BTreeMap::default(),
236 _marker: PhantomData,
237 }
238 }
239
240 fn __save_storage(&mut self, field: &StoragePath) {
242 self.save(field.get_path().to_vec(), false);
243 }
244}
245
246#[derive(Clone)]
248pub(crate) enum UpdateOperation<T> {
249 Insert(T, bool),
251 Delete
252}
253
254#[derive(BorshSerialize, BorshDeserialize)]
256struct Cell {
257 edition: u32,
259 data: Option<Vec<u8>>
262}
263
264pub trait Insertable : BorshSerialize + BorshDeserialize {
267 fn edition(key: &[u8]) -> u32 {
268 storage::get(key).map_or(0, |bytes|{
269 Cell::deserialize(&mut bytes.as_slice()).map_or(0, |c| c.edition)
270 })
271 }
272
273 fn load(key: Vec<u8>) -> Option<Self> {
274 let bytes = storage::get(&key)?;
275 let c = Cell::deserialize(&mut bytes.as_slice()).ok()?;
276 let data = c.data?;
277 Self::deserialize(&mut data.as_slice()).map_or(None, |s| Some(s))
278 }
279
280 fn save(&mut self, key: Vec<u8>, is_new: bool) {
281 let edition = storage::get(&key).map_or(0, |bytes| {
282 Cell::deserialize(&mut bytes.as_slice()).map_or(0, |c|
283 c.edition + u32::from(is_new)
284 )
285 });
286 let c = Cell { edition, data: Some(self.try_to_vec().unwrap()) };
287 storage::set(&key, c.try_to_vec().unwrap().as_slice());
288 }
289
290 fn delete(key: Vec<u8>) {
291 let edition = storage::get(&key).map_or(0, |bytes| {
292 Cell::deserialize(&mut bytes.as_slice()).map_or(0, |c| c.edition + 1)
293 });
294 let c = Cell { edition, data: None };
295 storage::set(&key, c.try_to_vec().unwrap().as_slice());
296 }
297}
298
299
300macro_rules! define_primitives {
303 ($($t:ty),*) => {
304 $(
305 impl Insertable for $t {}
306 )*
307 }
308}
309define_primitives!(
310 u8, u16, u32, u64, u128,
311 i8, i16, i32, i64, i128,
312 String, bool, usize,
313 [u8;32]
314);
315impl<T> Insertable for Option<T> where T: BorshSerialize + BorshDeserialize {}
316impl<T> Insertable for Vec<T> where T: BorshSerialize + BorshDeserialize {}
317macro_rules! impl_tuple {
318 ($($idx:tt $name:ident)+) => {
319 impl<$($name),+> Insertable for ($($name),+)
320 where $($name: BorshSerialize + BorshDeserialize,)+
321 {}
322 };
323}
324impl_tuple!(0 T0 1 T1);
325impl_tuple!(0 T0 1 T1 2 T2);
326impl_tuple!(0 T0 1 T1 2 T2 3 T3);
327impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4);
328impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5);
329impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6);
330impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7);
331impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8);
332impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9);
333impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9 10 T10);
334impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9 10 T10 11 T11);
335impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9 10 T10 11 T11 12 T12);
336impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9 10 T10 11 T11 12 T12 13 T13);
337impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9 10 T10 11 T11 12 T12 13 T13 14 T14);
338impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9 10 T10 11 T11 12 T12 13 T13 14 T14 15 T15);
339impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9 10 T10 11 T11 12 T12 13 T13 14 T14 15 T15 16 T16);
340impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9 10 T10 11 T11 12 T12 13 T13 14 T14 15 T15 16 T16 17 T17);
341impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9 10 T10 11 T11 12 T12 13 T13 14 T14 15 T15 16 T16 17 T17 18 T18);
342impl_tuple!(0 T0 1 T1 2 T2 3 T3 4 T4 5 T5 6 T6 7 T7 8 T8 9 T9 10 T10 11 T11 12 T12 13 T13 14 T14 15 T15 16 T16 17 T17 18 T18 19 T19);