feat: Implement caching layer for LazyOption

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

@@ -1,6 +1,6 @@
 use borsh::{BorshDeserialize, BorshSerialize};
 
-use super::Lazy;
+use super::{Lazy, LazyOption};
 
 impl<T> Drop for Lazy<T>
 where
@@ -89,3 +89,32 @@ where
         Self::get_mut(self)
     }
 }
+
+impl<T> Drop for LazyOption<T>
+where
+    T: BorshSerialize,
+{
+    fn drop(&mut self) {
+        self.flush()
+    }
+}
+
+impl<T> core::ops::Deref for LazyOption<T>
+where
+    T: BorshSerialize + BorshDeserialize,
+{
+    type Target = Option<T>;
+
+    fn deref(&self) -> &Self::Target {
+        Self::get(self)
+    }
+}
+
+impl<T> core::ops::DerefMut for LazyOption<T>
+where
+    T: BorshSerialize + BorshDeserialize,
+{
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        Self::get_mut(self)
+    }
+}

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

@@ -144,6 +144,137 @@ where
     }
 }
 
+/// An persistent lazily loaded option, 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::store::LazyOption;
+///
+///# near_sdk::test_utils::test_env::setup();
+/// let mut a = LazyOption::new(b"a", None);
+/// assert_eq!(*a, None);
+///
+/// *a = Some("new value".to_owned());
+/// assert_eq!(a.get(), &Some("new value".to_owned()));
+/// ```
+#[derive(BorshSerialize, BorshDeserialize)]
+pub struct LazyOption<T>
+where
+    T: BorshSerialize,
+{
+    /// Key bytes to index the contract's storage.
+    storage_key: Vec<u8>,
+
+    /// Cached value which is lazily loaded and deserialized from storage.
+    #[borsh_skip]
+    cache: OnceCell<CacheEntry<T>>,
+}
+
+impl<T> LazyOption<T>
+where
+    T: BorshSerialize,
+{
+    /// Returns `true` if the value is present in the storage.
+    pub fn is_some(&self) -> bool {
+        self.cache.get().map_or(false, |cache| cache.value().is_some())
+    }
+
+    /// Returns `true` if the value is not present in the storage.
+    pub fn is_none(&self) -> bool {
+        !self.is_some()
+    }
+
+    /// Create a new lazy option with the given `storage_key` and the initial value.
+    pub fn new<S>(storage_key: S, value: Option<T>) -> Self
+    where
+        S: IntoStorageKey,
+    {
+        let cache = match value {
+            Some(value) => CacheEntry::new_modified(Some(value)),
+            None => CacheEntry::new_cached(None),
+        };
+
+        Self { storage_key: storage_key.into_storage_key(), cache: OnceCell::from(cache) }
+    }
+
+    /// Removes the value from storage without reading it, and returns whether the value was present.
+    pub fn remove(&mut self) -> bool {
+        self.take().is_some()
+    }
+
+    /// Replaces the value in the storage and returns the previous value as an option.
+    pub fn replace(&mut self, value: T) -> Option<T> {
+        self.cache.get_mut().map_or(None, |cache| cache.replace(Some(value)))
+    }
+
+    /// Removes the value from storage without reading it, and returning cached value.
+    pub fn take(&mut self) -> Option<T> {
+        self.cache.get_mut().map_or(None, |cache| cache.replace(None))
+    }
+
+    /// Updates the value with a new value. This does not load the current value from storage.
+    /// Returns whether the previous value was present.
+    pub fn set(&mut self, value: T) -> bool {
+        if let Some(v) = self.cache.get_mut() {
+            *v.value_mut() = Some(value);
+            true
+        } else {
+            self.cache
+                .set(CacheEntry::new_modified(Some(value)))
+                .ok()
+                .expect("cache is checked to not be filled above");
+            false
+        }
+    }
+
+    /// 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() {
+                return;
+            }
+
+            match v.value().as_ref() {
+                Some(value) => serialize_and_store(&self.storage_key, value),
+                None => {
+                    env::storage_remove(&self.storage_key);
+                }
+            }
+
+            // 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> LazyOption<T>
+where
+    T: BorshSerialize + BorshDeserialize,
+{
+    /// Returns a reference to the lazily loaded optional.
+    /// The load from storage only happens once, and if the value is already cached, it will not
+    /// be reloaded.
+    pub fn get(&self) -> &Option<T> {
+        let entry = self.cache.get_or_init(|| load_and_deserialize(&self.storage_key));
+        entry.value()
+    }
+
+    /// Returns a reference to the lazily loaded optional.
+    /// The load from storage only happens once, and if the value is already cached, it will not
+    /// be reloaded.
+    pub fn get_mut(&mut self) -> &mut Option<T> {
+        self.cache.get_or_init(|| load_and_deserialize(&self.storage_key));
+        let entry = self.cache.get_mut().expect("cell should be filled above");
+        entry.value_mut()
+    }
+}
+
 #[cfg(not(target_arch = "wasm32"))]
 #[cfg(test)]
 mod tests {
@@ -182,4 +313,41 @@ mod tests {
         // be checked for equality.
         assert_eq!(lazy_loaded, b);
     }
+
+    #[test]
+    pub fn test_lazy_option_1() {
+        test_env::setup();
+        let mut a = LazyOption::new(b"a", None);
+        assert!(a.is_none());
+        assert!(!env::storage_has_key(b"a"));
+
+        // Check value has been set in via cache:
+        a.set(42u32);
+        assert!(a.is_some());
+        assert_eq!(a.get(), &Some(42));
+
+        // Flushing, then check if storage has been set:
+        a.flush();
+        assert!(env::storage_has_key(b"a"));
+        assert_eq!(u32::try_from_slice(&env::storage_read(b"a").unwrap()).unwrap(), 42);
+
+        // New value is set
+        *a = Some(49u32);
+        assert!(a.is_some());
+        assert_eq!(a.get(), &Some(49));
+
+        // Testing `replace`
+        let old = a.replace(69u32);
+        assert!(a.is_some());
+        assert_eq!(old, Some(49));
+
+        // Testing `take` deletes from internal storage
+        let taken = a.take();
+        assert!(a.is_none());
+        assert_eq!(taken, Some(69));
+
+        // `flush`/`drop` after `take` should remove from storage:
+        drop(a);
+        assert!(!env::storage_has_key(b"a"));
+    }
 }

near-sdk/src/store/mod.rs

@@ -1,2 +1,2 @@
 mod lazy;
-pub use lazy::Lazy;
+pub use lazy::{Lazy, LazyOption};