Merge pull request #43 from gavofyork/state

Account::storage_at, Account::ensure_cached and tests.

src/account.rs

@@ -83,10 +83,24 @@ impl Account {
 		self.code_cache = 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.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 {
+		match self.storage_overlay.get(&key) {
+			Some(x) => { return x.clone() },
+			_ => {}
+		}
+		// fetch - cannot be done in match because of the borrow rules.
+		let t = TrieDB::new_existing(db, &mut self.storage_root);
+		let r = H256::from_slice(t.at(key.bytes()).unwrap_or(&[0u8;32][..]));
+		self.storage_overlay.insert(key, r.clone());
+		r
+	}
+
 	/// return the balance associated with this account.
 	pub fn balance(&self) -> &U256 { &self.balance }
 
@@ -109,17 +123,39 @@ impl Account {
 	}
 
 	/// Provide a byte array which hashes to the `code_hash`. returns the hash as a result.
-	pub fn note_code(&mut self, code: Bytes) -> Result<H256, H256> {
+	pub fn note_code(&mut self, code: Bytes) -> Result<(), H256> {
 		let h = code.sha3();
 		match self.code_hash {
 			Some(ref i) if h == *i => {
 				self.code_cache = code;
-				Ok(h)
+				Ok(())
 			},
 			_ => Err(h)
 		}
 	}
 
+	/// Is `code_cache` valid; such that code is going to return Some?
+	pub fn is_cached(&self) -> bool {
+		!self.code_cache.is_empty() || (self.code_cache.is_empty() && self.code_hash == Some(SHA3_EMPTY))
+	}
+
+	/// Provide a database to lookup `code_hash`. Should not be called if it is a contract without code.
+	pub fn ensure_cached(&mut self, db: &HashDB) -> bool {
+		// TODO: fill out self.code_cache;
+/*		return !self.is_cached() ||
+			match db.lookup(&self.code_hash.unwrap()) {	// why doesn't this work? unwrap causes move?!
+				Some(x) => { self.code_cache = x.to_vec(); true },
+				_ => { false }
+			}*/
+		if self.is_cached() { return true; }
+		return if let Some(ref h) = self.code_hash {
+			match db.lookup(&h) {
+				Some(x) => { self.code_cache = x.to_vec(); true },
+				_ => { false }
+			}
+		} else { false }
+	}
+
 	/// return the storage root associated with this account.
 	pub fn base_root(&self) -> &H256 { &self.storage_root }
 	
@@ -151,8 +187,10 @@ impl Account {
 
 	/// 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) {
-		if self.code_hash.is_none() && !self.code_cache.is_empty() {
-			self.code_hash = Some(db.insert(&self.code_cache));
+		match (self.code_hash.is_none(), self.code_cache.is_empty()) {
+			(true, true) => self.code_hash = Some(self.code_cache.sha3()),
+			(true, false) => self.code_hash = Some(db.insert(&self.code_cache)),
+			(false, _) => {},
 		}
 	}
 
@@ -180,14 +218,46 @@ use util::uint::*;
 use util::overlaydb::*;
 
 #[test]
-fn playpen() {
+fn storage_at() {
+	let mut db = OverlayDB::new_temp();
+	let rlp = {
+		let mut a = Account::new_contract(U256::from(69u8));
+		a.set_storage(H256::from(&U256::from(0x00u64)), H256::from(&U256::from(0x1234u64)));
+		a.commit_storage(&mut db);
+		a.set_code(vec![]);
+		a.commit_code(&mut db);
+		a.rlp()
+	};
+
+	let mut 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(0x01u64))), H256::new());
+}
+
+#[test]
+fn note_code() {
+	let mut db = OverlayDB::new_temp();
+
+	let rlp = {
+		let mut a = Account::new_contract(U256::from(69u8));
+		a.set_code(vec![0x55, 0x44, 0xffu8]);
+		a.commit_code(&mut db);
+		a.rlp()
+	};
+
+	let mut a = Account::from_rlp(&rlp);
+	assert_eq!(a.ensure_cached(&db), true);
+
+	let mut a = Account::from_rlp(&rlp);
+	assert_eq!(a.note_code(vec![0x55, 0x44, 0xffu8]), Ok(()));
 }
 
 #[test]
 fn commit_storage() {
 	let mut a = Account::new_contract(U256::from(69u8));
 	let mut db = OverlayDB::new_temp();
-	a.set_storage(From::from(&U256::from(0x00u64)), From::from(&U256::from(0x1234u64)));
+	a.set_storage(H256::from(&U256::from(0x00u64)), H256::from(&U256::from(0x1234u64)));
 	assert_eq!(a.storage_root(), None);
 	a.commit_storage(&mut db);
 	assert_eq!(a.storage_root().unwrap().hex(), "3541f181d6dad5c504371884684d08c29a8bad04926f8ceddf5e279dbc3cc769");

src/state.rs

@@ -63,9 +63,8 @@ impl State {
 		&mut self.db
 	}
 
-	/// Commit accounts to TrieDB. This is simplified version of
-	/// cpp-ethereum's dev::eth::commit.
-	/// accounts mutable because we may need to commit the code or storage and record that.
+	/// Commit accounts to TrieDB. 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 {
 		// first, commit the sub trees.
 		// TODO: is this necessary or can we dispense with the `ref mut a` for just `a`?
@@ -96,4 +95,56 @@ impl State {
 		let r = self.root.clone();	// would prefer not to do this, really. 
 		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.
+	/// `force_create` creates a new, empty basic account if there is not currently an active account.
+	fn ensure_cached(&mut self, a: &Address, require_code: bool, force_create: bool) {
+		if self.cache.get(a).is_none() {
+			// load from trie.
+			self.cache.insert(a.clone(), TrieDB::new(&mut self.db, &mut self.root).at(&a).map(|rlp| Account::from_rlp(rlp)));
+		}
+
+		if self.cache.get(a).unwrap().is_none() {
+			if !force_create { return; }
+			// create a new account
+			self.cache.insert(a.clone(), Some(Account::new_basic(U256::from(0u8))));
+		}
+
+		if require_code {
+			if let &mut Some(ref mut account) = self.cache.get_mut(a).unwrap() {
+				account.ensure_cached(&self.db);
+			}
+		}
+	}
 }
+
+#[cfg(test)]
+mod tests {
+
+use super::*;
+use util::hash::*;
+use util::trie::*;
+use util::rlp::*;
+use std::str::FromStr;
+
+#[test]
+fn playpen() {
+}
+
+#[test]
+fn ensure_cached() {
+	let mut s = State::new_temp();
+	let a = Address::from_str("0000000000000000000000000000000000000000").unwrap();
+	s.ensure_cached(&a, false, true);
+	s.commit();
+	assert_eq!(s.root().hex(), "ec68b85fa2e0526dc0e821a5b33135459114f19173ce0479f5c09b21cc25b9a4");
+}
+
+#[test]
+fn create_empty() {
+	let mut s = State::new_temp();
+	s.commit();
+	assert_eq!(s.root().hex(), "56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421");
+}
+
+}