abstract_core/objects/pool/
unique_pool_id.rs1use std::{array::TryFromSliceError, convert::TryInto, fmt::Display};
2
3use cosmwasm_std::{StdError, StdResult};
4use cw_storage_plus::{IntKey, KeyDeserialize, Prefixer, PrimaryKey};
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7
8#[derive(
9 Deserialize, Serialize, Clone, Debug, PartialEq, Eq, JsonSchema, PartialOrd, Ord, Copy,
10)]
11#[cfg_attr(not(target_arch = "wasm32"), derive(Hash))]
13pub struct UniquePoolId(u64);
14
15impl UniquePoolId {
16 pub const fn new(id: u64) -> Self {
17 Self(id)
18 }
19 pub fn as_u64(&self) -> u64 {
20 self.0
21 }
22 pub fn increment(&mut self) {
23 self.0 += 1;
24 }
25}
26
27impl From<u64> for UniquePoolId {
28 fn from(id: u64) -> Self {
29 Self::new(id)
30 }
31}
32
33impl Display for UniquePoolId {
34 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35 write!(f, "{}", self.0)
36 }
37}
38
39impl<'a> PrimaryKey<'a> for UniquePoolId {
40 type Prefix = ();
41 type SubPrefix = ();
42 type Suffix = Self;
43 type SuperSuffix = Self;
44
45 fn key(&self) -> Vec<cw_storage_plus::Key> {
46 self.0.key()
47 }
48}
49
50impl<'a> Prefixer<'a> for UniquePoolId {
51 fn prefix(&self) -> Vec<cw_storage_plus::Key> {
52 self.0.prefix()
53 }
54}
55
56impl KeyDeserialize for UniquePoolId {
57 type Output = Self;
58 #[inline(always)]
59 fn from_vec(value: Vec<u8>) -> StdResult<Self::Output> {
60 Ok(Self::from_cw_bytes(value.as_slice().try_into().map_err(
61 |err: TryFromSliceError| StdError::generic_err(err.to_string()),
62 )?))
63 }
64}
65
66impl IntKey for UniquePoolId {
67 type Buf = [u8; std::mem::size_of::<u64>()];
68
69 #[inline]
70 fn to_cw_bytes(&self) -> Self::Buf {
71 self.0.to_be_bytes()
72 }
73
74 #[inline]
75 fn from_cw_bytes(bytes: Self::Buf) -> Self {
76 Self(u64::from_be_bytes(bytes))
77 }
78}
79
80#[cfg(test)]
81mod test {
82 use cosmwasm_std::{testing::mock_dependencies, Addr, Order};
83 use cw_storage_plus::Map;
84
85 use super::*;
86
87 fn mock_key() -> UniquePoolId {
88 UniquePoolId::new(1)
89 }
90
91 fn _mock_keys() -> (UniquePoolId, UniquePoolId, UniquePoolId) {
92 (
93 UniquePoolId::new(1),
94 UniquePoolId::new(2),
95 UniquePoolId::new(3),
96 )
97 }
98
99 #[test]
100 fn storage_key_works() {
101 let mut deps = mock_dependencies();
102 let key = mock_key();
103 let map: Map<UniquePoolId, u64> = Map::new("map");
104
105 map.save(deps.as_mut().storage, key, &42069).unwrap();
106
107 assert_eq!(map.load(deps.as_ref().storage, key).unwrap(), 42069);
108
109 let items = map
110 .range(deps.as_ref().storage, None, None, Order::Ascending)
111 .map(|item| item.unwrap())
112 .collect::<Vec<_>>();
113
114 assert_eq!(items.len(), 1);
115 assert_eq!(items[0], (key, 42069));
116 }
117
118 #[test]
119 fn composite_key_works() {
120 let mut deps = mock_dependencies();
121 let key = mock_key();
122 let map: Map<(UniquePoolId, Addr), u64> = Map::new("map");
123
124 map.save(
125 deps.as_mut().storage,
126 (key, Addr::unchecked("larry")),
127 &42069,
128 )
129 .unwrap();
130
131 map.save(
132 deps.as_mut().storage,
133 (key, Addr::unchecked("jake")),
134 &69420,
135 )
136 .unwrap();
137
138 let items = map
139 .prefix(key)
140 .range(deps.as_ref().storage, None, None, Order::Ascending)
141 .map(|item| item.unwrap())
142 .collect::<Vec<_>>();
143
144 assert_eq!(items.len(), 2);
145 assert_eq!(items[0], (Addr::unchecked("jake"), 69420));
146 assert_eq!(items[1], (Addr::unchecked("larry"), 42069));
147 }
148
149 #[test]
150 fn naked_64key_works() {
151 let k: UniquePoolId = 4242u64.into();
152 let path = k.key();
153 assert_eq!(1, path.len());
154 assert_eq!(4242u64.to_cw_bytes(), path[0].as_ref());
155 }
156}