From f1fc3c9ea241860ec81e15840a61213c75c10ec1 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 30 Jan 2023 04:00:38 +0000 Subject: [PATCH 1/2] TrieDB: Use im::OrdMap instead of im::HashMap --- crates/phala-trie-storage/src/lib.rs | 18 +++--- crates/phala-trie-storage/src/memdb.rs | 88 ++++++++++++++------------ crates/phala-trie-storage/src/ser.rs | 21 +++++- 3 files changed, 75 insertions(+), 52 deletions(-) diff --git a/crates/phala-trie-storage/src/lib.rs b/crates/phala-trie-storage/src/lib.rs index b0fb14de93..2d04cfddf6 100644 --- a/crates/phala-trie-storage/src/lib.rs +++ b/crates/phala-trie-storage/src/lib.rs @@ -29,11 +29,13 @@ pub type StorageCollection = Vec<(StorageKey, Option)>; pub type ChildStorageCollection = Vec<(StorageKey, StorageCollection)>; pub type InMemoryBackend = TrieBackend, H>; -pub struct TrieStorage(InMemoryBackend); +pub struct TrieStorage(InMemoryBackend) +where + H::Out: Ord; impl Default for TrieStorage where - H::Out: Codec, + H::Out: Codec + Ord, { fn default() -> Self { Self(TrieBackendBuilder::new(Default::default(), Default::default()).build()) @@ -44,7 +46,7 @@ pub fn load_trie_backend( pairs: impl Iterator, impl AsRef<[u8]>)>, ) -> TrieBackend, H> where - H::Out: Codec, + H::Out: Codec + Ord, { let mut root = Default::default(); let mut mdb = Default::default(); @@ -65,11 +67,11 @@ pub fn serialize_trie_backend( serializer: S, ) -> Result where - H::Out: Codec + Serialize, + H::Out: Codec + Serialize + Ord, S: Serializer, { let root = trie.root(); - let kvs: im::HashMap<_, _> = trie.backend_storage().clone().drain(); + let kvs = trie.backend_storage().clone().drain(); (root, kvs).serialize(serializer) } @@ -78,10 +80,10 @@ pub fn deserialize_trie_backend<'de, H: Hasher, De>( deserializer: De, ) -> Result, H>, De::Error> where - H::Out: Codec + Deserialize<'de>, + H::Out: Codec + Deserialize<'de> + Ord, De: Deserializer<'de>, { - let (root, kvs): (H::Out, im::HashMap<_, (Vec, i32)>) = + let (root, kvs): (H::Out, crate::memdb::Map<_, (Vec, i32)>) = Deserialize::deserialize(deserializer)?; let mdb = MemoryDB::from_inner(kvs); let backend = TrieBackendBuilder::new(mdb, root).build(); @@ -92,7 +94,7 @@ pub fn clone_trie_backend( trie: &TrieBackend, H>, ) -> TrieBackend, H> where - H::Out: Codec, + H::Out: Codec + Ord, { let root = trie.root(); let mdb = trie.backend_storage().clone(); diff --git a/crates/phala-trie-storage/src/memdb.rs b/crates/phala-trie-storage/src/memdb.rs index efdecc683b..e75ae39819 100644 --- a/crates/phala-trie-storage/src/memdb.rs +++ b/crates/phala-trie-storage/src/memdb.rs @@ -17,7 +17,7 @@ use hash_db::{ AsHashDB, AsPlainDB, HashDB, HashDBRef, Hasher as KeyHasher, PlainDB, PlainDBRef, Prefix, }; -use im::{hashmap::Entry, HashMap}; +pub(crate) use im::ordmap::{Entry, OrdMap as Map}; use std::{borrow::Borrow, cmp::Eq, hash, marker::PhantomData, mem}; use sp_state_machine::{backend::Consolidate, DefaultError, TrieBackendStorage}; @@ -28,13 +28,19 @@ impl MaybeDebug for T {} pub type GenericMemoryDB = MemoryDB, trie_db::DBValue>; -impl Consolidate for GenericMemoryDB { +impl Consolidate for GenericMemoryDB +where + H::Out: Ord, +{ fn consolidate(&mut self, other: Self) { MemoryDB::consolidate(self, other) } } -impl TrieBackendStorage for GenericMemoryDB { +impl TrieBackendStorage for GenericMemoryDB +where + H::Out: Ord, +{ type Overlay = Self; fn get( @@ -58,7 +64,7 @@ where H: KeyHasher, KF: KeyFunction, { - data: HashMap, + data: Map, hashed_null_node: H::Out, null_node_data: T, _kf: PhantomData, @@ -80,34 +86,6 @@ where } } -impl PartialEq> for MemoryDB -where - H: KeyHasher, - KF: KeyFunction, - >::Key: Eq + MaybeDebug, - T: Eq + MaybeDebug, -{ - fn eq(&self, other: &MemoryDB) -> bool { - for a in self.data.iter() { - match other.data.get(a.0) { - Some(v) if v != a.1 => return false, - None => return false, - _ => (), - } - } - true - } -} - -impl Eq for MemoryDB -where - H: KeyHasher, - KF: KeyFunction, - >::Key: Eq + MaybeDebug, - T: Eq + MaybeDebug, -{ -} - pub trait KeyFunction { type Key: Send + Sync + Clone + hash::Hash + Eq; @@ -181,6 +159,7 @@ where H: KeyHasher, T: for<'a> From<&'a [u8]> + Clone, KF: KeyFunction, + KF::Key: Ord, { fn default() -> Self { Self::from_null_node(&[0u8][..], [0u8][..].into()) @@ -193,6 +172,7 @@ where H: KeyHasher, T: Default + Clone, KF: KeyFunction, + KF::Key: Ord, { /// Remove an element and delete it from storage if reference count reaches zero. /// If the value was purged, return the old value. @@ -233,11 +213,12 @@ where H: KeyHasher, T: for<'a> From<&'a [u8]> + Clone, KF: KeyFunction, + KF::Key: Ord, { /// Create a new `MemoryDB` from a given null key/data pub fn from_null_node(null_key: &[u8], null_node_data: T) -> Self { MemoryDB { - data: HashMap::default(), + data: Map::default(), hashed_null_node: H::hash(null_key), null_node_data, _kf: Default::default(), @@ -245,7 +226,7 @@ where } /// Create a new `MemoryDB` from a given inner hash map. - pub fn from_inner(data: HashMap) -> Self { + pub fn from_inner(data: Map) -> Self { MemoryDB { data, ..Default::default() @@ -271,13 +252,8 @@ where self.data.clear(); } - /// Purge all zero-referenced data from the database. - pub fn purge(&mut self) { - self.data.retain(|_, (_, rc)| *rc != 0); - } - - /// Return the internal key-value HashMap, clearing the current state. - pub fn drain(&mut self) -> HashMap { + /// Return the internal key-value Map, clearing the current state. + pub fn drain(&mut self) -> Map { mem::take(&mut self.data) } @@ -317,7 +293,7 @@ where } /// Get the keys in the database together with number of underlying references. - pub fn keys(&self) -> HashMap { + pub fn keys(&self) -> Map { self.data .iter() .filter_map(|(k, v)| { @@ -331,12 +307,35 @@ where } } +impl MemoryDB +where + H: KeyHasher, + T: for<'a> From<&'a [u8]> + Clone, + KF: KeyFunction, + KF::Key: Ord, + T: serde::Serialize, +{ + pub fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + use serde::ser::SerializeSeq; + + let mut seq = serializer.serialize_seq(Some(self.data.len()))?; + for (_k, v) in self.data.iter() { + seq.serialize_element(&v)?; + } + seq.end() + } +} + impl PlainDB for MemoryDB where H: KeyHasher, T: Default + PartialEq + for<'a> From<&'a [u8]> + Clone + Send + Sync, KF: Send + Sync + KeyFunction, KF::Key: Borrow<[u8]> + for<'a> From<&'a [u8]>, + KF::Key: Ord, { fn get(&self, key: &H::Out) -> Option { match self.data.get(key.as_ref()) { @@ -390,6 +389,7 @@ where T: Default + PartialEq + for<'a> From<&'a [u8]> + Clone + Send + Sync, KF: Send + Sync + KeyFunction, KF::Key: Borrow<[u8]> + for<'a> From<&'a [u8]>, + KF::Key: Ord, { fn get(&self, key: &H::Out) -> Option { PlainDB::get(self, key) @@ -404,6 +404,7 @@ where H: KeyHasher, T: Default + PartialEq + AsRef<[u8]> + for<'a> From<&'a [u8]> + Clone + Send + Sync, KF: KeyFunction + Send + Sync, + KF::Key: Ord, { fn get(&self, key: &H::Out, prefix: Prefix) -> Option { if key == &self.hashed_null_node { @@ -486,6 +487,7 @@ where H: KeyHasher, T: Default + PartialEq + AsRef<[u8]> + for<'a> From<&'a [u8]> + Clone + Send + Sync, KF: KeyFunction + Send + Sync, + KF::Key: Ord, { fn get(&self, key: &H::Out, prefix: Prefix) -> Option { HashDB::get(self, key, prefix) @@ -501,6 +503,7 @@ where T: Default + PartialEq + for<'a> From<&'a [u8]> + Clone + Send + Sync, KF: KeyFunction + Send + Sync, KF::Key: Borrow<[u8]> + for<'a> From<&'a [u8]>, + KF::Key: Ord, { fn as_plain_db(&self) -> &dyn PlainDB { self @@ -515,6 +518,7 @@ where H: KeyHasher, T: Default + PartialEq + AsRef<[u8]> + for<'a> From<&'a [u8]> + Clone + Send + Sync, KF: KeyFunction + Send + Sync, + KF::Key: Ord, { fn as_hash_db(&self) -> &dyn HashDB { self diff --git a/crates/phala-trie-storage/src/ser.rs b/crates/phala-trie-storage/src/ser.rs index 3ece47ae60..75067562d2 100644 --- a/crates/phala-trie-storage/src/ser.rs +++ b/crates/phala-trie-storage/src/ser.rs @@ -1,6 +1,7 @@ -use serde::{Deserialize, Serialize}; -use parity_scale_codec::{Encode, Decode}; +use hash_db::Hasher; +use parity_scale_codec::{Decode, Encode}; use scale_info::TypeInfo; +use serde::{Deserialize, Serialize}; use super::{ChildStorageCollection, StorageCollection}; @@ -16,3 +17,19 @@ pub struct StorageChanges { pub struct StorageData { pub inner: Vec<(Vec, Vec)>, } + +pub struct SerAsSeq<'a, H: Hasher>(pub &'a crate::MemoryDB) +where + H::Out: Ord; + +impl Serialize for SerAsSeq<'_, H> +where + H::Out: Ord, +{ + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + self.0.serialize(serializer) + } +} From fabcda273c3adf52e0221bf811be53f5a2d6fca9 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 31 Jan 2023 11:37:49 +0800 Subject: [PATCH 2/2] Serializing Map as Seq --- crates/phala-trie-storage/src/lib.rs | 13 ++++++++----- crates/phala-trie-storage/src/memdb.rs | 15 +++------------ 2 files changed, 11 insertions(+), 17 deletions(-) diff --git a/crates/phala-trie-storage/src/lib.rs b/crates/phala-trie-storage/src/lib.rs index 2d04cfddf6..8748c318d3 100644 --- a/crates/phala-trie-storage/src/lib.rs +++ b/crates/phala-trie-storage/src/lib.rs @@ -71,8 +71,8 @@ where S: Serializer, { let root = trie.root(); - let kvs = trie.backend_storage().clone().drain(); - (root, kvs).serialize(serializer) + let kvs = trie.backend_storage(); + (root, ser::SerAsSeq(kvs)).serialize(serializer) } #[cfg(feature = "serde")] @@ -83,9 +83,12 @@ where H::Out: Codec + Deserialize<'de> + Ord, De: Deserializer<'de>, { - let (root, kvs): (H::Out, crate::memdb::Map<_, (Vec, i32)>) = - Deserialize::deserialize(deserializer)?; - let mdb = MemoryDB::from_inner(kvs); + let (root, kvs): (H::Out, Vec<(Vec, i32)>) = Deserialize::deserialize(deserializer)?; + let mdb = MemoryDB::from_inner( + kvs.into_iter() + .map(|(data, rc)| (H::hash(data.as_ref()), (data, rc))) + .collect(), + ); let backend = TrieBackendBuilder::new(mdb, root).build(); Ok(backend) } diff --git a/crates/phala-trie-storage/src/memdb.rs b/crates/phala-trie-storage/src/memdb.rs index e75ae39819..e6c9baadf4 100644 --- a/crates/phala-trie-storage/src/memdb.rs +++ b/crates/phala-trie-storage/src/memdb.rs @@ -18,7 +18,7 @@ use hash_db::{ AsHashDB, AsPlainDB, HashDB, HashDBRef, Hasher as KeyHasher, PlainDB, PlainDBRef, Prefix, }; pub(crate) use im::ordmap::{Entry, OrdMap as Map}; -use std::{borrow::Borrow, cmp::Eq, hash, marker::PhantomData, mem}; +use std::{borrow::Borrow, cmp::Eq, hash, marker::PhantomData}; use sp_state_machine::{backend::Consolidate, DefaultError, TrieBackendStorage}; use trie_db::DBValue; @@ -252,11 +252,6 @@ where self.data.clear(); } - /// Return the internal key-value Map, clearing the current state. - pub fn drain(&mut self) -> Map { - mem::take(&mut self.data) - } - /// Grab the raw information associated with a key. Returns None if the key /// doesn't exist. /// @@ -272,8 +267,8 @@ where } /// Consolidate all the entries of `other` into `self`. - pub fn consolidate(&mut self, mut other: Self) { - for (key, (value, rc)) in other.drain() { + pub fn consolidate(&mut self, other: Self) { + for (key, (value, rc)) in other.data { match self.data.entry(key) { Entry::Occupied(mut entry) => { if entry.get().1 < 0 { @@ -542,12 +537,8 @@ mod tests { let mut m = MemoryDB::, Vec>::default(); m.remove(&hello_key, EMPTY_PREFIX); assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, -1); - m.purge(); - assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, -1); m.insert(EMPTY_PREFIX, hello_bytes); assert_eq!(m.raw(&hello_key, EMPTY_PREFIX), None); - m.purge(); - assert_eq!(m.raw(&hello_key, EMPTY_PREFIX), None); let mut m = MemoryDB::, Vec>::default(); assert!(m.remove_and_purge(&hello_key, EMPTY_PREFIX).is_none());