Merge pull request godwokenrises#4 from quake/quake/customize-smt-ser…

…ializing

perf: customize smt node/branch serializing

crates/store/src/smt/mem_pool_smt_store.rs

@@ -9,9 +9,11 @@ use gw_common::{
  },
  H256,
 };
-use gw_types::{packed, prelude::*};
 
-use super::Columns;
+use super::{
+ serde::{branch_key_to_vec, branch_node_to_vec, slice_to_branch_node},
+ Columns,
+};
 
 const DELETED_FLAG: u8 = 0;
 
@@ -43,7 +45,7 @@ impl<'a> MemPoolSMTStore<'a> {
 
 impl<'a> Store<H256> for MemPoolSMTStore<'a> {
  fn get_branch(&self, branch_key: &BranchKey) -> Result<Option<BranchNode>, SMTError> {
- let branch_key: packed::SMTBranchKey = branch_key.pack();
+ let branch_key = branch_key_to_vec(branch_key);
  let opt = self
  .inner_store()
  .get(self.mem_pool_columns.branch_col, branch_key.as_slice())
@@ -52,10 +54,7 @@ impl<'a> Store<H256> for MemPoolSMTStore<'a> {
  self.inner_store()
  .get(self.under_layer_columns.branch_col, branch_key.as_slice())
  })
- .map(|slice| {
- let branch = packed::SMTBranchNodeReader::from_slice_should_be_ok(slice.as_ref());
- branch.to_entity().unpack()
- });
+ .map(|slice| slice_to_branch_node(slice.as_ref()));
  Ok(opt)
  }
 
@@ -79,8 +78,8 @@ impl<'a> Store<H256> for MemPoolSMTStore<'a> {
  }
 
  fn insert_branch(&mut self, branch_key: BranchKey, branch: BranchNode) -> Result<(), SMTError> {
- let branch_key: packed::SMTBranchKey = branch_key.pack();
- let branch: packed::SMTBranchNode = branch.pack();
+ let branch_key = branch_key_to_vec(&branch_key);
+ let branch = branch_node_to_vec(&branch);
 
  self.store
  .insert_raw(
@@ -106,7 +105,7 @@ impl<'a> Store<H256> for MemPoolSMTStore<'a> {
  }
 
  fn remove_branch(&mut self, branch_key: &BranchKey) -> Result<(), SMTError> {
- let branch_key: packed::SMTBranchKey = branch_key.pack();
+ let branch_key = branch_key_to_vec(branch_key);
 
  self.store
  .insert_raw(

crates/store/src/smt/mod.rs

@@ -2,6 +2,7 @@ use gw_db::schema::Col;
 
 pub mod mem_pool_smt_store;
 pub mod mem_smt_store;
+pub mod serde;
 pub mod smt_store;
 
 pub struct Columns {

crates/store/src/smt/serde.rs

@@ -0,0 +1,143 @@
+use std::convert::TryInto;
+
+use gw_common::sparse_merkle_tree::{
+ merge::MergeValue,
+ tree::{BranchKey, BranchNode},
+};
+
+pub fn branch_key_to_vec(key: &BranchKey) -> Vec<u8> {
+ let mut ret = Vec::with_capacity(33);
+ ret.extend_from_slice(key.node_key.as_slice());
+ ret.extend_from_slice(&[key.height]);
+ ret
+}
+
+pub fn branch_node_to_vec(node: &BranchNode) -> Vec<u8> {
+ match (&node.left, &node.right) {
+ (MergeValue::Value(left), MergeValue::Value(right)) => {
+ let mut ret = Vec::with_capacity(33);
+ ret.extend_from_slice(&[0]);
+ ret.extend_from_slice(left.as_slice());
+ ret.extend_from_slice(right.as_slice());
+ ret
+ }
+ (
+ MergeValue::Value(left),
+ MergeValue::MergeWithZero {
+ base_node,
+ zero_bits,
+ zero_count,
+ },
+ ) => {
+ let mut ret = Vec::with_capacity(98);
+ ret.extend_from_slice(&[1]);
+ ret.extend_from_slice(left.as_slice());
+ ret.extend_from_slice(base_node.as_slice());
+ ret.extend_from_slice(zero_bits.as_slice());
+ ret.extend_from_slice(&[*zero_count]);
+ ret
+ }
+ (
+ MergeValue::MergeWithZero {
+ base_node,
+ zero_bits,
+ zero_count,
+ },
+ MergeValue::Value(right),
+ ) => {
+ let mut ret = Vec::with_capacity(98);
+ ret.extend_from_slice(&[2]);
+ ret.extend_from_slice(base_node.as_slice());
+ ret.extend_from_slice(zero_bits.as_slice());
+ ret.extend_from_slice(&[*zero_count]);
+ ret.extend_from_slice(right.as_slice());
+ ret
+ }
+ (
+ MergeValue::MergeWithZero {
+ base_node: l_base_node,
+ zero_bits: l_zero_bits,
+ zero_count: l_zero_count,
+ },
+ MergeValue::MergeWithZero {
+ base_node: r_base_node,
+ zero_bits: r_zero_bits,
+ zero_count: r_zero_count,
+ },
+ ) => {
+ let mut ret = Vec::with_capacity(131);
+ ret.extend_from_slice(&[3]);
+ ret.extend_from_slice(l_base_node.as_slice());
+ ret.extend_from_slice(l_zero_bits.as_slice());
+ ret.extend_from_slice(&[*l_zero_count]);
+ ret.extend_from_slice(r_base_node.as_slice());
+ ret.extend_from_slice(r_zero_bits.as_slice());
+ ret.extend_from_slice(&[*r_zero_count]);
+ ret
+ }
+ }
+}
+
+pub fn slice_to_branch_node(slice: &[u8]) -> BranchNode {
+ match slice[0] {
+ 0 => {
+ let left: [u8; 32] = slice[1..33].try_into().expect("checked slice");
+ let right: [u8; 32] = slice[33..65].try_into().expect("checked slice");
+ BranchNode {
+ left: MergeValue::Value(left.into()),
+ right: MergeValue::Value(right.into()),
+ }
+ }
+ 1 => {
+ let left: [u8; 32] = slice[1..33].try_into().expect("checked slice");
+ let base_node: [u8; 32] = slice[33..65].try_into().expect("checked slice");
+ let zero_bits: [u8; 32] = slice[65..97].try_into().expect("checked slice");
+ let zero_count = slice[97];
+ BranchNode {
+ left: MergeValue::Value(left.into()),
+ right: MergeValue::MergeWithZero {
+ base_node: base_node.into(),
+ zero_bits: zero_bits.into(),
+ zero_count,
+ },
+ }
+ }
+ 2 => {
+ let base_node: [u8; 32] = slice[1..33].try_into().expect("checked slice");
+ let zero_bits: [u8; 32] = slice[33..65].try_into().expect("checked slice");
+ let zero_count = slice[65];
+ let right: [u8; 32] = slice[66..98].try_into().expect("checked slice");
+ BranchNode {
+ left: MergeValue::MergeWithZero {
+ base_node: base_node.into(),
+ zero_bits: zero_bits.into(),
+ zero_count,
+ },
+ right: MergeValue::Value(right.into()),
+ }
+ }
+ 3 => {
+ let l_base_node: [u8; 32] = slice[1..33].try_into().expect("checked slice");
+ let l_zero_bits: [u8; 32] = slice[33..65].try_into().expect("checked slice");
+ let l_zero_count = slice[65];
+ let r_base_node: [u8; 32] = slice[66..98].try_into().expect("checked slice");
+ let r_zero_bits: [u8; 32] = slice[98..130].try_into().expect("checked slice");
+ let r_zero_count = slice[130];
+ BranchNode {
+ left: MergeValue::MergeWithZero {
+ base_node: l_base_node.into(),
+ zero_bits: l_zero_bits.into(),
+ zero_count: l_zero_count,
+ },
+ right: MergeValue::MergeWithZero {
+ base_node: r_base_node.into(),
+ zero_bits: r_zero_bits.into(),
+ zero_count: r_zero_count,
+ },
+ }
+ }
+ _ => {
+ unreachable!()
+ }
+ }
+}

crates/store/src/smt/smt_store.rs

@@ -10,7 +10,8 @@ use gw_common::{
  H256,
 };
 use gw_db::schema::Col;
-use gw_types::{packed, prelude::*};
+
+use super::serde::{branch_key_to_vec, branch_node_to_vec, slice_to_branch_node};
 
 pub struct SMTStore<'a, DB: KVStore> {
  leaf_col: Col,
@@ -34,12 +35,11 @@ impl<'a, DB: KVStore> SMTStore<'a, DB> {
 
 impl<'a, DB: KVStore> Store<H256> for SMTStore<'a, DB> {
  fn get_branch(&self, branch_key: &BranchKey) -> Result<Option<BranchNode>, SMTError> {
- let branch_key: packed::SMTBranchKey = branch_key.pack();
- match self.store.get(self.branch_col, branch_key.as_slice()) {
- Some(slice) => {
- let branch = packed::SMTBranchNodeReader::from_slice_should_be_ok(slice.as_ref());
- Ok(Some(branch.to_entity().unpack()))
- }
+ match self
+ .store
+ .get(self.branch_col, &branch_key_to_vec(branch_key))
+ {
+ Some(slice) => Ok(Some(slice_to_branch_node(&slice))),
  None => Ok(None),
  }
  }
@@ -57,11 +57,12 @@ impl<'a, DB: KVStore> Store<H256> for SMTStore<'a, DB> {
  }
 
  fn insert_branch(&mut self, branch_key: BranchKey, branch: BranchNode) -> Result<(), SMTError> {
- let branch_key: packed::SMTBranchKey = branch_key.pack();
- let branch: packed::SMTBranchNode = branch.pack();
-
  self.store
- .insert_raw(self.branch_col, branch_key.as_slice(), branch.as_slice())
+ .insert_raw(
+ self.branch_col,
+ &branch_key_to_vec(&branch_key),
+ &branch_node_to_vec(&branch),
+ )
  .map_err(|err| SMTError::Store(format!("insert error {}", err)))?;
 
  Ok(())
@@ -76,10 +77,8 @@ impl<'a, DB: KVStore> Store<H256> for SMTStore<'a, DB> {
  }
 
  fn remove_branch(&mut self, branch_key: &BranchKey) -> Result<(), SMTError> {
- let branch_key: packed::SMTBranchKey = branch_key.pack();
-
  self.store
- .delete(self.branch_col, branch_key.as_slice())
+ .delete(self.branch_col, &branch_key_to_vec(branch_key))
  .map_err(|err| SMTError::Store(format!("delete error {}", err)))?;
 
  Ok(())