1use serde::de::DeserializeOwned;
2use serde::Serialize;
3
4use cosmwasm_std::{StdError, StdResult, Storage};
5
6use crate::bound::PrefixBound;
7use crate::de::KeyDeserialize;
8use crate::iter_helpers::deserialize_kv;
9use crate::keys::PrimaryKey;
10use crate::map::Map;
11use crate::namespace::Namespace;
12use crate::path::Path;
13use crate::prefix::{namespaced_prefix_range, Prefix};
14use crate::snapshot::{ChangeSet, Snapshot};
15use crate::{Bound, Prefixer, Strategy};
16
17pub struct SnapshotMap<K, T> {
21 primary: Map<K, T>,
22 snapshots: Snapshot<K, T>,
23}
24
25impl<K, T> SnapshotMap<K, T> {
26 pub const fn new(
43 pk: &'static str,
44 checkpoints: &'static str,
45 changelog: &'static str,
46 strategy: Strategy,
47 ) -> Self {
48 SnapshotMap {
49 primary: Map::new(pk),
50 snapshots: Snapshot::new(checkpoints, changelog, strategy),
51 }
52 }
53
54 pub fn new_dyn(
74 pk: impl Into<Namespace>,
75 checkpoints: impl Into<Namespace>,
76 changelog: impl Into<Namespace>,
77 strategy: Strategy,
78 ) -> Self {
79 SnapshotMap {
80 primary: Map::new_dyn(pk),
81 snapshots: Snapshot::new_dyn(checkpoints, changelog, strategy),
82 }
83 }
84
85 pub fn changelog(&self) -> &Map<(K, u64), ChangeSet<T>> {
86 &self.snapshots.changelog
87 }
88}
89
90impl<'a, K, T> SnapshotMap<K, T>
91where
92 T: Serialize + DeserializeOwned + Clone,
93 K: PrimaryKey<'a> + Prefixer<'a>,
94{
95 pub fn add_checkpoint(&self, store: &mut dyn Storage, height: u64) -> StdResult<()> {
96 self.snapshots.add_checkpoint(store, height)
97 }
98
99 pub fn remove_checkpoint(&self, store: &mut dyn Storage, height: u64) -> StdResult<()> {
100 self.snapshots.remove_checkpoint(store, height)
101 }
102}
103
104impl<'a, K, T> SnapshotMap<K, T>
105where
106 T: Serialize + DeserializeOwned + Clone,
107 K: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize,
108{
109 pub fn key(&self, k: K) -> Path<T> {
110 self.primary.key(k)
111 }
112
113 fn no_prefix_raw(&self) -> Prefix<Vec<u8>, T, K> {
114 self.primary.no_prefix_raw()
115 }
116
117 fn write_change(&self, store: &mut dyn Storage, k: K, height: u64) -> StdResult<()> {
119 if self.snapshots.has_changelog(store, k.clone(), height)? {
121 return Ok(());
122 }
123 let old = self.primary.may_load(store, k.clone())?;
125 self.snapshots.write_changelog(store, k, height, old)
126 }
127
128 pub fn save(&self, store: &mut dyn Storage, k: K, data: &T, height: u64) -> StdResult<()> {
129 if self.snapshots.should_checkpoint(store, &k)? {
130 self.write_change(store, k.clone(), height)?;
131 }
132 self.primary.save(store, k, data)
133 }
134
135 pub fn remove(&self, store: &mut dyn Storage, k: K, height: u64) -> StdResult<()> {
136 if self.snapshots.should_checkpoint(store, &k)? {
137 self.write_change(store, k.clone(), height)?;
138 }
139 self.primary.remove(store, k);
140 Ok(())
141 }
142
143 pub fn load(&self, store: &dyn Storage, k: K) -> StdResult<T> {
145 self.primary.load(store, k)
146 }
147
148 pub fn may_load(&self, store: &dyn Storage, k: K) -> StdResult<Option<T>> {
151 self.primary.may_load(store, k)
152 }
153
154 pub fn may_load_at_height(
155 &self,
156 store: &dyn Storage,
157 k: K,
158 height: u64,
159 ) -> StdResult<Option<T>> {
160 let snapshot = self
161 .snapshots
162 .may_load_at_height(store, k.clone(), height)?;
163
164 if let Some(r) = snapshot {
165 Ok(r)
166 } else {
167 self.may_load(store, k)
169 }
170 }
171
172 pub fn assert_checkpointed(&self, store: &dyn Storage, height: u64) -> StdResult<()> {
173 self.snapshots.assert_checkpointed(store, height)
174 }
175
176 pub fn update<A, E>(
183 &self,
184 store: &mut dyn Storage,
185 k: K,
186 height: u64,
187 action: A,
188 ) -> Result<T, E>
189 where
190 A: FnOnce(Option<T>) -> Result<T, E>,
191 E: From<StdError>,
192 {
193 let input = self.may_load(store, k.clone())?;
194 let output = action(input)?;
195 self.save(store, k, &output, height)?;
196 Ok(output)
197 }
198}
199
200impl<'a, K, T> SnapshotMap<K, T>
202where
203 T: Serialize + DeserializeOwned + Clone,
204 K: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize,
205{
206 pub fn range_raw<'c>(
209 &self,
210 store: &'c dyn Storage,
211 min: Option<Bound<'a, K>>,
212 max: Option<Bound<'a, K>>,
213 order: cosmwasm_std::Order,
214 ) -> Box<dyn Iterator<Item = StdResult<cosmwasm_std::Record<T>>> + 'c>
215 where
216 T: 'c,
217 {
218 self.no_prefix_raw().range_raw(store, min, max, order)
219 }
220
221 pub fn keys_raw<'c>(
222 &self,
223 store: &'c dyn Storage,
224 min: Option<Bound<'a, K>>,
225 max: Option<Bound<'a, K>>,
226 order: cosmwasm_std::Order,
227 ) -> Box<dyn Iterator<Item = Vec<u8>> + 'c>
228 where
229 T: 'c,
230 {
231 self.no_prefix_raw().keys_raw(store, min, max, order)
232 }
233}
234
235#[cfg(feature = "iterator")]
236impl<'a, K, T> SnapshotMap<K, T>
237where
238 T: Serialize + DeserializeOwned,
239 K: PrimaryKey<'a> + KeyDeserialize,
240{
241 pub fn prefix_range<'c>(
248 &self,
249 store: &'c dyn Storage,
250 min: Option<PrefixBound<'a, K::Prefix>>,
251 max: Option<PrefixBound<'a, K::Prefix>>,
252 order: cosmwasm_std::Order,
253 ) -> Box<dyn Iterator<Item = StdResult<(K::Output, T)>> + 'c>
254 where
255 T: 'c,
256 'a: 'c,
257 K: 'c,
258 K::Output: 'static,
259 {
260 let mapped =
261 namespaced_prefix_range(store, self.primary.namespace_bytes(), min, max, order)
262 .map(deserialize_kv::<K, T>);
263 Box::new(mapped)
264 }
265
266 pub fn range<'c>(
267 &self,
268 store: &'c dyn Storage,
269 min: Option<Bound<'a, K>>,
270 max: Option<Bound<'a, K>>,
271 order: cosmwasm_std::Order,
272 ) -> Box<dyn Iterator<Item = StdResult<(K::Output, T)>> + 'c>
273 where
274 T: 'c,
275 K::Output: 'static,
276 {
277 self.no_prefix().range(store, min, max, order)
278 }
279
280 pub fn keys<'c>(
281 &self,
282 store: &'c dyn Storage,
283 min: Option<Bound<'a, K>>,
284 max: Option<Bound<'a, K>>,
285 order: cosmwasm_std::Order,
286 ) -> Box<dyn Iterator<Item = StdResult<K::Output>> + 'c>
287 where
288 T: 'c,
289 K::Output: 'static,
290 {
291 self.no_prefix().keys(store, min, max, order)
292 }
293
294 pub fn prefix(&self, p: K::Prefix) -> Prefix<K::Suffix, T, K::Suffix> {
295 Prefix::new(self.primary.namespace_bytes(), &p.prefix())
296 }
297
298 pub fn sub_prefix(&self, p: K::SubPrefix) -> Prefix<K::SuperSuffix, T, K::SuperSuffix> {
299 Prefix::new(self.primary.namespace_bytes(), &p.prefix())
300 }
301
302 fn no_prefix(&self) -> Prefix<K, T, K> {
303 Prefix::new(self.primary.namespace_bytes(), &[])
304 }
305}
306
307#[cfg(test)]
308mod tests {
309 use super::*;
310 use cosmwasm_std::testing::MockStorage;
311
312 type TestMap = SnapshotMap<&'static str, u64>;
313 type TestMapCompositeKey = SnapshotMap<(&'static str, &'static str), u64>;
314
315 const NEVER: TestMap =
316 SnapshotMap::new("never", "never__check", "never__change", Strategy::Never);
317 const EVERY: TestMap = SnapshotMap::new(
318 "every",
319 "every__check",
320 "every__change",
321 Strategy::EveryBlock,
322 );
323 const EVERY_COMPOSITE_KEY: TestMapCompositeKey = SnapshotMap::new(
324 "every",
325 "every__check",
326 "every__change",
327 Strategy::EveryBlock,
328 );
329 const SELECT: TestMap = SnapshotMap::new(
330 "select",
331 "select__check",
332 "select__change",
333 Strategy::Selected,
334 );
335
336 fn init_data(map: &TestMap, storage: &mut dyn Storage) {
346 map.save(storage, "A", &5, 1).unwrap();
347 map.save(storage, "B", &7, 2).unwrap();
348
349 map.add_checkpoint(storage, 3).unwrap();
351
352 map.save(storage, "C", &1, 3).unwrap();
354 map.update(storage, "A", 3, |_| -> StdResult<u64> { Ok(8) })
355 .unwrap();
356
357 map.remove(storage, "B", 4).unwrap();
358 map.save(storage, "C", &13, 4).unwrap();
359
360 map.add_checkpoint(storage, 5).unwrap();
362 map.remove(storage, "A", 5).unwrap();
363 map.update(storage, "D", 5, |_| -> StdResult<u64> { Ok(22) })
364 .unwrap();
365 map.remove_checkpoint(storage, 5).unwrap();
367 }
368
369 const FINAL_VALUES: &[(&str, Option<u64>)] =
370 &[("A", None), ("B", None), ("C", Some(13)), ("D", Some(22))];
371
372 const VALUES_START_3: &[(&str, Option<u64>)] =
373 &[("A", Some(5)), ("B", Some(7)), ("C", None), ("D", None)];
374
375 const VALUES_START_5: &[(&str, Option<u64>)] =
376 &[("A", Some(8)), ("B", None), ("C", Some(13)), ("D", None)];
377
378 fn init_data_composite_key(map: &TestMapCompositeKey, storage: &mut dyn Storage) {
380 map.save(storage, ("A", "B"), &5, 1).unwrap();
381 map.save(storage, ("B", "A"), &7, 2).unwrap();
382
383 map.add_checkpoint(storage, 3).unwrap();
385
386 map.save(storage, ("B", "B"), &1, 3).unwrap();
388 map.update(storage, ("A", "B"), 3, |_| -> StdResult<u64> { Ok(8) })
389 .unwrap();
390
391 map.remove(storage, ("B", "A"), 4).unwrap();
392 map.save(storage, ("B", "B"), &13, 4).unwrap();
393
394 map.add_checkpoint(storage, 5).unwrap();
396 map.remove(storage, ("A", "B"), 5).unwrap();
397 map.update(storage, ("C", "A"), 5, |_| -> StdResult<u64> { Ok(22) })
398 .unwrap();
399 map.remove_checkpoint(storage, 5).unwrap();
401 }
402
403 fn assert_final_values(map: &TestMap, storage: &dyn Storage) {
404 for (k, v) in FINAL_VALUES.iter().cloned() {
405 assert_eq!(v, map.may_load(storage, k).unwrap());
406 }
407 }
408
409 fn assert_values_at_height(
410 map: &TestMap,
411 storage: &dyn Storage,
412 height: u64,
413 values: &[(&str, Option<u64>)],
414 ) {
415 for (k, v) in values.iter().cloned() {
416 assert_eq!(v, map.may_load_at_height(storage, k, height).unwrap());
417 }
418 }
419
420 fn assert_missing_checkpoint(map: &TestMap, storage: &dyn Storage, height: u64) {
421 for k in &["A", "B", "C", "D"] {
422 assert!(map.may_load_at_height(storage, *k, height).is_err());
423 }
424 }
425
426 #[test]
427 fn never_works_like_normal_map() {
428 let mut storage = MockStorage::new();
429 init_data(&NEVER, &mut storage);
430 assert_final_values(&NEVER, &storage);
431
432 assert_missing_checkpoint(&NEVER, &storage, 3);
434 assert_missing_checkpoint(&NEVER, &storage, 5);
435 }
436
437 #[test]
438 fn every_blocks_stores_present_and_past() {
439 let mut storage = MockStorage::new();
440 init_data(&EVERY, &mut storage);
441 assert_final_values(&EVERY, &storage);
442
443 assert_values_at_height(&EVERY, &storage, 3, VALUES_START_3);
445 assert_values_at_height(&EVERY, &storage, 5, VALUES_START_5);
446 }
447
448 #[test]
449 fn selected_shows_3_not_5() {
450 let mut storage = MockStorage::new();
451 init_data(&SELECT, &mut storage);
452 assert_final_values(&SELECT, &storage);
453
454 assert_values_at_height(&SELECT, &storage, 3, VALUES_START_3);
456 assert_missing_checkpoint(&NEVER, &storage, 1);
458 assert_missing_checkpoint(&NEVER, &storage, 5);
460 }
461
462 #[test]
463 fn handle_multiple_writes_in_one_block() {
464 let mut storage = MockStorage::new();
465
466 println!("SETUP");
467 EVERY.save(&mut storage, "A", &5, 1).unwrap();
468 EVERY.save(&mut storage, "B", &7, 2).unwrap();
469 EVERY.save(&mut storage, "C", &2, 2).unwrap();
470
471 EVERY
473 .update(&mut storage, "A", 3, |_| -> StdResult<u64> { Ok(9) })
474 .unwrap();
475 EVERY.save(&mut storage, "A", &12, 3).unwrap();
476 assert_eq!(Some(5), EVERY.may_load_at_height(&storage, "A", 2).unwrap());
477 assert_eq!(Some(5), EVERY.may_load_at_height(&storage, "A", 3).unwrap());
478 assert_eq!(
479 Some(12),
480 EVERY.may_load_at_height(&storage, "A", 4).unwrap()
481 );
482
483 EVERY.save(&mut storage, "B", &17, 4).unwrap();
485 EVERY.remove(&mut storage, "B", 4).unwrap();
486 assert_eq!(Some(7), EVERY.may_load_at_height(&storage, "B", 3).unwrap());
487 assert_eq!(Some(7), EVERY.may_load_at_height(&storage, "B", 4).unwrap());
488 assert_eq!(None, EVERY.may_load_at_height(&storage, "B", 5).unwrap());
489
490 EVERY.remove(&mut storage, "C", 5).unwrap();
492 EVERY
493 .update(&mut storage, "C", 5, |_| -> StdResult<u64> { Ok(16) })
494 .unwrap();
495 assert_eq!(Some(2), EVERY.may_load_at_height(&storage, "C", 4).unwrap());
496 assert_eq!(Some(2), EVERY.may_load_at_height(&storage, "C", 5).unwrap());
497 assert_eq!(
498 Some(16),
499 EVERY.may_load_at_height(&storage, "C", 6).unwrap()
500 );
501 }
502
503 #[test]
504 #[cfg(feature = "iterator")]
505 fn changelog_range_works() {
506 use cosmwasm_std::Order;
507
508 let mut store = MockStorage::new();
509
510 EVERY.save(&mut store, "A", &5, 1).unwrap();
512 EVERY.save(&mut store, "B", &7, 2).unwrap();
513 EVERY
514 .update(&mut store, "A", 3, |_| -> StdResult<u64> { Ok(8) })
515 .unwrap();
516 EVERY.remove(&mut store, "B", 4).unwrap();
517
518 let all: StdResult<Vec<_>> = EVERY
520 .changelog()
521 .range(&store, None, None, Order::Ascending)
522 .collect();
523 let all = all.unwrap();
524 assert_eq!(4, all.len());
525 assert_eq!(
526 all,
527 vec![
528 (("A".into(), 1), ChangeSet { old: None }),
529 (("A".into(), 3), ChangeSet { old: Some(5) }),
530 (("B".into(), 2), ChangeSet { old: None }),
531 (("B".into(), 4), ChangeSet { old: Some(7) })
532 ]
533 );
534
535 let all: StdResult<Vec<_>> = EVERY
537 .changelog()
538 .prefix("B")
539 .range(&store, None, None, Order::Ascending)
540 .collect();
541 let all = all.unwrap();
542 assert_eq!(2, all.len());
543 assert_eq!(
544 all,
545 vec![
546 (2, ChangeSet { old: None }),
547 (4, ChangeSet { old: Some(7) })
548 ]
549 );
550
551 let all: StdResult<Vec<_>> = EVERY
553 .changelog()
554 .prefix("A")
555 .range(&store, Some(Bound::inclusive(3u64)), None, Order::Ascending)
556 .collect();
557 let all = all.unwrap();
558 assert_eq!(1, all.len());
559 assert_eq!(all, vec![(3, ChangeSet { old: Some(5) }),]);
560 }
561
562 #[test]
563 #[cfg(feature = "iterator")]
564 fn range_simple_string_key() {
565 use cosmwasm_std::Order;
566
567 let mut store = MockStorage::new();
568 init_data(&EVERY, &mut store);
569
570 let all: StdResult<Vec<_>> = EVERY.range(&store, None, None, Order::Ascending).collect();
572 let all = all.unwrap();
573 assert_eq!(2, all.len());
574 assert_eq!(all, vec![("C".into(), 13), ("D".into(), 22)]);
575
576 let all: StdResult<Vec<_>> = EVERY
578 .range(&store, Some(Bound::inclusive("C")), None, Order::Ascending)
579 .collect();
580 let all = all.unwrap();
581 assert_eq!(2, all.len());
582 assert_eq!(all, vec![("C".into(), 13), ("D".into(), 22)]);
583
584 let all: StdResult<Vec<_>> = EVERY
586 .range(&store, Some(Bound::inclusive("D")), None, Order::Ascending)
587 .collect();
588 let all = all.unwrap();
589 assert_eq!(1, all.len());
590 assert_eq!(all, vec![("D".into(), 22)]);
591 }
592
593 #[test]
594 #[cfg(feature = "iterator")]
595 fn range_composite_key() {
596 use cosmwasm_std::Order;
597
598 let mut store = MockStorage::new();
599 init_data_composite_key(&EVERY_COMPOSITE_KEY, &mut store);
600
601 let all: StdResult<Vec<_>> = EVERY_COMPOSITE_KEY
603 .range(&store, None, None, Order::Ascending)
604 .collect();
605 let all = all.unwrap();
606 assert_eq!(2, all.len());
607 assert_eq!(
608 all,
609 vec![
610 (("B".into(), "B".into()), 13),
611 (("C".into(), "A".into()), 22)
612 ]
613 );
614 }
615
616 #[test]
617 #[cfg(feature = "iterator")]
618 fn prefix_range_composite_key() {
619 use cosmwasm_std::Order;
620
621 let mut store = MockStorage::new();
622 init_data_composite_key(&EVERY_COMPOSITE_KEY, &mut store);
623
624 let all: StdResult<Vec<_>> = EVERY_COMPOSITE_KEY
626 .prefix_range(
627 &store,
628 None,
629 Some(PrefixBound::exclusive("C")),
630 Order::Descending,
631 )
632 .collect();
633 let all = all.unwrap();
634 assert_eq!(1, all.len());
635 assert_eq!(all, vec![(("B".into(), "B".into()), 13)]);
636 }
637
638 #[test]
639 #[cfg(feature = "iterator")]
640 fn prefix_composite_key() {
641 use cosmwasm_std::Order;
642
643 let mut store = MockStorage::new();
644 init_data_composite_key(&EVERY_COMPOSITE_KEY, &mut store);
645
646 let all: StdResult<Vec<_>> = EVERY_COMPOSITE_KEY
648 .prefix("C")
649 .range(&store, None, None, Order::Ascending)
650 .collect();
651 let all = all.unwrap();
652 assert_eq!(1, all.len());
653 assert_eq!(all, vec![("A".into(), 22),]);
654 }
655
656 #[test]
657 #[cfg(feature = "iterator")]
658 fn sub_prefix_composite_key() {
659 use cosmwasm_std::Order;
660
661 let mut store = MockStorage::new();
662 init_data_composite_key(&EVERY_COMPOSITE_KEY, &mut store);
663
664 let all: StdResult<Vec<_>> = EVERY_COMPOSITE_KEY
668 .sub_prefix(())
669 .range(&store, None, None, Order::Ascending)
670 .collect();
671 let all = all.unwrap();
672 assert_eq!(2, all.len());
673 assert_eq!(
674 all,
675 vec![
676 (("B".into(), "B".into()), 13),
677 (("C".into(), "A".into()), 22)
678 ]
679 );
680 }
681}