fork tree: rebalance fork tree by max fork depth (#4605)

* fork-tree: rebalance fork tree by maximum branch height

* fork-tree: add docs for tree rebalancing

* fork-tree: add test

* babe: rebalance epoch changes tree on deserialization

* fork-tree: fix doc

* fork-tree: fix test

substrate/client/consensus/babe/src/aux_schema.rs

@@ -59,6 +59,12 @@ pub(crate) fn load_epoch_changes<Block: BlockT, B: AuxStore>(
 			SharedEpochChanges::new()
 		});
 
+	// rebalance the tree after deserialization. this isn't strictly necessary
+	// since the tree is now rebalanced on every update operation. but since the
+	// tree wasn't rebalanced initially it's useful to temporarily leave it here
+	// to avoid having to wait until an import for rebalancing.
+	epoch_changes.lock().rebalance();
+
 	Ok(epoch_changes)
 }
 

substrate/client/consensus/babe/src/epoch_changes.rs

@@ -185,6 +185,12 @@ impl<Hash, Number> EpochChanges<Hash, Number> where
 		EpochChanges { inner: ForkTree::new() }
 	}
 
+	/// Rebalances the tree of epoch changes so that it is sorted by length of
+	/// fork (longest fork first).
+	pub fn rebalance(&mut self) {
+		self.inner.rebalance()
+	}
+
 	/// Prune out finalized epochs, except for the ancestor of the finalized
 	/// block. The given slot should be the slot number at which the finalized
 	/// block was authored.