Merge pull request #45 from gavofyork/state

State::code and State::storage_at + tests.

src/account.rs

@@ -6,6 +6,7 @@ use util::bytes::*;
 use util::trie::*;
 use util::rlp::*;
 use util::uint::*;
+use std::cell::*;
 
 pub const SHA3_EMPTY: H256 = H256( [0xc5, 0xd2, 0x46, 0x01, 0x86, 0xf7, 0x23, 0x3c, 0x92, 0x7e, 0x7d, 0xb2, 0xdc, 0xc7, 0x03, 0xc0, 0xe5, 0x00, 0xb6, 0x53, 0xca, 0x82, 0x27, 0x3b, 0x7b, 0xfa, 0xd8, 0x04, 0x5d, 0x85, 0xa4, 0x70] );
 
@@ -19,7 +20,7 @@ pub struct Account {
 	// Trie-backed storage.
 	storage_root: H256,
 	// Overlay on trie-backed storage.
-	storage_overlay: HashMap<H256, H256>,
+	storage_overlay: RefCell<HashMap<H256, H256>>,
 	// Code hash of the account. If None, means that it's a contract whose code has not yet been set.
 	code_hash: Option<H256>,
 	// Code cache of the account.
@@ -33,7 +34,7 @@ impl Account {
 			balance: balance,
 			nonce: nonce,
 			storage_root: SHA3_NULL_RLP,
-			storage_overlay: storage,
+			storage_overlay: RefCell::new(storage),
 			code_hash: Some(code.sha3()),
 			code_cache: code
 		}
@@ -45,7 +46,7 @@ impl Account {
 			balance: balance,
 			nonce: U256::from(0u8),
 			storage_root: SHA3_NULL_RLP,
-			storage_overlay: HashMap::new(),
+			storage_overlay: RefCell::new(HashMap::new()),
 			code_hash: Some(SHA3_EMPTY),
 			code_cache: vec![],
 		}
@@ -58,7 +59,7 @@ impl Account {
 			nonce: r.val_at(0),
 			balance: r.val_at(1),
 			storage_root: r.val_at(2),
-			storage_overlay: HashMap::new(),
+			storage_overlay: RefCell::new(HashMap::new()),
 			code_hash: Some(r.val_at(3)),
 			code_cache: vec![],
 		}
@@ -71,7 +72,7 @@ impl Account {
 			balance: balance,
 			nonce: U256::from(0u8),
 			storage_root: SHA3_NULL_RLP,
-			storage_overlay: HashMap::new(),
+			storage_overlay: RefCell::new(HashMap::new()),
 			code_hash: None,
 			code_cache: vec![],
 		}
@@ -86,20 +87,14 @@ impl Account {
 
 	/// 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.insert(key, value);
+		self.storage_overlay.borrow_mut().insert(key, value);
 	}
 
 	/// Get (and cache) the contents of the trie's storage at `key`.
-	pub fn storage_at(&mut self, db: &mut HashDB, key: H256) -> H256 {
+	pub fn storage_at(&self, db: &HashDB, key: &H256) -> H256 {
-		match self.storage_overlay.get(&key) {
+		self.storage_overlay.borrow_mut().entry(key.clone()).or_insert_with(||{
-			Some(x) => { return x.clone() },
+			H256::from_slice(TrieDB::new(db, &self.storage_root).get(key.bytes()).unwrap_or(&[0u8;32][..]))
-			_ => {}
+		}).clone()
-		}
-		// fetch - cannot be done in match because of the borrow rules.
-		let t = TrieDBMut::new_existing(db, &mut self.storage_root);
-		let r = H256::from_slice(t.get(key.bytes()).unwrap_or(&[0u8;32][..]));
-		self.storage_overlay.insert(key, r.clone());
-		r
 	}
 
 	/// return the balance associated with this account.
@@ -119,6 +114,7 @@ impl Account {
 		match self.code_hash {
 			Some(SHA3_EMPTY) | None if self.code_cache.is_empty() => Some(&self.code_cache),
 			Some(_) if !self.code_cache.is_empty() => Some(&self.code_cache),
+			None => Some(&self.code_cache),
 			_ => None,
 		}
 	}
@@ -161,10 +157,10 @@ impl Account {
 	pub fn base_root(&self) -> &H256 { &self.storage_root }
 	
 	/// return the storage root associated with this account or None if it has been altered via the overlay.
-	pub fn storage_root(&self) -> Option<&H256> { if self.storage_overlay.is_empty() {Some(&self.storage_root)} else {None} }
+	pub fn storage_root(&self) -> Option<&H256> { if self.storage_overlay.borrow().is_empty() {Some(&self.storage_root)} else {None} }
 	
 	/// rturn the storage overlay.
-	pub fn storage_overlay(&self) -> &HashMap<H256, H256> { &self.storage_overlay }
+	pub fn storage_overlay(&self) -> Ref<HashMap<H256, H256>> { self.storage_overlay.borrow() }
 
 	/// Increment the nonce of the account by one.
 	pub fn inc_nonce(&mut self) { self.nonce = self.nonce + U256::from(1u8); }
@@ -178,19 +174,23 @@ impl Account {
 	/// Commit the `storage_overlay` to the backing DB and update `storage_root`.
 	pub fn commit_storage(&mut self, db: &mut HashDB) {
 		let mut t = TrieDBMut::new(db, &mut self.storage_root);
-		for (k, v) in self.storage_overlay.iter() {
+		for (k, v) in self.storage_overlay.borrow().iter() {
 			// cast key and value to trait type,
 			// so we can call overloaded `to_bytes` method
 			t.insert(k, v);
 		}
-		self.storage_overlay.clear();
+		self.storage_overlay.borrow_mut().clear();
 	}
 
 	/// Commit any unsaved code. `code_hash` will always return the hash of the `code_cache` after this.
 	pub fn commit_code(&mut self, db: &mut HashDB) {
+		println!("Commiting code of {:?} - {:?}, {:?}", self, self.code_hash.is_none(), self.code_cache.is_empty());
 		match (self.code_hash.is_none(), self.code_cache.is_empty()) {
-			(true, true) => self.code_hash = Some(self.code_cache.sha3()),
+			(true, true) => self.code_hash = Some(SHA3_EMPTY),
-			(true, false) => self.code_hash = Some(db.insert(&self.code_cache)),
+			(true, false) => {
+				println!("Writing into DB {:?}", self.code_cache);
+				self.code_hash = Some(db.insert(&self.code_cache));
+			},
 			(false, _) => {},
 		}
 	}
@@ -213,7 +213,6 @@ use super::*;
 use std::collections::HashMap;
 use util::hash::*;
 use util::bytes::*;
-use util::trie::*;
 use util::rlp::*;
 use util::uint::*;
 use util::overlaydb::*;
@@ -230,10 +229,10 @@ fn storage_at() {
 		a.rlp()
 	};
 
-	let mut a = Account::from_rlp(&rlp);
+	let a = Account::from_rlp(&rlp);
 	assert_eq!(a.storage_root().unwrap().hex(), "3541f181d6dad5c504371884684d08c29a8bad04926f8ceddf5e279dbc3cc769");
-	assert_eq!(a.storage_at(&mut db, H256::from(&U256::from(0x00u64))), H256::from(&U256::from(0x1234u64)));
+	assert_eq!(a.storage_at(&mut db, &H256::from(&U256::from(0x00u64))), H256::from(&U256::from(0x1234u64)));
-	assert_eq!(a.storage_at(&mut db, H256::from(&U256::from(0x01u64))), H256::new());
+	assert_eq!(a.storage_at(&mut db, &H256::from(&U256::from(0x01u64))), H256::new());
 }
 
 #[test]

src/state.rs

@@ -5,7 +5,6 @@ use util::overlaydb::*;
 use util::trie::*;
 use util::rlp::*;
 use util::uint::*;
-use std::mem;
 //use std::cell::*;
 //use std::ops::*;
 use account::Account;
@@ -61,7 +60,7 @@ impl State {
 	pub fn new_existing(mut db: OverlayDB, mut root: H256, account_start_nonce: U256) -> State {
 		{
 			// trie should panic! if root does not exist
-			let _ = TrieDBMut::new_existing(&mut db, &mut root);
+			let _ = TrieDB::new(&mut db, &mut root);
 		}
 
 		State {
@@ -82,15 +81,9 @@ impl State {
 		&self.root
 	}
 
-	/// Desttroy the current database and return it.
-	/// WARNING: the struct should be dropped immediately following this.
-	pub fn take_db(&mut self) -> OverlayDB {
-		mem::replace(&mut self.db, OverlayDB::new_temp())
-	}
-
 	/// Destroy the current object and return root and database.
-	pub fn drop(mut self) -> (H256, OverlayDB) {
+	pub fn drop(self) -> (H256, OverlayDB) {
-		(mem::replace(&mut self.root, H256::new()), mem::replace(&mut self.db, OverlayDB::new_temp()))
+		(self.root, self.db)
 	}
 
 	/// Expose the underlying database; good to use for calling `state.db().commit()`.
@@ -125,6 +118,23 @@ impl State {
 		self.require(a, false).inc_nonce()
 	}
 
+	/// Mutate storage of account `a` so that it is `value` for `key`.
+	pub fn storage_at(&mut self, a: &Address, key: &H256) -> H256 {
+		self.ensure_cached(a, false);
+		self.try_get(a).map(|a|a.storage_at(&self.db, key)).unwrap_or(H256::new())	
+	}
+
+	/// Mutate storage of account `a` so that it is `value` for `key`.
+	pub fn code(&mut self, a: &Address) -> Option<&[u8]> {
+		self.ensure_cached(a, true);
+		self.try_get(a).map(|a|a.code()).unwrap_or(None)
+	}
+
+	/// Mutate storage of account `a` so that it is `value` for `key`.
+	pub fn set_storage(&mut self, a: &Address, key: H256, value: H256) {
+		self.require(a, false).set_storage(key, value);
+	}
+
 	/// Commit accounts to TrieDBMut. This is similar to cpp-ethereum's dev::eth::commit.
 	/// `accounts` is mutable because we may need to commit the code or storage and record that.
 	pub fn commit_into(db: &mut HashDB, mut root: H256, accounts: &mut HashMap<Address, Option<Account>>) -> H256 {
@@ -158,35 +168,57 @@ impl State {
 		self.root = Self::commit_into(&mut self.db, r, &mut self.cache);
 	}
 
-	/// Pull account `a` in our cache from the trie DB. `require_code` requires that the code be cached, too.
+	/// Pull account `a` in our cache from the trie DB and return it.
-	/// `force_create` creates a new, empty basic account if there is not currently an active account.
+	/// `require_code` requires that the code be cached, too.
-	// TODO: make immutable.
+	// TODO: make immutable through returning an Option<Ref<Account>>
 	fn get(&mut self, a: &Address, require_code: bool) -> Option<&Account> {
+		self.ensure_cached(a, require_code);
+		self.try_get(a)
+	}
+
+	/// Return account `a` from our cache, or None if it doesn't exist in the cache or
+	/// the account is empty.
+	/// Call `ensure_cached` before if you want to avoid the "it doesn't exist in the cache"
+	/// possibility.
+	fn try_get(&self, a: &Address) -> Option<&Account> {
+		self.cache.get(a).map(|x| x.as_ref()).unwrap_or(None)
+	}
+
+	/// Ensure account `a` exists in our cache.
+	/// `require_code` requires that the code be cached, too.
+	fn ensure_cached(&mut self, a: &Address, require_code: bool) {
 		if self.cache.get(a).is_none() {
 			// load from trie.
-			let t = TrieDBMut::new_existing(&mut self.db, &mut self.root);
+			let act = TrieDB::new(&self.db, &self.root).get(&a).map(|rlp| Account::from_rlp(rlp));
-			self.cache.insert(a.clone(), t.get(&a).map(|rlp| { println!("RLP: {:?}", rlp); Account::from_rlp(rlp) }));
+			println!("Loaded {:?} from trie: {:?}", a, act);
+			self.cache.insert(a.clone(), act);
 		}
-
 		let db = &self.db;
-		self.cache.get_mut(a).unwrap().as_mut().map(|account| {
+		if require_code {
-			if require_code {
+			if let Some(ref mut account) = self.cache.get_mut(a).unwrap().as_mut() {
+				println!("Caching code");
 				account.cache_code(db);
+				println!("Now: {:?}", account);
 			}
-			account as &Account
+		}
-		})
 	}
 
 	/// Pull account `a` in our cache from the trie DB. `require_code` requires that the code be cached, too.
 	/// `force_create` creates a new, empty basic account if there is not currently an active account.
 	fn require(&mut self, a: &Address, require_code: bool) -> &mut Account {
+		self.require_or_from(a, require_code, || Account::new_basic(U256::from(0u8)))
+	}
+
+	/// Pull account `a` in our cache from the trie DB. `require_code` requires that the code be cached, too.
+	/// `force_create` creates a new, empty basic account if there is not currently an active account.
+	fn require_or_from<F: FnOnce() -> Account>(&mut self, a: &Address, require_code: bool, default: F) -> &mut Account {
 		if self.cache.get(a).is_none() {
 			// load from trie.
-			self.cache.insert(a.clone(), TrieDBMut::new(&mut self.db, &mut self.root).get(&a).map(|rlp| Account::from_rlp(rlp)));
+			self.cache.insert(a.clone(), TrieDB::new(&self.db, &self.root).get(&a).map(|rlp| Account::from_rlp(rlp)));
 		}
 
 		if self.cache.get(a).unwrap().is_none() {
-			self.cache.insert(a.clone(), Some(Account::new_basic(U256::from(0u8))));
+			self.cache.insert(a.clone(), Some(default()));
 		}
 
 		let db = &self.db;
@@ -208,6 +240,37 @@ use util::trie::*;
 use util::rlp::*;
 use util::uint::*;
 use std::str::FromStr;
+use account::*;
+
+#[test]
+fn code_from_database() {
+	let a = Address::from_str("0000000000000000000000000000000000000000").unwrap();
+	let (r, db) = {
+		let mut s = State::new_temp();
+		s.require_or_from(&a, false, ||Account::new_contract(U256::from(42u32))).set_code(vec![1, 2, 3]);
+		assert_eq!(s.code(&a), Some(&[1u8, 2, 3][..]));
+		s.commit();
+		assert_eq!(s.code(&a), Some(&[1u8, 2, 3][..]));
+		s.drop()
+	};
+
+	let mut s = State::new_existing(db, r, U256::from(0u8));
+	assert_eq!(s.code(&a), Some(&[1u8, 2, 3][..]));
+}
+
+#[test]
+fn storage_at_from_database() {
+	let a = Address::from_str("0000000000000000000000000000000000000000").unwrap();
+	let (r, db) = {
+		let mut s = State::new_temp();
+		s.set_storage(&a, H256::from(&U256::from(01u64)), H256::from(&U256::from(69u64)));
+		s.commit();
+		s.drop()
+	};
+
+	let mut s = State::new_existing(db, r, U256::from(0u8));
+	assert_eq!(s.storage_at(&a, &H256::from(&U256::from(01u64))), H256::from(&U256::from(69u64)));
+}
 
 #[test]
 fn get_from_database() {
@@ -218,7 +281,7 @@ fn get_from_database() {
 		s.add_balance(&a, &U256::from(69u64));
 		s.commit();
 		assert_eq!(s.balance(&a), U256::from(69u64));
-		(s.root().clone(), s.take_db())
+		s.drop()
 	};
 
 	let mut s = State::new_existing(db, r, U256::from(0u8));