From 3a3a389b5ad492cf221011923e3ea9ddbffc960d Mon Sep 17 00:00:00 2001
From: Aten <jincxmain@gmail.com>
Date: Tue, 22 Jun 2021 19:32:33 +0800
Subject: [PATCH] support revert all state data when revert block (#83)

---
 client/api/src/statekv.rs |  6 ++++++
 client/db/src/lib.rs      | 36 ++++++++++++++++++++++++++++++++++++
 client/rpc/src/lib.rs     | 16 ++++++++++++++--
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/client/api/src/statekv.rs b/client/api/src/statekv.rs
index 28ff59e..b698292 100644
--- a/client/api/src/statekv.rs
+++ b/client/api/src/statekv.rs
@@ -55,6 +55,8 @@ pub trait StateKv<Block: BlockT>: Send + Sync {
 	fn get_contract_tracing(&self, number: NumberFor<Block>, index: u32) -> Option<String>;
 	fn remove_contract_tracing(&self, number: NumberFor<Block>, index: u32) -> error::Result<()>;
 	fn remove_contract_tracings_by_number(&self, number: NumberFor<Block>) -> error::Result<()>;
+
+	fn revert_all(&self, number: NumberFor<Block>) -> error::Result<()>;
 }
 
 pub trait StateKvTransaction {
@@ -173,4 +175,8 @@ impl<Block: BlockT, T: StateKv<Block>> StateKv<Block> for Arc<T> {
 	fn remove_contract_tracings_by_number(&self, number: NumberFor<Block>) -> error::Result<()> {
 		(&**self).remove_contract_tracings_by_number(number)
 	}
+
+	fn revert_all(&self, number: NumberFor<Block>) -> error::Result<()> {
+		(&**self).remove_contract_tracings_by_number(number)
+	}
 }
diff --git a/client/db/src/lib.rs b/client/db/src/lib.rs
index 45b6fdc..dd5ae99 100644
--- a/client/db/src/lib.rs
+++ b/client/db/src/lib.rs
@@ -430,6 +430,42 @@ impl<B: BlockT> ec_client_api::statekv::StateKv<B> for StateKv {
 			t.delete_prefix(columns::TRACING, prefix);
 		})
 	}
+
+	fn revert_all(&self, number: NumberFor<B>) -> error::Result<()> {
+		let hash = <Self as ec_client_api::statekv::StateKv<B>>::get_hash(self, number).ok_or(
+			error::DatabaseError(format!("No hash for this number:{}", number).into()),
+		)?;
+		// state
+		<Self as ec_client_api::statekv::StateKv<B>>::delete_kvs_by_hash(self, hash)?;
+
+		// child state
+		let prefix = hash.as_ref();
+		let hash_len = prefix.len();
+		let mut lookup_key = Vec::with_capacity(hash_len + 1);
+		lookup_key.extend(prefix);
+		lookup_key.push(SEPARATOR);
+
+		let mut t = DBTransaction::with_capacity(1);
+		t.delete_prefix(columns::STATE_CHILD_KV, &lookup_key);
+		self.state_kv_db
+			.write(t)
+			.map_err(|e| error::DatabaseError(Box::new(e)))?;
+
+		// extrinsic changes
+		let num_u64: u64 = number.saturated_into::<u64>();
+		let prefix = &num_u64.to_le_bytes()[..];
+		let mut t = DBTransaction::with_capacity(1);
+		t.delete_prefix(columns::EXTRINSIC_CHANGES, &prefix);
+		self.state_kv_db
+			.write(t)
+			.map_err(|e| error::DatabaseError(Box::new(e)))?;
+
+		// contract tracing
+		<Self as ec_client_api::statekv::StateKv<B>>::remove_contract_tracings_by_number(
+			self, number,
+		)?;
+		Ok(())
+	}
 }
 
 fn tracing_key(number: u64, index: u32) -> Vec<u8> {
diff --git a/client/rpc/src/lib.rs b/client/rpc/src/lib.rs
index ee9b693..ed965d6 100644
--- a/client/rpc/src/lib.rs
+++ b/client/rpc/src/lib.rs
@@ -12,13 +12,16 @@ use jsonrpc_derive::rpc;
 
 use sp_blockchain::HeaderBackend;
 use sp_core::Bytes;
-use sp_runtime::traits::{Block as BlockT, BlockIdTo, Header};
+use sp_runtime::{
+	generic::BlockId,
+	traits::{Block as BlockT, BlockIdTo, Header},
+	SaturatedConversion,
+};
 use sp_utils::mpsc::{tracing_unbounded, TracingUnboundedReceiver, TracingUnboundedSender};
 
 use ec_client_api::statekv;
 
 use error::EuropaRpcError;
-use sp_runtime::generic::BlockId;
 
 pub enum Message<B: BlockT> {
 	Forward(NumberOf<B>),
@@ -128,6 +131,15 @@ where
 		self.backend
 			.revert(diff, true)
 			.map_err(error::client_err::<B>)?;
+		let state_kv = self.client.state_kv();
+		let mut current = best;
+		while current != height {
+			state_kv
+				.revert_all(current)
+				.map_err(|e| error::client_err::<B>(e.into()))?;
+			current -= 1_u64.saturated_into();
+		}
+
 		Ok(())
 	}
 	fn state_kvs(