Merge pull request #409 from near/austin/col/lazy

feat(collections): Lazy storage value

near-sdk/Cargo.toml

@@ -29,6 +29,9 @@ near-primitives-core = "=0.4.0"
 # Export dependencies for contracts
 wee_alloc = { version = "0.4.5", default-features = false, features = [] }
 
+# Used for caching, might be worth porting only functionality needed.
+once_cell = { version = "1.7.2", optional = true, default-features = false }
+
 [dev-dependencies]
 rand = "0.7.2"
 trybuild = "1.0"
@@ -38,3 +41,4 @@ quickcheck = "0.9.2"
 
 [features]
 expensive-debug = []
+unstable = ["once_cell"]

near-sdk/src/lib.rs

@@ -6,6 +6,9 @@ pub use near_sdk_macros::{
     serializer, BorshStorageKey, PanicOnDefault,
 };
 
+#[cfg(feature = "unstable")]
+pub mod store;
+
 pub mod collections;
 mod environment;
 pub use environment::env;

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

@@ -0,0 +1,91 @@
+use borsh::{BorshDeserialize, BorshSerialize};
+
+use super::Lazy;
+
+impl<T> Drop for Lazy<T>
+where
+    T: BorshSerialize,
+{
+    fn drop(&mut self) {
+        self.flush()
+    }
+}
+
+impl<T> core::ops::Deref for Lazy<T>
+where
+    T: BorshSerialize + BorshDeserialize,
+{
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        Self::get(self)
+    }
+}
+
+impl<T> core::ops::DerefMut for Lazy<T>
+where
+    T: BorshSerialize + BorshDeserialize,
+{
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        Self::get_mut(self)
+    }
+}
+
+impl<T> core::cmp::PartialEq for Lazy<T>
+where
+    T: PartialEq + BorshSerialize + BorshDeserialize,
+{
+    fn eq(&self, other: &Self) -> bool {
+        PartialEq::eq(self.get(), other.get())
+    }
+}
+
+impl<T> core::cmp::Eq for Lazy<T> where T: Eq + BorshSerialize + BorshDeserialize {}
+
+impl<T> core::cmp::PartialOrd for Lazy<T>
+where
+    T: PartialOrd + BorshSerialize + BorshDeserialize,
+{
+    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
+        PartialOrd::partial_cmp(self.get(), other.get())
+    }
+    fn lt(&self, other: &Self) -> bool {
+        PartialOrd::lt(self.get(), other.get())
+    }
+    fn le(&self, other: &Self) -> bool {
+        PartialOrd::le(self.get(), other.get())
+    }
+    fn ge(&self, other: &Self) -> bool {
+        PartialOrd::ge(self.get(), other.get())
+    }
+    fn gt(&self, other: &Self) -> bool {
+        PartialOrd::gt(self.get(), other.get())
+    }
+}
+
+impl<T> core::cmp::Ord for Lazy<T>
+where
+    T: core::cmp::Ord + BorshSerialize + BorshDeserialize,
+{
+    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
+        Ord::cmp(self.get(), other.get())
+    }
+}
+
+impl<T> core::convert::AsRef<T> for Lazy<T>
+where
+    T: BorshSerialize + BorshDeserialize,
+{
+    fn as_ref(&self) -> &T {
+        Self::get(self)
+    }
+}
+
+impl<T> core::convert::AsMut<T> for Lazy<T>
+where
+    T: BorshSerialize + BorshDeserialize,
+{
+    fn as_mut(&mut self) -> &mut T {
+        Self::get_mut(self)
+    }
+}

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

