cw_storage_plus/indexes/
prefix.rs

1#![cfg(feature = "iterator")]
2use core::fmt;
3use serde::de::DeserializeOwned;
4use serde::Serialize;
5use std::fmt::Debug;
6
7use cosmwasm_std::{Order, Record, StdResult, Storage};
8use std::ops::Deref;
9
10use crate::de::KeyDeserialize;
11use crate::iter_helpers::{deserialize_kv, deserialize_v};
12use crate::keys::Key;
13use crate::{Bound, PrimaryKey};
14
15type DeserializeVFn<T> = fn(&dyn Storage, &[u8], Record) -> StdResult<Record<T>>;
16
17type DeserializeKvFn<K, T> =
18    fn(&dyn Storage, &[u8], Record) -> StdResult<(<K as KeyDeserialize>::Output, T)>;
19
20pub fn default_deserializer_v<T: DeserializeOwned>(
21    _: &dyn Storage,
22    _: &[u8],
23    raw: Record,
24) -> StdResult<Record<T>> {
25    deserialize_v(raw)
26}
27
28pub fn default_deserializer_kv<K: KeyDeserialize, T: DeserializeOwned>(
29    _: &dyn Storage,
30    _: &[u8],
31    raw: Record,
32) -> StdResult<(K::Output, T)> {
33    deserialize_kv::<K, T>(raw)
34}
35
36#[derive(Clone)]
37pub struct IndexPrefix<K, T, B = Vec<u8>>
38where
39    K: KeyDeserialize,
40    T: Serialize + DeserializeOwned,
41{
42    inner: crate::prefix::Prefix<K, T, B>,
43    pk_name: Vec<u8>,
44    de_fn_kv: DeserializeKvFn<K, T>,
45    de_fn_v: DeserializeVFn<T>,
46}
47
48impl<K, T> Debug for IndexPrefix<K, T>
49where
50    K: KeyDeserialize,
51    T: Serialize + DeserializeOwned,
52{
53    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
54        f.debug_struct("IndexPrefix")
55            .field("storage_prefix", &self.inner.storage_prefix)
56            .field("pk_name", &self.pk_name)
57            .finish_non_exhaustive()
58    }
59}
60
61impl<K, T> Deref for IndexPrefix<K, T>
62where
63    K: KeyDeserialize,
64    T: Serialize + DeserializeOwned,
65{
66    type Target = [u8];
67
68    fn deref(&self) -> &[u8] {
69        &self.inner.storage_prefix
70    }
71}
72
73impl<K, T, B> IndexPrefix<K, T, B>
74where
75    K: KeyDeserialize,
76    T: Serialize + DeserializeOwned,
77{
78    pub fn new(top_name: &[u8], sub_names: &[Key]) -> Self {
79        IndexPrefix::with_deserialization_functions(
80            top_name,
81            sub_names,
82            &[],
83            default_deserializer_kv::<K, T>,
84            default_deserializer_v,
85        )
86    }
87
88    pub fn with_deserialization_functions(
89        top_name: &[u8],
90        sub_names: &[Key],
91        pk_name: &[u8],
92        de_fn_kv: DeserializeKvFn<K, T>,
93        de_fn_v: DeserializeVFn<T>,
94    ) -> Self {
95        IndexPrefix {
96            inner: crate::prefix::Prefix::new(top_name, sub_names),
97            pk_name: pk_name.to_vec(),
98            de_fn_kv,
99            de_fn_v,
100        }
101    }
102}
103
104impl<'b, K, T, B> IndexPrefix<K, T, B>
105where
106    B: PrimaryKey<'b>,
107    K: KeyDeserialize,
108    T: Serialize + DeserializeOwned,
109{
110    pub fn range_raw<'a>(
111        &self,
112        store: &'a dyn Storage,
113        min: Option<Bound<'b, B>>,
114        max: Option<Bound<'b, B>>,
115        order: Order,
116    ) -> Box<dyn Iterator<Item = StdResult<Record<T>>> + 'a>
117    where
118        T: 'a,
119    {
120        let de_fn = self.de_fn_v;
121        let pk_name = self.pk_name.clone();
122        let mapped = crate::prefix::range_with_prefix(
123            store,
124            &self.inner.storage_prefix,
125            min.map(|b| b.to_raw_bound()),
126            max.map(|b| b.to_raw_bound()),
127            order,
128        )
129        .map(move |kv| (de_fn)(store, &pk_name, kv));
130        Box::new(mapped)
131    }
132
133    pub fn keys_raw<'a>(
134        &self,
135        store: &'a dyn Storage,
136        min: Option<Bound<'b, B>>,
137        max: Option<Bound<'b, B>>,
138        order: Order,
139    ) -> Box<dyn Iterator<Item = Vec<u8>> + 'a> {
140        crate::prefix::keys_with_prefix(
141            store,
142            &self.inner.storage_prefix,
143            min.map(|b| b.to_raw_bound()),
144            max.map(|b| b.to_raw_bound()),
145            order,
146        )
147    }
148
149    /// Clears the prefix, removing the first `limit` elements (or all if `limit == None`).
150    pub fn clear(&self, store: &mut dyn Storage, limit: Option<usize>) {
151        self.inner.clear(store, limit);
152    }
153
154    /// Returns `true` if the prefix is empty.
155    pub fn is_empty(&self, store: &dyn Storage) -> bool {
156        crate::prefix::keys_full(
157            store,
158            &self.inner.storage_prefix,
159            None,
160            None,
161            Order::Ascending,
162        )
163        .next()
164        .is_none()
165    }
166
167    pub fn range<'a>(
168        &self,
169        store: &'a dyn Storage,
170        min: Option<Bound<'b, B>>,
171        max: Option<Bound<'b, B>>,
172        order: Order,
173    ) -> Box<dyn Iterator<Item = StdResult<(K::Output, T)>> + 'a>
174    where
175        T: 'a,
176        K::Output: 'static,
177    {
178        let de_fn = self.de_fn_kv;
179        let pk_name = self.pk_name.clone();
180        let mapped = crate::prefix::range_with_prefix(
181            store,
182            &self.inner.storage_prefix,
183            min.map(|b| b.to_raw_bound()),
184            max.map(|b| b.to_raw_bound()),
185            order,
186        )
187        .map(move |kv| (de_fn)(store, &pk_name, kv));
188        Box::new(mapped)
189    }
190
191    pub fn keys<'a>(
192        &self,
193        store: &'a dyn Storage,
194        min: Option<Bound<'b, B>>,
195        max: Option<Bound<'b, B>>,
196        order: Order,
197    ) -> Box<dyn Iterator<Item = StdResult<K::Output>> + 'a>
198    where
199        T: 'a,
200        K::Output: 'static,
201    {
202        let de_fn = self.de_fn_kv;
203        let pk_name = self.pk_name.clone();
204        let mapped = crate::prefix::range_with_prefix(
205            store,
206            &self.inner.storage_prefix,
207            min.map(|b| b.to_raw_bound()),
208            max.map(|b| b.to_raw_bound()),
209            order,
210        )
211        .map(move |kv| (de_fn)(store, &pk_name, kv).map(|(k, _)| k));
212        Box::new(mapped)
213    }
214}
215
216#[cfg(test)]
217mod test {
218    use super::*;
219
220    use std::marker::PhantomData;
221
222    use cosmwasm_std::testing::MockStorage;
223
224    #[test]
225    fn ensure_proper_range_bounds() {
226        let mut store = MockStorage::new();
227        // manually create this - not testing nested prefixes here
228        let prefix: IndexPrefix<Vec<u8>, u64> = IndexPrefix {
229            inner: crate::prefix::Prefix {
230                storage_prefix: b"foo".to_vec(),
231                data: PhantomData::<(u64, _, _)>,
232            },
233            pk_name: vec![],
234            de_fn_kv: |_, _, kv| deserialize_kv::<Vec<u8>, u64>(kv),
235            de_fn_v: |_, _, kv| deserialize_v(kv),
236        };
237
238        // set some data, we care about "foo" prefix
239        store.set(b"foobar", b"1");
240        store.set(b"foora", b"2");
241        store.set(b"foozi", b"3");
242        // these shouldn't match
243        store.set(b"foply", b"100");
244        store.set(b"font", b"200");
245
246        let expected = vec![
247            (b"bar".to_vec(), 1u64),
248            (b"ra".to_vec(), 2u64),
249            (b"zi".to_vec(), 3u64),
250        ];
251        let expected_reversed: Vec<(Vec<u8>, u64)> = expected.iter().rev().cloned().collect();
252
253        // let's do the basic sanity check
254        let res: StdResult<Vec<_>> = prefix
255            .range_raw(&store, None, None, Order::Ascending)
256            .collect();
257        assert_eq!(&expected, &res.unwrap());
258        let res: StdResult<Vec<_>> = prefix
259            .range_raw(&store, None, None, Order::Descending)
260            .collect();
261        assert_eq!(&expected_reversed, &res.unwrap());
262
263        // now let's check some ascending ranges
264        let res: StdResult<Vec<_>> = prefix
265            .range_raw(
266                &store,
267                Some(Bound::inclusive(b"ra".to_vec())),
268                None,
269                Order::Ascending,
270            )
271            .collect();
272        assert_eq!(&expected[1..], res.unwrap().as_slice());
273        // skip excluded
274        let res: StdResult<Vec<_>> = prefix
275            .range_raw(
276                &store,
277                Some(Bound::exclusive(b"ra".to_vec())),
278                None,
279                Order::Ascending,
280            )
281            .collect();
282        assert_eq!(&expected[2..], res.unwrap().as_slice());
283        // if we exclude something a little lower, we get matched
284        let res: StdResult<Vec<_>> = prefix
285            .range_raw(
286                &store,
287                Some(Bound::exclusive(b"r".to_vec())),
288                None,
289                Order::Ascending,
290            )
291            .collect();
292        assert_eq!(&expected[1..], res.unwrap().as_slice());
293
294        // now let's check some descending ranges
295        let res: StdResult<Vec<_>> = prefix
296            .range_raw(
297                &store,
298                None,
299                Some(Bound::inclusive(b"ra".to_vec())),
300                Order::Descending,
301            )
302            .collect();
303        assert_eq!(&expected_reversed[1..], res.unwrap().as_slice());
304        // skip excluded
305        let res: StdResult<Vec<_>> = prefix
306            .range_raw(
307                &store,
308                None,
309                Some(Bound::exclusive(b"ra".to_vec())),
310                Order::Descending,
311            )
312            .collect();
313        assert_eq!(&expected_reversed[2..], res.unwrap().as_slice());
314        // if we exclude something a little higher, we get matched
315        let res: StdResult<Vec<_>> = prefix
316            .range_raw(
317                &store,
318                None,
319                Some(Bound::exclusive(b"rb".to_vec())),
320                Order::Descending,
321            )
322            .collect();
323        assert_eq!(&expected_reversed[1..], res.unwrap().as_slice());
324
325        // now test when both sides are set
326        let res: StdResult<Vec<_>> = prefix
327            .range_raw(
328                &store,
329                Some(Bound::inclusive(b"ra".to_vec())),
330                Some(Bound::exclusive(b"zi".to_vec())),
331                Order::Ascending,
332            )
333            .collect();
334        assert_eq!(&expected[1..2], res.unwrap().as_slice());
335        // and descending
336        let res: StdResult<Vec<_>> = prefix
337            .range_raw(
338                &store,
339                Some(Bound::inclusive(b"ra".to_vec())),
340                Some(Bound::exclusive(b"zi".to_vec())),
341                Order::Descending,
342            )
343            .collect();
344        assert_eq!(&expected[1..2], res.unwrap().as_slice());
345        // Include both sides
346        let res: StdResult<Vec<_>> = prefix
347            .range_raw(
348                &store,
349                Some(Bound::inclusive(b"ra".to_vec())),
350                Some(Bound::inclusive(b"zi".to_vec())),
351                Order::Descending,
352            )
353            .collect();
354        assert_eq!(&expected_reversed[..2], res.unwrap().as_slice());
355        // Exclude both sides
356        let res: StdResult<Vec<_>> = prefix
357            .range_raw(
358                &store,
359                Some(Bound::exclusive(b"ra".to_vec())),
360                Some(Bound::exclusive(b"zi".to_vec())),
361                Order::Ascending,
362            )
363            .collect();
364        assert_eq!(res.unwrap().as_slice(), &[]);
365    }
366
367    #[test]
368    fn prefix_debug() {
369        let prefix: IndexPrefix<String, String> = IndexPrefix::new(b"lol", &[Key::Val8([8; 1])]);
370        assert_eq!(
371            format!("{prefix:?}"),
372            "IndexPrefix { storage_prefix: [0, 3, 108, 111, 108, 0, 1, 8], pk_name: [], .. }"
373        );
374    }
375
376    #[test]
377    fn is_empty_works() {
378        // manually create this - not testing nested prefixes here
379        let prefix: IndexPrefix<Vec<u8>, u64> = IndexPrefix {
380            inner: crate::prefix::Prefix {
381                storage_prefix: b"foo".to_vec(),
382                data: PhantomData::<(u64, _, _)>,
383            },
384            pk_name: vec![],
385            de_fn_kv: |_, _, kv| deserialize_kv::<Vec<u8>, u64>(kv),
386            de_fn_v: |_, _, kv| deserialize_v(kv),
387        };
388
389        let mut storage = MockStorage::new();
390
391        assert!(prefix.is_empty(&storage));
392
393        storage.set(b"fookey1", b"1");
394        storage.set(b"fookey2", b"2");
395
396        assert!(!prefix.is_empty(&storage));
397    }
398}
cw_storage_plus/indexes/prefix.rs

cw_storage_plus/indexes/
prefix.rs
