cw_storey/containers/
key_set.rs

1use cosmwasm_std::{Addr, Int128, Int256, Int512, Int64, Uint128, Uint256, Uint512, Uint64};
2use storey::containers::map::key::{DynamicKey, FixedSizeKey, NumericKeyDecodeError};
3use storey::containers::map::{Key, OwnedKey};
4
5/// The CosmWasm key set for use with storey's [`Map`](storey::containers::Map).
6///
7/// This key set includes the usual standard library types (like `u32` or `String`) as well as `cosmwasm_std` types (like `Addr` and `Uint128`).
8///
9/// For more information about key sets, take a look at the [`storey::containers::map::Key`] trait.
10pub struct CwKeySet;
11
12impl Key<CwKeySet> for Addr {
13    type Kind = DynamicKey;
14
15    fn encode(&self) -> Vec<u8> {
16        self.as_str().as_bytes().to_vec()
17    }
18}
19
20impl OwnedKey<CwKeySet> for Addr {
21    type Error = <String as OwnedKey>::Error;
22
23    fn from_bytes(bytes: &[u8]) -> Result<Self, Self::Error>
24    where
25        Self: Sized,
26    {
27        <String as OwnedKey>::from_bytes(bytes).map(Addr::unchecked)
28    }
29}
30
31macro_rules! cosmwasm_std_uints1 {
32    ($($ty:ty => $size:expr, $stdty:ty),*) => {
33        $(
34            impl Key<CwKeySet> for $ty {
35                type Kind = FixedSizeKey<$size>;
36
37                fn encode(&self) -> Vec<u8> {
38                    self.to_be_bytes().to_vec()
39                }
40            }
41
42            impl OwnedKey<CwKeySet> for $ty {
43                type Error = NumericKeyDecodeError;
44
45                fn from_bytes(bytes: &[u8]) -> Result<Self, Self::Error>
46                where
47                    Self: Sized,
48                {
49                    let array: [u8; $size] = bytes.try_into().map_err(|_| NumericKeyDecodeError::InvalidLength)?;
50                    Ok(<$stdty>::from_be_bytes(array).into())
51                }
52            }
53        )*
54    }
55}
56
57cosmwasm_std_uints1!(
58    Uint64 => 8, u64,
59    Uint128 => 16, u128
60);
61
62macro_rules! cosmwasm_std_uints2 {
63    ($($ty:ty => $size:expr),*) => {
64        $(
65            impl Key<CwKeySet> for $ty {
66                type Kind = FixedSizeKey<$size>;
67
68                fn encode(&self) -> Vec<u8> {
69                    self.to_be_bytes().to_vec()
70                }
71            }
72
73            impl OwnedKey<CwKeySet> for $ty {
74                type Error = NumericKeyDecodeError;
75
76                fn from_bytes(bytes: &[u8]) -> Result<Self, Self::Error>
77                where
78                    Self: Sized,
79                {
80                    let array: [u8; $size] = bytes.try_into().map_err(|_| NumericKeyDecodeError::InvalidLength)?;
81                    Ok(<$ty>::from_be_bytes(array))
82                }
83            }
84        )*
85    }
86}
87
88cosmwasm_std_uints2!(
89    Uint256 => 32,
90    Uint512 => 64
91);
92
93macro_rules! cosmwasm_std_ints {
94    ($($ty:ty => $size:expr),*) => {
95        $(
96            impl Key<CwKeySet> for $ty {
97                type Kind = FixedSizeKey<$size>;
98
99                fn encode(&self) -> Vec<u8> {
100                    let mut bytes = self.to_be_bytes();
101                    bytes[0] ^= 0x80;
102
103                    bytes.to_vec()
104                }
105            }
106
107            impl OwnedKey<CwKeySet> for $ty {
108                type Error = NumericKeyDecodeError;
109
110                fn from_bytes(bytes: &[u8]) -> Result<Self, Self::Error>
111                where
112                    Self: Sized,
113                {
114                    let mut array: [u8; $size] = bytes.try_into().map_err(|_| NumericKeyDecodeError::InvalidLength)?;
115                    array[0] ^= 0x80;
116
117                    Ok(<$ty>::from_be_bytes(array))
118                }
119            }
120        )*
121    }
122}
123
124cosmwasm_std_ints!(Int64 => 8, Int128 => 16, Int256 => 32, Int512 => 64);
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn unsigned_ints_1() {
132        let test_vector = [
133            (Uint64::from(0u64), [0, 0, 0, 0, 0, 0, 0, 0]),
134            (Uint64::from(1u64), [0, 0, 0, 0, 0, 0, 0, 1]),
135            (
136                Uint64::from(0x1234567890abcdefu64),
137                [0x12, 0x34, 0x56, 0x78, 0x90, 0xab, 0xcd, 0xef],
138            ),
139        ];
140
141        for (num, expected) in test_vector.iter() {
142            let encoded = num.encode();
143            assert_eq!(encoded, *expected);
144        }
145
146        for (expected, bytes) in test_vector.iter() {
147            let decoded = Uint64::from_bytes(bytes).unwrap();
148            assert_eq!(decoded, *expected);
149        }
150    }
151
152    #[test]
153    fn unsigned_ints_2() {
154        let test_vector = [
155            (Uint256::from(0u64), [0; 32]),
156            (
157                Uint256::new([
158                    0x12, 0x34, 0x56, 0x78, 0x90, 0xab, 0xcd, 0xef, 0x12, 0x34, 0x56, 0x78, 0x90,
159                    0xab, 0xcd, 0xef, 0x12, 0x34, 0x56, 0x78, 0x90, 0xab, 0xcd, 0xef, 0x12, 0x34,
160                    0x56, 0x78, 0x90, 0xab, 0xcd, 0xef,
161                ]),
162                [
163                    0x12, 0x34, 0x56, 0x78, 0x90, 0xab, 0xcd, 0xef, 0x12, 0x34, 0x56, 0x78, 0x90,
164                    0xab, 0xcd, 0xef, 0x12, 0x34, 0x56, 0x78, 0x90, 0xab, 0xcd, 0xef, 0x12, 0x34,
165                    0x56, 0x78, 0x90, 0xab, 0xcd, 0xef,
166                ],
167            ),
168        ];
169
170        for (num, expected) in test_vector.iter() {
171            let encoded = num.encode();
172            assert_eq!(encoded, *expected);
173        }
174
175        for (expected, bytes) in test_vector.iter() {
176            let decoded = Uint256::from_bytes(bytes).unwrap();
177            assert_eq!(decoded, *expected);
178        }
179    }
180
181    #[test]
182    fn signed_ints() {
183        let nums = [
184            Int256::from(-542),
185            Int256::from(-111),
186            Int256::from(0),
187            Int256::from(121),
188            Int256::from(342),
189        ];
190
191        let mut byte_nums = nums.iter().map(|n| n.encode()).collect::<Vec<_>>();
192        byte_nums.sort();
193
194        let result = byte_nums
195            .iter()
196            .map(|bytes| Int256::from_bytes(bytes).unwrap())
197            .collect::<Vec<_>>();
198
199        assert_eq!(result, nums);
200    }
201}