diff --git a/ethcore/src/json_tests/trie.rs b/ethcore/src/json_tests/trie.rs index 065985b0a..9107675d5 100644 --- a/ethcore/src/json_tests/trie.rs +++ b/ethcore/src/json_tests/trie.rs @@ -33,7 +33,7 @@ fn test_trie(json: &[u8], trie: TrieSpec) -> Vec { let key: Vec = key.into(); let value: Vec = 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() { diff --git a/util/src/trie/fatdbmut.rs b/util/src/trie/fatdbmut.rs index ceed14bab..1af476358 100644 --- a/util/src/trie/fatdbmut.rs +++ b/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> { 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> { let hash = key.sha3(); - self.raw.db_mut().remove(&Self::to_aux_key(&hash)); - self.raw.remove(&hash) + let out = 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) } } diff --git a/util/src/trie/mod.rs b/util/src/trie/mod.rs index 151760c8c..01c351fc1 100644 --- a/util/src/trie/mod.rs +++ b/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> where 'a: 'key; - /// Insert a `key`/`value` pair into the trie. An `empty` value is equivalent to removing - /// `key` from the trie. - fn insert(&mut self, key: &[u8], value: &[u8]) -> Result<()>; + /// Insert a `key`/`value` pair into the trie. An empty value is equivalent to removing + /// `key` from the trie. Returns the old value associated with this key, if it existed. + fn insert(&mut self, key: &[u8], value: &[u8]) -> Result>; /// Remove a `key` from the trie. Equivalent to making it equal to the empty - /// value. - fn remove(&mut self, key: &[u8]) -> Result<()>; + /// value. Returns the old value associated with this key, if it existed. + fn remove(&mut self, key: &[u8]) -> Result>; } /// A trie iterator that also supports random access. diff --git a/util/src/trie/sectriedbmut.rs b/util/src/trie/sectriedbmut.rs index ca14b2df6..ab695921f 100644 --- a/util/src/trie/sectriedbmut.rs +++ b/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> { self.raw.insert(&key.sha3(), value) } - fn remove(&mut self, key: &[u8]) -> super::Result<()> { + fn remove(&mut self, key: &[u8]) -> super::Result> { self.raw.remove(&key.sha3()) } } diff --git a/util/src/trie/triedbmut.rs b/util/src/trie/triedbmut.rs index b062cd70b..90d5867dd 100644 --- a/util/src/trie/triedbmut.rs +++ b/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) + -> 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 { + fn insert_inspector(&mut self, node: Node, partial: NibbleSlice, value: DBValue, old_val: &mut Option) + -> super::Result + { 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> { + fn remove_at(&mut self, handle: NodeHandle, partial: NibbleSlice, old_val: &mut Option) + -> super::Result> + { 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 { + fn remove_inspector(&mut self, node: Node, partial: NibbleSlice, old_val: &mut Option) -> super::Result { 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<()> { - if value.is_empty() { - return self.remove(key); - } + fn insert(&mut self, key: &[u8], value: &[u8]) -> super::Result> { + if value.is_empty() { 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> { 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()); + } + } }