fetch known code from the database during restoration

ethcore/src/snapshot/account.rs

@@ -19,11 +19,11 @@
 use account_db::{AccountDB, AccountDBMut};
 use snapshot::Error;
 
-use util::{U256, FixedHash, H256, Bytes, HashDB, DBValue, SHA3_EMPTY, SHA3_NULL_RLP};
+use util::{U256, FixedHash, H256, Bytes, HashDB, SHA3_EMPTY, SHA3_NULL_RLP};
 use util::trie::{TrieDB, Trie};
 use rlp::{Rlp, RlpStream, Stream, UntrustedRlp, View};
 
-use std::collections::{HashMap, HashSet};
+use std::collections::HashSet;
 
 // An empty account -- these are replaced with RLP null data for a space optimization.
 const ACC_EMPTY: Account = Account {
@@ -150,7 +150,6 @@ impl Account {
 	pub fn from_fat_rlp(
 		acct_db: &mut AccountDBMut,
 		rlp: UntrustedRlp,
-		code_map: &HashMap<H256, Bytes>,
 	) -> Result<(Self, Option<Bytes>), Error> {
 		use util::{TrieDBMut, TrieMut};
 
@@ -177,9 +176,6 @@ impl Account {
 			}
 			CodeState::Hash => {
 				let code_hash = try!(rlp.val_at(3));
-				if let Some(code) = code_map.get(&code_hash) {
-					acct_db.emplace(code_hash.clone(), DBValue::from_slice(code));
-				}
 
 				(code_hash, None)
 			}
@@ -229,7 +225,7 @@ mod tests {
 	use util::{Address, FixedHash, H256, HashDB, DBValue};
 	use rlp::{UntrustedRlp, View};
 
-	use std::collections::{HashSet, HashMap};
+	use std::collections::HashSet;
 
 	use super::{ACC_EMPTY, Account};
 
@@ -250,7 +246,7 @@ mod tests {
 
 		let fat_rlp = account.to_fat_rlp(&AccountDB::new(db.as_hashdb(), &addr), &mut Default::default()).unwrap();
 		let fat_rlp = UntrustedRlp::new(&fat_rlp);
-		assert_eq!(Account::from_fat_rlp(&mut AccountDBMut::new(db.as_hashdb_mut(), &addr), fat_rlp, &Default::default()).unwrap().0, account);
+		assert_eq!(Account::from_fat_rlp(&mut AccountDBMut::new(db.as_hashdb_mut(), &addr), fat_rlp).unwrap().0, account);
 	}
 
 	#[test]
@@ -275,7 +271,7 @@ mod tests {
 
 		let fat_rlp = account.to_fat_rlp(&AccountDB::new(db.as_hashdb(), &addr), &mut Default::default()).unwrap();
 		let fat_rlp = UntrustedRlp::new(&fat_rlp);
-		assert_eq!(Account::from_fat_rlp(&mut AccountDBMut::new(db.as_hashdb_mut(), &addr), fat_rlp, &Default::default()).unwrap().0, account);
+		assert_eq!(Account::from_fat_rlp(&mut AccountDBMut::new(db.as_hashdb_mut(), &addr), fat_rlp).unwrap().0, account);
 	}
 
 	#[test]
@@ -318,12 +314,11 @@ mod tests {
 		let fat_rlp1 = UntrustedRlp::new(&fat_rlp1);
 		let fat_rlp2 = UntrustedRlp::new(&fat_rlp2);
 
-		let code_map = HashMap::new();
-		let (acc, maybe_code) = Account::from_fat_rlp(&mut AccountDBMut::new(db.as_hashdb_mut(), &addr2), fat_rlp2, &code_map).unwrap();
+		let (acc, maybe_code) = Account::from_fat_rlp(&mut AccountDBMut::new(db.as_hashdb_mut(), &addr2), fat_rlp2).unwrap();
 		assert!(maybe_code.is_none());
 		assert_eq!(acc, account2);
 
-		let (acc, maybe_code) = Account::from_fat_rlp(&mut AccountDBMut::new(db.as_hashdb_mut(), &addr1), fat_rlp1, &code_map).unwrap();
+		let (acc, maybe_code) = Account::from_fat_rlp(&mut AccountDBMut::new(db.as_hashdb_mut(), &addr1), fat_rlp1).unwrap();
 		assert_eq!(maybe_code, Some(b"this is definitely code".to_vec()));
 		assert_eq!(acc, account1);
 	}
@@ -332,9 +327,8 @@ mod tests {
 	fn encoding_empty_acc() {
 		let mut db = get_temp_state_db();
 		let mut used_code = HashSet::new();
-		let code_map = HashMap::new();
 
 		assert_eq!(ACC_EMPTY.to_fat_rlp(&AccountDB::new(db.as_hashdb(), &Address::default()), &mut used_code).unwrap(), ::rlp::NULL_RLP.to_vec());
-		assert_eq!(Account::from_fat_rlp(&mut AccountDBMut::new(db.as_hashdb_mut(), &Address::default()), UntrustedRlp::new(&::rlp::NULL_RLP), &code_map).unwrap(), (ACC_EMPTY, None));
+		assert_eq!(Account::from_fat_rlp(&mut AccountDBMut::new(db.as_hashdb_mut(), &Address::default()), UntrustedRlp::new(&::rlp::NULL_RLP)).unwrap(), (ACC_EMPTY, None));
 	}
 }

ethcore/src/snapshot/mod.rs

@@ -389,7 +389,7 @@ pub fn chunk_state<'a>(db: &HashDB, root: &H256, writer: &Mutex<SnapshotWriter +
 pub struct StateRebuilder {
 	db: Box<JournalDB>,
 	state_root: H256,
-	code_map: HashMap<H256, Bytes>, // maps code hashes to code itself.
+	known_code: HashMap<H256, H256>, // code hashes mapped to first account with this code.
 	missing_code: HashMap<H256, Vec<H256>>, // maps code hashes to lists of accounts missing that code.
 	bloom: Bloom,
 }
@@ -400,7 +400,7 @@ impl StateRebuilder {
 		StateRebuilder {
 			db: journaldb::new(db.clone(), pruning, ::db::COL_STATE),
 			state_root: SHA3_NULL_RLP,
-			code_map: HashMap::new(),
+			known_code: HashMap::new(),
 			missing_code: HashMap::new(),
 			bloom: StateDB::load_bloom(&*db),
 		}
@@ -419,7 +419,7 @@ impl StateRebuilder {
 			self.db.as_hashdb_mut(),
 			rlp,
 			&mut pairs,
-			&self.code_map,
+			&self.known_code,
 			flag
 		));
 
@@ -428,13 +428,13 @@ impl StateRebuilder {
 		}
 
 		// patch up all missing code. must be done after collecting all new missing code entries.
-		for (code_hash, code) in status.new_code {
+		for (code_hash, code, first_with) in status.new_code {
 			for addr_hash in self.missing_code.remove(&code_hash).unwrap_or_else(Vec::new) {
 				let mut db = AccountDBMut::from_hash(self.db.as_hashdb_mut(), addr_hash);
 				db.emplace(code_hash, DBValue::from_slice(&code));
 			}
 
-			self.code_map.insert(code_hash, code);
+			self.known_code.insert(code_hash, first_with);
 		}
 
 		let backing = self.db.backing().clone();
@@ -482,7 +482,8 @@ impl StateRebuilder {
 
 #[derive(Default)]
 struct RebuiltStatus {
-	new_code: Vec<(H256, Bytes)>, // new code that's become available.
+	// new code that's become available. (code_hash, code, addr_hash)
+	new_code: Vec<(H256, Bytes, H256)>,
 	missing_code: Vec<(H256, H256)>, // accounts that are missing code.
 }
 
@@ -492,10 +493,9 @@ fn rebuild_accounts(
 	db: &mut HashDB,
 	account_fat_rlps: UntrustedRlp,
 	out_chunk: &mut [(H256, Bytes)],
-	code_map: &HashMap<H256, Bytes>,
-	abort_flag: &AtomicBool
-) -> Result<RebuiltStatus, ::error::Error>
-{
+	known_code: &HashMap<H256, H256>,
+	abort_flag: &AtomicBool,
+) -> Result<RebuiltStatus, ::error::Error> {
 	let mut status = RebuiltStatus::default();
 	for (account_rlp, out) in account_fat_rlps.into_iter().zip(out_chunk) {
 		if !abort_flag.load(Ordering::SeqCst) { return Err(Error::RestorationAborted.into()) }
@@ -504,17 +504,33 @@ fn rebuild_accounts(
 		let fat_rlp = try!(account_rlp.at(1));
 
 		let thin_rlp = {
-			let mut acct_db = AccountDBMut::from_hash(db, hash);
 
 			// fill out the storage trie and code while decoding.
-			let (acc, maybe_code) = try!(Account::from_fat_rlp(&mut acct_db, fat_rlp, code_map));
+			let (acc, maybe_code) = {
+				let mut acct_db = AccountDBMut::from_hash(db, hash);
+				try!(Account::from_fat_rlp(&mut acct_db, fat_rlp))
+			};
 
 			let code_hash = acc.code_hash().clone();
 			match maybe_code {
-				Some(code) => status.new_code.push((code_hash, code)),
+				// new inline code
+				Some(code) => status.new_code.push((code_hash, code, hash)),
 				None => {
-					if code_hash != ::util::SHA3_EMPTY && !code_map.contains_key(&code_hash) {
-						status.missing_code.push((hash, code_hash));
+					if code_hash != ::util::SHA3_EMPTY {
+						// see if this code has already been included inline
+						match known_code.get(&code_hash) {
+							Some(&first_with) => {
+								// if so, load it from the database.
+								let code = try!(AccountDB::from_hash(db, first_with)
+									.get(&code_hash)
+									.ok_or_else(|| Error::MissingCode(vec![first_with])));
+
+								// and write it again under a different mangled key
+								AccountDBMut::from_hash(db, hash).emplace(code_hash, code);
+							}
+							// if not, queue it up to be filled later
+							None => status.missing_code.push((hash, code_hash)),
+						}
 					}
 				}
 			}

ethcore/src/snapshot/tests/state.rs

@@ -17,6 +17,7 @@
 //! State snapshotting tests.
 
 use snapshot::{chunk_state, Error as SnapshotError, Progress, StateRebuilder};
+use snapshot::account::Account;
 use snapshot::io::{PackedReader, PackedWriter, SnapshotReader, SnapshotWriter};
 use super::helpers::{compare_dbs, StateProducer};
 
@@ -30,6 +31,8 @@ use util::memorydb::MemoryDB;
 use util::Mutex;
 use devtools::RandomTempPath;
 
+use util::sha3::SHA3_NULL_RLP;
+
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 
@@ -88,6 +91,58 @@ fn snap_and_restore() {
 	compare_dbs(&old_db, new_db.as_hashdb());
 }
 
+#[test]
+fn get_code_from_prev_chunk() {
+	use std::collections::HashSet;
+	use rlp::{RlpStream, Stream};
+	use util::{HashDB, H256, FixedHash, U256, Hashable};
+
+	use account_db::{AccountDBMut, AccountDB};
+
+	let code = b"this is definitely code";
+	let mut used_code = HashSet::new();
+	let mut acc_stream = RlpStream::new_list(4);
+	acc_stream.append(&U256::default())
+		.append(&U256::default())
+		.append(&SHA3_NULL_RLP)
+		.append(&code.sha3());
+
+	let (h1, h2) = (H256::random(), H256::random());
+
+	// two accounts with the same code, one per chunk.
+	// first one will have code inlined,
+	// second will just have its hash.
+	let thin_rlp = acc_stream.out();
+	let acc1 = Account::from_thin_rlp(&thin_rlp);
+	let acc2 = Account::from_thin_rlp(&thin_rlp);
+
+	let mut make_chunk = |acc: Account, hash| {
+		let mut db = MemoryDB::new();
+		AccountDBMut::from_hash(&mut db, hash).insert(&code[..]);
+
+		let fat_rlp = acc.to_fat_rlp(&AccountDB::from_hash(&db, hash), &mut used_code).unwrap();
+
+		let mut stream = RlpStream::new_list(1);
+		stream.begin_list(2).append(&hash).append_raw(&fat_rlp, 1);
+		stream.out()
+	};
+
+	let chunk1 = make_chunk(acc1, h1);
+	let chunk2 = make_chunk(acc2, h2);
+
+	let db_path = RandomTempPath::create_dir();
+	let db_cfg = DatabaseConfig::with_columns(::db::NUM_COLUMNS);
+	let new_db = Arc::new(Database::open(&db_cfg, &db_path.to_string_lossy()).unwrap());
+
+	let mut rebuilder = StateRebuilder::new(new_db, Algorithm::Archive);
+	let flag = AtomicBool::new(true);
+
+	rebuilder.feed(&chunk1, &flag).unwrap();
+	rebuilder.feed(&chunk2, &flag).unwrap();
+
+	rebuilder.check_missing().unwrap();
+}
+
 #[test]
 fn checks_flag() {
 	let mut producer = StateProducer::new();