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#[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 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#[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#[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}