refactor: impl Trie for MPTTrie (#1363)

refactor: impl Trie for MPTTrie

core/api/src/adapter.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use protocol::traits::{
     APIAdapter, Context, Executor, ExecutorReadOnlyAdapter, MemPool, Network, ReadOnlyStorage,
 };
+use protocol::trie::Trie as _;
 use protocol::types::{
     Account, BigEndianHash, Block, BlockNumber, Bytes, CkbRelatedInfo, ExecutorContext, Hash,
     Header, Metadata, Proposal, Receipt, SignedTransaction, TxResp, H160, H256,
@@ -231,6 +232,7 @@ where
         let hash: Hash = BigEndianHash::from_uint(&position);
         storage_mpt_tree
             .get(hash.as_bytes())?
+            .map(Into::into)
             .ok_or_else(|| APIError::Adapter("Can't find this position".to_string()).into())
     }
 

core/executor/benches/bench_transfer.rs

@@ -9,6 +9,7 @@ use common_crypto::{
 };
 use protocol::codec::{hex_decode, ProtocolCodec};
 use protocol::traits::Executor;
+use protocol::trie::Trie as _;
 use protocol::types::{
     public_to_address, Account, Address, Eip1559Transaction, ExecutorContext, Hash, Public,
     SignedTransaction, TransactionAction, UnsignedTransaction, UnverifiedTransaction, NIL_DATA,
@@ -52,8 +53,8 @@ impl BenchAdapter {
         };
 
         mpt.insert(
-            DISTRIBUTE_ADDRESS.as_slice(),
-            distribute_account.encode().unwrap().as_ref(),
+            DISTRIBUTE_ADDRESS.as_slice().to_vec(),
+            distribute_account.encode().unwrap().to_vec(),
         )
         .unwrap();
 

core/executor/benches/bench_vm.rs

@@ -9,7 +9,7 @@ use core_executor::{AxonExecutor, AxonExecutorApplyAdapter, MPTTrie};
 use protocol::{
     codec::ProtocolCodec,
     traits::{Executor, Storage},
-    trie,
+    trie::{self, Trie as _},
     types::{Account, Address, ExecutorContext},
 };
 
@@ -59,8 +59,11 @@ where
         let db = Arc::new(db);
         let mut mpt = MPTTrie::new(Arc::clone(&db));
 
-        mpt.insert(addr.as_slice(), init_account.encode().unwrap().as_ref())
-            .unwrap();
+        mpt.insert(
+            addr.as_slice().to_vec(),
+            init_account.encode().unwrap().to_vec(),
+        )
+        .unwrap();
 
         let state_root = mpt.commit().unwrap();
         AxonExecutorApplyAdapter::from_root(state_root, db, Arc::new(storage), exec_ctx).unwrap()

core/executor/benches/revm_adapter.rs

@@ -7,6 +7,7 @@ use revm::{AccountInfo, Bytecode, Database, DatabaseCommit};
 use common_merkle::TrieMerkle;
 use core_executor::{code_address, MPTTrie};
 use protocol::traits::{Context, Storage};
+use protocol::trie::Trie as _;
 use protocol::types::{
     Account, Address, Bytes, ExecResp, ExecutorContext, Hasher, SignedTransaction,
     TransactionAction, TxResp, H160, H256, NIL_DATA, RLP_NULL, U256,
@@ -140,8 +141,10 @@ where
                 MPTTrie::from_root(storage_root, Arc::clone(&self.db)).unwrap()
             };
             change.storage.into_iter().for_each(|(k, v)| {
-                let _ =
-                    storage_trie.insert(u256_to_u8_slice(&k), u256_to_u8_slice(&v.present_value()));
+                let _ = storage_trie.insert(
+                    u256_to_u8_slice(&k).to_vec(),
+                    u256_to_u8_slice(&v.present_value()).to_vec(),
+                );
             });
 
             let code_hash = if let Some(code) = change.info.code {
@@ -171,7 +174,7 @@ where
 
             let account_bytes = new_account.encode().unwrap();
             self.trie
-                .insert(addr.as_bytes(), account_bytes.as_ref())
+                .insert(addr.as_bytes().to_vec(), account_bytes.to_vec())
                 .unwrap();
         });
     }
