abstract_os/objects/
contract_entry.rs

1use crate::constants::ATTRIBUTE_DELIMITER;
2use cosmwasm_std::{StdError, StdResult};
3use cw_storage_plus::{Key, KeyDeserialize, Prefixer, PrimaryKey};
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6use std::{
7    convert::{TryFrom, TryInto},
8    fmt::Display,
9};
10
11/// Key to get the Address of a contract
12#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq, JsonSchema, PartialOrd, Ord)]
13pub struct UncheckedContractEntry {
14    pub protocol: String,
15    pub contract: String,
16}
17
18impl UncheckedContractEntry {
19    pub fn new<T: ToString, R: ToString>(protocol: T, contract: R) -> Self {
20        Self {
21            protocol: protocol.to_string(),
22            contract: contract.to_string(),
23        }
24    }
25    pub fn check(self) -> ContractEntry {
26        ContractEntry {
27            contract: self.contract.to_ascii_lowercase(),
28            protocol: self.protocol.to_ascii_lowercase(),
29        }
30    }
31}
32
33impl From<ContractEntry> for UncheckedContractEntry {
34    fn from(contract_entry: ContractEntry) -> Self {
35        Self {
36            protocol: contract_entry.protocol,
37            contract: contract_entry.contract,
38        }
39    }
40}
41
42impl TryFrom<String> for UncheckedContractEntry {
43    type Error = StdError;
44    /// Try from a string like "protocol:contract_name"
45    fn try_from(entry: String) -> Result<Self, Self::Error> {
46        let composite: Vec<&str> = entry.split(ATTRIBUTE_DELIMITER).collect();
47        if composite.len() != 2 {
48            return Err(StdError::generic_err(
49                "contract entry should be formatted as \"protocol:contract_name\".",
50            ));
51        }
52        Ok(Self::new(composite[0], composite[1]))
53    }
54}
55
56/// Key to get the Address of a contract
57/// Use [`UncheckedContractEntry`] to construct this type.  
58#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, JsonSchema, Eq, PartialOrd, Ord)]
59pub struct ContractEntry {
60    pub protocol: String,
61    pub contract: String,
62}
63
64impl From<UncheckedContractEntry> for ContractEntry {
65    fn from(entry: UncheckedContractEntry) -> Self {
66        entry.check()
67    }
68}
69
70impl Display for ContractEntry {
71    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72        write!(f, "{}:{}", self.protocol, self.contract)
73    }
74}
75
76impl<'a> PrimaryKey<'a> for &ContractEntry {
77    type Prefix = String;
78
79    type SubPrefix = ();
80
81    type Suffix = String;
82
83    type SuperSuffix = Self;
84
85    fn key(&self) -> Vec<cw_storage_plus::Key> {
86        let mut keys = self.protocol.key();
87        keys.extend(self.contract.key());
88        keys
89    }
90}
91
92impl<'a> Prefixer<'a> for &ContractEntry {
93    fn prefix(&self) -> Vec<Key> {
94        let mut res = self.protocol.prefix();
95        res.extend(self.contract.prefix().into_iter());
96        res
97    }
98}
99
100impl KeyDeserialize for &ContractEntry {
101    type Output = ContractEntry;
102
103    #[inline(always)]
104    fn from_vec(mut value: Vec<u8>) -> StdResult<Self::Output> {
105        let mut tu = value.split_off(2);
106        let t_len = parse_length(&value)?;
107        let u = tu.split_off(t_len);
108
109        Ok(ContractEntry {
110            protocol: String::from_vec(tu)?,
111            contract: String::from_vec(u)?,
112        })
113    }
114}
115
116#[inline(always)]
117fn parse_length(value: &[u8]) -> StdResult<usize> {
118    Ok(u16::from_be_bytes(
119        value
120            .try_into()
121            .map_err(|_| StdError::generic_err("Could not read 2 byte length"))?,
122    )
123    .into())
124}
125
126//--------------------------------------------------------------------------------------------------
127// Tests
128//--------------------------------------------------------------------------------------------------
129
130#[cfg(test)]
131mod test {
132    use super::*;
133    use cosmwasm_std::{testing::mock_dependencies, Addr, Order};
134    use cw_storage_plus::Map;
135
136    mod key {
137        use super::*;
138
139        fn mock_key() -> ContractEntry {
140            ContractEntry {
141                protocol: "abstract".to_string(),
142                contract: "rocket-ship".to_string(),
143            }
144        }
145
146        fn mock_keys() -> (ContractEntry, ContractEntry, ContractEntry) {
147            (
148                ContractEntry {
149                    protocol: "abstract".to_string(),
150                    contract: "sailing-ship".to_string(),
151                },
152                ContractEntry {
153                    protocol: "abstract".to_string(),
154                    contract: "rocket-ship".to_string(),
155                },
156                ContractEntry {
157                    protocol: "shitcoin".to_string(),
158                    contract: "pump'n dump".to_string(),
159                },
160            )
161        }
162
163        #[test]
164        fn storage_key_works() {
165            let mut deps = mock_dependencies();
166            let key = mock_key();
167            let map: Map<&ContractEntry, u64> = Map::new("map");
168
169            map.save(deps.as_mut().storage, &key, &42069).unwrap();
170
171            assert_eq!(map.load(deps.as_ref().storage, &key).unwrap(), 42069);
172
173            let items = map
174                .range(deps.as_ref().storage, None, None, Order::Ascending)
175                .map(|item| item.unwrap())
176                .collect::<Vec<_>>();
177
178            assert_eq!(items.len(), 1);
179            assert_eq!(items[0], (key, 42069));
180        }
181
182        #[test]
183        fn composite_key_works() {
184            let mut deps = mock_dependencies();
185            let key = mock_key();
186            let map: Map<(&ContractEntry, Addr), u64> = Map::new("map");
187
188            map.save(
189                deps.as_mut().storage,
190                (&key, Addr::unchecked("larry")),
191                &42069,
192            )
193            .unwrap();
194
195            map.save(
196                deps.as_mut().storage,
197                (&key, Addr::unchecked("jake")),
198                &69420,
199            )
200            .unwrap();
201
202            let items = map
203                .prefix(&key)
204                .range(deps.as_ref().storage, None, None, Order::Ascending)
205                .map(|item| item.unwrap())
206                .collect::<Vec<_>>();
207
208            assert_eq!(items.len(), 2);
209            assert_eq!(items[0], (Addr::unchecked("jake"), 69420));
210            assert_eq!(items[1], (Addr::unchecked("larry"), 42069));
211        }
212
213        #[test]
214        fn partial_key_works() {
215            let mut deps = mock_dependencies();
216            let (key1, key2, key3) = mock_keys();
217            let map: Map<&ContractEntry, u64> = Map::new("map");
218
219            map.save(deps.as_mut().storage, &key1, &42069).unwrap();
220
221            map.save(deps.as_mut().storage, &key2, &69420).unwrap();
222
223            map.save(deps.as_mut().storage, &key3, &999).unwrap();
224
225            let items = map
226                .prefix("abstract".to_string())
227                .range(deps.as_ref().storage, None, None, Order::Ascending)
228                .map(|item| item.unwrap())
229                .collect::<Vec<_>>();
230
231            assert_eq!(items.len(), 2);
232            assert_eq!(items[0], ("rocket-ship".to_string(), 69420));
233            assert_eq!(items[1], ("sailing-ship".to_string(), 42069));
234        }
235    }
236}