cw_item_set/
lib.rs

1#![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/README.md"))]
2
3use std::marker::PhantomData;
4
5#[cfg(feature = "iterator")]
6use cosmwasm_std::Order;
7#[cfg(feature = "counter")]
8use cosmwasm_std::StdError;
9use cosmwasm_std::{Empty, StdResult, Storage};
10#[cfg(feature = "counter")]
11use cw_storage_plus::Item;
12#[cfg(feature = "iterator")]
13use cw_storage_plus::{Bound, Prefix, Prefixer};
14use cw_storage_plus::{Key, KeyDeserialize, Path, PrimaryKey};
15
16/// A set of non-duplicate items.
17///
18/// On a high level, a `Set<T>` is equivalent to a `Map<T, Empty>`, but offers
19/// a more intuitive API similar to native HashSet and BTreeSet.
20pub struct Set<'a, T> {
21    namespace: &'a [u8],
22
23    #[cfg(feature = "counter")]
24    counter: Item<u64>,
25
26    item_type: PhantomData<T>,
27}
28
29#[cfg(not(feature = "counter"))]
30impl<'a, T> Set<'a, T> {
31    /// Create a new instance of the item set with the given namespace.
32    pub const fn new(namespace: &'a str) -> Self {
33        Set {
34            namespace: namespace.as_bytes(),
35            item_type: PhantomData,
36        }
37    }
38}
39
40#[cfg(feature = "counter")]
41impl<'a, T> Set<'a, T> {
42    /// Create a new instance of the item set with the given map and counter namespaces.
43    pub const fn new(namespace: &'a str, counter_namespace: &'static str) -> Self {
44        Set {
45            namespace: namespace.as_bytes(),
46            counter: Item::new(counter_namespace),
47            item_type: PhantomData,
48        }
49    }
50
51    /// Return the total amount of items in the set.
52    pub fn count(&self, store: &dyn Storage) -> StdResult<u64> {
53        Ok(self.counter.may_load(store)?.unwrap_or(0))
54    }
55
56    /// Increase the item count by 1.
57    fn increment_count(&self, store: &mut dyn Storage) -> StdResult<()> {
58        let mut count = self.counter.may_load(store)?.unwrap_or(0);
59        count += 1;
60        self.counter.save(store, &count)
61    }
62
63    /// Reduce the item count by 1; throw error if the current count is zero.
64    fn reduce_count(&self, store: &mut dyn Storage) -> StdResult<()> {
65        match self.counter.may_load(store)? {
66            None | Some(0) => {
67                Err(StdError::generic_err("[cw-item-set]: count cannot be reduced below zero"))
68            },
69            Some(mut count) => {
70                count -= 1;
71                self.counter.save(store, &count)
72            },
73        }
74    }
75}
76
77impl<'a, T> Set<'a, T>
78where
79    T: PrimaryKey<'a> + KeyDeserialize,
80{
81    /// Returns the key for storing an item
82    ///
83    /// This is copied from
84    /// https://github.com/CosmWasm/cw-plus/blob/v0.14.0/packages/storage-plus/src/map.rs#L47-L52
85    fn key(&self, item: T) -> Path<Empty> {
86        Path::new(self.namespace, &item.key().iter().map(Key::as_ref).collect::<Vec<_>>())
87    }
88
89    /// Returns `true` if the set contains an item
90    pub fn contains(&self, store: &dyn Storage, item: T) -> bool {
91        self.key(item).has(store)
92    }
93
94    /// Adds an item to the set. Returns whether the item was newly added.
95    pub fn insert(&self, store: &mut dyn Storage, item: T) -> StdResult<bool> {
96        let key = self.key(item);
97        if key.has(store) {
98            Ok(false)
99        } else {
100            key.save(store, &Empty {})?;
101
102            #[cfg(feature = "counter")]
103            self.increment_count(store)?;
104
105            Ok(true)
106        }
107    }
108
109    /// Remove an item from the set. Returns whether the item was present in the set.
110    pub fn remove(&self, store: &mut dyn Storage, item: T) -> StdResult<bool> {
111        let key = self.key(item);
112        if key.has(store) {
113            key.remove(store);
114
115            #[cfg(feature = "counter")]
116            self.reduce_count(store)?;
117
118            Ok(true)
119        } else {
120            Ok(false)
121        }
122    }
123}
124
125#[cfg(feature = "iterator")]
126impl<'a, T> Set<'a, T>
127where
128    T: PrimaryKey<'a> + KeyDeserialize,
129{
130    /// Copied from cw-storage-plus:
131    /// https://github.com/CosmWasm/cw-storage-plus/blob/v0.16.0/src/map.rs#L55-L57
132    fn no_prefix_raw(&self) -> Prefix<Vec<u8>, Empty, T> {
133        Prefix::new(self.namespace, &[])
134    }
135
136    /// Access items in the set under the given prefix.\
137    ///
138    /// Copied from cw-storage-plus:
139    /// https://github.com/CosmWasm/cw-plus/blob/v0.14.0/packages/storage-plus/src/map.rs#L124-126
140    pub fn prefix(&self, p: T::Prefix) -> Prefix<T::Suffix, Empty, T::Suffix> {
141        Prefix::new(self.namespace, &p.prefix())
142    }
143
144    /// Iterates items in the set with the specified bounds and ordering.
145    pub fn items<'c>(
146        &self,
147        store: &'c dyn Storage,
148        min: Option<Bound<'a, T>>,
149        max: Option<Bound<'a, T>>,
150        order: Order,
151    ) -> Box<dyn Iterator<Item = StdResult<T::Output>> + 'c>
152    where
153        T::Output: 'static,
154    {
155        Prefix::<T, Empty, T>::new(self.namespace, &[]).keys(store, min, max, order)
156    }
157
158    /// Delete all elements from the set.
159    ///
160    /// Copied from cw-storage-plus:
161    /// https://github.com/CosmWasm/cw-storage-plus/blob/v0.16.0/src/map.rs#L115-L132
162    pub fn clear(&self, store: &mut dyn Storage) {
163        const TAKE: usize = 10;
164        let mut cleared = false;
165
166        while !cleared {
167            let paths = self
168                .no_prefix_raw()
169                .keys_raw(store, None, None, Order::Ascending)
170                .map(|raw_key| Path::<Empty>::new(self.namespace, &[raw_key.as_slice()]))
171                .take(TAKE)
172                .collect::<Vec<_>>();
173
174            for path in &paths {
175                store.remove(path);
176            }
177
178            cleared = paths.len() < TAKE;
179        }
180
181        #[cfg(feature = "counter")]
182        self.counter.remove(store);
183    }
184}
185
186//--------------------------------------------------------------------------------------------------
187// Tests
188//--------------------------------------------------------------------------------------------------
189
190#[cfg(test)]
191mod tests {
192    #[cfg(feature = "iterator")]
193    use std::ops::Range;
194
195    use cosmwasm_std::testing::MockStorage;
196
197    use super::*;
198
199    const NAMESPACE: &str = "names";
200
201    #[cfg(not(feature = "counter"))]
202    const NAMES: Set<&str> = Set::new(NAMESPACE);
203    #[cfg(feature = "counter")]
204    const NAMES: Set<&str> = Set::new(NAMESPACE, "names__counter");
205
206    #[cfg(all(feature = "iterator", not(feature = "counter")))]
207    const TUPLES: Set<(u64, &str)> = Set::new("tuples");
208    #[cfg(all(feature = "iterator", feature = "counter"))]
209    const TUPLES: Set<(u64, &str)> = Set::new("tuples", "tuples__counter");
210
211    /// Returns the raw storage key for a given name.
212    /// The key is: length of namespace (2 bytes) + namespace + the name
213    fn key(name: &str) -> Vec<u8> {
214        let length_bytes = (NAMESPACE.len() as u32).to_be_bytes();
215        let mut out = Vec::with_capacity(2 + NAMESPACE.len() + name.len());
216        out.extend_from_slice(&[length_bytes[2], length_bytes[3]]);
217        out.extend_from_slice(NAMESPACE.as_bytes());
218        out.extend_from_slice(name.as_bytes());
219        out
220    }
221
222    /// Return a list of mockup names for use in testing.
223    #[cfg(feature = "iterator")]
224    fn mock_names(indexes: Range<usize>) -> Vec<String> {
225        let mut names = indexes.map(|i| format!("test-name-{i}")).collect::<Vec<_>>();
226        names.sort();
227        names
228    }
229
230    /// Insert mock names into a set.
231    #[cfg(feature = "iterator")]
232    fn insert_mock_names(set: Set<&str>, store: &mut dyn Storage) {
233        for name in mock_names(1..100) {
234            set.insert(store, &name).unwrap();
235        }
236    }
237
238    #[test]
239    fn containing() {
240        let mut store = MockStorage::default();
241
242        NAMES.insert(&mut store, "larry").unwrap();
243        assert!(NAMES.contains(&store, "larry"));
244        assert!(!NAMES.contains(&store, "jake"));
245
246        NAMES.insert(&mut store, "jake").unwrap();
247        assert!(NAMES.contains(&store, "larry"));
248        assert!(NAMES.contains(&store, "jake"));
249    }
250
251    #[test]
252    fn inserting() {
253        let mut store = MockStorage::default();
254
255        let new = NAMES.insert(&mut store, "larry").unwrap();
256        assert!(new);
257        assert_eq!(store.get(&key("larry")), Some(b"{}".to_vec()));
258        assert_eq!(store.get(&key("jake")), None);
259
260        let new = NAMES.insert(&mut store, "larry").unwrap();
261        assert!(!new);
262        assert_eq!(store.get(&key("larry")), Some(b"{}".to_vec()));
263        assert_eq!(store.get(&key("jake")), None);
264
265        let new = NAMES.insert(&mut store, "jake").unwrap();
266        assert!(new);
267        assert_eq!(store.get(&key("larry")), Some(b"{}".to_vec()));
268        assert_eq!(store.get(&key("jake")), Some(b"{}".to_vec()));
269    }
270
271    #[test]
272    fn removing() {
273        let mut store = MockStorage::default();
274
275        NAMES.insert(&mut store, "larry").unwrap();
276
277        let existed = NAMES.remove(&mut store, "larry").unwrap();
278        assert!(existed);
279        assert_eq!(store.get(&key("larry")), None);
280
281        let existed = NAMES.remove(&mut store, "jake").unwrap();
282        assert!(!existed);
283        assert_eq!(store.get(&key("jake")), None);
284    }
285
286    #[cfg(feature = "counter")]
287    #[test]
288    fn counting() {
289        let mut store = MockStorage::default();
290
291        let count = NAMES.count(&store).unwrap();
292        assert_eq!(count, 0);
293
294        NAMES.insert(&mut store, "larry").unwrap();
295        assert_eq!(NAMES.count(&store).unwrap(), 1);
296
297        NAMES.insert(&mut store, "jake").unwrap();
298        assert_eq!(NAMES.count(&store).unwrap(), 2);
299
300        NAMES.insert(&mut store, "pumpkin").unwrap();
301        assert_eq!(NAMES.count(&store).unwrap(), 3);
302
303        #[cfg(feature = "iterator")]
304        {
305            NAMES.clear(&mut store);
306            assert_eq!(NAMES.count(&store).unwrap(), 0);
307        }
308    }
309
310    #[cfg(feature = "iterator")]
311    #[test]
312    fn iterating() {
313        let mut store = MockStorage::default();
314
315        insert_mock_names(NAMES, &mut store);
316
317        let names = NAMES
318            .items(&store, None, None, Order::Ascending)
319            .collect::<StdResult<Vec<_>>>()
320            .unwrap();
321        assert_eq!(names, mock_names(1..100));
322
323        let start_after = Bound::ExclusiveRaw(b"test-name-2".to_vec());
324        let names = NAMES
325            .items(&store, Some(start_after), None, Order::Ascending)
326            .take(10)
327            .collect::<StdResult<Vec<_>>>()
328            .unwrap();
329        assert_eq!(names, mock_names(20..30));
330    }
331
332    #[cfg(feature = "iterator")]
333    #[test]
334    fn clearing() {
335        let mut store = MockStorage::default();
336
337        insert_mock_names(NAMES, &mut store);
338
339        NAMES.clear(&mut store);
340
341        let names = NAMES
342            .items(&store, None, None, Order::Ascending)
343            .collect::<StdResult<Vec<_>>>()
344            .unwrap();
345        assert_eq!(names.len(), 0);
346    }
347
348    #[cfg(feature = "iterator")]
349    #[test]
350    fn prefixes() {
351        let mut store = MockStorage::default();
352
353        let tuples = vec![(1u64, "larry"), (1u64, "jake"), (2u64, "pumpkin")];
354
355        for tuple in &tuples {
356            TUPLES.insert(&mut store, *tuple).unwrap();
357        }
358
359        let names = TUPLES
360            .prefix(1)
361            .keys(&store, None, None, Order::Ascending)
362            .collect::<StdResult<Vec<_>>>()
363            .unwrap();
364        assert_eq!(names, vec!["jake", "larry"]);
365
366        let names = TUPLES
367            .prefix(2)
368            .keys(&store, None, None, Order::Ascending)
369            .collect::<StdResult<Vec<_>>>()
370            .unwrap();
371        assert_eq!(names, vec!["pumpkin"]);
372    }
373}
cw_item_set/lib.rs

cw_item_set/
lib.rs
