diff --git a/substrate/core/client/src/leaves.rs b/substrate/core/client/src/leaves.rs index a41f91933c07680e0bb100623c7e945cb4df143d..bc1551e6fadc60fd590fa7cc8d732acfa841a5ec 100644 --- a/substrate/core/client/src/leaves.rs +++ b/substrate/core/client/src/leaves.rs @@ -64,6 +64,8 @@ pub struct DisplacedLeaf<H, N> { #[derive(Debug, Clone, PartialEq, Eq)] pub struct LeafSet<H, N> { storage: BTreeSet<LeafSetItem<H, N>>, + pending_added: Vec<LeafSetItem<H, N>>, + pending_removed: Vec<H>, } impl<H, N> LeafSet<H, N> where @@ -73,7 +75,9 @@ impl<H, N> LeafSet<H, N> where /// Construct a new, blank leaf set. pub fn new() -> Self { Self { - storage: BTreeSet::new() + storage: BTreeSet::new(), + pending_added: Vec::new(), + pending_removed: Vec::new(), } } @@ -94,7 +98,11 @@ impl<H, N> LeafSet<H, N> where }; storage.insert(LeafSetItem { hash, number }); } - Ok(Self { storage }) + Ok(Self { + storage, + pending_added: Vec::new(), + pending_removed: Vec::new(), + }) } /// update the leaf list on import. returns a displaced leaf if there was one. @@ -102,12 +110,13 @@ impl<H, N> LeafSet<H, N> where // avoid underflow for genesis. let displaced = if number != N::zero() { let displaced = LeafSetItem { - hash: parent_hash, + hash: parent_hash.clone(), number: number.clone() - N::one(), }; let was_displaced = self.storage.remove(&displaced); if was_displaced { + self.pending_removed.push(parent_hash); Some(DisplacedLeaf { new_hash: hash.clone(), displaced, @@ -119,7 +128,9 @@ impl<H, N> LeafSet<H, N> where None }; - self.storage.insert(LeafSetItem { hash, number }); + let item = LeafSetItem { hash, number }; + self.storage.insert(item.clone()); + self.pending_added.push(item); displaced } @@ -128,6 +139,8 @@ impl<H, N> LeafSet<H, N> where let new_number = displaced.displaced.number.clone() + N::one(); self.storage.remove(&LeafSetItem { hash: displaced.new_hash, number: new_number }); self.storage.insert(displaced.displaced); + self.pending_added.clear(); + self.pending_removed.clear(); } /// currently since revert only affects the canonical chain @@ -148,13 +161,18 @@ impl<H, N> LeafSet<H, N> where } /// Write the leaf list to the database transaction. - pub fn prepare_transaction(&self, tx: &mut DBTransaction, column: Option<u32>, prefix: &[u8]) { + pub fn prepare_transaction(&mut self, tx: &mut DBTransaction, column: Option<u32>, prefix: &[u8]) { let mut buf = prefix.to_vec(); - for &LeafSetItem { ref hash, ref number } in &self.storage { + for LeafSetItem { hash, number } in self.pending_added.drain(..) { hash.using_encoded(|s| buf.extend(s)); tx.put_vec(column, &buf[..], number.encode()); buf.truncate(prefix.len()); // reuse allocation. } + for hash in self.pending_removed.drain(..) { + hash.using_encoded(|s| buf.extend(s)); + tx.delete(column, &buf[..]); + buf.truncate(prefix.len()); // reuse allocation. + } } }