Make sure displaced leaves are removed from the db (#1338)

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.
+		}
 	}
 }