1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

use crate::arithmetic::{OverflowingAdd, OverflowingSub};
use crate::casper_types::{
    bytesrepr::{FromBytes, ToBytes},
    CLTyped
};
use crate::module::{ModuleComponent, ModulePrimitive};
use crate::prelude::*;
use crate::{
    module::{Module, SubModule},
    var::Var,
    ContractEnv, UnwrapOrRevert
};
use core::fmt::Debug;

/// Data structure for storing key-value pairs.
pub struct Mapping<K, V> {
    parent_env: Rc<ContractEnv>,
    phantom: core::marker::PhantomData<(K, V)>,
    index: u8
}

impl<K: ToBytes, V> ModuleComponent for Mapping<K, V> {
    /// Creates a new instance of `Mapping` with the given environment and index.
    fn instance(env: Rc<ContractEnv>, index: u8) -> Self {
        Self {
            parent_env: env,
            phantom: core::marker::PhantomData,
            index
        }
    }
}

impl<K: ToBytes, V> ModulePrimitive for Mapping<K, V> {}

impl<K: ToBytes, V> Mapping<K, V> {
    fn env_for_key(&self, key: &K) -> ContractEnv {
        let mut env = (*self.parent_env).clone();
        let key = key.to_bytes().unwrap_or_default();
        env.add_to_mapping_data(&key);
        env
    }
}

impl<K: ToBytes, V: FromBytes + CLTyped> Mapping<K, V> {
    /// Retrieves the value associated with the given key.
    ///
    /// Returns an `Option<V>` representing the value associated with the key, or `None` if the key is not found.
    pub fn get(&self, key: &K) -> Option<V> {
        let env = self.env_for_key(key);
        Var::<V>::instance(Rc::new(env), self.index).get()
    }
}

impl<K: ToBytes, V: FromBytes + CLTyped + Default> Mapping<K, V> {
    /// Retrieves the value associated with the given key from the mapping.
    /// If the key does not exist, returns the default value of type `V`.
    pub fn get_or_default(&self, key: &K) -> V {
        let env = self.env_for_key(key);
        Var::<V>::instance(Rc::new(env), self.index).get_or_default()
    }
}

impl<K: ToBytes, V: ToBytes + CLTyped> Mapping<K, V> {
    /// Sets the value associated with the given key in the mapping.
    pub fn set(&mut self, key: &K, value: V) {
        let env = self.env_for_key(key);
        Var::<V>::instance(Rc::new(env), self.index).set(value)
    }
}

impl<K: ToBytes, V: Module> Mapping<K, V> {
    /// Retrieves the module associated with the given key.
    ///
    /// A [`SubModule`] instance containing the module associated with the key.
    pub fn module(&self, key: &K) -> SubModule<V> {
        let env = self.env_for_key(key);
        SubModule::instance(Rc::new(env), self.index)
    }
}

impl<K: ToBytes, V: ToBytes + FromBytes + CLTyped + OverflowingAdd + Default> Mapping<K, V> {
    /// Utility function that gets the current value and adds the passed `value`
    /// and sets the new value to the storage.
    ///
    /// If the operation fails due to overflow, the currently executing contract reverts.
    pub fn add(&mut self, key: &K, value: V) {
        let current_value = self.get_or_default(key);
        let new_value = current_value
            .overflowing_add(value)
            .unwrap_or_revert(&self.env_for_key(key));
        self.set(key, new_value);
    }
}

impl<
        K: ToBytes,
        V: ToBytes + FromBytes + CLTyped + OverflowingSub + Default + Debug + PartialOrd
    > Mapping<K, V>
{
    /// Utility function that gets the current value and subtracts the passed `value`
    /// and sets the new value to the storage.
    ///
    /// If the operation fails due to overflow, the currently executing contract reverts.
    pub fn subtract(&mut self, key: &K, value: V) {
        let current_value = self.get_or_default(key);
        let new_value = current_value
            .overflowing_sub(value)
            .unwrap_or_revert(&self.env_for_key(key));
        self.set(key, new_value);
    }
}