cw_storage_plus/indexes/
multi.rs

1// this module requires iterator to be useful at all
2#![cfg(feature = "iterator")]
3
4use cosmwasm_std::storage_keys::namespace_with_key;
5use serde::de::DeserializeOwned;
6use serde::Serialize;
7
8use cosmwasm_std::{from_json, Order, Record, StdError, StdResult, Storage};
9
10use crate::bound::PrefixBound;
11use crate::de::KeyDeserialize;
12use crate::indexes::IndexPrefix;
13use crate::iter_helpers::deserialize_kv;
14use crate::map::Map;
15use crate::prefix::namespaced_prefix_range;
16use crate::{Bound, Index, Prefixer, PrimaryKey};
17use std::marker::PhantomData;
18
19/// MultiIndex stores (namespace, index_name, idx_value, pk) -> b"pk_len".
20/// Allows many values per index, and references pk.
21/// The associated primary key value is stored in the main (pk_namespace) map,
22/// which stores (namespace, pk_namespace, pk) -> value.
23///
24/// The stored pk_len is used to recover the pk from the index namespace, and perform
25/// the secondary load of the associated value from the main map.
26///
27/// The PK type defines the type of Primary Key, both for deserialization, and
28/// more important, as the type-safe bound key type.
29/// This type must match the encompassing `IndexedMap` primary key type,
30/// or its owned variant.
31pub struct MultiIndex<'a, IK, T, PK> {
32    index: fn(&[u8], &T) -> IK,
33    idx_namespace: &'a [u8],
34    // note, we collapse the ik - combining everything under the namespace - and concatenating the pk
35    idx_map: Map<Vec<u8>, u32>,
36    pk_namespace: &'a [u8],
37    phantom: PhantomData<PK>,
38}
39
40impl<'a, IK, T, PK> MultiIndex<'a, IK, T, PK>
41where
42    T: Serialize + DeserializeOwned + Clone,
43{
44    /// Create a new MultiIndex
45    ///
46    /// idx_fn - lambda creating index key from value
47    /// pk_namespace - prefix for the primary key
48    /// idx_namespace - prefix for the index value
49    ///
50    /// ## Example:
51    ///
52    /// ```rust
53    /// use cw_storage_plus::MultiIndex;
54    /// use serde::{Deserialize, Serialize};
55    ///
56    /// #[derive(Deserialize, Serialize, Clone)]
57    /// struct Data {
58    ///     pub name: String,
59    ///     pub age: u32,
60    /// }
61    ///
62    /// let index: MultiIndex<_, _, String> = MultiIndex::new(
63    ///     |_pk: &[u8], d: &Data| d.age,
64    ///     "age",
65    ///     "age__owner",
66    /// );
67    /// ```
68    pub const fn new(
69        idx_fn: fn(&[u8], &T) -> IK,
70        pk_namespace: &'a str,
71        idx_namespace: &'static str,
72    ) -> Self {
73        MultiIndex {
74            index: idx_fn,
75            idx_namespace: idx_namespace.as_bytes(),
76            idx_map: Map::new(idx_namespace),
77            pk_namespace: pk_namespace.as_bytes(),
78            phantom: PhantomData,
79        }
80    }
81}
82
83fn deserialize_multi_v<T: DeserializeOwned>(
84    store: &dyn Storage,
85    pk_namespace: &[u8],
86    kv: Record,
87) -> StdResult<Record<T>> {
88    let (key, pk_len) = kv;
89
90    // Deserialize pk_len
91    let pk_len = from_json::<u32>(pk_len.as_slice())?;
92
93    // Recover pk from last part of k
94    let offset = key.len() - pk_len as usize;
95    let pk = &key[offset..];
96
97    let full_key = namespace_with_key(&[pk_namespace], pk);
98
99    let v = store
100        .get(&full_key)
101        .ok_or_else(|| StdError::msg("pk not found"))?;
102    let v = from_json::<T>(&v)?;
103
104    Ok((pk.to_vec(), v))
105}
106
107fn deserialize_multi_kv<K: KeyDeserialize, T: DeserializeOwned>(
108    store: &dyn Storage,
109    pk_namespace: &[u8],
110    kv: Record,
111) -> StdResult<(K::Output, T)> {
112    let (key, pk_len) = kv;
113
114    // Deserialize pk_len
115    let pk_len = from_json::<u32>(pk_len.as_slice())?;
116
117    // Recover pk from last part of k
118    let offset = key.len() - pk_len as usize;
119    let pk = &key[offset..];
120
121    let full_key = namespace_with_key(&[pk_namespace], pk);
122
123    let v = store
124        .get(&full_key)
125        .ok_or_else(|| StdError::msg("pk not found"))?;
126    let v = from_json::<T>(&v)?;
127
128    // We return deserialized `pk` here for consistency
129    Ok((K::from_slice(pk)?, v))
130}
131
132impl<'a, IK, T, PK> Index<T> for MultiIndex<'a, IK, T, PK>
133where
134    T: Serialize + DeserializeOwned + Clone,
135    IK: PrimaryKey<'a>,
136{
137    fn save(&self, store: &mut dyn Storage, pk: &[u8], data: &T) -> StdResult<()> {
138        let idx = (self.index)(pk, data).joined_extra_key(pk);
139        self.idx_map.save(store, idx, &(pk.len() as u32))
140    }
141
142    fn remove(&self, store: &mut dyn Storage, pk: &[u8], old_data: &T) -> StdResult<()> {
143        let idx = (self.index)(pk, old_data).joined_extra_key(pk);
144        self.idx_map.remove(store, idx);
145        Ok(())
146    }
147}
148
149impl<'a, IK, T, PK> MultiIndex<'a, IK, T, PK>
150where
151    T: Serialize + DeserializeOwned + Clone,
152    IK: PrimaryKey<'a> + Prefixer<'a>,
153{
154    fn no_prefix_raw(&self) -> IndexPrefix<Vec<u8>, T, (IK, PK)> {
155        IndexPrefix::with_deserialization_functions(
156            self.idx_namespace,
157            &[],
158            self.pk_namespace,
159            deserialize_multi_v,
160            deserialize_multi_v,
161        )
162    }
163}
164
165impl<'a, IK, T, PK> MultiIndex<'a, IK, T, PK>
166where
167    PK: PrimaryKey<'a> + KeyDeserialize,
168    T: Serialize + DeserializeOwned + Clone,
169    IK: PrimaryKey<'a> + Prefixer<'a>,
170{
171    pub fn index_key(&self, k: IK) -> Vec<u8> {
172        k.joined_extra_key(b"")
173    }
174
175    #[cfg(test)]
176    pub fn count(&self, store: &dyn Storage, p: IK) -> usize {
177        let prefix = self.prefix(p);
178        prefix.keys_raw(store, None, None, Order::Ascending).count()
179    }
180
181    #[cfg(test)]
182    pub fn all_pks(&self, store: &dyn Storage, p: IK) -> Vec<Vec<u8>> {
183        let prefix = self.prefix(p);
184        prefix
185            .keys_raw(store, None, None, Order::Ascending)
186            .collect::<Vec<Vec<u8>>>()
187    }
188
189    #[cfg(test)]
190    pub fn all_items(&self, store: &dyn Storage, p: IK) -> StdResult<Vec<Record<T>>> {
191        let prefix = self.prefix(p);
192        prefix
193            .range_raw(store, None, None, Order::Ascending)
194            .collect()
195    }
196}
197
198// short-cut for simple keys, rather than .prefix(()).range_raw(...)
199impl<'a, IK, T, PK> MultiIndex<'a, IK, T, PK>
200where
201    T: Serialize + DeserializeOwned + Clone,
202    IK: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize,
203    PK: PrimaryKey<'a> + KeyDeserialize,
204{
205    // I would prefer not to copy code from Prefix, but no other way
206    // with lifetimes (create Prefix inside function and return ref = no no)
207    pub fn range_raw<'c>(
208        &'c self,
209        store: &'c dyn Storage,
210        min: Option<Bound<'a, (IK, PK)>>,
211        max: Option<Bound<'a, (IK, PK)>>,
212        order: Order,
213    ) -> Box<dyn Iterator<Item = StdResult<Record<T>>> + 'c>
214    where
215        T: 'c,
216    {
217        self.no_prefix_raw().range_raw(store, min, max, order)
218    }
219
220    pub fn keys_raw<'c>(
221        &'c self,
222        store: &'c dyn Storage,
223        min: Option<Bound<'a, (IK, PK)>>,
224        max: Option<Bound<'a, (IK, PK)>>,
225        order: Order,
226    ) -> Box<dyn Iterator<Item = Vec<u8>> + 'c> {
227        self.no_prefix_raw().keys_raw(store, min, max, order)
228    }
229
230    /// While `range_raw` over a `prefix` fixes the prefix to one element and iterates over the
231    /// remaining, `prefix_range_raw` accepts bounds for the lowest and highest elements of the
232    /// `Prefix` itself, and iterates over those (inclusively or exclusively, depending on
233    /// `PrefixBound`).
234    /// There are some issues that distinguish these two, and blindly casting to `Vec<u8>` doesn't
235    /// solve them.
236    pub fn prefix_range_raw<'c>(
237        &'c self,
238        store: &'c dyn Storage,
239        min: Option<PrefixBound<'a, IK>>,
240        max: Option<PrefixBound<'a, IK>>,
241        order: cosmwasm_std::Order,
242    ) -> Box<dyn Iterator<Item = StdResult<cosmwasm_std::Record<T>>> + 'c>
243    where
244        T: 'c,
245        'a: 'c,
246    {
247        let mapped = namespaced_prefix_range(store, self.idx_namespace, min, max, order)
248            .map(move |kv| (deserialize_multi_v)(store, self.pk_namespace, kv));
249        Box::new(mapped)
250    }
251}
252
253#[cfg(feature = "iterator")]
254impl<'a, IK, T, PK> MultiIndex<'a, IK, T, PK>
255where
256    PK: PrimaryKey<'a> + KeyDeserialize,
257    T: Serialize + DeserializeOwned + Clone,
258    IK: PrimaryKey<'a> + Prefixer<'a>,
259{
260    pub fn prefix(&self, p: IK) -> IndexPrefix<PK, T, PK> {
261        IndexPrefix::with_deserialization_functions(
262            self.idx_namespace,
263            &p.prefix(),
264            self.pk_namespace,
265            deserialize_multi_kv::<PK, T>,
266            deserialize_multi_v,
267        )
268    }
269
270    pub fn sub_prefix(&self, p: IK::Prefix) -> IndexPrefix<PK, T, (IK::Suffix, PK)> {
271        IndexPrefix::with_deserialization_functions(
272            self.idx_namespace,
273            &p.prefix(),
274            self.pk_namespace,
275            deserialize_multi_kv::<PK, T>,
276            deserialize_multi_v,
277        )
278    }
279}
280
281#[cfg(feature = "iterator")]
282impl<'a, IK, T, PK> MultiIndex<'a, IK, T, PK>
283where
284    PK: PrimaryKey<'a> + KeyDeserialize,
285    T: Serialize + DeserializeOwned + Clone,
286    IK: PrimaryKey<'a> + KeyDeserialize + Prefixer<'a>,
287{
288    /// While `range` over a `prefix` fixes the prefix to one element and iterates over the
289    /// remaining, `prefix_range` accepts bounds for the lowest and highest elements of the
290    /// `Prefix` itself, and iterates over those (inclusively or exclusively, depending on
291    /// `PrefixBound`).
292    /// There are some issues that distinguish these two, and blindly casting to `Vec<u8>` doesn't
293    /// solve them.
294    pub fn prefix_range<'c>(
295        &self,
296        store: &'c dyn Storage,
297        min: Option<PrefixBound<'a, IK>>,
298        max: Option<PrefixBound<'a, IK>>,
299        order: cosmwasm_std::Order,
300    ) -> Box<dyn Iterator<Item = StdResult<(PK::Output, T)>> + 'c>
301    where
302        T: 'c,
303        'a: 'c,
304        IK: 'c,
305        PK: 'c,
306        PK::Output: 'static,
307    {
308        let mapped = namespaced_prefix_range(store, self.idx_namespace, min, max, order)
309            .map(deserialize_kv::<PK, T>);
310        Box::new(mapped)
311    }
312
313    pub fn range<'c>(
314        &self,
315        store: &'c dyn Storage,
316        min: Option<Bound<'a, (IK, PK)>>,
317        max: Option<Bound<'a, (IK, PK)>>,
318        order: cosmwasm_std::Order,
319    ) -> Box<dyn Iterator<Item = StdResult<(PK::Output, T)>> + 'c>
320    where
321        T: 'c,
322        PK::Output: 'static,
323    {
324        self.no_prefix().range(store, min, max, order)
325    }
326
327    pub fn keys<'c>(
328        &self,
329        store: &'c dyn Storage,
330        min: Option<Bound<'a, (IK, PK)>>,
331        max: Option<Bound<'a, (IK, PK)>>,
332        order: cosmwasm_std::Order,
333    ) -> Box<dyn Iterator<Item = StdResult<PK::Output>> + 'c>
334    where
335        T: 'c,
336        PK::Output: 'static,
337    {
338        self.no_prefix().keys(store, min, max, order)
339    }
340
341    fn no_prefix(&self) -> IndexPrefix<PK, T, (IK, PK)> {
342        IndexPrefix::with_deserialization_functions(
343            self.idx_namespace,
344            &[],
345            self.pk_namespace,
346            deserialize_multi_kv::<PK, T>,
347            deserialize_multi_v,
348        )
349    }
350}
cw_storage_plus/indexes/multi.rs

cw_storage_plus/indexes/
multi.rs
