return old trie values on insert and remove (#4053)

* return old trie values on insert and remove

* fix json tests

ethcore/src/json_tests/trie.rs

@@ -33,7 +33,7 @@ fn test_trie(json: &[u8], trie: TrieSpec) -> Vec<String> {
 			let key: Vec<u8> = key.into();
 			let value: Vec<u8> = value.map_or_else(Vec::new, Into::into);
 			t.insert(&key, &value)
-				.expect(&format!("Trie test '{:?}' failed due to internal error", name))
+				.expect(&format!("Trie test '{:?}' failed due to internal error", name));
 		}
 
 		if *t.root() != test.root.into() {

util/src/trie/fatdbmut.rs

@@ -76,18 +76,28 @@ impl<'db> TrieMut for FatDBMut<'db> {
 		self.raw.get(&key.sha3())
 	}
 
-	fn insert(&mut self, key: &[u8], value: &[u8]) -> super::Result<()> {
+	fn insert(&mut self, key: &[u8], value: &[u8]) -> super::Result<Option<DBValue>> {
 		let hash = key.sha3();
-		self.raw.insert(&hash, value)?;
+		let out = self.raw.insert(&hash, value)?;
 		let db = self.raw.db_mut();
-		db.emplace(Self::to_aux_key(&hash), DBValue::from_slice(key));
+
-		Ok(())
+		// don't insert if it doesn't exist.
+		if out.is_none() {
+			db.emplace(Self::to_aux_key(&hash), DBValue::from_slice(key));
+		}
+		Ok(out)
 	}
 
-	fn remove(&mut self, key: &[u8]) -> super::Result<()> {
+	fn remove(&mut self, key: &[u8]) -> super::Result<Option<DBValue>> {
 		let hash = key.sha3();
-		self.raw.db_mut().remove(&Self::to_aux_key(&hash));
+		let out = self.raw.remove(&hash)?;
-		self.raw.remove(&hash)
+
+		// don't remove if it already exists.
+		if out.is_some() {
+			self.raw.db_mut().remove(&Self::to_aux_key(&hash));
+		}
+
+		Ok(out)
 	}
 }
 

util/src/trie/mod.rs

@@ -119,13 +119,13 @@ pub trait TrieMut {
 	/// What is the value of the given key in this trie?
 	fn get<'a, 'key>(&'a self, key: &'key [u8]) -> Result<Option<DBValue>> where 'a: 'key;
 
-	/// Insert a `key`/`value` pair into the trie. An `empty` value is equivalent to removing
+	/// Insert a `key`/`value` pair into the trie. An empty value is equivalent to removing
-	/// `key` from the trie.
+	/// `key` from the trie. Returns the old value associated with this key, if it existed.
-	fn insert(&mut self, key: &[u8], value: &[u8]) -> Result<()>;
+	fn insert(&mut self, key: &[u8], value: &[u8]) -> Result<Option<DBValue>>;
 
 	/// Remove a `key` from the trie. Equivalent to making it equal to the empty
-	/// value.
+	/// value. Returns the old value associated with this key, if it existed.
-	fn remove(&mut self, key: &[u8]) -> Result<()>;
+	fn remove(&mut self, key: &[u8]) -> Result<Option<DBValue>>;
 }
 
 /// A trie iterator that also supports random access.

util/src/trie/sectriedbmut.rs

@@ -68,11 +68,11 @@ impl<'db> TrieMut for SecTrieDBMut<'db> {
 		self.raw.get(&key.sha3())
 	}
 
-	fn insert(&mut self, key: &[u8], value: &[u8]) -> super::Result<()> {
+	fn insert(&mut self, key: &[u8], value: &[u8]) -> super::Result<Option<DBValue>> {
 		self.raw.insert(&key.sha3(), value)
 	}
 
-	fn remove(&mut self, key: &[u8]) -> super::Result<()> {
+	fn remove(&mut self, key: &[u8]) -> super::Result<Option<DBValue>> {
 		self.raw.remove(&key.sha3())
 	}
 }

util/src/trie/triedbmut.rs

@@ -452,14 +452,16 @@ impl<'a> TrieDBMut<'a> {
 	}
 
 	/// insert a key, value pair into the trie, creating new nodes if necessary.
-	fn insert_at(&mut self, handle: NodeHandle, partial: NibbleSlice, value: DBValue) -> super::Result<(StorageHandle, bool)> {
+	fn insert_at(&mut self, handle: NodeHandle, partial: NibbleSlice, value: DBValue, old_val: &mut Option<DBValue>)
+		-> super::Result<(StorageHandle, bool)>
+	{
 		let h = match handle {
 			NodeHandle::InMemory(h) => h,
 			NodeHandle::Hash(h) => self.cache(h)?,
 		};
 		let stored = self.storage.destroy(h);
 		let (new_stored, changed) = self.inspect(stored, move |trie, stored| {
-			trie.insert_inspector(stored, partial, value).map(|a| a.into_action())
+			trie.insert_inspector(stored, partial, value, old_val).map(|a| a.into_action())
 		})?.expect("Insertion never deletes.");
 
 		Ok((self.storage.alloc(new_stored), changed))
@@ -467,7 +469,9 @@ impl<'a> TrieDBMut<'a> {
 
 	/// the insertion inspector.
 	#[cfg_attr(feature = "dev", allow(cyclomatic_complexity))]
-	fn insert_inspector(&mut self, node: Node, partial: NibbleSlice, value: DBValue) -> super::Result<InsertAction> {
+	fn insert_inspector(&mut self, node: Node, partial: NibbleSlice, value: DBValue, old_val: &mut Option<DBValue>)
+		-> super::Result<InsertAction>
+	{
 		trace!(target: "trie", "augmented (partial: {:?}, value: {:?})", partial, value.pretty());
 
 		Ok(match node {
@@ -481,6 +485,8 @@ impl<'a> TrieDBMut<'a> {
 				if partial.is_empty() {
 					let unchanged = stored_value.as_ref() == Some(&value);
 					let branch = Node::Branch(children, Some(value));
+					*old_val = stored_value;
+
 					match unchanged {
 						true => InsertAction::Restore(branch),
 						false => InsertAction::Replace(branch),
@@ -490,7 +496,7 @@ impl<'a> TrieDBMut<'a> {
 					let partial = partial.mid(1);
 					if let Some(child) = children[idx].take() {
 						// original had something there. recurse down into it.
-						let (new_child, changed) = self.insert_at(child, partial, value)?;
+						let (new_child, changed) = self.insert_at(child, partial, value, old_val)?;
 						children[idx] = Some(new_child.into());
 						if !changed {
 							// the new node we composed didn't change. that means our branch is untouched too.
@@ -511,7 +517,10 @@ impl<'a> TrieDBMut<'a> {
 				if cp == existing_key.len() && cp == partial.len() {
 					trace!(target: "trie", "equivalent-leaf: REPLACE");
 					// equivalent leaf.
-					match stored_value == value {
+					let unchanged = stored_value == value;
+					*old_val = Some(stored_value);
+
+					match unchanged {
 						// unchanged. restore
 						true => InsertAction::Restore(Node::Leaf(encoded.clone(), value)),
 						false => InsertAction::Replace(Node::Leaf(encoded.clone(), value)),
@@ -533,7 +542,7 @@ impl<'a> TrieDBMut<'a> {
 					};
 
 					// always replace because whatever we get out here is not the branch we started with.
-					let branch_action = self.insert_inspector(branch, partial, value)?.unwrap_node();
+					let branch_action = self.insert_inspector(branch, partial, value, old_val)?.unwrap_node();
 					InsertAction::Replace(branch_action)
 				} else if cp == existing_key.len() {
 				    trace!(target: "trie", "complete-prefix (cp={:?}): AUGMENT-AT-END", cp);
@@ -542,7 +551,7 @@ impl<'a> TrieDBMut<'a> {
 					// make a stub branch and an extension.
 					let branch = Node::Branch(empty_children(), Some(stored_value));
 					// augment the new branch.
-					let branch = self.insert_inspector(branch, partial.mid(cp), value)?.unwrap_node();
+					let branch = self.insert_inspector(branch, partial.mid(cp), value, old_val)?.unwrap_node();
 
 					// always replace since we took a leaf and made an extension.
 					let branch_handle = self.storage.alloc(Stored::New(branch)).into();
@@ -553,9 +562,10 @@ impl<'a> TrieDBMut<'a> {
 					// partially-shared prefix for an extension.
 					// start by making a leaf.
 					let low = Node::Leaf(existing_key.mid(cp).encoded(true), stored_value);
+
 					// augment it. this will result in the Leaf -> cp == 0 routine,
 					// which creates a branch.
-					let augmented_low = self.insert_inspector(low, partial.mid(cp), value)?.unwrap_node();
+					let augmented_low = self.insert_inspector(low, partial.mid(cp), value, old_val)?.unwrap_node();
 
 					// make an extension using it. this is a replacement.
 					InsertAction::Replace(Node::Extension(
@@ -586,7 +596,7 @@ impl<'a> TrieDBMut<'a> {
 					};
 
 					// continue inserting.
-					let branch_action = self.insert_inspector(Node::Branch(children, None), partial, value)?.unwrap_node();
+					let branch_action = self.insert_inspector(Node::Branch(children, None), partial, value, old_val)?.unwrap_node();
 					InsertAction::Replace(branch_action)
 				} else if cp == existing_key.len() {
 					trace!(target: "trie", "complete-prefix (cp={:?}): AUGMENT-AT-END", cp);
@@ -594,7 +604,7 @@ impl<'a> TrieDBMut<'a> {
 					// fully-shared prefix.
 
 					// insert into the child node.
-					let (new_child, changed) = self.insert_at(child_branch, partial.mid(cp), value)?;
+					let (new_child, changed) = self.insert_at(child_branch, partial.mid(cp), value, old_val)?;
 					let new_ext = Node::Extension(existing_key.encoded(false), new_child.into());
 
 					// if the child branch wasn't changed, meaning this extension remains the same.
@@ -608,7 +618,7 @@ impl<'a> TrieDBMut<'a> {
 					// partially-shared.
 					let low = Node::Extension(existing_key.mid(cp).encoded(false), child_branch);
 					// augment the extension. this will take the cp == 0 path, creating a branch.
-					let augmented_low = self.insert_inspector(low, partial.mid(cp), value)?.unwrap_node();
+					let augmented_low = self.insert_inspector(low, partial.mid(cp), value, old_val)?.unwrap_node();
 
 					// always replace, since this extension is not the one we started with.
 					// this is known because the partial key is only the common prefix.
@@ -622,7 +632,9 @@ impl<'a> TrieDBMut<'a> {
 	}
 
 	/// Remove a node from the trie based on key.
-	fn remove_at(&mut self, handle: NodeHandle, partial: NibbleSlice) -> super::Result<Option<(StorageHandle, bool)>> {
+	fn remove_at(&mut self, handle: NodeHandle, partial: NibbleSlice, old_val: &mut Option<DBValue>)
+		-> super::Result<Option<(StorageHandle, bool)>>
+	{
 		let stored = match handle {
 			NodeHandle::InMemory(h) => self.storage.destroy(h),
 			NodeHandle::Hash(h) => {
@@ -631,17 +643,18 @@ impl<'a> TrieDBMut<'a> {
 			}
 		};
 
-		let opt = self.inspect(stored, move |trie, node| trie.remove_inspector(node, partial))?;
+		let opt = self.inspect(stored, move |trie, node| trie.remove_inspector(node, partial, old_val))?;
 
 		Ok(opt.map(|(new, changed)| (self.storage.alloc(new), changed)))
 	}
 
 	/// the removal inspector
-	fn remove_inspector(&mut self, node: Node, partial: NibbleSlice) -> super::Result<Action> {
+	fn remove_inspector(&mut self, node: Node, partial: NibbleSlice, old_val: &mut Option<DBValue>) -> super::Result<Action> {
 		Ok(match (node, partial.is_empty()) {
 			(Node::Empty, _) => Action::Delete,
 			(Node::Branch(c, None), true) => Action::Restore(Node::Branch(c, None)),
-			(Node::Branch(children, _), true) => {
+			(Node::Branch(children, Some(val)), true) => {
+				*old_val = Some(val);
 				// always replace since we took the value out.
 				Action::Replace(self.fix(Node::Branch(children, None))?)
 			}
@@ -649,7 +662,7 @@ impl<'a> TrieDBMut<'a> {
 				let idx = partial.at(0) as usize;
 				if let Some(child) = children[idx].take() {
 					trace!(target: "trie", "removing value out of branch child, partial={:?}", partial);
-					match self.remove_at(child, partial.mid(1))? {
+					match self.remove_at(child, partial.mid(1), old_val)? {
 						Some((new, changed)) => {
 							children[idx] = Some(new.into());
 							let branch = Node::Branch(children, value);
@@ -675,6 +688,7 @@ impl<'a> TrieDBMut<'a> {
 			(Node::Leaf(encoded, value), _) => {
 				if NibbleSlice::from_encoded(&encoded).0 == partial {
 					// this is the node we were looking for. Let's delete it.
+					*old_val = Some(value);
 					Action::Delete
 				} else {
 					// leaf the node alone.
@@ -690,7 +704,7 @@ impl<'a> TrieDBMut<'a> {
 				if cp == existing_len {
 					// try to remove from the child branch.
 					trace!(target: "trie", "removing from extension child, partial={:?}", partial);
-					match self.remove_at(child_branch, partial.mid(cp))? {
+					match self.remove_at(child_branch, partial.mid(cp), old_val)? {
 						Some((new_child, changed)) => {
 							let new_child = new_child.into();
 
@@ -907,28 +921,35 @@ impl<'a> TrieMut for TrieDBMut<'a> {
 	}
 
 
-	fn insert(&mut self, key: &[u8], value: &[u8]) -> super::Result<()> {
+	fn insert(&mut self, key: &[u8], value: &[u8]) -> super::Result<Option<DBValue>> {
-		if value.is_empty() {
+		if value.is_empty() { return self.remove(key) }
-			return self.remove(key);
+
-		}
+		let mut old_val = None;
 
 		trace!(target: "trie", "insert: key={:?}, value={:?}", key.pretty(), value.pretty());
 
 		let root_handle = self.root_handle();
-		let (new_handle, changed) = self.insert_at(root_handle, NibbleSlice::new(key), DBValue::from_slice(value))?;
+		let (new_handle, changed) = self.insert_at(
+			root_handle,
+			NibbleSlice::new(key),
+			DBValue::from_slice(value),
+			&mut old_val,
+		)?;
 
 		trace!(target: "trie", "insert: altered trie={}", changed);
 		self.root_handle = NodeHandle::InMemory(new_handle);
 
-		Ok(())
+		Ok(old_val)
 	}
 
-	fn remove(&mut self, key: &[u8]) -> super::Result<()> {
+	fn remove(&mut self, key: &[u8]) -> super::Result<Option<DBValue>> {
 		trace!(target: "trie", "remove: key={:?}", key.pretty());
 
 		let root_handle = self.root_handle();
 		let key = NibbleSlice::new(key);
-		match self.remove_at(root_handle, key)? {
+		let mut old_val = None;
+
+		match self.remove_at(root_handle, key, &mut old_val)? {
 			Some((handle, changed)) => {
 				trace!(target: "trie", "remove: altered trie={}", changed);
 				self.root_handle = NodeHandle::InMemory(handle);
@@ -938,9 +959,9 @@ impl<'a> TrieMut for TrieDBMut<'a> {
 				self.root_handle = NodeHandle::Hash(SHA3_NULL_RLP);
 				*self.root = SHA3_NULL_RLP;
 			}
-		};
+		}
 
-		Ok(())
+		Ok(old_val)
 	}
 }
 
@@ -1287,4 +1308,29 @@ mod tests {
 		assert!(t.is_empty());
 		assert_eq!(*t.root(), SHA3_NULL_RLP);
 	}
+
+	#[test]
+	fn return_old_values() {
+		let mut seed = H256::new();
+		let x = StandardMap {
+				alphabet: Alphabet::Custom(b"@QWERTYUIOPASDFGHJKLZXCVBNM[/]^_".to_vec()),
+				min_key: 5,
+				journal_key: 0,
+				value_mode: ValueMode::Index,
+				count: 4,
+		}.make_with(&mut seed);
+
+		let mut db = MemoryDB::new();
+		let mut root = H256::new();
+		let mut t = TrieDBMut::new(&mut db, &mut root);
+		for &(ref key, ref value) in &x {
+			assert!(t.insert(key, value).unwrap().is_none());
+			assert_eq!(t.insert(key, value).unwrap(), Some(DBValue::from_slice(value)));
+		}
+
+		for (key, value) in x {
+			assert_eq!(t.remove(&key).unwrap(), Some(DBValue::from_slice(&value)));
+			assert!(t.remove(&key).unwrap().is_none());
+		}
+	}
 }