fvm_std/collections/

hyper_map.rs

use crate::{runtime};
use crate::spread::SpreadLayout;
use core::hash::Hash;
use core::borrow::Borrow;
use crate::prelude::{Vec};
use crate::collections::entry_state::{ValueWrapper, EntryState};
use scale::{Decode, Encode};
use alloc::collections::BTreeMap;
use alloc::borrow::ToOwned;
use alloc::collections::btree_map::Entry;

/// An non-iterable implementation of a map that stores its content directly on the trie.
pub struct HyperMap<K, V> {
    // the size of the map values
    origin_size: u64,
    size: u64,
    // when deploy: storage_prefix is `_`, it will always empty of read result
    // when invoke: storage_prefix is `field_name` from the framework
    storage_prefix: Vec<u8>,
    // 为每一个元素记录是否有被修改过
    cache_values: BTreeMap<K, ValueWrapper<V>>,
}

impl<K, V> HyperMap<K, V> {
    /// Make a new, empty `HyperMap`,
    ///
    /// It should be called when contract init.
    ///
    /// # Examples
    ///
    /// Basic usage:
    /// ``` ignore
    /// use fvm_std::collections::hyper_map::HyperMap;
    /// use fvm_macros::contract;
    /// use fvm_macros::storage;
    ///
    /// #[storage]
    /// pub struct HyperMapDemo {
    ///     map: HyperMap<String, String>,
    /// }
    ///
    /// #[contract]
    /// impl HyperMapDemo {
    ///     /// used in construct method
    ///     fn new() -> Self {
    ///         Self { map: HyperMap::new() }
    ///     }
    /// }
    /// ```
    pub fn new() -> Self {
        Self {
            origin_size: 0,
            size: 0,
            storage_prefix: "_".as_bytes().to_vec(),
            cache_values: BTreeMap::<K, ValueWrapper<V>>::new(),
        }
    }
}

