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

* return old trie values on insert and remove

* fix json tests
This commit is contained in:
Robert Habermeier 2017-01-05 21:15:43 +01:00 committed by Gav Wood
parent 20c1d37b59
commit b4e713efdc
5 changed files with 98 additions and 42 deletions

View File

@ -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() {

View File

@ -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));
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)
}
}

View File

@ -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
/// `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<Option<DBValue>>;
/// 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<Option<DBValue>>;
}
/// A trie iterator that also supports random access.

View File

@ -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())
}
}

View File

@ -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<()> {
if value.is_empty() {
return self.remove(key);
}
fn insert(&mut self, key: &[u8], value: &[u8]) -> super::Result<Option<DBValue>> {
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<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());
}
}
}