cw_storage_plus/
keys.rs

1use cosmwasm_std::{storage_keys::namespace_with_key, Addr};
2use cosmwasm_std::{Int128, Int64, Uint128, Uint64};
3
4use crate::de::KeyDeserialize;
5use crate::int_key::IntKey;
6
7#[derive(Debug)]
8pub enum Key<'a> {
9    Ref(&'a [u8]),
10    Val8([u8; 1]),
11    Val16([u8; 2]),
12    Val32([u8; 4]),
13    Val64([u8; 8]),
14    Val128([u8; 16]),
15}
16
17impl AsRef<[u8]> for Key<'_> {
18    fn as_ref(&self) -> &[u8] {
19        match self {
20            Key::Ref(r) => r,
21            Key::Val8(v) => v,
22            Key::Val16(v) => v,
23            Key::Val32(v) => v,
24            Key::Val64(v) => v,
25            Key::Val128(v) => v,
26        }
27    }
28}
29
30impl PartialEq<&[u8]> for Key<'_> {
31    fn eq(&self, other: &&[u8]) -> bool {
32        self.as_ref() == *other
33    }
34}
35
36/// `PrimaryKey` needs to be implemented for types that want to be a `Map` (or `Map`-like) key,
37/// or part of a key.
38///
39/// In particular, it defines a series of types that help iterating over parts of a (composite) key:
40///
41/// `Prefix`: Prefix is eager. That is, except for empty keys, it's always "one less" than the full key.
42/// `Suffix`: Suffix is the complement of prefix.
43/// `SubPrefix`: Sub-prefix is "one less" than prefix.
44/// `SuperSuffix`: Super-suffix is "one more" than suffix. The complement of sub-prefix.
45///
46/// By example, for a 2-tuple `(T, U)`:
47///
48/// `T`: Prefix.
49/// `U`: Suffix.
50/// `()`: Sub-prefix.
51/// `(T, U)`: Super-suffix.
52///
53/// `SubPrefix` and `SuperSuffix` only make real sense in the case of triples. Still, they need to be
54/// consistently defined for all types.
55pub trait PrimaryKey<'a>: Clone {
56    /// These associated types need to implement `Prefixer`, so that they can be useful arguments
57    /// for `prefix()`, `sub_prefix()`, and their key-deserializable variants.
58    type Prefix: Prefixer<'a>;
59    type SubPrefix: Prefixer<'a>;
60
61    /// These associated types need to implement `KeyDeserialize`, so that they can be returned from
62    /// `range_de()` and friends.
63    type Suffix: KeyDeserialize;
64    type SuperSuffix: KeyDeserialize;
65
66    /// returns a slice of key steps, which can be optionally combined
67    fn key(&self) -> Vec<Key<'_>>;
68
69    fn joined_key(&self) -> Vec<u8> {
70        let keys = self.key();
71        let l = keys.len();
72        namespace_with_key(
73            &keys[0..l - 1].iter().map(Key::as_ref).collect::<Vec<_>>(),
74            keys[l - 1].as_ref(),
75        )
76    }
77
78    fn joined_extra_key(&self, key: &[u8]) -> Vec<u8> {
79        let keys = self.key();
80        namespace_with_key(&keys.iter().map(Key::as_ref).collect::<Vec<_>>(), key)
81    }
82}
83
84// Empty / no primary key
85impl PrimaryKey<'_> for () {
86    type Prefix = Self;
87    type SubPrefix = Self;
88    type Suffix = Self;
89    type SuperSuffix = Self;
90
91    fn key(&self) -> Vec<Key<'_>> {
92        vec![]
93    }
94}
95
96impl<'a> PrimaryKey<'a> for &'a [u8] {
97    type Prefix = ();
98    type SubPrefix = ();
99    type Suffix = Self;
100    type SuperSuffix = Self;
101
102    fn key(&self) -> Vec<Key<'_>> {
103        // this is simple, we don't add more prefixes
104        vec![Key::Ref(self)]
105    }
106}
107
108impl<const N: usize> PrimaryKey<'_> for [u8; N] {
109    type Prefix = ();
110    type SubPrefix = ();
111    type Suffix = Self;
112    type SuperSuffix = Self;
113
114    fn key(&self) -> Vec<Key<'_>> {
115        // this is simple, we don't add more prefixes
116        vec![Key::Ref(self.as_slice())]
117    }
118}
119
120// Provide a string version of this to raw encode strings
121impl<'a> PrimaryKey<'a> for &'a str {
122    type Prefix = ();
123    type SubPrefix = ();
124    type Suffix = Self;
125    type SuperSuffix = Self;
126
127    fn key(&self) -> Vec<Key<'_>> {
128        // this is simple, we don't add more prefixes
129        vec![Key::Ref(self.as_bytes())]
130    }
131}
132
133// use generics for combining there - so we can use &[u8], Vec<u8>, or IntKey
134impl<'a, T: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize, U: PrimaryKey<'a> + KeyDeserialize>
135    PrimaryKey<'a> for (T, U)
136{
137    type Prefix = T;
138    type SubPrefix = ();
139    type Suffix = U;
140    type SuperSuffix = Self;
141
142    fn key(&self) -> Vec<Key<'_>> {
143        let mut keys = self.0.key();
144        keys.extend(self.1.key());
145        keys
146    }
147}
148
149// implements PrimaryKey for all &T where T implements PrimaryKey.
150impl<'a, T> PrimaryKey<'a> for &'a T
151where
152    T: PrimaryKey<'a>,
153{
154    type Prefix = <T as PrimaryKey<'a>>::Prefix;
155    type SubPrefix = <T as PrimaryKey<'a>>::SubPrefix;
156    type Suffix = T::Suffix;
157    type SuperSuffix = T::SuperSuffix;
158
159    fn key(&self) -> Vec<Key<'_>> {
160        <T as PrimaryKey<'a>>::key(self)
161    }
162}
163
164// use generics for combining there - so we can use &[u8], Vec<u8>, or IntKey
165impl<
166        'a,
167        T: PrimaryKey<'a> + Prefixer<'a>,
168        U: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize,
169        V: PrimaryKey<'a> + KeyDeserialize,
170    > PrimaryKey<'a> for (T, U, V)
171{
172    type Prefix = (T, U);
173    type SubPrefix = T;
174    type Suffix = V;
175    type SuperSuffix = (U, V);
176
177    fn key(&self) -> Vec<Key<'_>> {
178        let mut keys = self.0.key();
179        keys.extend(self.1.key());
180        keys.extend(self.2.key());
181        keys
182    }
183}
184
185pub trait Prefixer<'a> {
186    /// returns 0 or more namespaces that should be length-prefixed and concatenated for range searches
187    fn prefix(&self) -> Vec<Key<'_>>;
188
189    fn joined_prefix(&self) -> Vec<u8> {
190        let prefixes = self.prefix();
191        namespace_with_key(&prefixes.iter().map(Key::as_ref).collect::<Vec<_>>(), &[])
192    }
193}
194
195impl Prefixer<'_> for () {
196    fn prefix(&self) -> Vec<Key<'_>> {
197        vec![]
198    }
199}
200
201impl<'a> Prefixer<'a> for &'a [u8] {
202    fn prefix(&self) -> Vec<Key<'_>> {
203        vec![Key::Ref(self)]
204    }
205}
206
207impl<'a, T: Prefixer<'a>, U: Prefixer<'a>> Prefixer<'a> for (T, U) {
208    fn prefix(&self) -> Vec<Key<'_>> {
209        let mut res = self.0.prefix();
210        res.extend(self.1.prefix());
211        res
212    }
213}
214
215impl<'a, T: Prefixer<'a>, U: Prefixer<'a>, V: Prefixer<'a>> Prefixer<'a> for (T, U, V) {
216    fn prefix(&self) -> Vec<Key<'_>> {
217        let mut res = self.0.prefix();
218        res.extend(self.1.prefix());
219        res.extend(self.2.prefix());
220        res
221    }
222}
223
224impl<'a, T> Prefixer<'a> for &'a T
225where
226    T: Prefixer<'a>,
227{
228    fn prefix(&self) -> Vec<Key<'_>> {
229        <T as Prefixer<'a>>::prefix(self)
230    }
231}
232
233// Provide a string version of this to raw encode strings
234impl<'a> Prefixer<'a> for &'a str {
235    fn prefix(&self) -> Vec<Key<'_>> {
236        vec![Key::Ref(self.as_bytes())]
237    }
238}
239
240impl PrimaryKey<'_> for Vec<u8> {
241    type Prefix = ();
242    type SubPrefix = ();
243    type Suffix = Self;
244    type SuperSuffix = Self;
245
246    fn key(&self) -> Vec<Key<'_>> {
247        vec![Key::Ref(self)]
248    }
249}
250
251impl Prefixer<'_> for Vec<u8> {
252    fn prefix(&self) -> Vec<Key<'_>> {
253        vec![Key::Ref(self.as_ref())]
254    }
255}
256
257impl PrimaryKey<'_> for String {
258    type Prefix = ();
259    type SubPrefix = ();
260    type Suffix = Self;
261    type SuperSuffix = Self;
262
263    fn key(&self) -> Vec<Key<'_>> {
264        vec![Key::Ref(self.as_bytes())]
265    }
266}
267
268impl Prefixer<'_> for String {
269    fn prefix(&self) -> Vec<Key<'_>> {
270        vec![Key::Ref(self.as_bytes())]
271    }
272}
273
274/// owned variant.
275impl PrimaryKey<'_> for Addr {
276    type Prefix = ();
277    type SubPrefix = ();
278    type Suffix = Self;
279    type SuperSuffix = Self;
280
281    fn key(&self) -> Vec<Key<'_>> {
282        // this is simple, we don't add more prefixes
283        vec![Key::Ref(self.as_bytes())]
284    }
285}
286
287impl Prefixer<'_> for Addr {
288    fn prefix(&self) -> Vec<Key<'_>> {
289        vec![Key::Ref(self.as_bytes())]
290    }
291}
292
293macro_rules! integer_key {
294    (for $($t:ty, $v:tt),+) => {
295        $(impl<'a> PrimaryKey<'a> for $t {
296            type Prefix = ();
297            type SubPrefix = ();
298            type Suffix = Self;
299            type SuperSuffix = Self;
300
301            fn key(&self) -> Vec<Key<'_>> {
302                vec![Key::$v(self.to_cw_bytes())]
303            }
304        })*
305    }
306}
307
308integer_key!(for i8, Val8, u8, Val8, i16, Val16, u16, Val16, i32, Val32, u32, Val32, i64, Val64, u64, Val64, i128, Val128, u128, Val128, Uint64, Val64, Uint128, Val128, Int64, Val64, Int128, Val128);
309macro_rules! integer_prefix {
310    (for $($t:ty, $v:tt),+) => {
311        $(impl<'a> Prefixer<'a> for $t {
312            fn prefix(&self) -> Vec<Key<'_>> {
313                vec![Key::$v(self.to_cw_bytes())]
314            }
315        })*
316    }
317}
318
319integer_prefix!(for i8, Val8, u8, Val8, i16, Val16, u16, Val16, i32, Val32, u32, Val32, i64, Val64, u64, Val64, i128, Val128, u128, Val128, Uint64, Val64, Uint128, Val128, Int64, Val64, Int128, Val128);
320
321#[cfg(test)]
322mod test {
323    use cosmwasm_std::{Uint256, Uint512};
324
325    use super::*;
326
327    #[test]
328    fn naked_8key_works() {
329        let k: u8 = 42u8;
330        let path = k.key();
331        assert_eq!(1, path.len());
332        assert_eq!(42u8.to_cw_bytes(), path[0].as_ref());
333
334        let k: i8 = 42i8;
335        let path = k.key();
336        assert_eq!(1, path.len());
337        assert_eq!(42i8.to_cw_bytes(), path[0].as_ref());
338    }
339
340    #[test]
341    fn naked_16key_works() {
342        let k: u16 = 4242u16;
343        let path = k.key();
344        assert_eq!(1, path.len());
345        assert_eq!(4242u16.to_cw_bytes(), path[0].as_ref());
346
347        let k: i16 = 4242i16;
348        let path = k.key();
349        assert_eq!(1, path.len());
350        assert_eq!(4242i16.to_cw_bytes(), path[0].as_ref());
351    }
352
353    #[test]
354    fn naked_32key_works() {
355        let k: u32 = 4242u32;
356        let path = k.key();
357        assert_eq!(1, path.len());
358        assert_eq!(4242u32.to_cw_bytes(), path[0].as_ref());
359
360        let k: i32 = 4242i32;
361        let path = k.key();
362        assert_eq!(1, path.len());
363        assert_eq!(4242i32.to_cw_bytes(), path[0].as_ref());
364    }
365
366    #[test]
367    fn naked_64key_works() {
368        let k: u64 = 4242u64;
369        let path = k.key();
370        assert_eq!(1, path.len());
371        assert_eq!(4242u64.to_cw_bytes(), path[0].as_ref());
372
373        let k: i64 = 4242i64;
374        let path = k.key();
375        assert_eq!(1, path.len());
376        assert_eq!(4242i64.to_cw_bytes(), path[0].as_ref());
377    }
378
379    #[test]
380    fn naked_128key_works() {
381        let k: u128 = 4242u128;
382        let path = k.key();
383        assert_eq!(1, path.len());
384        assert_eq!(4242u128.to_cw_bytes(), path[0].as_ref());
385
386        let k: i128 = 4242i128;
387        let path = k.key();
388        assert_eq!(1, path.len());
389        assert_eq!(4242i128.to_cw_bytes(), path[0].as_ref());
390    }
391
392    #[test]
393    fn std_uint64_key_works() {
394        let k: Uint64 = Uint64::from(4242u64);
395        let path = k.key();
396        assert_eq!(1, path.len());
397        assert_eq!(4242u64.to_cw_bytes(), path[0].as_ref());
398    }
399
400    #[test]
401    fn std_uint128_key_works() {
402        let k: Uint128 = Uint128::from(4242u128);
403        let path = k.key();
404        assert_eq!(1, path.len());
405        assert_eq!(4242u128.to_cw_bytes(), path[0].as_ref());
406    }
407
408    #[test]
409    fn std_int64_key_works() {
410        let k: Int64 = Int64::from(-4242i64);
411        let path = k.key();
412        assert_eq!(1, path.len());
413        assert_eq!((-4242i64).to_cw_bytes(), path[0].as_ref());
414    }
415
416    #[test]
417    fn std_int128_key_works() {
418        let k: Int128 = Int128::from(-4242i128);
419        let path = k.key();
420        assert_eq!(1, path.len());
421        assert_eq!((-4242i128).to_cw_bytes(), path[0].as_ref());
422    }
423
424    #[test]
425    fn str_key_works() {
426        type K<'a> = &'a str;
427
428        let k: K = "hello";
429        let path = k.key();
430        assert_eq!(1, path.len());
431        assert_eq!(b"hello", path[0].as_ref());
432
433        let joined = k.joined_key();
434        assert_eq!(joined, b"hello")
435    }
436
437    #[test]
438    fn string_key_works() {
439        type K = String;
440
441        let k: K = "hello".to_string();
442        let path = k.key();
443        assert_eq!(1, path.len());
444        assert_eq!(b"hello", path[0].as_ref());
445
446        let joined = k.joined_key();
447        assert_eq!(joined, b"hello")
448    }
449
450    #[test]
451    fn fixed_size_bytes_key_works() {
452        type K = [u8; 32];
453
454        let k: K = Uint256::MAX.to_be_bytes();
455        let path = k.key();
456        assert_eq!(1, path.len());
457        assert_eq!(k, path[0].as_ref());
458
459        let joined = k.joined_key();
460        assert_eq!(joined, k);
461
462        // ref also works
463        type K2<'a> = &'a [u8; 64];
464
465        let k: K2 = &Uint512::MAX.to_be_bytes();
466        let path = k.key();
467        assert_eq!(1, path.len());
468        assert_eq!(k, path[0].as_ref());
469
470        let joined = k.joined_key();
471        assert_eq!(joined, k);
472    }
473
474    #[test]
475    fn nested_str_key_works() {
476        type K<'a> = (&'a str, &'a [u8]);
477
478        let k: K = ("hello", b"world");
479        let path = k.key();
480        assert_eq!(2, path.len());
481        assert_eq!(b"hello", path[0].as_ref());
482        assert_eq!(b"world", path[1].as_ref());
483    }
484
485    #[test]
486    fn composite_byte_key() {
487        let k: (&[u8], &[u8]) = ("foo".as_bytes(), b"bar");
488        let path = k.key();
489        assert_eq!(2, path.len());
490        assert_eq!(path, vec!["foo".as_bytes(), b"bar"],);
491    }
492
493    #[test]
494    fn naked_composite_int_key() {
495        let k: (u32, u64) = (123, 87654);
496        let path = k.key();
497        assert_eq!(2, path.len());
498        assert_eq!(4, path[0].as_ref().len());
499        assert_eq!(8, path[1].as_ref().len());
500        assert_eq!(path[0].as_ref(), 123u32.to_cw_bytes());
501        assert_eq!(path[1].as_ref(), 87654u64.to_cw_bytes());
502    }
503
504    #[test]
505    fn nested_composite_keys() {
506        // use this to ensure proper type-casts below
507        let first: &[u8] = b"foo";
508        // this function tests how well the generics extend to "edge cases"
509        let k: ((&[u8], &[u8]), &[u8]) = ((first, b"bar"), b"zoom");
510        let path = k.key();
511        assert_eq!(3, path.len());
512        assert_eq!(path, vec![first, b"bar", b"zoom"]);
513
514        // ensure prefix also works
515        let dir = k.0.prefix();
516        assert_eq!(2, dir.len());
517        assert_eq!(dir, vec![first, b"bar"]);
518    }
519
520    #[test]
521    fn naked_8bit_prefixes() {
522        let pair: (u8, &[u8]) = (123, b"random");
523        let one: Vec<u8> = vec![123];
524        let two: Vec<u8> = b"random".to_vec();
525        assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
526
527        let pair: (i8, &[u8]) = (123, b"random");
528        // Signed int keys are "sign-flipped"
529        let one: Vec<u8> = vec![123 ^ 0x80];
530        let two: Vec<u8> = b"random".to_vec();
531        assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
532    }
533
534    #[test]
535    fn naked_16bit_prefixes() {
536        let pair: (u16, &[u8]) = (12345, b"random");
537        let one: Vec<u8> = vec![48, 57];
538        let two: Vec<u8> = b"random".to_vec();
539        assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
540
541        let pair: (i16, &[u8]) = (12345, b"random");
542        // Signed int keys are "sign-flipped"
543        let one: Vec<u8> = vec![48 ^ 0x80, 57];
544        let two: Vec<u8> = b"random".to_vec();
545        assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
546    }
547
548    #[test]
549    fn naked_64bit_prefixes() {
550        let pair: (u64, &[u8]) = (12345, b"random");
551        let one: Vec<u8> = vec![0, 0, 0, 0, 0, 0, 48, 57];
552        let two: Vec<u8> = b"random".to_vec();
553        assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
554
555        let pair: (i64, &[u8]) = (12345, b"random");
556        // Signed int keys are "sign-flipped"
557        #[allow(clippy::identity_op)]
558        let one: Vec<u8> = vec![0 ^ 0x80, 0, 0, 0, 0, 0, 48, 57];
559        let two: Vec<u8> = b"random".to_vec();
560        assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
561    }
562
563    #[test]
564    fn naked_proper_prefixes() {
565        let pair: (u32, &[u8]) = (12345, b"random");
566        let one: Vec<u8> = vec![0, 0, 48, 57];
567        let two: Vec<u8> = b"random".to_vec();
568        assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
569
570        let triple: (&str, u32, &[u8]) = ("begin", 12345, b"end");
571        let one: Vec<u8> = b"begin".to_vec();
572        let two: Vec<u8> = vec![0, 0, 48, 57];
573        let three: Vec<u8> = b"end".to_vec();
574        assert_eq!(
575            triple.prefix(),
576            vec![one.as_slice(), two.as_slice(), three.as_slice()]
577        );
578
579        // same works with owned variants (&str -> String, &[u8] -> Vec<u8>)
580        let owned_triple: (String, u32, Vec<u8>) = ("begin".to_string(), 12345, b"end".to_vec());
581        assert_eq!(
582            owned_triple.prefix(),
583            vec![one.as_slice(), two.as_slice(), three.as_slice()]
584        );
585    }
586}
cw_storage_plus/keys.rs

cw_storage_plus/
keys.rs
