abstract_core/objects/pool/
unique_pool_id.rs

1use 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// Need hash for ans scraper
12#[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}