cw_storage_plus/
prefix.rs

1#![cfg(feature = "iterator")]
2use core::fmt;
3use cosmwasm_std::storage_keys::to_length_prefixed_nested;
4use serde::de::DeserializeOwned;
5use serde::Serialize;
6use std::fmt::Debug;
7use std::marker::PhantomData;
8
9use cosmwasm_std::{Order, Record, StdResult, Storage};
10use std::ops::Deref;
11
12use crate::bound::{PrefixBound, RawBound};
13use crate::de::KeyDeserialize;
14use crate::iter_helpers::{concat, deserialize_kv, deserialize_v, trim};
15use crate::keys::Key;
16use crate::{Bound, Prefixer, PrimaryKey};
17
18#[derive(Clone)]
19pub struct Prefix<K, T, B = Vec<u8>>
20where
21    K: KeyDeserialize,
22    T: Serialize + DeserializeOwned,
23{
24    /// all namespaces prefixes and concatenated with the key
25    pub(crate) storage_prefix: Vec<u8>,
26    // see https://doc.rust-lang.org/std/marker/struct.PhantomData.html#unused-type-parameters for why this is needed
27    pub(crate) data: PhantomData<(T, K, B)>,
28}
29
30impl<K, T> Debug for Prefix<K, T>
31where
32    K: KeyDeserialize,
33    T: Serialize + DeserializeOwned,
34{
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        f.debug_struct("Prefix")
37            .field("storage_prefix", &self.storage_prefix)
38            .finish_non_exhaustive()
39    }
40}
41
42impl<K, T> Deref for Prefix<K, T>
43where
44    K: KeyDeserialize,
45    T: Serialize + DeserializeOwned,
46{
47    type Target = [u8];
48
49    fn deref(&self) -> &[u8] {
50        &self.storage_prefix
51    }
52}
53
54impl<K, T, B> Prefix<K, T, B>
55where
56    K: KeyDeserialize,
57    T: Serialize + DeserializeOwned,
58{
59    pub fn new(top_name: &[u8], sub_names: &[Key]) -> Self {
60        let calculated_len = 1 + sub_names.len();
61        let mut combined: Vec<&[u8]> = Vec::with_capacity(calculated_len);
62        combined.push(top_name);
63        combined.extend(sub_names.iter().map(|sub_name| sub_name.as_ref()));
64        debug_assert_eq!(calculated_len, combined.len()); // as long as we calculate correctly, we don't need to reallocate
65        let storage_prefix = to_length_prefixed_nested(&combined);
66        Prefix {
67            storage_prefix,
68            data: PhantomData,
69        }
70    }
71}
72
73impl<'b, K, T, B> Prefix<K, T, B>
74where
75    B: PrimaryKey<'b>,
76    K: KeyDeserialize,
77    T: Serialize + DeserializeOwned,
78{
79    pub fn range_raw<'a>(
80        &self,
81        store: &'a dyn Storage,
82        min: Option<Bound<'b, B>>,
83        max: Option<Bound<'b, B>>,
84        order: Order,
85    ) -> Box<dyn Iterator<Item = StdResult<Record<T>>> + 'a>
86    where
87        T: 'a,
88    {
89        let mapped = range_with_prefix(
90            store,
91            &self.storage_prefix,
92            min.map(|b| b.to_raw_bound()),
93            max.map(|b| b.to_raw_bound()),
94            order,
95        )
96        .map(deserialize_v);
97        Box::new(mapped)
98    }
99
100    pub fn keys_raw<'a>(
101        &self,
102        store: &'a dyn Storage,
103        min: Option<Bound<'b, B>>,
104        max: Option<Bound<'b, B>>,
105        order: Order,
106    ) -> Box<dyn Iterator<Item = Vec<u8>> + 'a> {
107        keys_with_prefix(
108            store,
109            &self.storage_prefix,
110            min.map(|b| b.to_raw_bound()),
111            max.map(|b| b.to_raw_bound()),
112            order,
113        )
114    }
115
116    /// Clears the prefix, removing the first `limit` elements (or all if `limit == None`).
117    pub fn clear(&self, store: &mut dyn Storage, limit: Option<usize>) {
118        const TAKE: usize = 10;
119        let mut cleared = false;
120
121        let mut left_to_clear = limit.unwrap_or(usize::MAX);
122
123        while !cleared {
124            // Take just TAKE elements to prevent possible heap overflow if the prefix is big,
125            // but don't take more than we want to clear.
126            let take = TAKE.min(left_to_clear);
127
128            let paths = keys_full(store, &self.storage_prefix, None, None, Order::Ascending)
129                .take(take)
130                .collect::<Vec<_>>();
131
132            for path in &paths {
133                store.remove(path);
134            }
135            left_to_clear -= paths.len();
136
137            cleared = paths.len() < take || left_to_clear == 0;
138        }
139    }
140
141    /// Returns `true` if the prefix is empty.
142    pub fn is_empty(&self, store: &dyn Storage) -> bool {
143        keys_full(store, &self.storage_prefix, None, None, Order::Ascending)
144            .next()
145            .is_none()
146    }
147
148    pub fn range<'a>(
149        &self,
150        store: &'a dyn Storage,
151        min: Option<Bound<'b, B>>,
152        max: Option<Bound<'b, B>>,
153        order: Order,
154    ) -> Box<dyn Iterator<Item = StdResult<(K::Output, T)>> + 'a>
155    where
156        T: 'a,
157        K::Output: 'static,
158    {
159        let mapped = range_with_prefix(
160            store,
161            &self.storage_prefix,
162            min.map(|b| b.to_raw_bound()),
163            max.map(|b| b.to_raw_bound()),
164            order,
165        )
166        .map(|kv| deserialize_kv::<K, T>(kv));
167        Box::new(mapped)
168    }
169
170    pub fn keys<'a>(
171        &self,
172        store: &'a dyn Storage,
173        min: Option<Bound<'b, B>>,
174        max: Option<Bound<'b, B>>,
175        order: Order,
176    ) -> Box<dyn Iterator<Item = StdResult<K::Output>> + 'a>
177    where
178        T: 'a,
179        K::Output: 'static,
180    {
181        let mapped = keys_with_prefix(
182            store,
183            &self.storage_prefix,
184            min.map(|b| b.to_raw_bound()),
185            max.map(|b| b.to_raw_bound()),
186            order,
187        )
188        .map(|k| K::from_vec(k));
189        Box::new(mapped)
190    }
191}
192
193/// Returns an iterator through all records in storage with the given prefix and
194/// within the given bounds, yielding the key without prefix and value.
195pub fn range_with_prefix<'a>(
196    storage: &'a dyn Storage,
197    namespace: &[u8],
198    start: Option<RawBound>,
199    end: Option<RawBound>,
200    order: Order,
201) -> Box<dyn Iterator<Item = Record> + 'a> {
202    // make a copy for the closure to handle lifetimes safely
203    let prefix = namespace.to_vec();
204    let mapped =
205        range_full(storage, namespace, start, end, order).map(move |(k, v)| (trim(&prefix, &k), v));
206    Box::new(mapped)
207}
208
209/// Returns an iterator through all keys in storage with the given prefix and
210/// within the given bounds, yielding the key without the prefix.
211pub fn keys_with_prefix<'a>(
212    storage: &'a dyn Storage,
213    namespace: &[u8],
214    start: Option<RawBound>,
215    end: Option<RawBound>,
216    order: Order,
217) -> Box<dyn Iterator<Item = Vec<u8>> + 'a> {
218    // make a copy for the closure to handle lifetimes safely
219    let prefix = namespace.to_vec();
220    let mapped = keys_full(storage, namespace, start, end, order).map(move |k| trim(&prefix, &k));
221    Box::new(mapped)
222}
223
224/// Returns an iterator through all records in storage within the given bounds,
225/// yielding the full key (including the prefix) and value.
226pub(crate) fn range_full<'a>(
227    store: &'a dyn Storage,
228    namespace: &[u8],
229    start: Option<RawBound>,
230    end: Option<RawBound>,
231    order: Order,
232) -> impl Iterator<Item = Record> + 'a {
233    let start = calc_start_bound(namespace, start);
234    let end = calc_end_bound(namespace, end);
235
236    // get iterator from storage
237    store.range(Some(&start), Some(&end), order)
238}
239
240/// Returns an iterator through all keys in storage within the given bounds,
241/// yielding the full key including the prefix.
242pub(crate) fn keys_full<'a>(
243    store: &'a dyn Storage,
244    namespace: &[u8],
245    start: Option<RawBound>,
246    end: Option<RawBound>,
247    order: Order,
248) -> impl Iterator<Item = Vec<u8>> + 'a {
249    let start = calc_start_bound(namespace, start);
250    let end = calc_end_bound(namespace, end);
251
252    // get iterator from storage
253    store.range_keys(Some(&start), Some(&end), order)
254}
255
256fn calc_start_bound(namespace: &[u8], bound: Option<RawBound>) -> Vec<u8> {
257    match bound {
258        None => namespace.to_vec(),
259        // this is the natural limits of the underlying Storage
260        Some(RawBound::Inclusive(limit)) => concat(namespace, &limit),
261        Some(RawBound::Exclusive(limit)) => concat(namespace, &extend_one_byte(&limit)),
262    }
263}
264
265fn calc_end_bound(namespace: &[u8], bound: Option<RawBound>) -> Vec<u8> {
266    match bound {
267        None => increment_last_byte(namespace),
268        // this is the natural limits of the underlying Storage
269        Some(RawBound::Exclusive(limit)) => concat(namespace, &limit),
270        Some(RawBound::Inclusive(limit)) => concat(namespace, &extend_one_byte(&limit)),
271    }
272}
273
274pub fn namespaced_prefix_range<'a, 'c, K: Prefixer<'a>>(
275    storage: &'c dyn Storage,
276    namespace: &[u8],
277    start: Option<PrefixBound<'a, K>>,
278    end: Option<PrefixBound<'a, K>>,
279    order: Order,
280) -> Box<dyn Iterator<Item = Record> + 'c> {
281    let prefix = to_length_prefixed_nested(&[namespace]);
282    let start = calc_prefix_start_bound(&prefix, start);
283    let end = calc_prefix_end_bound(&prefix, end);
284
285    // get iterator from storage
286    let base_iterator = storage.range(Some(&start), Some(&end), order);
287
288    // make a copy for the closure to handle lifetimes safely
289    let mapped = base_iterator.map(move |(k, v)| (trim(&prefix, &k), v));
290    Box::new(mapped)
291}
292
293fn calc_prefix_start_bound<'a, K: Prefixer<'a>>(
294    namespace: &[u8],
295    bound: Option<PrefixBound<'a, K>>,
296) -> Vec<u8> {
297    match bound.map(|b| b.to_raw_bound()) {
298        None => namespace.to_vec(),
299        // this is the natural limits of the underlying Storage
300        Some(RawBound::Inclusive(limit)) => concat(namespace, &limit),
301        Some(RawBound::Exclusive(limit)) => concat(namespace, &increment_last_byte(&limit)),
302    }
303}
304
305fn calc_prefix_end_bound<'a, K: Prefixer<'a>>(
306    namespace: &[u8],
307    bound: Option<PrefixBound<'a, K>>,
308) -> Vec<u8> {
309    match bound.map(|b| b.to_raw_bound()) {
310        None => increment_last_byte(namespace),
311        // this is the natural limits of the underlying Storage
312        Some(RawBound::Exclusive(limit)) => concat(namespace, &limit),
313        Some(RawBound::Inclusive(limit)) => concat(namespace, &increment_last_byte(&limit)),
314    }
315}
316
317pub(crate) fn extend_one_byte(limit: &[u8]) -> Vec<u8> {
318    let mut v = limit.to_vec();
319    v.push(0);
320    v
321}
322
323/// Returns a new vec of same length and last byte incremented by one
324/// If last bytes are 255, we handle overflow up the chain.
325/// If all bytes are 255, this returns wrong data - but that is never possible as a namespace
326fn increment_last_byte(input: &[u8]) -> Vec<u8> {
327    let mut copy = input.to_vec();
328    // zero out all trailing 255, increment first that is not such
329    for i in (0..input.len()).rev() {
330        if copy[i] == 255 {
331            copy[i] = 0;
332        } else {
333            copy[i] += 1;
334            break;
335        }
336    }
337    copy
338}
339
340#[cfg(test)]
341mod test {
342    use super::*;
343    use cosmwasm_std::testing::MockStorage;
344
345    #[test]
346    fn ensure_proper_range_bounds() {
347        let mut store = MockStorage::new();
348        // manually create this - not testing nested prefixes here
349        let prefix: Prefix<Vec<u8>, u64> = Prefix {
350            storage_prefix: b"foo".to_vec(),
351            data: PhantomData,
352        };
353
354        // set some data, we care about "foo" prefix
355        store.set(b"foobar", b"1");
356        store.set(b"foora", b"2");
357        store.set(b"foozi", b"3");
358        // these shouldn't match
359        store.set(b"foply", b"100");
360        store.set(b"font", b"200");
361
362        let expected = vec![
363            (b"bar".to_vec(), 1u64),
364            (b"ra".to_vec(), 2u64),
365            (b"zi".to_vec(), 3u64),
366        ];
367        let expected_reversed: Vec<(Vec<u8>, u64)> = expected.iter().rev().cloned().collect();
368
369        // let's do the basic sanity check
370        let res: StdResult<Vec<_>> = prefix
371            .range_raw(&store, None, None, Order::Ascending)
372            .collect();
373        assert_eq!(&expected, &res.unwrap());
374        let res: StdResult<Vec<_>> = prefix
375            .range_raw(&store, None, None, Order::Descending)
376            .collect();
377        assert_eq!(&expected_reversed, &res.unwrap());
378
379        // now let's check some ascending ranges
380        let res: StdResult<Vec<_>> = prefix
381            .range_raw(
382                &store,
383                Some(Bound::inclusive(b"ra".to_vec())),
384                None,
385                Order::Ascending,
386            )
387            .collect();
388        assert_eq!(&expected[1..], res.unwrap().as_slice());
389        // skip excluded
390        let res: StdResult<Vec<_>> = prefix
391            .range_raw(
392                &store,
393                Some(Bound::exclusive(b"ra".to_vec())),
394                None,
395                Order::Ascending,
396            )
397            .collect();
398        assert_eq!(&expected[2..], res.unwrap().as_slice());
399        // if we exclude something a little lower, we get matched
400        let res: StdResult<Vec<_>> = prefix
401            .range_raw(
402                &store,
403                Some(Bound::exclusive(b"r".to_vec())),
404                None,
405                Order::Ascending,
406            )
407            .collect();
408        assert_eq!(&expected[1..], res.unwrap().as_slice());
409
410        // now let's check some descending ranges
411        let res: StdResult<Vec<_>> = prefix
412            .range_raw(
413                &store,
414                None,
415                Some(Bound::inclusive(b"ra".to_vec())),
416                Order::Descending,
417            )
418            .collect();
419        assert_eq!(&expected_reversed[1..], res.unwrap().as_slice());
420        // skip excluded
421        let res: StdResult<Vec<_>> = prefix
422            .range_raw(
423                &store,
424                None,
425                Some(Bound::exclusive(b"ra".to_vec())),
426                Order::Descending,
427            )
428            .collect();
429        assert_eq!(&expected_reversed[2..], res.unwrap().as_slice());
430        // if we exclude something a little higher, we get matched
431        let res: StdResult<Vec<_>> = prefix
432            .range_raw(
433                &store,
434                None,
435                Some(Bound::exclusive(b"rb".to_vec())),
436                Order::Descending,
437            )
438            .collect();
439        assert_eq!(&expected_reversed[1..], res.unwrap().as_slice());
440
441        // now test when both sides are set
442        let res: StdResult<Vec<_>> = prefix
443            .range_raw(
444                &store,
445                Some(Bound::inclusive(b"ra".to_vec())),
446                Some(Bound::exclusive(b"zi".to_vec())),
447                Order::Ascending,
448            )
449            .collect();
450        assert_eq!(&expected[1..2], res.unwrap().as_slice());
451        // and descending
452        let res: StdResult<Vec<_>> = prefix
453            .range_raw(
454                &store,
455                Some(Bound::inclusive(b"ra".to_vec())),
456                Some(Bound::exclusive(b"zi".to_vec())),
457                Order::Descending,
458            )
459            .collect();
460        assert_eq!(&expected[1..2], res.unwrap().as_slice());
461        // Include both sides
462        let res: StdResult<Vec<_>> = prefix
463            .range_raw(
464                &store,
465                Some(Bound::inclusive(b"ra".to_vec())),
466                Some(Bound::inclusive(b"zi".to_vec())),
467                Order::Descending,
468            )
469            .collect();
470        assert_eq!(&expected_reversed[..2], res.unwrap().as_slice());
471        // Exclude both sides
472        let res: StdResult<Vec<_>> = prefix
473            .range_raw(
474                &store,
475                Some(Bound::exclusive(b"ra".to_vec())),
476                Some(Bound::exclusive(b"zi".to_vec())),
477                Order::Ascending,
478            )
479            .collect();
480        assert_eq!(res.unwrap().as_slice(), &[]);
481    }
482
483    #[test]
484    fn prefix_debug() {
485        let prefix: Prefix<String, String> = Prefix::new(b"lol", &[Key::Val8([8; 1])]);
486        assert_eq!(
487            format!("{:?}", prefix),
488            "Prefix { storage_prefix: [0, 3, 108, 111, 108, 0, 1, 8], .. }"
489        );
490    }
491
492    #[test]
493    fn prefix_clear_limited() {
494        let mut store = MockStorage::new();
495        // manually create this - not testing nested prefixes here
496        let prefix: Prefix<Vec<u8>, u64> = Prefix {
497            storage_prefix: b"foo".to_vec(),
498            data: PhantomData,
499        };
500
501        // set some data, we care about "foo" prefix
502        for i in 0..100u32 {
503            store.set(format!("foo{}", i).as_bytes(), b"1");
504        }
505
506        // clearing less than `TAKE` should work
507        prefix.clear(&mut store, Some(1));
508        assert_eq!(
509            prefix.range(&store, None, None, Order::Ascending).count(),
510            99
511        );
512
513        // clearing more than `TAKE` should work
514        prefix.clear(&mut store, Some(12));
515        assert_eq!(
516            prefix.range(&store, None, None, Order::Ascending).count(),
517            99 - 12
518        );
519
520        // clearing an exact multiple of `TAKE` should work
521        prefix.clear(&mut store, Some(20));
522        assert_eq!(
523            prefix.range(&store, None, None, Order::Ascending).count(),
524            99 - 12 - 20
525        );
526
527        // clearing more than available should work
528        prefix.clear(&mut store, Some(1000));
529        assert_eq!(
530            prefix.range(&store, None, None, Order::Ascending).count(),
531            0
532        );
533    }
534
535    #[test]
536    fn prefix_clear_unlimited() {
537        let mut store = MockStorage::new();
538        // manually create this - not testing nested prefixes here
539        let prefix: Prefix<Vec<u8>, u64> = Prefix {
540            storage_prefix: b"foo".to_vec(),
541            data: PhantomData,
542        };
543
544        // set some data, we care about "foo" prefix
545        for i in 0..1000u32 {
546            store.set(format!("foo{}", i).as_bytes(), b"1");
547        }
548
549        // clearing all should work
550        prefix.clear(&mut store, None);
551        assert_eq!(
552            prefix.range(&store, None, None, Order::Ascending).count(),
553            0
554        );
555
556        // set less data
557        for i in 0..5u32 {
558            store.set(format!("foo{}", i).as_bytes(), b"1");
559        }
560
561        // clearing all should work
562        prefix.clear(&mut store, None);
563        assert_eq!(
564            prefix.range(&store, None, None, Order::Ascending).count(),
565            0
566        );
567    }
568
569    #[test]
570    fn is_empty_works() {
571        // manually create this - not testing nested prefixes here
572        let prefix: Prefix<Vec<u8>, u64> = Prefix {
573            storage_prefix: b"foo".to_vec(),
574            data: PhantomData,
575        };
576
577        let mut storage = MockStorage::new();
578
579        assert!(prefix.is_empty(&storage));
580
581        storage.set(b"fookey1", b"1");
582        storage.set(b"fookey2", b"2");
583
584        assert!(!prefix.is_empty(&storage));
585    }
586
587    #[test]
588    fn keys_raw_works() {
589        // manually create this - not testing nested prefixes here
590        let prefix: Prefix<Vec<u8>, u64> = Prefix {
591            storage_prefix: b"foo".to_vec(),
592            data: PhantomData,
593        };
594
595        let mut storage = MockStorage::new();
596        storage.set(b"fookey1", b"1");
597        storage.set(b"fookey2", b"2");
598
599        let keys: Vec<_> = prefix
600            .keys_raw(&storage, None, None, Order::Ascending)
601            .collect();
602        assert_eq!(keys, vec![b"key1", b"key2"]);
603
604        let keys: Vec<_> = prefix
605            .keys_raw(
606                &storage,
607                Some(Bound::exclusive("key1")),
608                None,
609                Order::Ascending,
610            )
611            .collect();
612        assert_eq!(keys, vec![b"key2"]);
613    }
614}
cw_storage_plus/prefix.rs

cw_storage_plus/
prefix.rs