impl<K, V> HyperMap<K, V>
    where
        K: Ord + Hash + Encode + Decode,
        V: Encode + Decode
{
    /// Returns a reference to the value corresponding to the key.
    ///
    /// The key may be any borrowed form of the map's key type, but the ordering
    /// on the borrowed form *must* match the ordering on the key type.
    pub fn get<Q: ?Sized>(&mut self, k: &Q) -> Option<&V>
        where
            K: Borrow<Q> + Ord,
            Q: Hash + Ord + Encode + ToOwned<Owned=K>,
    {
        match self.cache_values.entry(k.to_owned()) {
            Entry::Vacant(entry) => {
                // get from ledger
                let storage_key = k.encode();
                let value = runtime::storage_read(storage_key.as_slice(), self.storage_prefix.as_slice());
                match value {
                    Some(v) => {
                        let input = &mut &v[..];
                        let res: V = V::decode(input).unwrap();
                        let wrapper = ValueWrapper::new(res, EntryState::NoChanged);
                        // push it to cache
                        let va = entry.insert(wrapper);
                        va.value.as_ref()
                    }
                    None => {
                        None
                    }
                }
            }
            Entry::Occupied(entry) => {
                let ent = entry.into_mut();
                if ent.state == EntryState::Deleted {
                    return None;
                }
                // return from cache
                ent.value.as_ref()
            }
        }
    }

    /// Returns a mutable reference to the value corresponding to the key.
    ///
    /// The key may be any borrowed form of the map's key type, but the ordering
    /// on the borrowed form *must* match the ordering on the key type.
    ///
    /// All `get_mut` result we wil think it to be changed.
    pub fn get_mut<Q: ?Sized>(&mut self, k: &Q) -> Option<&mut V>
        where
            K: Borrow<Q> + Ord,
            Q: Hash + Ord + Encode + ToOwned<Owned=K>,
    {
        match self.cache_values.entry(k.to_owned()) {
            Entry::Vacant(entry) => {
                // get from ledger
                let storage_key = k.encode();
                let value = runtime::storage_read(storage_key.as_slice(), self.storage_prefix.as_slice());
                match value {
                    Some(v) => {
                        let input = &mut &v[..];
                        let res: V = V::decode(input).unwrap();
                        let wrapper = ValueWrapper::new(res, EntryState::Changed);
                        // push it to cache
                        let va = entry.insert(wrapper);
                        va.value.as_mut()
                    }
                    None => {
                        None
                    }
                }
            }
            Entry::Occupied(entry) => {
                // return from cache
                let ent = entry.into_mut();
                if ent.state == EntryState::Deleted {
                    return None;
                }
                ent.state = EntryState::Changed;
                return ent.value.as_mut();
            }
        }
    }


    /// Inserts a key-value pair into the map.
    ///
    /// If the map did not have this key present, [`None`] is returned.
    ///
    /// If the map did have this key present, the value is updated, and the old
    /// value is returned. The key is not updated, though; this matters for
    /// types that can be `==` without being identical. See the [module-level
    /// documentation] for more.
    pub fn insert(&mut self, key: K, value: V) -> Option<V> where K: Clone {
        // insert before get
        match self.get_mut(&key) {
            Some(v) => {
                // change the value and return old value
                let old_value = core::mem::replace(v, value);
                Some(old_value)
            }
            None => {
                // insert to cache first and make it `Changed`
                let wrapper = ValueWrapper::new(value, EntryState::Changed);
                self.cache_values.insert(key, wrapper);
                self.size += 1;
                None
            }
        }
    }

    /// Removes a key from the map, returning the stored key and value if the
    /// key was previously in the map.
    ///
    /// The key may be any borrowed form of the map's key type, but the ordering
    /// on the borrowed form *must* match the ordering on the key type.
    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
        where
            K: Borrow<Q>,
            Q: Hash + Eq + Ord + Encode + ToOwned<Owned=K>,
    {
        // remove before get
        match self.cache_values.entry(k.to_owned()) {
            Entry::Vacant(entry) => {
                // get from ledger
                let storage_key = k.encode();
                let value = runtime::storage_read(storage_key.as_slice(), self.storage_prefix.as_slice());
                match value {
                    Some(v) => {
                        let input = &mut &v[..];
                        let res: V = V::decode(input).unwrap();
                        let wrapper = ValueWrapper::new_empty_value(EntryState::Deleted);
                        // push it to cache
                        entry.insert(wrapper);
                        self.size -= 1;  // change the size of the map
                        Some(res)
                    }
                    None => {
                        None
                    }
                }
            }
            Entry::Occupied(mut entry) => {
                // return from cache
                let old = entry.insert(ValueWrapper::new_empty_value(EntryState::Deleted));
                if old.state == EntryState::Deleted {
                    return None;
                }
                self.size -= 1;  // change the size of the map
                return old.value;
            }
        }
    }

    /// Returns the number of elements in the map.
    ///
    pub fn len(&self) -> u64 {
        self.size
    }

    /// Returns `true` if the map contains no elements.
    ///
    pub fn is_empty(&self) -> bool {
        self.size == 0
    }

    /// Returns the key-value pair corresponding to the supplied key.
    ///
    /// The key may be any borrowed form of the map's key type, but the ordering
    /// on the borrowed form *must* match the ordering on the key type.
    pub fn get_key_value<Q: ?Sized>(&mut self, k: &Q) -> Option<(&K, &V)>
        where
            K: Borrow<Q> + Ord,
            Q: Hash + Ord + Encode + ToOwned<Owned=K>,
    {
        // use the
        match self.get(k) {
            None => { None }
            Some(_) => {
                let (key, value) = self.cache_values.get_key_value(k).unwrap();
                Some((key, &value.value.as_ref().unwrap()))
            }
        }
    }

    /// Returns `true` if the map contains a value for the specified key.
    ///
    /// The key may be any borrowed form of the map's key type, but the ordering
    /// on the borrowed form *must* match the ordering on the key type.
    pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
        where
            K: Borrow<Q> + Ord,
            Q: Hash + Ord + Encode + ToOwned<Owned=K>,
    {
        // form cache first
        if let Some(v) = self.cache_values.get(k) {
            return v.state != EntryState::Deleted;
        }

        // from ledger
        let storage_key = k.encode();
        let value = runtime::storage_read(storage_key.as_slice(), self.storage_prefix.as_slice());

        value.is_some()
    }
}

