diff --git a/crates/store/src/smt/mem_pool_smt_store.rs b/crates/store/src/smt/mem_pool_smt_store.rs index b7968610a..a13c8b8b0 100644 --- a/crates/store/src/smt/mem_pool_smt_store.rs +++ b/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 for MemPoolSMTStore<'a> { fn get_branch(&self, branch_key: &BranchKey) -> Result, 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 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 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 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( diff --git a/crates/store/src/smt/mod.rs b/crates/store/src/smt/mod.rs index 32532490f..575589bfa 100644 --- a/crates/store/src/smt/mod.rs +++ b/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 { diff --git a/crates/store/src/smt/serde.rs b/crates/store/src/smt/serde.rs new file mode 100644 index 000000000..c32ace05c --- /dev/null +++ b/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 { + 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 { + 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!() + } + } +} diff --git a/crates/store/src/smt/smt_store.rs b/crates/store/src/smt/smt_store.rs index a838a4548..9dabf8797 100644 --- a/crates/store/src/smt/smt_store.rs +++ b/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 for SMTStore<'a, DB> { fn get_branch(&self, branch_key: &BranchKey) -> Result, 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 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 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(())