slot.rs

//! Managed storage slots
//!
//! Makes it easy to create and manage storage keys and avoid unnecessary
//! writes to contract storage. This reduces transaction IO  and saves on gas.
use std::{marker::PhantomData, ops::Deref};

use near_sdk::{
    borsh::{self, BorshDeserialize, BorshSerialize},
    env, near, IntoStorageKey,
};

use crate::utils::prefix_key;

#[derive(Clone, Debug)]
#[near]
/// A storage slot, composed of a storage location (key) and a data type
pub struct Slot<T> {
    /// The storage key this slot controls.
    pub key: Vec<u8>,
    #[borsh(skip)]
    _marker: PhantomData<T>,
}

impl Slot<()> {
    /// A placeholder slot. Useful for creating namespaced fields.
    pub fn root<K: IntoStorageKey>(key: K) -> Self {
        Self {
            key: key.into_storage_key(),
            _marker: PhantomData,
        }
    }
}

impl<T> Slot<T> {
    /// Creates a new [`Slot`] that controls the given storage key
    pub fn new(key: impl IntoStorageKey) -> Self {
        Self {
            key: key.into_storage_key(),
            _marker: PhantomData,
        }
    }

    /// Creates a new [`Slot`] that controls the given key namespaced (prefixed)
    /// by the parent key, to be used as a namespace for another subfield.
    pub fn ns(&self, key: impl IntoStorageKey) -> Slot<()> {
        Slot {
            key: prefix_key(&self.key, &key.into_storage_key()),
            _marker: PhantomData,
        }
    }

    /// Creates a new [`Slot`] that controls the given key namespaced (prefixed)
    /// by the parent key.
    pub fn field<U>(&self, key: impl IntoStorageKey) -> Slot<U> {
        Slot {
            key: prefix_key(&self.key, &key.into_storage_key()),
            _marker: PhantomData,
        }
    }

    /// Creates a [`Slot`] that tries to parse a different data type from the same
    /// storage slot.
    ///
    /// # Warning
    ///
    /// If the data in the slot is not parsable into the new type, methods like
    /// [`Slot::read`] and [`Slot::take`] will panic.
    #[must_use]
    pub fn transmute<U>(&self) -> Slot<U> {
        Slot {
            key: self.key.clone(),
            _marker: PhantomData,
        }
    }

    /// Write raw bytes into the storage slot. No type checking.
    pub fn write_raw(&mut self, value: &[u8]) -> bool {
        env::storage_write(&self.key, value)
    }

    /// Read raw bytes from the slot. No type checking or parsing.
    #[must_use]
    pub fn read_raw(&self) -> Option<Vec<u8>> {
        env::storage_read(&self.key)
    }

    /// Returns `true` if this slot's key is currently present in the smart
    /// contract storage, `false` otherwise.
    #[must_use]
    pub fn exists(&self) -> bool {
        env::storage_has_key(&self.key)
    }

    /// Removes the managed key from storage.
    pub fn remove(&mut self) -> bool {
        env::storage_remove(&self.key)
    }
}

impl<T: BorshSerialize> Slot<T> {
    /// Writes a value to the managed storage slot.
    ///
    /// # Panics
    ///
    /// If Borsh serialization fails.
    pub fn write(&mut self, value: &T) -> bool {
        self.write_raw(&borsh::to_vec(value).unwrap())
    }

    /// Writes a value to the managed storage slot which is dereferenced from
    /// the target type.
    ///
    /// # Panics
    ///
    /// If Borsh serialization fails.
    pub fn write_deref<U: BorshSerialize + ?Sized>(&mut self, value: &U) -> bool
    where
        T: Deref<Target = U>,
    {
        self.write_raw(&borsh::to_vec(value).unwrap())
    }

    /// If the given value is `Some(T)`, writes `T` to storage. Otherwise,
    /// removes the key from storage.
    pub fn set(&mut self, value: Option<&T>) -> bool {
        match value {
            Some(value) => self.write(value),
            _ => self.remove(),
        }
    }
}

impl<T: BorshDeserialize> Slot<T> {
    /// Reads a value from storage, if present.
    ///
    /// # Panics
    ///
    /// If Borsh deserialization fails.
    #[must_use]
    pub fn read(&self) -> Option<T> {
        self.read_raw().map(|v| T::try_from_slice(&v).unwrap())
    }

    /// Removes a value from storage and returns it if present.
    ///
    /// # Panics
    ///
    /// If Borsh deserialization fails.
    #[must_use]
    pub fn take(&mut self) -> Option<T> {
        if self.remove() {
            // unwrap should be safe if remove returns true
            Some(T::try_from_slice(&env::storage_get_evicted().unwrap()).unwrap())
        } else {
            None
        }
    }
}

impl<T: BorshSerialize + BorshDeserialize> Slot<T> {
    /// Writes a value to storage and returns the evicted value, if present.
    ///
    /// # Panics
    ///
    /// If Borsh serialization fails.
    #[must_use]
    pub fn swap(&mut self, value: &T) -> Option<T> {
        let v = borsh::to_vec(&value).unwrap();

        if self.write_raw(&v) {
            // unwrap should be safe because write_raw returned true
            Some(T::try_from_slice(&env::storage_get_evicted().unwrap()).unwrap())
        } else {
            None
        }
    }
}

impl<T> IntoStorageKey for Slot<T> {
    fn into_storage_key(self) -> Vec<u8> {
        self.key
    }
}

impl<T, U> PartialEq<Slot<U>> for Slot<T> {
    fn eq(&self, other: &Slot<U>) -> bool {
        self.key == other.key
    }
}

#[cfg(test)]
mod tests {
    use super::Slot;

    #[test]
    fn partialeq() {
        let a1 = Slot::<u32>::new(b"a");
        let a2 = Slot::<i32>::new(b"a");
        assert_eq!(a1, a2);
        let b = Slot::<u32>::new(b"b");
        assert_ne!(a1, b);
    }
}