impl<K, V> Default for HyperMap<K, V> {
    fn default() -> Self {
        HyperMap::new()
    }
}

impl<K, V> SpreadLayout for HyperMap<K, V>
    where
        K: Ord + Hash + Encode + Decode,
        V: Encode + Decode
{
    fn read_ledger(key_prefix: &[u8], _: &[u8]) -> Self {
        // 这里只是创建一个包含缓存的对象
        let mut map = HyperMap::<K, V>::new();
        map.storage_prefix = Vec::from(key_prefix);
        let size = runtime::storage_read(key_prefix, "@HyperMap".as_bytes()).unwrap();
        map.size = <u64 as scale::Decode>::decode(&mut &size[..]).unwrap();
        map.origin_size = map.size;
        map
    }

    fn write_ledger(&self, key_prefix: &[u8], deploy_flag: &[u8]) {
        if deploy_flag == "deploy".as_bytes() {
            // if deploy, write a basic size 0 of the map, real key is: @HyperMap@{fieldName}
            runtime::storage_write(key_prefix, "@HyperMap".as_bytes(), &[0, 0, 0, 0, 0, 0, 0, 0])
        }
        // write all cached k->v
        self.cache_values.iter().for_each(|(key, value)| {
            match value.state {
                EntryState::Changed => {
                    let sto_key = key.encode();
                    let sto_value = value.value.as_ref().unwrap().encode();
                    runtime::storage_write(sto_key.as_slice(), key_prefix, sto_value.as_slice());
                }
                EntryState::Deleted => {
                    let sto_key = key.encode();
                    runtime::storage_delete(sto_key.as_slice(), key_prefix);
                }
                _ => {}
            }
        });
        // update size
        if self.size != self.origin_size {
            runtime::storage_write(key_prefix, "@HyperMap".as_bytes(), self.size.encode().as_slice())
        }
    }
}

#[cfg(test)]
mod test {
    use crate::collections::hyper_map::HyperMap;
    use std::string::String;
    use crate::spread::SpreadLayout;
    use crate::runtime;
    use scale::Encode;
    use scale::Decode;

    #[test]
    fn use_map_normal() {
        let _mock = fvm_mock::build_runtime();
        let mut map: HyperMap<String, String> = HyperMap::new();
        map.insert(String::from("123"), String::from("456"));
        let res = map.get(&String::from("123")).unwrap();
        assert_eq!(*res, String::from("456"))
    }

    #[test]
    fn map_get() {
        let mock = fvm_mock::build_runtime();
        let mut map: HyperMap<i32, i64> = Default::default();
        assert_eq!(map.get(&1), None);

        mock.storage_write(1.encode().as_slice(), "_".as_bytes(), 100i64.encode());
        assert_eq!(map.get(&1), Some(&100i64));
    }


    #[test]
    fn map_get_insert() {
        let _mock = fvm_mock::build_runtime();
        let mut map: HyperMap<i32, i32> = Default::default();
        map.insert(1, 100);
        map.insert(2, 200);
        assert_eq!(map.len(), 2);
        assert_eq!(map.get(&1), Some(&100));
        map.write_ledger("map".as_bytes(), "invoke".as_bytes());
        // check result
        assert_eq!(runtime::storage_read(1.encode().as_slice(), "map".as_bytes()), Some(100.encode()));
        assert_eq!(runtime::storage_read(2.encode().as_slice(), "map".as_bytes()), Some(200.encode()));
        assert_eq!(runtime::storage_read(3.encode().as_slice(), "map".as_bytes()), None);
        assert_eq!(runtime::storage_read("map".as_bytes(), "@HyperMap".as_bytes()), Some(2u64.encode()));
    }

