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:
parent
20c1d37b59
commit
b4e713efdc
@ -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() {
|
||||
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -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.
|
||||
|
@ -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())
|
||||
}
|
||||
}
|
||||
|
@ -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());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user