Preserve cache on reverting the snapshot

ethcore/src/executive.rs

@@ -265,7 +265,7 @@ impl<'a> Executive<'a> {
 			let cost = self.engine.cost_of_builtin(&params.code_address, data);
 			if cost <= params.gas {
 				self.engine.execute_builtin(&params.code_address, data, &mut output);
-				self.state.clear_snapshot();
+				self.state.discard_snapshot();
 
 				// trace only top level calls to builtins to avoid DDoS attacks
 				if self.depth == 0 {
@@ -285,7 +285,7 @@ impl<'a> Executive<'a> {
 				Ok(params.gas - cost)
 			} else {
 				// just drain the whole gas
-				self.state.revert_snapshot();
+				self.state.revert_to_snapshot();
 
 				tracer.trace_failed_call(trace_info, vec![], evm::Error::OutOfGas.into());
 
@@ -331,7 +331,7 @@ impl<'a> Executive<'a> {
 				res
 			} else {
 				// otherwise it's just a basic transaction, only do tracing, if necessary.
-				self.state.clear_snapshot();
+				self.state.discard_snapshot();
 
 				tracer.trace_call(trace_info, U256::zero(), trace_output, vec![]);
 				Ok(params.gas)
@@ -473,10 +473,10 @@ impl<'a> Executive<'a> {
 				| Err(evm::Error::BadInstruction {.. })
 				| Err(evm::Error::StackUnderflow {..})
 				| Err(evm::Error::OutOfStack {..}) => {
-					self.state.revert_snapshot();
+					self.state.revert_to_snapshot();
 			},
 			Ok(_) | Err(evm::Error::Internal) => {
-				self.state.clear_snapshot();
+				self.state.discard_snapshot();
 				substate.accrue(un_substate);
 			}
 		}

ethcore/src/state/account.rs

@@ -394,7 +394,7 @@ impl Account {
 	/// Replace self with the data from other account merging storage cache.
 	/// Basic account data and all modifications are overwritten
 	/// with new values.
-	pub fn merge_with(&mut self, other: Account) {
+	pub fn overwrite_with(&mut self, other: Account) {
 		self.balance = other.balance;
 		self.nonce = other.nonce;
 		self.storage_root = other.storage_root;

ethcore/src/state/mod.rs

@@ -93,21 +93,35 @@ impl AccountEntry {
 		}
 	}
 
-	// Create a new account entry and mark it as dirty
+	// Create a new account entry and mark it as dirty.
 	fn new_dirty(account: Option<Account>) -> AccountEntry {
 		AccountEntry {
 			account: account,
 			state: AccountState::Dirty,
 		}
 	}
-	
-	// Create a new account entry and mark it as clean
+
+	// Create a new account entry and mark it as clean.
 	fn new_clean(account: Option<Account>) -> AccountEntry {
 		AccountEntry {
 			account: account,
 			state: AccountState::Clean,
 		}
 	}
+
+	// Replace data with another entry but preserve storage cache.
+	fn overwrite_with(&mut self, other: AccountEntry) {
+		self.state = other.state;
+		match other.account {
+			Some(acc) => match self.account {
+				Some(ref mut ours) => {
+					ours.overwrite_with(acc);
+				},
+				None => {},
+			},
+			None => self.account = None,
+		}
+	}
 }
 
 /// Representation of the entire state of all accounts in the system.
@@ -142,10 +156,24 @@ impl AccountEntry {
 ///
 /// Upon destruction all the local cache data merged into the global cache.
 /// The merge might be rejected if current state is non-canonical.
+///
+/// State snapshotting.
+///
+/// A new snapshot can be created with `snapshot()`. Snapshots can be
+/// created in a hierarchy.
+/// When a snapshot is active all changes are applied directly into
+/// `cache` and the original value is copied into an active snapshot.
+/// Reverting a snapshot with `revert_to_snapshot` involves copying
+/// original values from the latest snapshot back into `cache`. The code
+/// takes care not to overwrite cached storage while doing that.
+/// Snapshot can be discateded with `discard_snapshot`. All of the orignal
+/// backed-up values are moved into a parent snapshot (if any).
+///
 pub struct State {
 	db: StateDB,
 	root: H256,
 	cache: RefCell<HashMap<Address, AccountEntry>>,
+	// The original account is preserved in
 	snapshots: RefCell<Vec<HashMap<Address, Option<AccountEntry>>>>,
 	account_start_nonce: U256,
 	factories: Factories,
@@ -199,31 +227,44 @@ impl State {
 		Ok(state)
 	}
 
-	/// Create a recoverable snaphot of this state
+	/// Create a recoverable snaphot of this state.
 	pub fn snapshot(&mut self) {
 		self.snapshots.borrow_mut().push(HashMap::new());
 	}
 
-	/// Merge last snapshot with previous
-	pub fn clear_snapshot(&mut self) {
+	/// Merge last snapshot with previous.
+	pub fn discard_snapshot(&mut self) {
 		// merge with previous snapshot
 		let last = self.snapshots.borrow_mut().pop();
 		if let Some(mut snapshot) = last {
 			if let Some(ref mut prev) = self.snapshots.borrow_mut().last_mut() {
-				for (k, v) in snapshot.drain() {
-					prev.entry(k).or_insert(v);
+				if prev.is_empty() {
+					**prev = snapshot;
+				} else {
+					for (k, v) in snapshot.drain() {
+						prev.entry(k).or_insert(v);
+					}
 				}
 			}
 		}
 	}
 
-	/// Revert to snapshot
-	pub fn revert_snapshot(&mut self) {
+	/// Revert to the last snapshot and discard it.
+	pub fn revert_to_snapshot(&mut self) {
 		if let Some(mut snapshot) = self.snapshots.borrow_mut().pop() {
 			for (k, v) in snapshot.drain() {
 				match v {
 					Some(v) => {
-						self.cache.borrow_mut().insert(k, v);
+						match self.cache.borrow_mut().entry(k) {
+							Entry::Occupied(mut e) => {
+								// Merge snapshotted changes back into the main account
+								// storage preserving the cache.
+								e.get_mut().overwrite_with(v);
+							},
+							Entry::Vacant(e) => {
+								e.insert(v);
+							}
+						}
 					},
 					None => {
 						match self.cache.borrow_mut().entry(k) {
@@ -1717,12 +1758,12 @@ fn snapshot_basic() {
 	state.snapshot();
 	state.add_balance(&a, &U256::from(69u64));
 	assert_eq!(state.balance(&a), U256::from(69u64));
-	state.clear_snapshot();
+	state.discard_snapshot();
 	assert_eq!(state.balance(&a), U256::from(69u64));
 	state.snapshot();
 	state.add_balance(&a, &U256::from(1u64));
 	assert_eq!(state.balance(&a), U256::from(70u64));
-	state.revert_snapshot();
+	state.revert_to_snapshot();
 	assert_eq!(state.balance(&a), U256::from(69u64));
 }
 
@@ -1735,9 +1776,9 @@ fn snapshot_nested() {
 	state.snapshot();
 	state.add_balance(&a, &U256::from(69u64));
 	assert_eq!(state.balance(&a), U256::from(69u64));
-	state.clear_snapshot();
+	state.discard_snapshot();
 	assert_eq!(state.balance(&a), U256::from(69u64));
-	state.revert_snapshot();
+	state.revert_to_snapshot();
 	assert_eq!(state.balance(&a), U256::from(0));
 }
 

ethcore/src/state_db.rs

@@ -191,7 +191,7 @@ impl StateDB {
 		for (address, account) in self.cache_overlay.drain(..) {
 			if let Some(&mut Some(ref mut existing)) = cache.accounts.get_mut(&address) {
 				if let Some(new) = account {
-					existing.merge_with(new);
+					existing.overwrite_with(new);
 					continue;
 				}
 			}