fix(lib/blocktree): fix setting leaves after blocktree pruning (Chain…

…Safe#1605)
timwu20 · May 26, 2021 · 58c0854 · 58c0854
1 parent f195204
commit 58c0854
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git 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
@@ -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) {