cw_storage_plus/
keys.rs

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

cw_storage_plus/
keys.rs
