diff --git a/substrate/state-machine/src/ext.rs b/substrate/state-machine/src/ext.rs index cd177b19845f6..5a8ba096f5753 100644 --- a/substrate/state-machine/src/ext.rs +++ b/substrate/state-machine/src/ext.rs @@ -121,6 +121,7 @@ impl<'a, B: 'a> Externalities for Ext<'a, B> fn clear_prefix(&mut self, prefix: &[u8]) { self.mark_dirty(); + self.overlay.clear_prefix(prefix); self.backend.for_keys_with_prefix(prefix, |key| { self.overlay.set_storage(key.to_vec(), None); }); diff --git a/substrate/state-machine/src/lib.rs b/substrate/state-machine/src/lib.rs index 5cc10027c6284..17fd85852b6ac 100644 --- a/substrate/state-machine/src/lib.rs +++ b/substrate/state-machine/src/lib.rs @@ -68,10 +68,37 @@ impl OverlayedChanges { .map(|x| x.as_ref().map(AsRef::as_ref)) } + /// Inserts the given key-value pair into the prospective change set. + /// + /// `None` can be used to delete a value specified by the given key. fn set_storage(&mut self, key: Vec, val: Option>) { self.prospective.insert(key, val); } + /// Removes all key-value pairs which keys share the given prefix. + /// + /// NOTE that this doesn't take place immediately but written into the prospective + /// change set, and still can be reverted by [`discard_prospective`]. + /// + /// [`discard_prospective`]: #method.discard_prospective + fn clear_prefix(&mut self, prefix: &[u8]) { + // Iterate over all prospective and mark all keys that share + // the given prefix as removed (None). + for (key, value) in self.prospective.iter_mut() { + if key.starts_with(prefix) { + *value = None; + } + } + + // Then do the same with keys from commited changes. + // NOTE that we are making changes in the prospective change set. + for key in self.committed.keys() { + if key.starts_with(prefix) { + self.prospective.insert(key.to_owned(), None); + } + } + } + /// Discard prospective changes to state. pub fn discard_prospective(&mut self) { self.prospective.clear(); @@ -540,4 +567,44 @@ mod tests { assert_eq!(remote_result, vec![66]); assert_eq!(remote_result, local_result); } + + #[test] + fn clear_prefix_in_ext_works() { + let initial: HashMap<_, _> = map![ + b"aaa".to_vec() => b"0".to_vec(), + b"abb".to_vec() => b"1".to_vec(), + b"abc".to_vec() => b"2".to_vec(), + b"bbb".to_vec() => b"3".to_vec() + ]; + let backend = InMemory::from(initial).try_into_trie_backend().unwrap(); + let mut overlay = OverlayedChanges { + committed: map![ + b"aba".to_vec() => Some(b"1312".to_vec()), + b"bab".to_vec() => Some(b"228".to_vec()) + ], + prospective: map![ + b"abd".to_vec() => Some(b"69".to_vec()), + b"bbd".to_vec() => Some(b"42".to_vec()) + ], + }; + + { + let mut ext = Ext::new(&mut overlay, &backend); + ext.clear_prefix(b"ab"); + } + overlay.commit_prospective(); + + assert_eq!( + overlay.committed, + map![ + b"abb".to_vec() => None, + b"abc".to_vec() => None, + b"aba".to_vec() => None, + b"abd".to_vec() => None, + + b"bab".to_vec() => Some(b"228".to_vec()), + b"bbd".to_vec() => Some(b"42".to_vec()) + ], + ); + } }