@@ -187,8 +190,8 @@ where
         let distribute_account = account;
 
         mpt.insert(
-            addr.as_slice(),
-            distribute_account.encode().unwrap().as_ref(),
+            addr.as_slice().to_vec(),
+            distribute_account.encode().unwrap().to_vec(),
         )
         .unwrap();
 

core/executor/src/adapter/backend/apply.rs

@@ -6,6 +6,7 @@ use protocol::traits::{
     ApplyBackend, Backend, Context, ExecutorAdapter, ExecutorReadOnlyAdapter, ReadOnlyStorage,
     Storage,
 };
+use protocol::trie::Trie;
 use protocol::types::{
     Account, Bytes, ExecutorContext, Hasher, Log, MerkleRoot, H160, H256, NIL_DATA, RLP_NULL, U256,
 };
@@ -130,7 +131,10 @@ where
     fn save_account(&mut self, address: &H160, account: &Account) {
         self.inner
             .trie
-            .insert(address.as_bytes(), &account.encode().unwrap())
+            .insert(
+                address.as_bytes().to_vec(),
+                account.encode().unwrap().to_vec(),
+            )
             .unwrap();
     }
 }
@@ -171,7 +175,7 @@ where
         };
 
         storage.into_iter().for_each(|(k, v)| {
-            let _ = storage_trie.insert(k.as_bytes(), v.as_bytes());
+            let _ = storage_trie.insert(k.as_bytes().to_vec(), v.as_bytes().to_vec());
         });
 
         let storage_root = storage_trie
@@ -207,7 +211,7 @@ where
         {
             self.inner
                 .trie
-                .insert(address.as_bytes(), bytes.as_ref())
+                .insert(address.as_bytes().to_vec(), bytes.to_vec())
                 .unwrap();
         }
 

core/executor/src/adapter/backend/read_only.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use evm::backend::Basic;
 
 use protocol::traits::{Backend, Context, ExecutorReadOnlyAdapter, ReadOnlyStorage};
+use protocol::trie::Trie as _;
 use protocol::types::{
     Account, Bytes, ExecutorContext, MerkleRoot, H160, H256, NIL_DATA, RLP_NULL, U256,
 };