@@ -0,0 +1,185 @@
+//! A persistent lazy storage value. Stores a value for a given key.
+//! Example:
+//! If the underlying value is large, e.g. the contract needs to store an image, but it doesn't need
+//! to have access to this image at regular calls, then the contract can wrap this image into
+//! [`Lazy`] and it will not be deserialized until requested.
+
+mod impls;
+
+use borsh::{BorshDeserialize, BorshSerialize};
+use once_cell::unsync::OnceCell;
+
+use crate::env;
+use crate::utils::{CacheEntry, EntryState};
+use crate::IntoStorageKey;
+
+const ERR_VALUE_SERIALIZATION: &[u8] = b"Cannot serialize value with Borsh";
+const ERR_VALUE_DESERIALIZATION: &[u8] = b"Cannot deserialize value with Borsh";
+const ERR_NOT_FOUND: &[u8] = b"No value found for the given key";
+const ERR_DELETED: &[u8] = b"The Lazy cell's value has been deleted. Verify the key has not been\
+                            deleted manually.";
+
+fn expect_key_exists<T>(val: Option<T>) -> T {
+    val.unwrap_or_else(|| env::panic(ERR_NOT_FOUND))
+}
+
+fn expect_consistent_state<T>(val: Option<T>) -> T {
+    val.unwrap_or_else(|| env::panic(ERR_DELETED))
+}
+
+fn load_and_deserialize<T>(key: &[u8]) -> CacheEntry<T>
+where
+    T: BorshDeserialize,
+{
+    let bytes = expect_key_exists(env::storage_read(key));
+    let val = T::try_from_slice(&bytes).unwrap_or_else(|_| env::panic(ERR_VALUE_DESERIALIZATION));
+    CacheEntry::new_cached(Some(val))
+}
+
+fn serialize_and_store<T>(key: &[u8], value: &T)
+where
+    T: BorshSerialize,
+{
+    let serialized = value.try_to_vec().unwrap_or_else(|_| env::panic(ERR_VALUE_SERIALIZATION));
+    env::storage_write(&key, &serialized);
+}
+
+/// An persistent lazily loaded value, that stores a value in the storage.
+///
+/// This will only write to the underlying store if the value has changed, and will only read the
+/// existing value from storage once.
+///
+/// # Examples
+/// ```
+/// use near_sdk::collections::Lazy;
+///
+///# near_sdk::test_utils::test_env::setup();
+/// let mut a = Lazy::new(b"a", "test string".to_string());
+/// assert_eq!(*a, "test string");
+///
+/// *a = "new string".to_string();
+/// assert_eq!(a.get(), "new string");
+/// ```
+#[derive(BorshSerialize, BorshDeserialize, Debug)]
+pub struct Lazy<T>
+where
+    T: BorshSerialize,
+{
+    /// Key bytes to index the contract's storage.
+    storage_key: Vec<u8>,
+    #[borsh_skip]
+    /// Cached value which is lazily loaded and deserialized from storage.
+    cache: OnceCell<CacheEntry<T>>,
+}
+
+impl<T> Lazy<T>
+where
+    T: BorshSerialize,
+{
+    pub fn new<S>(key: S, value: T) -> Self
+    where
+        S: IntoStorageKey,
+    {
+        Self {
+            storage_key: key.into_storage_key(),
+            cache: OnceCell::from(CacheEntry::new_modified(Some(value))),
+        }
+    }
+
+    /// Updates the value with a new value. This does not load the current value from storage.
+    pub fn set(&mut self, value: T) {
+        if let Some(v) = self.cache.get_mut() {
+            *v.value_mut() = Some(value);
+        } else {
+            self.cache
+                .set(CacheEntry::new_modified(Some(value)))
+                .ok()
+                .expect("cache is checked to not be filled above");
+        }
+    }
+
+    /// Writes any changes to the value to storage. This will automatically be done when the
+    /// value is dropped through [`Drop`] so this should only be used when the changes need to be
+    /// reflected in the underlying storage before then.
+    pub fn flush(&mut self) {
+        if let Some(v) = self.cache.get_mut() {
+            if v.is_modified() {
+                // Value was modified, serialize and put the serialized bytes in storage.
+                let value = expect_consistent_state(v.value().as_ref());
+                serialize_and_store(&self.storage_key, value);
+
+                // Replaces cache entry state to cached because the value in memory matches the
+                // stored value. This avoids writing the same value twice.
+                v.replace_state(EntryState::Cached);
+            }
+        }
+    }
+}
+
+impl<T> Lazy<T>
+where
+    T: BorshSerialize + BorshDeserialize,
+{
+    /// Returns a reference to the lazily loaded storage value.
+    /// The load from storage only happens once, and if the value is already cached, it will not
+    /// be reloaded.
+    ///
+    /// This function will panic if the cache is not loaded and the value at the key does not exist.
+    pub fn get(&self) -> &T {
+        let entry = self.cache.get_or_init(|| load_and_deserialize(&self.storage_key));
+
+        expect_consistent_state(entry.value().as_ref())
+    }
+
+    /// Returns a reference to the lazily loaded storage value.
+    /// The load from storage only happens once, and if the value is already cached, it will not
+    /// be reloaded.
+    ///
+    /// This function will panic if the cache is not loaded and the value at the key does not exist.
+    pub fn get_mut(&mut self) -> &mut T {
+        self.cache.get_or_init(|| load_and_deserialize(&self.storage_key));
+        let entry = self.cache.get_mut().expect("cell should be filled above");
+
+        expect_consistent_state(entry.value_mut().as_mut())
+    }
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use crate::test_utils::test_env;
+
+    #[test]
+    pub fn test_lazy() {
+        test_env::setup();
+        let mut a = Lazy::new(b"a", 8u32);
+        assert_eq!(a.get(), &8);
+
+        assert!(!env::storage_has_key(b"a"));
+        a.flush();
+        assert_eq!(u32::try_from_slice(&env::storage_read(b"a").unwrap()).unwrap(), 8);
+
+        a.set(42);
+
+        // Value in storage will still be 8 until the value is flushed
+        assert_eq!(u32::try_from_slice(&env::storage_read(b"a").unwrap()).unwrap(), 8);
+        assert_eq!(*a, 42);
+
+        *a = 30;
+        let serialized = a.try_to_vec().unwrap();
+        drop(a);
+        assert_eq!(u32::try_from_slice(&env::storage_read(b"a").unwrap()).unwrap(), 30);
+
+        let lazy_loaded = Lazy::<u32>::try_from_slice(&serialized).unwrap();
+        assert!(lazy_loaded.cache.get().is_none());
+
+        let b = Lazy::new(b"b", 30);
+        assert!(!env::storage_has_key(b"b"));
+
+        // A value that is not stored in storage yet and one that has not been loaded yet can
+        // be checked for equality.
+        assert_eq!(lazy_loaded, b);
+    }
+}

