From 58c08541f113d524ef8868cc85f89f7e669ba570 Mon Sep 17 00:00:00 2001 From: noot <36753753+noot@users.noreply.github.com> Date: Wed, 26 May 2021 12:09:21 -0400 Subject: [PATCH] fix(lib/blocktree): fix setting leaves after blocktree pruning (#1605) --- lib/blocktree/blocktree.go | 5 ++++- lib/blocktree/blocktree_test.go | 6 ++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/lib/blocktree/blocktree.go b/lib/blocktree/blocktree.go index 274fe58a7e..9f5cc2d5e8 100644 --- a/lib/blocktree/blocktree.go +++ b/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 } diff --git a/lib/blocktree/blocktree_test.go b/lib/blocktree/blocktree_test.go index 8a8032d5ff..2e0b905613 100644 --- a/lib/blocktree/blocktree_test.go +++ b/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) {