@@ -32,7 +33,7 @@ where
     }
 
     fn get(&self, key: &[u8]) -> Option<Bytes> {
-        self.trie.get(key).ok().flatten()
+        self.trie.get(key).ok().flatten().map(Into::into)
     }
 
     fn get_account(&self, address: &H160) -> Account {

core/executor/src/adapter/trie/wrapped.rs

@@ -3,73 +3,64 @@ use std::sync::Arc;
 use hasher::HasherKeccak;
 
 use protocol::trie::{PatriciaTrie, Trie, TrieError, DB as TrieDB};
-use protocol::types::{Bytes, MerkleRoot};
-use protocol::{
-    codec::hex_encode, Display, From, ProtocolError, ProtocolErrorKind, ProtocolResult,
-};
+use protocol::types::MerkleRoot;
+use protocol::ProtocolResult;
 
 pub struct MPTTrie<DB: TrieDB>(PatriciaTrie<DB, HasherKeccak>);
 
-impl<DB: TrieDB> MPTTrie<DB> {
-    pub fn new(db: Arc<DB>) -> Self {
-        MPTTrie(PatriciaTrie::new(db, Arc::new(HasherKeccak::new())))
+impl<DB: TrieDB> Trie<DB, HasherKeccak> for MPTTrie<DB> {
+    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, TrieError> {
+        self.0.get(key)
     }
 
-    pub fn from_root(root: MerkleRoot, db: Arc<DB>) -> ProtocolResult<Self> {
-        Ok(MPTTrie(
-            PatriciaTrie::from(db, Arc::new(HasherKeccak::new()), root.as_bytes())
-                .map_err(MPTTrieError::from)?,
-        ))
+    fn contains(&self, key: &[u8]) -> Result<bool, TrieError> {
+        self.0.contains(key)
     }
 
-    pub fn get(&self, key: &[u8]) -> ProtocolResult<Option<Bytes>> {
-        Ok(self
-            .0
-            .get(key)
-            .map_err(MPTTrieError::from)?
-            .map(Bytes::from))
+    fn insert(&mut self, key: Vec<u8>, value: Vec<u8>) -> Result<(), TrieError> {
+        self.0.insert(key, value)
     }
 
-    pub fn contains(&self, key: &[u8]) -> ProtocolResult<bool> {
-        Ok(self.0.contains(key).map_err(MPTTrieError::from)?)
+    fn remove(&mut self, key: &[u8]) -> Result<bool, TrieError> {
+        self.0.remove(key)
     }
 
-    pub fn insert(&mut self, key: &[u8], value: &[u8]) -> ProtocolResult<()> {
-        self.0
-            .insert(key.to_vec(), value.to_vec())
-            .map_err(MPTTrieError::from)?;
-        Ok(())
+    fn root(&mut self) -> Result<Vec<u8>, TrieError> {
+        self.0.root()
     }
 
-    pub fn remove(&mut self, key: &[u8]) -> ProtocolResult<()> {
-        if self.0.remove(key).map_err(MPTTrieError::from)? {
-            Ok(())
-        } else {
-            Err(MPTTrieError::RemoveFailed(hex_encode(key)).into())
-        }
+    fn get_proof(&self, key: &[u8]) -> Result<Vec<Vec<u8>>, TrieError> {
+        self.0.get_proof(key)
     }
 
-    pub fn commit(&mut self) -> ProtocolResult<MerkleRoot> {
-        Ok(MerkleRoot::from_slice(
-            &self.0.root().map_err(MPTTrieError::from)?,
-        ))
+    fn verify_proof(
+        &self,
+        root_hash: &[u8],
+        key: &[u8],
+        proof: Vec<Vec<u8>>,
+    ) -> Result<Option<Vec<u8>>, TrieError> {
+        self.0.verify_proof(root_hash, key, proof)
     }
 }
 
-#[derive(Debug, Display, From)]
-pub enum MPTTrieError {
-    #[display(fmt = "Trie {:?}", _0)]
-    Trie(TrieError),
-
-    #[display(fmt = "Remove {:?} failed", _0)]
-    RemoveFailed(String),
-}
+impl<DB: TrieDB> MPTTrie<DB> {
+    pub fn new(db: Arc<DB>) -> Self {
+        MPTTrie(PatriciaTrie::new(db, Arc::new(HasherKeccak::new())))
+    }
 
-impl std::error::Error for MPTTrieError {}
+    pub fn from_root(root: MerkleRoot, db: Arc<DB>) -> ProtocolResult<Self> {
+        Ok(MPTTrie(PatriciaTrie::from(
+            db,
+            Arc::new(HasherKeccak::new()),
+            root.as_bytes(),
+        )?))
+    }
 
-impl From<MPTTrieError> for ProtocolError {
-    fn from(err: MPTTrieError) -> ProtocolError {
-        ProtocolError::new(ProtocolErrorKind::Executor, Box::new(err))
+    pub fn commit(&mut self) -> ProtocolResult<MerkleRoot> {
+        self.0
+            .root()
+            .map(|r| MerkleRoot::from_slice(&r))
+            .map_err(Into::into)
     }
 }
 
@@ -99,12 +90,12 @@ mod tests {
         let key_2 = rand_bytes(10);
         let val_2 = rand_bytes(20);
 
-        mpt.insert(&key_1, &val_1).unwrap();
-        mpt.insert(&key_2, &val_2).unwrap();
+        mpt.insert(key_1.clone(), val_1.clone()).unwrap();
+        mpt.insert(key_2.clone(), val_2.clone()).unwrap();
         mpt.commit().unwrap();
 
-        assert_eq!(mpt.get(&key_1).unwrap(), Some(Bytes::from(val_1)));
-        assert_eq!(mpt.get(&key_2).unwrap(), Some(Bytes::from(val_2)));
+        assert_eq!(mpt.get(&key_1).unwrap(), Some(val_1));
+        assert_eq!(mpt.get(&key_2).unwrap(), Some(val_2));
         assert!(mpt.remove(&key_1).is_ok());
         assert!(mpt.get(&key_1).unwrap().is_none());
 

core/executor/src/debugger/mod.rs

@@ -12,6 +12,7 @@ use common_config_parser::parse_file;
 use common_crypto::{PrivateKey, Secp256k1RecoverablePrivateKey, Signature};
 use protocol::codec::{hex_decode, ProtocolCodec};
 use protocol::traits::{Backend, Executor};
+use protocol::trie::Trie as _;
 use protocol::types::{
     Account, Eip1559Transaction, ExecResp, ExecutorContext, Hash, Hasher, RichBlock,
     SignedTransaction, TxResp, UnsignedTransaction, UnverifiedTransaction, H160, H256,
@@ -59,8 +60,8 @@ impl EvmDebugger {
             };
 
             mpt.insert(
-                distribute_address.as_bytes(),
-                distribute_account.encode().unwrap().as_ref(),
+                distribute_address.as_bytes().to_vec(),
+                distribute_account.encode().unwrap().to_vec(),
             )
             .unwrap();
         }

core/executor/src/system_contract/ckb_light_client/mod.rs

@@ -8,6 +8,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use ethers::abi::AbiDecode;
 
 use protocol::traits::{ApplyBackend, ExecutorAdapter};
+use protocol::trie::Trie as _;
 use protocol::types::{SignedTransaction, TxResp, H160, H256};
 use protocol::ProtocolResult;
 

core/executor/src/system_contract/ckb_light_client/store.rs

@@ -2,7 +2,8 @@ use std::sync::Arc;
 
 use ethers::abi::{AbiDecode, AbiEncode};
 
-use protocol::{types::H256, ProtocolResult};
+use protocol::trie::Trie as _;
+use protocol::{codec::hex_encode, types::H256, ProtocolResult};
 
 use crate::system_contract::{
     ckb_light_client::ckb_light_client_abi, error::SystemScriptError, HEADER_CELL_DB,
@@ -71,14 +72,23 @@ impl CkbLightClientStore {
 
     fn save_header(&mut self, header: &ckb_light_client_abi::Header) -> ProtocolResult<()> {
         self.trie
-            .insert(&header.block_hash, &header.clone().encode())
+            .insert(header.block_hash.to_vec(), header.clone().encode().to_vec())
             .map_err(|e| SystemScriptError::InsertHeader(e.to_string()).into())
     }
 
     fn remove_header(&mut self, block_hash: &[u8]) -> ProtocolResult<()> {
         self.trie
             .remove(block_hash)
-            .map_err(|e| SystemScriptError::RemoveHeader(e.to_string()).into())
+            .map_err(|e| SystemScriptError::RemoveHeader(e.to_string()))
+            .and_then(|removed| {
+                if removed {
+                    Ok(())
+                } else {
+                    let content = format!("remove header {} failed", hex_encode(block_hash));
+                    Err(SystemScriptError::RemoveHeader(content))
+                }
+            })
+            .map_err(Into::into)
     }
 
     pub fn commit(&mut self) -> ProtocolResult<()> {

core/executor/src/system_contract/image_cell/store.rs

@@ -3,9 +3,7 @@ use std::sync::Arc;
 use ckb_types::{bytes::Bytes, core::cell::CellMeta, packed, prelude::*};
 use rlp::{RlpDecodable, RlpEncodable};
 
-use protocol::ckb_blake2b_256;
-use protocol::types::H256;
-use protocol::ProtocolResult;
+use protocol::{ckb_blake2b_256, codec::hex_encode, trie::Trie as _, types::H256, ProtocolResult};
 
 use crate::system_contract::image_cell::{image_cell_abi, MPTTrie};
 use crate::system_contract::HEADER_CELL_DB;
@@ -164,14 +162,24 @@ impl ImageCellStore {
 
     pub fn insert_cell(&mut self, key: &CellKey, cell: &CellInfo) -> ProtocolResult<()> {
         self.trie
-            .insert(&key.encode(), &rlp::encode(cell))
+            .insert(key.encode().to_vec(), rlp::encode(cell).to_vec())
             .map_err(|e| SystemScriptError::InsertCell(e.to_string()).into())
     }
 
-    pub fn remove_cell(&mut self, key: &CellKey) -> ProtocolResult<()> {
+    pub fn remove_cell(&mut self, cell_key: &CellKey) -> ProtocolResult<()> {
+        let key = cell_key.encode();
         self.trie
-            .remove(&key.encode())
-            .map_err(|e| SystemScriptError::RemoveCell(e.to_string()).into())
+            .remove(&key)
+            .map_err(|e| SystemScriptError::RemoveCell(e.to_string()))
+            .and_then(|removed| {
+                if removed {
+                    Ok(())
+                } else {
+                    let content = format!("remove cell {} failed", hex_encode(&key));
+                    Err(SystemScriptError::RemoveCell(content))
+                }
+            })
+            .map_err(Into::into)
     }
 
     pub fn commit(&mut self) -> ProtocolResult<()> {