    #[test]
    fn map_delete_normal() {
        let mock = fvm_mock::build_runtime();
        let mut map: HyperMap<i32, i64> = Default::default();
        map.origin_size = 1;
        map.size = 1;
        assert_eq!(map.remove(&1), None);

        mock.storage_write(1.encode().as_slice(), "_".as_bytes(), 100i64.encode());
        assert_eq!(map.remove(&1), Some(100i64));
    }

    #[test]
    fn map_delete_complex() {
        #[derive(Encode, Decode, Eq, PartialEq, Debug)]
        struct Student {
            id: u32,
            name: String,
        }

        let _mock = fvm_mock::build_runtime();
        let mut map: HyperMap<i32, Student> = Default::default();
        map.insert(1, Student { id: 0, name: String::from("Xiao Ming") });
        map.insert(2, Student { id: 1, name: String::from("Xiao Fang") });
        assert_eq!(map.len(), 2);
        assert_eq!(map.get(&1), Some(&Student { id: 0, name: String::from("Xiao Ming") }));

        // begin delete
        let old_value = map.remove(&1);
        assert_eq!(old_value, Some(Student { id: 0, name: String::from("Xiao Ming") }));
        assert_eq!(map.get(&1), None);

        map.write_ledger("map".as_bytes(), "invoke".as_bytes());
        // check result
        assert_eq!(runtime::storage_read(1.encode().as_slice(), "map".as_bytes()), None);
        assert_eq!(runtime::storage_read(2.encode().as_slice(), "map".as_bytes()), Some(vec![1, 0, 0, 0, 36, 88, 105, 97, 111, 32, 70, 97, 110, 103]));
        assert_eq!(runtime::storage_read(3.encode().as_slice(), "map".as_bytes()), None);
        assert_eq!(runtime::storage_read("map".as_bytes(), "@HyperMap".as_bytes()), Some(1u64.encode()));
    }

    #[test]
    fn map_operate_many_times() {
        let _mock = fvm_mock::build_runtime();
        let mut map = HyperMap::new();
        // insert twice
        map.insert(String::from("abc"), String::from("def"));
        map.insert(String::from("abc"), String::from("ggg"));
        assert_eq!(map.get(&String::from("abc")), Some(&String::from("ggg")));

        // get twice
        assert_eq!(map.get(&String::from("abc")), Some(&String::from("ggg")));
        assert_eq!(map.get_mut(&String::from("abc")), Some(&mut String::from("ggg")));

        // delete twice
        let old = map.remove(&String::from("abc"));
        assert_eq!(old, Some(String::from("ggg")));
        let old = map.remove(&String::from("abc"));
        assert_eq!(old, None);

        // check result
        map.write_ledger("map".as_bytes(), "deploy".as_bytes());
        assert_eq!(runtime::storage_read("abc".encode().as_slice(), "map".as_bytes()), None);
        assert_eq!(runtime::storage_read("map".as_bytes(), "@HyperMap".as_bytes()), Some(0u64.encode()));
    }

    #[test]
    fn map_others_methods() {
        let _mock = fvm_mock::build_runtime();
        let mut map = HyperMap::new();
        assert_eq!(map.is_empty(), true);
        map.insert(b'A', 100);
        map.insert(9, 200);
        assert_eq!(map.is_empty(), false);
        assert_eq!(map.len(), 2);

        let key_value = map.get_key_value(&b'A');
        assert!(key_value.is_some());
        assert_eq!(key_value.unwrap().0, &b'A');
        assert_eq!(key_value.unwrap().1, &100);
    }
}