feat: caching LookupSet implementation

CHANGELOG.md

@@ -10,6 +10,7 @@
 - Add consts for near, yocto, and tgas. [PR 640](https://github.com/near/near-sdk-rs/pull/640).
   - `near_sdk::ONE_NEAR`, `near_sdk::ONE_YOCTO`, `near_sdk::Gas::ONE_TERA`
 - Update SDK dependencies for `nearcore` crates used for mocking (`0.10`) and `borsh` (`0.9`)
+- store: Implement caching `LookupSet` type. This is the new iteration of the previous version of `near_sdk::collections::LookupSet` that has an updated API, and is located at `near_sdk::store::LookupSet`. [PR 654](https://github.com/near/near-sdk-rs/pull/654).
 
 ## `4.0.0-pre.4` [10-15-2021]
 - Unpin `syn` dependency in macros from `=1.0.57` to be more composable with other crates. [PR 605](https://github.com/near/near-sdk-rs/pull/605)

near-sdk/src/store/lookup_map/mod.rs

@@ -95,7 +95,6 @@ where
     hasher: PhantomData<H>,
 }
 
-// #[derive(Default)]
 struct EntryAndHash<V> {
     value: OnceCell<CacheEntry<V>>,
     hash: OnceCell<[u8; 32]>,

near-sdk/src/store/lookup_set/impls.rs

@@ -0,0 +1,16 @@
+use super::LookupSet;
+use crate::crypto_hash::CryptoHasher;
+use borsh::BorshSerialize;
+
+impl<T, H> Extend<T> for LookupSet<T, H>
+where
+    T: BorshSerialize + Ord,
+    H: CryptoHasher<Digest = [u8; 32]>,
+{
+    fn extend<I>(&mut self, iter: I)
+    where
+        I: IntoIterator<Item = T>,
+    {
+        self.map.extend(iter.into_iter().map(|k| (k, ())))
+    }
+}

near-sdk/src/store/lookup_set/mod.rs

@@ -0,0 +1,288 @@
+mod impls;
+
+use crate::crypto_hash::{CryptoHasher, Sha256};
+use crate::store::LookupMap;
+use crate::IntoStorageKey;
+use borsh::{BorshDeserialize, BorshSerialize};
+use std::borrow::Borrow;
+use std::fmt;
+
+#[derive(BorshSerialize, BorshDeserialize)]
+pub struct LookupSet<T, H = Sha256>
+where
+    T: BorshSerialize + Ord,
+    H: CryptoHasher<Digest = [u8; 32]>,
+{
+    map: LookupMap<T, (), H>,
+}
+
+impl<T, H> Drop for LookupSet<T, H>
+where
+    T: BorshSerialize + Ord,
+    H: CryptoHasher<Digest = [u8; 32]>,
+{
+    fn drop(&mut self) {
+        self.flush()
+    }
+}
+
+impl<T, H> fmt::Debug for LookupSet<T, H>
+where
+    T: BorshSerialize + Ord,
+    H: CryptoHasher<Digest = [u8; 32]>,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("LookupSet").field("map", &self.map).finish()
+    }
+}
+
+impl<T> LookupSet<T, Sha256>
+where
+    T: BorshSerialize + Ord,
+{
+    #[inline]
+    pub fn new<S>(prefix: S) -> Self
+    where
+        S: IntoStorageKey,
+    {
+        Self::with_hasher(prefix)
+    }
+}
+
+impl<T, H> LookupSet<T, H>
+where
+    T: BorshSerialize + Ord,
+    H: CryptoHasher<Digest = [u8; 32]>,
+{
+    /// Initialize a [`LookupSet`] with a custom hash function.
+    ///
+    /// # Example
+    /// ```
+    /// use near_sdk::crypto_hash::Keccak256;
+    /// use near_sdk::store::LookupSet;
+    ///
+    /// let map = LookupSet::<String, Keccak256>::with_hasher(b"m");
+    /// ```
+    pub fn with_hasher<S>(prefix: S) -> Self
+    where
+        S: IntoStorageKey,
+    {
+        Self { map: LookupMap::with_hasher(prefix) }
+    }
+
+    /// Returns `true` if the set contains the specified value.
+    ///
+    /// The value may be any borrowed form of the set's value type, but
+    /// [`BorshSerialize`], [`ToOwned<Owned = T>`](ToOwned) and [`Ord`] on the borrowed form *must*
+    /// match those for the value type.
+    pub fn contains<Q: ?Sized>(&self, value: &Q) -> bool
+    where
+        T: Borrow<Q>,
+        Q: BorshSerialize + ToOwned<Owned = T> + Ord,
+    {
+        self.map.contains_key(value)
+    }
+
+    /// Adds a value to the set.
+    ///
+    /// If the set did not have this value present, true is returned.
+    ///
+    /// If the set did have this value present, false is returned.
+    pub fn insert(&mut self, value: T) -> bool
+    where
+        T: Clone,
+    {
+        self.map.insert(value, ()).is_none()
+    }
+
+    /// Removes a value from the set. Returns whether the value was present in the set.
+    ///
+    /// The value may be any borrowed form of the set's value type, but
+    /// [`BorshSerialize`], [`ToOwned<Owned = K>`](ToOwned) and [`Ord`] on the borrowed form *must*
+    /// match those for the value type.
+    pub fn remove<Q: ?Sized>(&mut self, value: &Q) -> bool
+    where
+        T: Borrow<Q>,
+        Q: BorshSerialize + ToOwned<Owned = T> + Ord,
+    {
+        self.map.remove(value).is_some()
+    }
+}
+
+impl<T, H> LookupSet<T, H>
+where
+    T: BorshSerialize + Ord,
+    H: CryptoHasher<Digest = [u8; 32]>,
+{
+    /// Flushes the intermediate values of the set before this is called when the structure is
+    /// [`Drop`]ed. This will write all modified values to storage but keep all cached values
+    /// in memory.
+    pub fn flush(&mut self) {
+        self.map.flush()
+    }
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+#[cfg(test)]
+mod tests {
+    use super::LookupSet;
+    use crate::crypto_hash::{Keccak256, Sha256};
+    use crate::test_utils::test_env::setup_free;
+    use arbitrary::{Arbitrary, Unstructured};
+    use rand::seq::SliceRandom;
+    use rand::RngCore;
+    use rand::{Rng, SeedableRng};
+    use std::collections::HashSet;
+
+    #[test]
+    fn test_insert_contains() {
+        let mut set = LookupSet::new(b"m");
+        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
+        let mut baseline = HashSet::new();
+        for _ in 0..100 {
+            let value = rng.gen::<u64>();
+            set.insert(value);
+            baseline.insert(value);
+        }
+        // Non existing
+        for _ in 0..100 {
+            let value = rng.gen::<u64>();
+            assert_eq!(set.contains(&value), baseline.contains(&value));
+        }
+        // Existing
+        for value in baseline.iter() {
+            assert!(set.contains(value));
+        }
+    }
+
+    #[test]
+    fn test_insert_remove() {
+        let mut set = LookupSet::new(b"m");
+        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(1);
+        let mut values = vec![];
+        for _ in 0..100 {
+            let value = rng.gen::<u64>();
+            values.push(value);
+            set.insert(value);
+        }
+        values.shuffle(&mut rng);
+        for value in values {
+            assert!(set.remove(&value));
+        }
+    }
+
+    #[test]
+    fn test_remove_last_reinsert() {
+        let mut set = LookupSet::new(b"m");
+        let value1 = 2u64;
+        set.insert(value1);
+        let value2 = 4u64;
+        set.insert(value2);
+
+        assert!(set.remove(&value2));
+        assert!(set.insert(value2));
+    }
+
+    #[test]
+    fn test_extend() {
+        let mut set = LookupSet::new(b"m");
+        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(4);
+        let mut values = vec![];
+        for _ in 0..100 {
+            let value = rng.gen::<u64>();
+            values.push(value);
+            set.insert(value);
+        }
+        for _ in 0..10 {
+            let mut tmp = vec![];
+            for _ in 0..=(rng.gen::<u64>() % 20 + 1) {
+                let value = rng.gen::<u64>();
+                tmp.push(value);
+            }
+            values.extend(tmp.iter().cloned());
+            set.extend(tmp.iter().cloned());
+        }
+
+        for value in values {
+            assert!(set.contains(&value));
+        }
+    }
+
+    #[test]
+    fn test_debug() {
+        let set = LookupSet::<u8, Sha256>::new(b"m");
+
+        assert_eq!(format!("{:?}", set), "LookupSet { map: LookupMap { prefix: [109] } }")
+    }
+
+    #[test]
+    fn test_flush_on_drop() {
+        let mut set = LookupSet::<_, Keccak256>::with_hasher(b"m");
+
+        // Set a value, which does not write to storage yet
+        set.insert(5u8);
+        assert!(set.contains(&5u8));
+
+        // Drop the set which should flush all data
+        drop(set);
+
+        // Create a duplicate set which references same data
+        let dup_set = LookupSet::<u8, Keccak256>::with_hasher(b"m");
+
+        // New map can now load the value
+        assert!(dup_set.contains(&5u8));
+    }
+
+    #[derive(Arbitrary, Debug)]
+    enum Op {
+        Insert(u8),
+        Remove(u8),
+        Flush,
+        Restore,
+        Contains(u8),
+    }
+
+    #[test]
+    fn test_arbitrary() {
+        setup_free();
+
+        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
+        let mut buf = vec![0; 4096];
+        for _ in 0..512 {
+            // Clear storage in-between runs
+            crate::mock::with_mocked_blockchain(|b| b.take_storage());
+            rng.fill_bytes(&mut buf);
+
+            let mut ls = LookupSet::new(b"l");
+            let mut hs = HashSet::new();
+            let u = Unstructured::new(&buf);
+            if let Ok(ops) = Vec::<Op>::arbitrary_take_rest(u) {
+                for op in ops {
+                    match op {
+                        Op::Insert(v) => {
+                            let r1 = ls.insert(v);
+                            let r2 = hs.insert(v);
+                            assert_eq!(r1, r2)
+                        }
+                        Op::Remove(v) => {
+                            let r1 = ls.remove(&v);
+                            let r2 = hs.remove(&v);
+                            assert_eq!(r1, r2)
+                        }
+                        Op::Flush => {
+                            ls.flush();
+                        }
+                        Op::Restore => {
+                            ls = LookupSet::new(b"l");
+                        }
+                        Op::Contains(v) => {
+                            let r1 = ls.contains(&v);
+                            let r2 = hs.contains(&v);
+                            assert_eq!(r1, r2)
+                        }
+                    }
+                }
+            }
+        }
+    }
+}

near-sdk/src/store/mod.rs

@@ -10,6 +10,9 @@ pub use vec::Vector;
 pub mod lookup_map;
 pub use self::lookup_map::LookupMap;
 
+mod lookup_set;
+pub use self::lookup_set::LookupSet;
+
 pub mod unordered_map;
 pub use self::unordered_map::UnorderedMap;