From eb808f7449be4cf4c6ea968481c458b0d3e4703d Mon Sep 17 00:00:00 2001
From: NikVolf <nikvolf@gmail.com>
Date: Tue, 29 Nov 2016 17:24:04 +0300
Subject: [PATCH] rollback function for storage

---
 db/src/storage.rs | 60 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/db/src/storage.rs b/db/src/storage.rs
index 9d6c0bb8..9c1c53c2 100644
--- a/db/src/storage.rs
+++ b/db/src/storage.rs
@@ -383,6 +383,34 @@ impl Storage {
 
 		Ok(Some(reorganization))
 	}
+
+	/// Decanonizes main chain until best block is the one with `hash`
+	pub fn rollback(&self, hash: &H256) -> Result<Vec<H256>, Error> {
+		if self.block_number(hash).is_none() {
+			return Err(Error::not_main(hash));
+		}
+
+		// lock will be held until the end of the routine
+		let mut best_block = self.best_block.write();
+
+		let mut context = UpdateContext::new(&self.database);
+		let mut result = Vec::new();
+		let mut best_number = try!(self.best_number().ok_or(Error::Consistency(ConsistencyError::NoBestBlock)));
+		loop {
+			let next = try!(self.block_hash(best_number).ok_or(Error::unknown_number(best_number)));
+			if &next == hash {
+				context.db_transaction.put(Some(COL_META), KEY_BEST_BLOCK_HASH, &**hash);
+				context.db_transaction.write_u32(Some(COL_META), KEY_BEST_BLOCK_NUMBER, best_number);
+				try!(context.apply(&self.database));
+				*best_block = Some(BestBlock { hash: next, number: best_number });
+
+				return Ok(result);
+			}
+			try!(self.decanonize_block(&mut context, &next));
+			best_number -= 1;
+			result.push(next);
+		}
+	}
 }
 
 impl BlockProvider for Storage {
@@ -1334,6 +1362,38 @@ mod tests {
 		});
 	}
 
+	#[test]
+	fn rollback() {
+		let path = RandomTempPath::create_dir();
+		let store = Storage::new(path.as_path()).unwrap();
+
+		let genesis = test_data::genesis();
+		store.insert_block(&genesis).unwrap();
+
+		let (main_hash1, main_block1) = test_data::block_hash_builder()
+			.block()
+				.header().parent(genesis.hash())
+					.nonce(1)
+					.build()
+				.build()
+			.build();
+		store.insert_block(&main_block1).expect("main block 1 should insert with no problems");
+
+		let (main_hash2, main_block2) = test_data::block_hash_builder()
+			.block()
+				.header().parent(main_hash1.clone())
+					.nonce(2)
+					.build()
+				.build()
+			.build();
+		store.insert_block(&main_block2).expect("main block 2 should insert with no problems");
+		assert_eq!(store.best_block().unwrap().hash, main_hash2);
+
+		store.rollback(&main_hash1).unwrap();
+
+		assert_eq!(store.best_block().unwrap().hash, main_hash1);
+	}
+
 	#[test]
 	fn fork_route() {
 		let path = RandomTempPath::create_dir();