near-sdk/src/store/mod.rs

@@ -0,0 +1,2 @@
+mod lazy;
+pub use lazy::Lazy;

near-sdk/src/utils/cache_entry.rs

@@ -0,0 +1,69 @@
+#[derive(Clone, Debug)]
+pub(crate) struct CacheEntry<T> {
+    value: Option<T>,
+    state: EntryState,
+}
+
+impl<T> CacheEntry<T> {
+    pub fn new(value: Option<T>, state: EntryState) -> Self {
+        Self { value, state }
+    }
+
+    pub fn new_cached(value: Option<T>) -> Self {
+        Self::new(value, EntryState::Cached)
+    }
+
+    pub fn new_modified(value: Option<T>) -> Self {
+        Self::new(value, EntryState::Modified)
+    }
+
+    pub fn value(&self) -> &Option<T> {
+        &self.value
+    }
+
+    pub fn value_mut(&mut self) -> &mut Option<T> {
+        self.state = EntryState::Modified;
+        &mut self.value
+    }
+
+    #[allow(dead_code)]
+    pub fn into_value(self) -> Option<T> {
+        self.value
+    }
+
+    /// Replaces the current value with a new one. This changes the state of the cell to mutated
+    /// if either the old or new value is [`Some<T>`].
+    #[allow(dead_code)]
+    pub fn replace(&mut self, value: Option<T>) -> Option<T> {
+        let old_value = core::mem::replace(&mut self.value, value);
+
+        if self.value.is_some() || old_value.is_some() {
+            // Set modified if both values are not `None`
+            self.state = EntryState::Modified;
+        }
+
+        old_value
+    }
+
+    /// Replaces the state of the cache entry and returns the previous value.
+    pub fn replace_state(&mut self, state: EntryState) -> EntryState {
+        core::mem::replace(&mut self.state, state)
+    }
+
+    /// Returns true if the entry has been modified
+    pub fn is_modified(&self) -> bool {
+        matches!(self.state, EntryState::Modified)
+    }
+
+    #[allow(dead_code)]
+    /// Returns true if the entry state has not been changed.
+    pub fn is_cached(&self) -> bool {
+        !self.is_modified()
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub(crate) enum EntryState {
+    Modified,
+    Cached,
+}