fix(lib/blocktree): fix setting leaves after blocktree pruning (ChainSafe#1605)

Author: noot

lib/blocktree/blocktree.go

@@ -187,8 +187,11 @@ func (bt *BlockTree) Prune(finalised Hash) (pruned []Hash) {
 
 	pruned = bt.head.prune(n, nil)
 	bt.head = n
+	leaves := n.getLeaves(nil)
 	bt.leaves = newEmptyLeafMap()
-	bt.leaves.store(n.hash, n)
+	for _, leaf := range leaves {
+		bt.leaves.store(leaf.hash, leaf)
+	}
 	return pruned
 }
 

lib/blocktree/blocktree_test.go

@@ -378,6 +378,12 @@ func TestBlockTree_Prune(t *testing.T) {
 			t.Fatal("pruned an ancestor of the finalised node!!")
 		}
 	}
+
+	require.NotEqual(t, 0, len(bt.leaves.nodes()))
+	for _, leaf := range bt.leaves.nodes() {
+		require.NotEqual(t, leaf.hash, finalised.hash)
+		require.True(t, leaf.isDescendantOf(finalised))
+	}
 }
 
 func TestBlockTree_DeepCopy(t *testing.T) {