From 98ae9cad6f042a69dafd53a27498f3fec7f7a0b0 Mon Sep 17 00:00:00 2001
From: Gav Wood <gavin@ethcore.io>
Date: Wed, 29 Jun 2016 16:29:04 +0200
Subject: [PATCH] Minor additions to allow resetting of code. (#1482)

* Minor additions to allow resetting of code.

* Add test.
---
 ethcore/src/account.rs | 23 +++++++++++++++++++++--
 ethcore/src/state.rs   |  7 ++++++-
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/ethcore/src/account.rs b/ethcore/src/account.rs
index d0628436d..2edbf87ae 100644
--- a/ethcore/src/account.rs
+++ b/ethcore/src/account.rs
@@ -108,6 +108,12 @@ impl Account {
 		self.code_cache = code;
 	}
 
+	/// Reset this account's code to the given code.
+	pub fn reset_code(&mut self, code: Bytes) {
+		self.code_hash = None;
+		self.init_code(code);
+	}
+
 	/// Set (and cache) the contents of the trie's storage at `key` to `value`.
 	pub fn set_storage(&mut self, key: H256, value: H256) {
 		self.storage_overlay.borrow_mut().insert(key, (Filth::Dirty, value));
@@ -336,6 +342,21 @@ mod tests {
 		assert_eq!(a.code_hash().hex(), "af231e631776a517ca23125370d542873eca1fb4d613ed9b5d5335a46ae5b7eb");
 	}
 
+	#[test]
+	fn reset_code() {
+		let mut a = Account::new_contract(69.into(), 0.into());
+		let mut db = MemoryDB::new();
+		let mut db = AccountDBMut::new(&mut db, &Address::new());
+		a.init_code(vec![0x55, 0x44, 0xffu8]);
+		assert_eq!(a.code_hash(), SHA3_EMPTY);
+		a.commit_code(&mut db);
+		assert_eq!(a.code_hash().hex(), "af231e631776a517ca23125370d542873eca1fb4d613ed9b5d5335a46ae5b7eb");
+		a.reset_code(vec![0x55]);
+		assert_eq!(a.code_hash(), SHA3_EMPTY);
+		a.commit_code(&mut db);
+		assert_eq!(a.code_hash().hex(), "37bf2238b11b68cdc8382cece82651b59d3c3988873b6e0f33d79694aa45f1be");
+	}
+
 	#[test]
 	fn rlpio() {
 		let a = Account::new(U256::from(69u8), U256::from(0u8), HashMap::new(), Bytes::new());
@@ -348,7 +369,6 @@ mod tests {
 
 	#[test]
 	fn new_account() {
-
 		let a = Account::new(U256::from(69u8), U256::from(0u8), HashMap::new(), Bytes::new());
 		assert_eq!(a.rlp().to_hex(), "f8448045a056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421a0c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470");
 		assert_eq!(a.balance(), &U256::from(69u8));
@@ -359,7 +379,6 @@ mod tests {
 
 	#[test]
 	fn create_account() {
-
 		let a = Account::new(U256::from(69u8), U256::from(0u8), HashMap::new(), Bytes::new());
 		assert_eq!(a.rlp().to_hex(), "f8448045a056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421a0c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470");
 	}
diff --git a/ethcore/src/state.rs b/ethcore/src/state.rs
index e7981abba..0c086ffc3 100644
--- a/ethcore/src/state.rs
+++ b/ethcore/src/state.rs
@@ -208,12 +208,17 @@ impl State {
 		self.require(a, false).set_storage(key, value)
 	}
 
-	/// Initialise the code of account `a` so that it is `value` for `key`.
+	/// Initialise the code of account `a` so that it is `code`.
 	/// NOTE: Account should have been created with `new_contract`.
 	pub fn init_code(&mut self, a: &Address, code: Bytes) {
 		self.require_or_from(a, true, || Account::new_contract(0.into(), self.account_start_nonce), |_|{}).init_code(code);
 	}
 
+	/// Reset the code of account `a` so that it is `code`.
+	pub fn reset_code(&mut self, a: &Address, code: Bytes) {
+		self.require_or_from(a, true, || Account::new_contract(0.into(), self.account_start_nonce), |_|{}).reset_code(code);
+	}	
+
 	/// Execute a given transaction.
 	/// This will change the state accordingly.
 	pub fn apply(&mut self, env_info: &EnvInfo, engine: &Engine, vm_factory: &EvmFactory, t: &SignedTransaction, tracing: bool) -> ApplyResult {